Port the magnetic exchange workflow from atomate1 (Heisenberg fit + Vampire Tc)

pyproject.toml

@@ -158,6 +158,20 @@ docs = [
     "jupyterlab==4.6.1",
 ]
 
+# The strict forcefield groups have mutually incompatible pins (e.g. numpy>=2
+# via mattersim vs numpy==1.26.4). They are meant to be installed one at a time.
+# uv builds a single universal lockfile, so these must be declared as conflicting
+# extras or `uv sync` fails to resolve them together.
+[tool.uv]
+conflicts = [
+    [
+        { extra = "strict-forcefields-generic" },
+        { extra = "strict-forcefields-torch-limited" },
+        { extra = "strict-forcefields-e3nn-limited" },
+        { extra = "strict-forcefields-numpy-limited" },
+    ],
+]
+
 [tool.setuptools.package-data]
 atomate2 = ["py.typed"]
 "atomate2.vasp.sets" = ["*.yaml"]
@@ -204,6 +218,9 @@ exclude_lines = [
 [tool.ruff]
 target-version = "py310"
 output-format = "concise"
+# vampire_caller.py is vendored verbatim from pymatgen (removed there in 2026.3.23);
+# exclude it from lint/format so it stays a faithful copy of the upstream source.
+extend-exclude = ["src/atomate2/vampire/vampire_caller.py"]
 
 [tool.ruff.lint]
 select = ["ALL"]

src/atomate2/common/flows/exchange.py

@@ -0,0 +1,164 @@
+"""Flow for fitting magnetic exchange parameters and estimating Tc.
+
+This is a post-processing workflow and runs no DFT itself. Given magnetic
+structures and their energies (e.g. the output of the magnetic-orderings
+workflow), it fits a classical Heisenberg Hamiltonian and optionally runs
+Vampire Monte-Carlo for the critical temperature.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from jobflow import Flow, Maker
+
+from atomate2.common.jobs.exchange import build_exchange_doc, heisenberg_mapping
+from atomate2.vampire.jobs.run_vampire import run_vampire
+
+if TYPE_CHECKING:
+    from pymatgen.core.structure import Structure
+
+    from atomate2.common.schemas.magnetism import MagneticOrderingsDocument
+
+
+__all__ = ["ExchangeMaker"]
+
+
+@dataclass
+class ExchangeMaker(Maker):
+    """Maker to fit Heisenberg exchange parameters from magnetic structures.
+
+    Given a set of magnetic structures and their energies (per atom), this fits a
+    classical Heisenberg Hamiltonian via pymatgen's ``HeisenbergMapper`` to extract
+    shell-resolved exchange parameters ``J_ij``, fit by least squares over all
+    supplied orderings.
+
+    This is a post-processing workflow and runs no DFT. It is code-agnostic: the
+    inputs are plain pymatgen structures and energies, typically taken from the
+    magnetic-orderings workflow (see :meth:`make_from_ordering_doc`).
+
+    Originally implemented in atomate (v1) for VASP as the ``ExchangeWF``.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flows produced by this Maker.
+    heisenberg_settings : dict
+        Keyword arguments for pymatgen's HeisenbergMapper, i.e. the nearest-neighbour
+        ``cutoff`` (Angstrom) and the distance-grouping ``tol``.
+    run_vampire : bool
+        Whether to run the Vampire Monte-Carlo step to estimate the critical
+        temperature. Requires a VAMPIRE binary on PATH - ``vampire-serial`` unless
+        ``ATOMATE2_VAMPIRE_CMD`` names another (e.g. ``srun -n 24 vampire-parallel``
+        for the MPI build); the step raises a clear error if it is missing.
+        Defaults to True (atomate1 parity).
+    mc_settings : dict | None
+        Keyword arguments for the Vampire Monte-Carlo run (e.g. ``mc_box_size``,
+        ``equil_timesteps``, ``mc_timesteps``). Only used if ``run_vampire``.
+    """
+
+    name: str = "exchange"
+    heisenberg_settings: dict = field(
+        default_factory=lambda: {"cutoff": 3.0, "tol": 0.04}
+    )
+    run_vampire: bool = True
+    mc_settings: dict | None = None
+
+    def make(
+        self,
+        structures: list[Structure],
+        energies: list[float],
+        parent: Structure | None = None,
+    ) -> Flow:
+        """Make a flow to fit Heisenberg exchange parameters.
+
+        Parameters
+        ----------
+        structures : list[Structure]
+            Magnetic structures, each carrying a "magmom" site property.
+        energies : list[float]
+            Energies **per atom** (eV) corresponding to each structure.
+
+        Returns
+        -------
+        Flow
+            The exchange-parameter fitting workflow.
+        """
+        if len(structures) != len(energies):
+            raise ValueError(
+                f"Got {len(structures)} structures but {len(energies)} energies; "
+                "these must be equal."
+            )
+
+        # sort so the ground state (lowest energy) is index 0
+        order = sorted(range(len(energies)), key=lambda i: energies[i])
+        structures = [structures[i] for i in order]
+        energies = [energies[i] for i in order]
+
+        # HeisenbergMapper requires a 'magmom' site property on every structure
+        for idx, structure in enumerate(structures):
+            if not structure.site_properties.get("magmom"):
+                raise ValueError(
+                    f"Structure {idx} is missing a 'magmom' site property, which is "
+                    "required to fit a Heisenberg model."
+                )
+
+        hmap = heisenberg_mapping(
+            structures, energies, parent, self.heisenberg_settings
+        )
+        jobs = [hmap]
+
+        vampire_output = None
+        if self.run_vampire:
+            vmc = run_vampire(hmap.output, self.mc_settings)
+            jobs.append(vmc)
+            vampire_output = vmc.output
+
+        doc = build_exchange_doc(
+            hmap.output,
+            parent_structure=parent,
+            vampire_output=vampire_output,
+            vampire_settings=self.mc_settings if self.run_vampire else None,
+        )
+        jobs.append(doc)
+
+        formula = structures[0].composition.reduced_formula
+        return Flow(
+            jobs=jobs,
+            output=doc.output,
+            name=f"{self.name} ({formula})",
+        )
+
+    def make_from_ordering_doc(self, doc: MagneticOrderingsDocument) -> Flow:
+        """Make an exchange flow from a (concrete) magnetic-orderings document.
+
+        This is a convenience constructor for chaining onto the magnetic-orderings
+        workflow: it pulls the structures and per-atom energies out of an
+        already-computed ``MagneticOrderingsDocument`` and forwards them to
+        :meth:`make`. Because the structures must be inspected (sorted, validated)
+        when the flow is built, ``doc`` must be a resolved document, not a jobflow
+        output reference.
+
+        Parameters
+        ----------
+        doc : MagneticOrderingsDocument
+            A computed magnetic-orderings document (e.g. from a finished
+            MagneticOrderingsMaker run).
+
+        Returns
+        -------
+        Flow
+            The exchange-parameter fitting workflow.
+        """
+        structures, energies = [], []
+
+        for output in doc.outputs:
+            structure = output.structure.copy()
+            # HeisenbergMapper needs magmoms; the output stores them separately
+            if not structure.site_properties.get("magmom"):
+                structure.add_site_property("magmom", output.magmoms)
+            structures.append(structure)
+            energies.append(output.energy_per_atom)
+
+        return self.make(structures, energies, doc.parent_structure)

src/atomate2/common/jobs/exchange.py

@@ -0,0 +1,95 @@
+"""Jobs for fitting Heisenberg exchange parameters and building exchange docs.
+
+These replace the atomate1 ``HeisenbergModelMapping``, ``HeisenbergModelToDb`` and
+``VampireToDb`` firetasks. Outputs flow between jobs through jobflow references
+rather than a MongoDB collection.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from jobflow import job
+from pymatgen.analysis.magnetism.heisenberg import HeisenbergMapper
+
+from atomate2.common.schemas.exchange import ExchangeDocument
+
+if TYPE_CHECKING:
+    from pymatgen.analysis.magnetism.heisenberg import HeisenbergModel
+    from pymatgen.core.structure import Structure
+
+    from atomate2.vampire.schemas.vampire_output import VampireOutput
+
+logger = logging.getLogger(__name__)
+
+
+@job(name="heisenberg mapping")
+def heisenberg_mapping(
+    structures: list[Structure],
+    energies: list[float],
+    parent: Structure | None = None,
+    heisenberg_settings: dict | None = None,
+) -> HeisenbergModel:
+    """Fit a classical Heisenberg Hamiltonian to magnetic structures and energies.
+
+    This wraps pymatgen's ``HeisenbergMapper`` to extract shell-resolved exchange
+    parameters ``J_ij``, fit by least squares over all supplied orderings.
+
+    Parameters
+    ----------
+    structures : list[Structure]
+        Magnetic structures, each carrying a "magmom" site property. These should be
+        ordered with the ground state first (index 0).
+    energies : list[float]
+        Energies **per atom** (eV) corresponding to each structure. These are
+        converted to total energies internally, as required by HeisenbergMapper.
+    heisenberg_settings : dict or None
+        Keyword arguments for HeisenbergMapper, e.g. ``{"cutoff": 3.0, "tol": 0.04}``.
+
+    Returns
+    -------
+    HeisenbergModel
+        The fitted, MSONable Heisenberg model.
+    """
+    heisenberg_settings = heisenberg_settings or {}
+    total_energies = [e * len(s) for s, e in zip(structures, energies, strict=True)]
+    hmapper = HeisenbergMapper(
+        structures, total_energies, parent, **heisenberg_settings
+    )
+    return hmapper.get_heisenberg_model()
+
+
+@job(name="build exchange doc")
+def build_exchange_doc(
+    heisenberg_model: HeisenbergModel,
+    parent_structure: Structure | None = None,
+    vampire_output: VampireOutput | None = None,
+    vampire_settings: dict | None = None,
+) -> ExchangeDocument:
+    """Assemble the final ExchangeDocument from a fitted model and optional Tc run.
+
+    Parameters
+    ----------
+    heisenberg_model : HeisenbergModel
+        The fitted Heisenberg model from :func:`heisenberg_mapping`.
+    parent_structure : Structure or None
+        The full parent structure from which the magnetic structures were derived.
+        This is used to store the final fitted exchange parameters in the context
+        of the original structure.
+    vampire_output : VampireOutput or None
+        The Vampire Monte-Carlo result, if the critical-temperature step was run.
+    vampire_settings : dict or None
+        The keyword arguments :func:`run_vampire` was called with, if it was run.
+
+    Returns
+    -------
+    ExchangeDocument
+        The final summary document.
+    """
+    return ExchangeDocument.from_model(
+        heisenberg_model,
+        parent_structure=parent_structure,
+        vampire_output=vampire_output,
+        vampire_settings=vampire_settings,
+    )