clean up magnetic ordering wf files

src/atomate2/common/builders/magnetism.py

@@ -9,10 +9,12 @@
 from monty.serialization import MontyDecoder
 from pymatgen.analysis.structure_matcher import StructureMatcher
 
+from atomate2.common.schemas.magnetism import MagneticOrderingsDocument
+
 if TYPE_CHECKING:
-    from maggma.core import Store
+    from collections.abc import Iterator
 
-    from atomate2.common.schemas.magnetism import MagneticOrderingsDocument
+    from maggma.core import Store
 
 
 class MagneticOrderingsBuilder(Builder):
@@ -21,13 +23,8 @@ class MagneticOrderingsBuilder(Builder):
     This job will process the output documents of the calculations and return new
     documents with relevant parameters (such as the total magnetization, whether the
     ordering changed, whether the particular ordering is the ground state, etc.). This
-    is especially useful for performing postprocessing of magnetic ordering calculations
-
-    .. Note::
-        This builder can be trivially implemented for your DFT code of choice by
-        adding functionality for constructing the static/relax outputs (see schema for
-        details) and implementing the 1) _build_relax_output, 2) _build_static_output,
-        and 3) _dft_code_query methods.
+    is especially useful for performing postprocessing of magnetic ordering
+    calculations.
 
     Parameters
     ----------
@@ -37,8 +34,12 @@ class MagneticOrderingsBuilder(Builder):
         Store for magnetic ordering documents.
     query : dict
         Dictionary query to limit tasks to be analyzed.
-    structure_match_tol : float
-        Numerical tolerance for structure equivalence.
+    structure_match_stol : float
+        Numerical site tolerance for structure equivalence. Default is 0.3.
+    structure_match_ltol : float
+        Numerical length tolerance for structure equivalence. Default is 0.3
+    structure_match_angle_tol : float
+        Numerical  angle tolerance in degrees for structure equivalence. Default is 5.
     **kwargs : dict
         Keyword arguments that will be passed to the Builder init.
     """
@@ -48,11 +49,11 @@ def __init__(
         tasks: Store,
         magnetic_orderings: Store,
         query: dict = None,
-        structure_match_stol: float = 0.2,
-        structure_match_ltol: float = 0.3,
+        structure_match_stol: float = 0.3,
+        structure_match_ltol: float = 0.2,
         structure_match_angle_tol: float = 5,
         **kwargs,
-    ):
+    ) -> None:
         self.tasks = tasks
         self.magnetic_orderings = magnetic_orderings
         self.query = query if query else {}
@@ -70,7 +71,7 @@ def ensure_indexes(self) -> None:
         self.tasks.ensure_index("last_updated")
         self.magnetic_orderings.ensure_index("last_updated")
 
-    def get_items(self):
+    def get_items(self) -> Iterator[list[dict]]:
         """Get all items to process into magnetic ordering documents.
 
         This step does a first grouping by formula (which is fast) and then the magnetic
@@ -87,7 +88,6 @@ def get_items(self):
         self.ensure_indexes()
 
         criteria = dict(self.query)
