Aims magnetic ordering (#922)

* Add tests for magnetism

The Base FHI-aims schemas will also have to be modified

* ADd magnetism job for FHI-aims

* Modify the Aims Schemas to run with magnetism workflows

* Add magnetic ordering support for FHI-aims

This is approximate as magnetic properties in aims are only done approximately

* Add magnetic ordering jobs

Facilitates calculations via making default InputSet Generators

* Update to pymatgen 2024.9.10 for new aims interface

* fix lint error

* Default phonon tests no longer needed with latest pymatge

* Remove properties in ASE jobs test

output.structure.properties is filled, fcc_ne_structure does not
the equality assertion is therefore false

* Fix ASE jobs test

* Update src/atomate2/aims/schemas/calculation.py

Co-authored-by: J. George <JaGeo@users.noreply.github.com>

---------

Co-authored-by: J. George <JaGeo@users.noreply.github.com>

src/atomate2/aims/jobs/magnetism.py

@@ -0,0 +1,78 @@
+"""Define Makers for Magnetic ordering flow in FHI-aims."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from pymatgen.io.aims.sets.magnetism import (
+    MagneticRelaxSetGenerator,
+    MagneticStaticSetGenerator,
+)
+
+if TYPE_CHECKING:
+    from pymatgen.io.aims.sets.base import AimsInputGenerator
+
+
+from atomate2.aims.jobs.core import RelaxMaker, StaticMaker
+
+
+@dataclass
+class MagneticStaticMaker(StaticMaker):
+    """Maker to create FHI-aims SCF jobs.
+
+    Parameters
+    ----------
+    calc_type: str
+        The type key for the calculation
+    name: str
+        The job name
+    input_set_generator: .AimsInputGenerator
+        The InputGenerator for the calculation
+    """
+
+    calc_type: str = "magnetic_scf"
+    name: str = "Magnetic SCF Calculation"
+    input_set_generator: AimsInputGenerator = field(
+        default_factory=MagneticStaticSetGenerator
+    )
+
+
+@dataclass
+class MagneticRelaxMaker(RelaxMaker):
+    """Maker to create relaxation calculations.
+
+    Parameters
+    ----------
+    calc_type: str
+        The type key for the calculation
+    name: str
+        The job name
+    input_set_generator: .AimsInputGenerator
+        The InputGenerator for the calculation
+    """
+
+    calc_type: str = "relax"
+    name: str = "Magnetic Relaxation calculation"
+    input_set_generator: AimsInputGenerator = field(
+        default_factory=MagneticRelaxSetGenerator
+    )
+
+    @classmethod
+    def fixed_cell_relaxation(cls, *args, **kwargs) -> RelaxMaker:
+        """Create a fixed cell relaxation maker."""
+        return cls(
+            input_set_generator=MagneticRelaxSetGenerator(
+                *args, relax_cell=False, **kwargs
+            ),
+            name=cls.name + " (fixed cell)",
+        )
+
+    @classmethod
+    def full_relaxation(cls, *args, **kwargs) -> RelaxMaker:
+        """Create a full relaxation maker."""
+        return cls(
+            input_set_generator=MagneticRelaxSetGenerator(
+                *args, relax_cell=True, **kwargs
+            )
+        )

src/atomate2/aims/schemas/calculation.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import json
 import os
 from collections.abc import Sequence
 from datetime import datetime, timezone
@@ -15,6 +16,7 @@
 from pymatgen.core import Molecule, Structure
 from pymatgen.core.trajectory import Trajectory
 from pymatgen.electronic_structure.dos import Dos
+from pymatgen.io.aims.inputs import AimsGeometryIn
 from pymatgen.io.aims.outputs import AimsOutput
 from pymatgen.io.common import VolumetricData
 from typing_extensions import Self
@@ -185,6 +187,25 @@ def from_aims_output(
         )
 
 
+class CalculationInput(BaseModel):
+    """The FHI-aims Calculation input doc.
+
+    Parameters
+    ----------
+    structure: Structure or Molecule
+        The input pymatgen Structure or Molecule of the system
+    parameters: dict[str, Any]
+        The parameters passed in the control.in file
+    """
+
+    structure: Union[Structure, Molecule] = Field(
+        None, description="The input structure object"
+    )
+    parameters: dict[str, Any] = Field(
+        {}, description="The input parameters for FHI-aims"
+    )
+
+
 class Calculation(BaseModel):
     """Full FHI-aims calculation inputs and outputs.
 
@@ -198,6 +219,8 @@ class Calculation(BaseModel):
         Whether FHI-aims completed the calculation successfully
     output: .CalculationOutput
         The FHI-aims calculation output
+    input: .CalculationInput
+        The FHI-aims calculation input
     completed_at: str
         Timestamp for when the calculation was completed
     output_file_paths: Dict[str, str]
@@ -214,6 +237,10 @@ class Calculation(BaseModel):
     has_aims_completed: TaskState = Field(
         None, description="Whether FHI-aims completed the calculation successfully"
     )
+    completed: bool = Field(
+        None, description="Whether FHI-aims completed the calculation successfully"
+    )
+    input: CalculationInput = Field(None, description="The FHI-aims calculation input")
     output: CalculationOutput = Field(
         None, description="The FHI-aims calculation output"
     )
@@ -236,7 +263,6 @@ def from_aims_files(
         parse_dos: str | bool = False,
         parse_bandstructure: str | bool = False,
         store_trajectory: bool = False,
-        # store_scf: bool = False,
         store_volumetric_data: Optional[Sequence[str]] = STORE_VOLUMETRIC_DATA,
     ) -> tuple[Self, dict[AimsObject, dict]]:
         """Create an FHI-aims calculation document from a directory and file paths.
@@ -289,6 +315,15 @@ def from_aims_files(
         aims_output_file = dir_name / aims_output_file
 
         volumetric_files = [] if volumetric_files is None else volumetric_files
+
+        aims_geo_in = AimsGeometryIn.from_file(dir_name / "geometry.in")
+        aims_parameters = {}
+        with open(str(dir_name / "parameters.json")) as pj:
+            aims_parameters = json.load(pj)
+
+        input_doc = CalculationInput(
+            structure=aims_geo_in.structure, parameters=aims_parameters
+        )
         aims_output = AimsOutput.from_outfile(aims_output_file)
 
         completed_at = str(
@@ -323,8 +358,10 @@ def from_aims_files(
             task_name=task_name,
             aims_version=aims_output.aims_version,
             has_aims_completed=has_aims_completed,
+            completed=has_aims_completed == TaskState.SUCCESS,
             completed_at=completed_at,
             output=output_doc,
+            input=input_doc,
             output_file_paths={k.name.lower(): v for k, v in output_file_paths.items()},
         )
 

src/atomate2/aims/schemas/task.py

@@ -11,7 +11,9 @@
 import numpy as np
 from emmet.core.math import Matrix3D, Vector3D
 from emmet.core.structure import MoleculeMetadata, StructureMetadata
+from emmet.core.task import BaseTaskDocument
 from emmet.core.tasks import get_uri
+from emmet.core.vasp.calc_types.enums import TaskType
 from pydantic import BaseModel, Field
 from pymatgen.core import Molecule, Structure
 from pymatgen.entries.computed_entries import ComputedEntry
@@ -78,56 +80,6 @@ def from_aims_calc_docs(cls, calc_docs: list[Calculation]) -> Self:
         )
 
 
-class Species(BaseModel):
-    """A representation of the most important information about each type of species.
-
-    Parameters
-    ----------
-    element: str
-        Element assigned to this atom kind
-    species_defaults: str
-        Basis set for this atom kind
-    """
-
-    element: str = Field(None, description="Element assigned to this atom kind")
-    species_defaults: str = Field(None, description="Basis set for this atom kind")
-
-
-class SpeciesSummary(BaseModel):
-    """A summary of species defaults.
-
-    Parameters
-    ----------
-    species_defaults: Dict[str, .Species]
-        Dictionary mapping atomic kind labels to their info
-    """
-
-    species_defaults: dict[str, Species] = Field(
-        None, description="Dictionary mapping atomic kind labels to their info"
-    )
-
-    @classmethod
-    def from_species_info(cls, species_info: dict[str, dict[str, Any]]) -> Self:
-        """Initialize from the atomic_kind_info dictionary.
-
-        Parameters
-        ----------
-        species_info: Dict[str, Dict[str, Any]]
-            The information for the basis set for the calculation
-
-        Returns
-        -------
-        The SpeciesSummary
-        """
-        dct: dict[str, dict[str, Any]] = {"species_defaults": {}}
-        for kind, info in species_info.items():
-            dct["species_defaults"][kind] = {
-                "element": info["element"],
-                "species_defaults": info["species_defaults"],
-            }
-        return cls(**dct)
-
-
 class InputDoc(BaseModel):
     """Summary of inputs for an FHI-aims calculation.
 
@@ -137,19 +89,24 @@ class InputDoc(BaseModel):
         The input pymatgen Structure or Molecule of the system
     species_info: .SpeciesSummary
         Summary of the species defaults used for each atom kind
+    parameters: dict[str, Any]
+        The parameters passed in the control.in file
     xc: str
         Exchange-correlation functional used if not the default
     """
 
     structure: Union[Structure, Molecule] = Field(
         None, description="The input structure object"
     )
-    species_info: SpeciesSummary = Field(
-        None, description="Summary of the species defaults used for each atom kind"
+    parameters: dict[str, Any] = Field(
+        {}, description="The input parameters for FHI-aims"
     )
     xc: str = Field(
         None, description="Exchange-correlation functional used if not the default"
     )
+    magnetic_moments: Optional[list[float]] = Field(
+        None, description="Magnetic moments for each atom"
+    )
 
     @classmethod
     def from_aims_calc_doc(cls, calc_doc: Calculation) -> Self:
@@ -165,12 +122,16 @@ def from_aims_calc_doc(cls, calc_doc: Calculation) -> Self:
         .InputDoc
             A summary of the input structure and parameters.
         """
-        summary = SpeciesSummary.from_species_info(calc_doc.input.species_info)
+        structure = calc_doc.input.structure
+        magnetic_moments = None
+        if "magmom" in structure.site_properties:
+            magnetic_moments = structure.site_properties["magmom"]
 
         return cls(
-            structure=calc_doc.input.structure,
-            atomic_kind_info=summary,
-            xc=str(calc_doc.run_type),
+            structure=structure,
+            parameters=calc_doc.input.parameters,
+            xc=calc_doc.input.parameters["xc"],
+            magnetic_moments=magnetic_moments,
         )
 
 
@@ -383,15 +344,19 @@ def from_data(
         )
 
 
-class AimsTaskDoc(StructureMetadata, MoleculeMetadata):
+class AimsTaskDoc(BaseTaskDocument, StructureMetadata, MoleculeMetadata):
     """Definition of FHI-aims task document.
 
     Parameters
     ----------
+    calc_code: str
+        The calculation code used to compute the task
     dir_name: str
         The directory for this FHI-aims task
     last_updated: str
         Timestamp for this task document was last updated
+    completed: bool
+        Whether this calculation completed
     completed_at: str
         Timestamp for when this task was completed
     input: .InputDoc
@@ -430,11 +395,13 @@ class AimsTaskDoc(StructureMetadata, MoleculeMetadata):
         Additional json loaded from the calculation directory
     """
 
+    calc_code: str = "aims"
     dir_name: str = Field(None, description="The directory for this FHI-aims task")
     last_updated: str = Field(
         default_factory=datetime_str,
         description="Timestamp for this task document was last updated",
     )
+    completed: bool = Field(None, description="Whether this calculation completed")
     completed_at: str = Field(
         None, description="Timestamp for when this task was completed"
     )
@@ -553,7 +520,9 @@ def from_directory(
             "calcs_reversed": calcs_reversed,
             "analysis": analysis,
             "tags": tags,
+            "completed": calcs_reversed[-1].completed,
             "completed_at": calcs_reversed[-1].completed_at,
+            "input": InputDoc.from_aims_calc_doc(calcs_reversed[-1]),
             "output": OutputDoc.from_aims_calc_doc(calcs_reversed[-1]),
             "state": _get_state(calcs_reversed, analysis),
             "entry": cls.get_entry(calcs_reversed),
@@ -597,6 +566,17 @@ def get_entry(
         }
         return ComputedEntry.from_dict(entry_dict)
 
+    # TARP: This is done because the mangnetism schema assume that VASP
+    #       TaskTypes are used. I think this should be changed, but that
+    #       would require modifications in emmet
+    @property
+    def task_type(self) -> TaskType:
+        """Get the task type of the calculation."""
+        if "Relaxation calculation" in self.task_label:
+            return TaskType("Structure Optimization")
+
+        return TaskType("Static")
+
 
 def _find_aims_files(
     path: Path | str,

src/atomate2/common/jobs/magnetism.py

@@ -79,7 +79,6 @@ def enumerate_magnetic_orderings(
         truncate_by_symmetry=truncate_by_symmetry,
         transformation_kwargs=transformation_kwargs,
     )
-
     return enumerator.ordered_structures, enumerator.ordered_structure_origins