"""Schemas for magnetic ordering calculations."""
from __future__ import annotations
from typing import Optional
import numpy as np
from emmet.core.tasks import TaskDoc
from pydantic import BaseModel, Field
from pymatgen.analysis.magnetism.analyzer import (
CollinearMagneticStructureAnalyzer,
Ordering,
)
from pymatgen.core.structure import Structure
[docs]
class MagneticOrderingRelaxation(BaseModel):
"""Defines the relaxation information for a magnetic ordering calculation.
This is embedded within the MagneticOrderingOutput.
"""
uuid: Optional[str] = Field(None, description="Unique ID of the calculation.")
dir_name: Optional[str] = Field(None, description="Directory of the calculation.")
input: Optional[MagneticOrderingInput] = Field(
None, description="Input ordering information."
)
structure: Optional[Structure] = Field(
None, description="Final structure from the calculation."
)
symmetry_changed: Optional[bool] = Field(
None, description="Whether or not symmetry changed."
)
ordering_changed: Optional[bool] = Field(
None,
description=(
"Specifies whether or not the magnetic ordering changed during the"
" calculation."
),
)
ordering: Optional[Ordering] = Field(
None, description="Final ordering from the calculation."
)
magmoms: Optional[list[float]] = Field(
None, description="Magnetic moments of the structure."
)
symmetry: Optional[str] = Field(None, description="Detected space group symbol.")
energy: Optional[float] = Field(
None, description="Final energy result from the calculation."
)
energy_per_atom: Optional[float] = Field(None, description="Final energy per atom.")
total_magnetization: Optional[float] = Field(
None,
description=(
"Total magnetization as a sum of individual atomic moments in "
"the calculated unit cell."
),
)
total_magnetization_per_formula_unit: Optional[float] = Field(
None, description="Total magnetization normalized to per formula unit."
)
total_magnetization_per_unit_volume: Optional[float] = Field(
None, description="Total magnetization normalized to per unit volume."
)
[docs]
@classmethod
def from_structures_and_energies(
cls,
input_structure: Structure,
output_structure: Structure,
output_energy: float,
uuid: str | None = None,
dir_name: str | None = None,
) -> MagneticOrderingRelaxation:
"""Construct a relaxation output doc from structures and energies."""
return cls(
uuid=uuid,
dir_name=dir_name,
structure=output_structure,
energy=output_energy,
energy_per_atom=output_energy / output_structure.num_sites,
**_compare_ordering_and_symmetry(input_structure, output_structure),
)
[docs]
@classmethod
def from_task_document(
cls, task_document: TaskDoc, uuid: str | None = None
) -> MagneticOrderingRelaxation:
"""Construct a MagneticOrderingRelaxation from a task document.
This does not include the uuid, which must be provided separately.
.. Warning:: Currently, the TaskDoc defined in emmet is VASP-specific.
Ensure that the TaskDoc provided contains an InputDoc with a magnetic
moments field.
"""
dir_name = task_document.dir_name
structure = task_document.structure
input_structure = task_document.input.structure
if not input_structure.site_properties.get("magmom"):
input_structure.add_site_property( # input struct likely has no magmoms
"magmom", task_document.input.magnetic_moments
)
energy = task_document.output.energy
return cls.from_structures_and_energies(
input_structure=input_structure,
output_structure=structure,
output_energy=energy,
uuid=uuid,
dir_name=dir_name,
)
[docs]
class MagneticOrderingOutput(BaseModel):
"""Defines the output for a *static* magnetic ordering calculation.
This is used within the construction of the MagneticOrderingDocument. If a
relaxation was performed, this information will be stored within the relax_output
field.
"""
uuid: Optional[str] = Field(None, description="Unique ID of the calculation.")
dir_name: Optional[str] = Field(None, description="Directory of the calculation.")
input: Optional[MagneticOrderingInput] = Field(
None, description="Input ordering information."
)
structure: Optional[Structure] = Field(
None, description="Final structure from the calculation."
)
ordering: Optional[Ordering] = Field(
None,
description=(
"The magnetic ordering of the output structure, "
"as defined in pymatgen.analysis.magnetism.analyzer."
),
)
magmoms: Optional[list[float]] = Field(
None, description="Magnetic moments of the structure."
)
symmetry: Optional[str] = Field(None, description="Detected space group symbol.")
energy: Optional[float] = Field(
None, description="Final energy result from the calculation."
)
energy_per_atom: Optional[float] = Field(None, description="Final energy per atom.")
total_magnetization: Optional[float] = Field(
None,
description=(
"Total magnetization as a sum of individual atomic moments in "
"the calculated unit cell."
),
)
total_magnetization_per_formula_unit: Optional[float] = Field(
None, description="Total magnetization normalized to per formula unit."
)
total_magnetization_per_unit_volume: Optional[float] = Field(
None, description="Total magnetization normalized to per unit volume."
)
ordering_changed: Optional[bool] = Field(
None,
description=(
"Specifies whether or not the magnetic ordering changed during the"
" calculation."
),
)
symmetry_changed: Optional[bool] = Field(
None,
description=(
"Specifies whether or not the symmetry changed during the calculation."
),
)
energy_above_ground_state_per_atom: Optional[float] = Field(
None, description="Energy per atom above the calculated ground state ordering."
)
relax_output: Optional[MagneticOrderingRelaxation] = Field(
None, description="Relaxation output, if relaxation performed."
)
energy_diff_relax_static: Optional[float] = Field(
None,
description=(
"Difference in energy between relaxation and final static calculation, if"
" relaxation performed (useful for benchmarking). Specifically, this is"
" calculated as energy[static] - energy[relax]."
),
)
[docs]
@classmethod
def from_structures_and_energies(
cls,
input_structure: Structure,
output_structure: Structure,
output_energy: float,
relax_output: MagneticOrderingRelaxation | None = None,
uuid: str | None = None,
dir_name: str | None = None,
ground_state_energy_per_atom: float | None = None,
) -> MagneticOrderingOutput:
"""Construct a MagneticOrderingOutput doc from structures and energies."""
energy_diff_relax_static = (
relax_output.energy - output_energy if relax_output else None
)
output_energy_per_atom = output_energy / output_structure.num_sites
energy_above_ground_state_per_atom = (
output_energy_per_atom - ground_state_energy_per_atom
if ground_state_energy_per_atom
else None
)
return cls(
uuid=uuid,
dir_name=dir_name,
structure=output_structure,
energy=output_energy,
relax_output=relax_output,
energy_diff_relax_static=energy_diff_relax_static,
energy_per_atom=output_energy_per_atom,
energy_above_ground_state_per_atom=energy_above_ground_state_per_atom,
**_compare_ordering_and_symmetry(input_structure, output_structure),
)
[docs]
@classmethod
def from_task_document(
cls,
task_document: TaskDoc,
uuid: str | None = None,
relax_output: MagneticOrderingRelaxation | None = None,
) -> MagneticOrderingOutput:
"""Construct a MagneticOrderingOutput from a task document.
This does not include the uuid, which must be set separately.
.. Warning:: Currently, the TaskDoc defined in emmet is VASP-specific.
Ensure that the TaskDoc provided contains an InputDoc with a
magnetic moments field.
"""
dir_name = task_document.dir_name
structure = task_document.structure
input_structure = task_document.input.structure
if not input_structure.site_properties.get("magmom"):
input_structure.add_site_property( # input struct likely has no magmoms
"magmom", task_document.input.magnetic_moments
)
energy = task_document.output.energy
return cls.from_structures_and_energies(
input_structure=input_structure,
output_structure=structure,
output_energy=energy,
uuid=uuid,
dir_name=dir_name,
relax_output=relax_output,
)
[docs]
class MagneticOrderingsDocument(BaseModel):
"""Final document containing information about calculated magnetic orderings.
Includes description of the ground state ordering. This document is returned by the
MagneticOrderingsBuilder corresponding to your DFT code.
"""
formula: Optional[str] = Field(
None,
description="Formula taken from pymatgen.core.structure.Structure.formula.",
)
formula_pretty: Optional[str] = Field(
None,
description="Cleaned representation of the formula",
)
parent_structure: Optional[Structure] = Field(
None,
description=(
"The parent structure from which individual magnetic "
"orderings are generated."
),
)
outputs: Optional[list[MagneticOrderingOutput]] = Field(
None,
description="All magnetic ordering calculation results for this structure.",
)
ground_state_uuid: Optional[str] = Field(
None, description="UUID of the ground state ordering."
)
ground_state_structure: Optional[Structure] = Field(
None, description="Ground state structure."
)
ground_state_ordering: Optional[Ordering] = Field(
None, description="Ground state magnetic ordering."
)
ground_state_energy: Optional[float] = Field(
None, description="Ground state energy."
)
ground_state_energy_per_atom: Optional[float] = Field(
None, description="Ground state energy, normalized per atom."
)
[docs]
@classmethod
def from_outputs(
cls,
outputs: list[MagneticOrderingOutput],
parent_structure: Structure,
) -> MagneticOrderingsDocument:
"""Construct a MagneticOrderingDocument from a list of output docs.
This is general and should not need to be implemented for a specific DFT code.
"""
formula = outputs[0].structure.formula
formula_pretty = outputs[0].structure.composition.reduced_formula
ground_state = min(outputs, key=lambda struct: struct.energy_per_atom)
ground_state_energy_per_atom = ground_state.energy_per_atom
for output in outputs:
output.energy_above_ground_state_per_atom = (
output.energy_per_atom - ground_state_energy_per_atom
)
return cls(
formula=formula,
formula_pretty=formula_pretty,
parent_structure=parent_structure,
outputs=outputs,
ground_state_uuid=ground_state.uuid,
ground_state_structure=ground_state.structure,
ground_state_ordering=ground_state.ordering,
ground_state_energy=ground_state.energy,
ground_state_energy_per_atom=ground_state.energy_per_atom,
)
[docs]
@classmethod
def from_tasks(cls, tasks: list[dict]) -> MagneticOrderingsDocument:
"""Construct a MagneticOrderingsDocument from a list of task dicts.
.. Note::
this function assumes the tasks contain the keys "output" and "metadata".
These keys are automatically constructed when jobflow stores its outputs;
however, you may need to put the data in this format if using this
manually (as in a postprocessing job).
"""
parent_structure = tasks[0]["metadata"]["parent_structure"]
relax_tasks, static_tasks = [], []
for task in tasks:
if task["output"].task_type.value.lower() == "structure optimization":
relax_tasks.append(task)
elif task["output"].task_type.value.lower() == "static":
static_tasks.append(task)
outputs = []
for task in static_tasks:
relax_output = None
for r_task in relax_tasks:
if r_task["uuid"] == task["metadata"]["parent_uuid"]:
relax_output = MagneticOrderingRelaxation.from_task_document(
r_task["output"],
uuid=r_task["uuid"],
)
break
output = MagneticOrderingOutput.from_task_document(
task["output"],
uuid=task["uuid"],
relax_output=relax_output,
)
outputs.append(output)
return cls.from_outputs(outputs, parent_structure=parent_structure)
def _compare_ordering_and_symmetry(
input_structure: Structure, output_structure: Structure
) -> dict:
"""Compare ordering and symmetry of input and output structures.
This is especially useful for debugging purposes.
"""
# process input structure
input_analyzer = CollinearMagneticStructureAnalyzer(input_structure, threshold=0.61)
input_ordering = input_analyzer.ordering
input_magmoms = input_analyzer.magmoms
input_symmetry = input_structure.get_space_group_info()[0]
# process output structure
output_analyzer = CollinearMagneticStructureAnalyzer(
output_structure, threshold=0.61
)
output_ordering = output_analyzer.ordering
output_magmoms = output_analyzer.magmoms
output_symmetry = output_structure.get_space_group_info()[0]
total_magnetization = output_analyzer.total_magmoms
num_formula_units = (
output_structure.composition.get_reduced_composition_and_factor()[1]
)
total_magnetization_per_formula_unit = total_magnetization / num_formula_units
total_magnetization_per_unit_volume = total_magnetization / output_structure.volume
# compare
ordering_changed = not np.array_equal(
np.sign(input_analyzer.magmoms), np.sign(output_magmoms)
)
symmetry_changed = output_symmetry != input_symmetry
input_doc = MagneticOrderingInput(
structure=input_structure,
ordering=input_ordering,
magmoms=list(input_magmoms),
symmetry=input_symmetry,
)
return {
"input": input_doc,
"magmoms": list(output_magmoms),
"ordering": output_ordering,
"symmetry": output_symmetry,
"ordering_changed": ordering_changed,
"symmetry_changed": symmetry_changed,
"total_magnetization": total_magnetization,
"total_magnetization_per_formula_unit": total_magnetization_per_formula_unit,
"total_magnetization_per_unit_volume": total_magnetization_per_unit_volume,
}