-        criteria.update(self._dft_code_query)
         criteria.update(
             {
                 "metadata.ordering": {"$exists": True},
@@ -102,7 +102,7 @@ def get_items(self):
         for n_formula, (keys, docs) in enumerate(results):
             formula = keys["output"]["formula_pretty"]
             self.logger.debug(
-                f"Getting {formula} (Formula {n_formula + 1} of {num_formulas})"
+                "Getting %s (Formula %d of %d)", formula, n_formula + 1, num_formulas
             )
             decoded_docs = MontyDecoder().process_decoded(docs)
             grouped_tasks = _group_orderings(
@@ -114,8 +114,11 @@ def get_items(self):
             n_groups = len(grouped_tasks)
             for n_group, group in enumerate(grouped_tasks):
                 self.logger.debug(
-                    f"Found {len(group)} tasks for {formula} (Parent structure"
-                    f" {n_group+1} of {n_groups})"
+                    "Found %d tasks for %s (Parent structure %d of %d)",
+                    len(group),
+                    formula,
+                    n_group + 1,
+                    n_groups,
                 )
                 yield group
 
@@ -137,13 +140,13 @@ def process_item(self, tasks: list[dict]) -> list[MagneticOrderingsDocument]:
             A list of magnetic ordering documents (one for each unique parent
             structure).
         """
-        self.logger.debug(f"Processing {tasks[0]['output'].formula_pretty}")
+        self.logger.debug("Processing %s", tasks[0]["output"].formula_pretty)
 
         if not tasks:
             return []
 
         return jsanitize(
-            self._build_doc_fn(tasks).dict(),
+            MagneticOrderingsDocument.from_tasks(tasks).model_dump(),
             allow_bson=True,
         )
 
@@ -155,31 +158,9 @@ def update_targets(self, items: list[MagneticOrderingsDocument]) -> None:
         items : list of .MagneticOrderingsDocument
             A list of magnetic ordering documents to add to the database.
         """
-        self.logger.info(f"Updating {len(items)} magnetic orderings documents")
+        self.logger.info("Updating %s magnetic orderings documents", len(items))
         self.magnetic_orderings.update(items, key="ground_state_uuid")
 
-    @staticmethod
-    def _build_doc_fn(tasks):
-        """Build a magnetic ordering document from a list of tasks.
-
-        This is to be implemented for the DFT code of choice. The implementation can be
-        completed by fully implementing the MagneticOrderingsDocument class for your DFT
-        code and then wrapping the MagneticOrderingsDocument.from_tasks method.
-        """
-        raise NotImplementedError
-
-    @property
-    def _dft_code_query(self):
-        """Return a query criterion for querying the tasks collection for tasks.
-
-        This is to be implemented for the DFT code of choice. This prevents mixing of
-        tasks from different DFT codes in the same builder.
-
-        For example, for VASP, this query is:
-            {"output.orig_inputs.incar": {"$exists": True}}
-        """
-        raise NotImplementedError
-
 
 def _group_orderings(
     tasks: list[dict], ltol: float, stol: float, angle_tol: float

src/atomate2/common/flows/magnetism.py

@@ -37,7 +37,15 @@ class MagneticOrderingsMaker(Maker, ABC):
     total energy. The lowest energy ordering is the predicted ground-state collinear
     ordering.
 
-    This workflow was benchmarked with VASP for a wide range of test materials. It was
+    This Maker can be trivially implemented for your DFT code of choice by
+    utilizing the appropriate static/relax makers. However, for postprocessing to
+    work correctly, one must ensure that the calculation outputs can be processed into a
+    MagneticOrderingsDocument. As the TaskDoc class is currently defined only for
+    VASP, one should ensure that the task document returned during their DFT runs
+    contains the necessary parameters (e.g., TaskDoc.input.magnetic_moments).
+    This warning will be removed once a universal TaskDoc is implemented (Issue #741).
+
+    This workflow was benchmarked with VASP for a wide range of test materials and
     originally implemented in atomate (v1) for VASP as the MagneticOrderingsWF. Please
     refer to the following paper for more information and cite appropriately:
 
@@ -46,15 +54,6 @@ class MagneticOrderingsMaker(Maker, ABC):
         Functional Theory. npj Computational Materials 5, 64 (2019).
         https://doi.org/10.1038/s41524-019-0199-7
 
-    .. Warning::
-        This Maker can be trivially implemented for your DFT code of choice by
-        utilizing the appropriate static/relax makers. However, for postprocessing to
-        work, one must ensure that the calculation outputs can be processed into a
-        MagneticOrderingsDocument. As the TaskDoc class is currently defined for VASP,
-        one should ensure that the corresponding fields are present in the output
-        documents of their DFT code. For example: the TaskDoc.input.magnetic_moments
-        field should be present.
-
     .. Note::
         Good performance of this workflow is ultimately dependent on an appropriate
         choice of Hubbard U, Hund J values and/or the functional. The defaults will work
@@ -65,11 +64,11 @@ class MagneticOrderingsMaker(Maker, ABC):
     name : str
         Name of the flows produced by this Maker.
     static_maker : Maker
-        Maker used to perform static calculations for total energy. The default DFT code
-        used is VASP with atomate2.vasp.jobs.StaticMaker.
+        Maker used to perform static calculations for total energy. VASP is selected as
+        the default DFT code (atomate2.vasp.jobs.StaticMaker).
     relax_maker : Maker | None
-        Maker used to perform relaxations of the enumerated structures. The default DFT
-        code used is VASP with atomate2.vasp.jobs.RelaxMaker. If this field is None,
+        Maker used to perform relaxations of the enumerated structures. VASP is selected
+        as the default DFT code (atomate2.vasp.jobs.RelaxMaker). If this field is None,
         relaxations will be skipped (i.e., only static calculations are performed).
     default_magmoms : dict | None
         Optional default mapping of magnetic elements to their initial magnetic moments
@@ -91,6 +90,7 @@ class MagneticOrderingsMaker(Maker, ABC):
         None.
     """
 
+    name: str = "magnetic_orderings"
     static_maker: Maker = field(
         default_factory=lambda: StaticMaker(
             input_set_generator=StaticSetGenerator(user_incar_settings={"EDIFF": 1e-7})
@@ -111,7 +111,6 @@ class MagneticOrderingsMaker(Maker, ABC):
     automatic: bool = True
     truncate_by_symmetry: bool = True
     transformation_kwargs: dict | None = None
-    name: str = "magnetic_orderings"
 
     def __post_init__(self) -> None:
         """Ensure that the static and relax makers come from the same base maker.