Add magnetic orderings workflow (#432)

* Initial sketch of `MagneticOrderingsMaker`

* Add schema to match original atomate `MagneticOrderingsToDb` format

* move magnetism jobs to common

* move magnetism flows to common

* clean up magnetic ordering documentation/typing

* make "prev calc dir" a code-dependent argument name for static maker

* clean up typing, imports, docstrings

* begin porting over magnetic workflow analysis code / schemas

* fix jobflow logic for magnetic enumerator flow

* rework for loop for orderings

* migrate magnetic orderings wf to builder format

* Improve naming

* remove output storing

* add magnetic ordering builder implementation

* remove unused imports in magnetism jobs

* fix bug in task creation in magnetism builder

* update magnetic ordering job documentation

* make sure name "magnetic orderings" is consistent

* make sure name "magnetic orderings" is consistent pt. 2

* move magnetic orderings builder into common

* ensure copying of magmoms between relax and static; linting

* fix copy magmom

* fix kwarg location (mistake previously)

* add final False to copy vasp magmoms

* debug test of magmoms

* fix bug in magmoms

* delete unused line

* move magnetic ordering & symmetry comparison logic to common

* linting

* linting pt 2

* add from_tasks() support; extract logic to common

* add and fix postprocessing job

* linting

* fix error in structure matching in builder

It's actually dangerous to match by lattice/site matrices; sometimes the structures are not identically defined despite being the same structure

* add test for magnetic ordering workflow

* Add warning if relax_maker not provided; stricter EDIFF settings

* Add enumlib install and Cobalt warning

* move Co warning

* add Optional to magnetism schemas

* ignore enumlib's tests

* make testing directory explicit

* regenerate test data for mag_ordering workflow

* fix broken test (due to new EDIFF=1e-7 data)

* remove the patch fix we had for prev_dir now that atomate2 is updated

* fix breaking python 3.8 test

* move imports and fix typing problem

* try making tests dir not explicit

* change dir in github action to move enumlib outside testing scope

* update @mattmcdermott contributor info

* attempt to fix coverage CI bug

* monty decoding for magnetism builder

* fully implement magnetic_orderings workflow into common

* add warning about TaskDoc

* bump emmet-core version

* clean up magnetic ordering wf files

* bump emmet-core again

---------

Co-authored-by: Matthew Horton <[email protected]>
Co-authored-by: J. George <[email protected]>

.github/workflows/testing.yml

@@ -40,6 +40,20 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
+      - name: Install enumlib
+        run: |
+          cd ..
+          git clone --recursive https://github.com/msg-byu/enumlib.git
+          cd enumlib/symlib/src
+          export F90=gfortran
+          make
+          cd ../../src
+          make enum.x
+          sudo mv enum.x /usr/local/bin/
+          cd ..
+          sudo cp aux_src/makeStr.py /usr/local/bin/
+        continue-on-error: true # This is not critical to succeed.
+
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip

docs/about/contributors.md

@@ -90,8 +90,8 @@ Lawrence Berkeley National Laboratory
 [0000-0003-3439-4856]: https://orcid.org/0000-0003-3439-4856
 
 **Matthew McDermott** [![gh]][mattmcdermott] [![orc]][0000-0002-4071-3000] \
-PhD student \
-University of California, Berkeley
+Postdoctoral Researcher \
+Lawrence Berkeley National Laboratory
 
 [mattmcdermott]: https://github.com/mattmcdermott
 [0000-0002-4071-3000]: https://orcid.org/0000-0002-4071-3000

pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
     "PyYAML",
     "click",
     "custodian>=2024.4.18",
-    "emmet-core>=0.78.0rc4",
+    "emmet-core==0.82.2",
     "jobflow>=0.1.11",
     "monty>=2024.2.2",
     "numpy",
@@ -131,7 +131,7 @@ ignore_missing_imports = true
 no_strict_optional = true
 
 [tool.pytest.ini_options]
-addopts = "-p no:warnings --import-mode=importlib"
+addopts = "-p no:warnings --import-mode=importlib --cov-config=pyproject.toml"
 filterwarnings = [
     "ignore:.*POTCAR.*:UserWarning",
     "ignore:.*input structure.*:UserWarning",

src/atomate2/common/builders/__init__.py

@@ -0,0 +1 @@
+"""DFT code agnostic builders."""

src/atomate2/common/builders/magnetism.py

@@ -0,0 +1,210 @@
+"""Module defining DFT code agnostic magnetic orderings builder."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from emmet.core.utils import jsanitize
+from maggma.builders import Builder
+from monty.serialization import MontyDecoder
+from pymatgen.analysis.structure_matcher import StructureMatcher
+
+from atomate2.common.schemas.magnetism import MagneticOrderingsDocument
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from maggma.core import Store
+
+
+class MagneticOrderingsBuilder(Builder):
+    """Builder to analyze the results of magnetic orderings calculations.
+
+    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.
+
+    Parameters
+    ----------
+    tasks : .Store
+        Store of task documents.
+    magnetic_orderings : .Store
+        Store for magnetic ordering documents.
+    query : dict
+        Dictionary query to limit tasks to be analyzed.
+    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.
+    """
+
+    def __init__(
+        self,
+        tasks: Store,
+        magnetic_orderings: Store,
+        query: dict = None,
+        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 {}
+        self.structure_match_stol = structure_match_stol
+        self.structure_match_ltol = structure_match_ltol
+        self.structure_match_angle_tol = structure_match_angle_tol
+
+        self.kwargs = kwargs
+
+        super().__init__(sources=[tasks], targets=[magnetic_orderings], **kwargs)
+
+    def ensure_indexes(self) -> None:
+        """Ensure indices on the tasks and magnetic orderings collections."""
+        self.tasks.ensure_index("output.formula_pretty")
+        self.tasks.ensure_index("last_updated")
+        self.magnetic_orderings.ensure_index("last_updated")
+
+    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
+        orderings are grouped by parent structure.
+
+        Yields
+        ------
+        list of dict
+            A list of magnetic ordering relaxation or static task outputs, grouped by
+            formula.
+        """
+        self.logger.info("Magnetic orderings builder started")
+        self.logger.debug("Adding/ensuring indices...")
+        self.ensure_indexes()
+
+        criteria = dict(self.query)
+        criteria.update(
+            {
+                "metadata.ordering": {"$exists": True},
+            }
+        )
+        self.logger.info("Grouping by formula...")
+        num_formulas = len(
+            self.tasks.distinct("output.formula_pretty", criteria=criteria)
+        )
+        results = self.tasks.groupby("output.formula_pretty", criteria=criteria)
+
+        for n_formula, (keys, docs) in enumerate(results):
+            formula = keys["output"]["formula_pretty"]
+            self.logger.debug(
+                "Getting %s (Formula %d of %d)", formula, n_formula + 1, num_formulas
+            )
+            decoded_docs = MontyDecoder().process_decoded(docs)
+            grouped_tasks = _group_orderings(
+                decoded_docs,
+                self.structure_match_ltol,
+                self.structure_match_stol,
+                self.structure_match_angle_tol,
+            )
+            n_groups = len(grouped_tasks)
+            for n_group, group in enumerate(grouped_tasks):
+                self.logger.debug(
+                    "Found %d tasks for %s (Parent structure %d of %d)",
+                    len(group),
+                    formula,
+                    n_group + 1,
+                    n_groups,
+                )
+                yield group
+
+    def process_item(self, tasks: list[dict]) -> list[MagneticOrderingsDocument]:
+        """Process magnetic ordering relaxation/static calculations into documents.
+
+        The magnetic ordering tasks will be grouped based on their parent structure
+        (i.e., the structure before the magnetic ordering transformation was applied).
+        See _group_orderings for more details.
+
+        Parameters
+        ----------
+        tasks : list[dict]
+            A list of magnetic ordering tasks grouped by same formula.
+
+        Returns
+        -------
+        list of .MagneticOrderingsDocument
+            A list of magnetic ordering documents (one for each unique parent
+            structure).
+        """
+        self.logger.debug("Processing %s", tasks[0]["output"].formula_pretty)
+
+        if not tasks:
+            return []
+
+        return jsanitize(
+            MagneticOrderingsDocument.from_tasks(tasks).model_dump(),
+            allow_bson=True,
+        )
+
+    def update_targets(self, items: list[MagneticOrderingsDocument]) -> None:
+        """Insert new magnetic orderings into the magnetic orderings Store.
+
+        Parameters
+        ----------
+        items : list of .MagneticOrderingsDocument
+            A list of magnetic ordering documents to add to the database.
+        """
+        self.logger.info("Updating %s magnetic orderings documents", len(items))
+        self.magnetic_orderings.update(items, key="ground_state_uuid")
+
+
+def _group_orderings(
+    tasks: list[dict], ltol: float, stol: float, angle_tol: float
+) -> list[list[dict]]:
+    """Group ordering tasks by their parent structure.
+
+    This is useful for distinguishing between different polymorphs (i.e., same formula).
+
+    Parameters
+    ----------
+    tasks : list[dict]
+        A list of ordering tasks.
+    tol : float
+        Numerical tolerance for structure equivalence.
+
+    Returns
+    -------
+    list[list[dict]]
+        The tasks grouped by their parent structure.
+    """
+    tasks = [dict(task) for task in tasks]
+
+    grouped_tasks = [[tasks[0]]]
+    sm = StructureMatcher(ltol=ltol, stol=stol, angle_tol=angle_tol)
+
+    for task in tasks[1:]:
+        parent_structure = MontyDecoder().process_decoded(
+            task["metadata"]["parent_structure"]
+        )
+
+        match = False
+        for group in grouped_tasks:
+            group_parent_structure = MontyDecoder().process_decoded(
+                group[0]["metadata"]["parent_structure"]
+            )
+
+            #  parent structure lattice/coords may be same but in different order
+            #  so we need to be more rigorous in checking equivalence
+            if sm.fit(parent_structure, group_parent_structure):
+                group.append(task)
+                match = True
+                break
+
+        if not match:
+            grouped_tasks.append([task])
+
+    return MontyDecoder().process_decoded(grouped_tasks)