ASE MD NPT bug fixes + housekeeping (#1255)

* Close #1254: Ensure ASE NPT MD runs correctly transform structure to upper triangular matrix
* Close #1249: Default MatPES flow reflects flow used in paper
* Close #1252: Openff tests use consistent values of the atomic masses

Author: esoteric-ephemera

pyproject.toml

@@ -121,7 +121,7 @@ strict-openff = [
     "monty==2025.3.3",
     "openmm-mdanalysis-reporter==0.1.0",
     "openmm==8.1.1",
-    "pymatgen==2025.6.14", # TODO: open ff is extremely sensitive to pymatgen version
+    "pymatgen==2025.6.14", # EXERCISE CAUTION WHEN UPDATING - open ff is extremely sensitive to pymatgen version
 ]
 strict-forcefields = [
     "calorine==3.1",

src/atomate2/ase/md.py

@@ -28,7 +28,6 @@
 from pymatgen.core.structure import Molecule, Structure
 from pymatgen.io.ase import AseAtomsAdaptor
 from scipy.interpolate import interp1d
-from scipy.linalg import schur
 
 from atomate2.ase.jobs import _ASE_DATA_OBJECTS, AseMaker
 from atomate2.ase.schemas import AseResult, AseTaskDoc
@@ -237,11 +236,31 @@ def _get_ensemble_defaults(self) -> None:
             self.ase_md_kwargs.pop("temperature_K", None)
             self.ase_md_kwargs.pop("externalstress", None)
         elif self.ensemble == MDEnsemble.nvt:
-            self.ase_md_kwargs["temperature_K"] = self.t_schedule[0]
+            self.ase_md_kwargs["temperature_K"] = self.ase_md_kwargs.get(
+                "temperature_K", self.t_schedule[0]
+            )
             self.ase_md_kwargs.pop("externalstress", None)
         elif self.ensemble == MDEnsemble.npt:
-            self.ase_md_kwargs["temperature_K"] = self.t_schedule[0]
-            self.ase_md_kwargs["externalstress"] = self.p_schedule[0] * 1e3 * units.bar
+            self.ase_md_kwargs["temperature_K"] = self.ase_md_kwargs.get(
+                "temperature_K", self.t_schedule[0]
+            )
+
+            # These use different kwargs for pressure
+            if (
+                isinstance(self.dynamics, DynamicsPresets)
+                and DynamicsPresets(self.dynamics) == DynamicsPresets.npt_berendsen
+            ) or (
+                isinstance(self.dynamics, type)
+                and issubclass(self.dynamics, MolecularDynamics)
+                and self.dynamics.__name__ == "NPTBerendsen"
+            ):
+                stress_kwarg = "pressure_au"
+            else:
+                stress_kwarg = "externalstress"
+
+            self.ase_md_kwargs[stress_kwarg] = self.ase_md_kwargs.get(
+                stress_kwarg, self.p_schedule[0] * 1e3 * units.bar
+            )
 
         if isinstance(self.dynamics, str) and self.dynamics.lower() == "langevin":
             self.ase_md_kwargs["friction"] = self.ase_md_kwargs.get(
@@ -338,8 +357,7 @@ def run_ase(
             # `isinstance(dynamics,NPT)` is False
 
             # ASE NPT implementation requires upper triangular cell
-            schur_decomp, _ = schur(atoms.get_cell(complete=True), output="complex")
-            atoms.set_cell(schur_decomp.real, scale_atoms=True)
+            atoms.set_cell(atoms.cell.standard_form(form="upper")[0])
 
         if initial_velocities:
             atoms.set_velocities(initial_velocities)

src/atomate2/openff/utils.py

@@ -11,10 +11,13 @@
 from emmet.core.openff import MoleculeSpec
 from pymatgen.core import Element, Molecule
 from pymatgen.io.openff import create_openff_mol
+from scipy.constants import Avogadro
 
 if TYPE_CHECKING:
     import pathlib
 
+DEFAULT_ATOMIC_MASSES = {el.Z: el.atomic_mass for el in Element}
+
 
 def create_mol_spec(
     smiles: str,
@@ -145,6 +148,7 @@ def calculate_elyte_composition(
     salts: dict[str, float],
     solvent_densities: dict = None,
     solvent_ratio_dimension: Literal["mass", "volume"] = "mass",
+    atomic_masses: dict[int, float] | None = None,
 ) -> dict[str, float]:
     """Calculate the normalized mass ratios of an electrolyte solution.
 
@@ -158,6 +162,9 @@ def calculate_elyte_composition(
         Dictionary of solvent SMILES strings and their densities (g/ml).
     solvent_ratio_dimension: optional, str
         Whether the solvents are included with a ratio of "mass" or "volume"
+    atomic_masses : dict[int,float] or None (Default)
+        The mass of each element in the species dict. Defaults to the
+        most recent data from pymatgen.core.periodic_table.Element
 
     Returns
     -------
@@ -186,11 +193,11 @@ def calculate_elyte_composition(
     }
 
     # Calculate the molecular weights of the solvent
-    masses = {el.Z: el.atomic_mass for el in Element}
+    atomic_masses = atomic_masses or DEFAULT_ATOMIC_MASSES
     salt_mws = {}
     for smile in salts:
         mol = tk.Molecule.from_smiles(smile, allow_undefined_stereo=True)
-        salt_mws[smile] = sum(masses[atom.atomic_number] for atom in mol.atoms)
+        salt_mws[smile] = sum(atomic_masses[atom.atomic_number] for atom in mol.atoms)
 
     # Convert salt mole ratios to mass ratios
     salt_mass_ratio = {
@@ -207,7 +214,9 @@ def calculate_elyte_composition(
     return {species: mass / total_mass for species, mass in combined_mass_ratio.items()}
 
 
-def counts_from_masses(species: dict[str, float], n_mol: int) -> dict[str, float]:
+def counts_from_masses(
+    species: dict[str, float], n_mol: int, atomic_masses: dict[int, float] | None = None
+) -> dict[str, float]:
     """Calculate the number of mols needed to yield a given mass ratio.
 
     Parameters
@@ -216,19 +225,21 @@ def counts_from_masses(species: dict[str, float], n_mol: int) -> dict[str, float
         Dictionary of species SMILES strings and their relative mass fractions.
     n_mol : float
         Total number of mols. Returned array will sum to near n_mol.
-
+    atomic_masses : dict[int,float] or None (Default)
+        The mass of each element in the species dict. Defaults to the
+        most recent data from pymatgen.core.periodic_table.Element
 
     Returns
     -------
     numpy.ndarray
         n_mols: Number of each SMILES needed for the given mass ratio.
     """
-    masses = {el.Z: el.atomic_mass for el in Element}
+    atomic_masses = atomic_masses or DEFAULT_ATOMIC_MASSES
 
     mol_weights = []
     for smile in species:
         mol = tk.Molecule.from_smiles(smile, allow_undefined_stereo=True)
-        mol_weights.append(sum(masses[atom.atomic_number] for atom in mol.atoms))
+        mol_weights.append(sum(atomic_masses[atom.atomic_number] for atom in mol.atoms))
 
     mol_ratio = np.array(list(species.values())) / np.array(mol_weights)
     mol_ratio /= sum(mol_ratio)
@@ -239,7 +250,10 @@ def counts_from_masses(species: dict[str, float], n_mol: int) -> dict[str, float
 
 
 def counts_from_box_size(
-    species: dict[str, float], side_length: float, density: float = 0.8
+    species: dict[str, float],
+    side_length: float,
+    density: float = 0.8,
+    atomic_masses: dict[int, float] | None = None,
 ) -> dict[str, float]:
     """Calculate the number of molecules needed to fill a box.
 
@@ -251,25 +265,27 @@ def counts_from_box_size(
         Side length of the cubic simulation box in nm.
     density : int, optional
         Density of the system in g/cm^3. Default is 1 g/cm^3.
+    atomic_masses : dict[int,float] or None (Default)
+        The mass of each element in the species dict. Defaults to the
+        most recent data from pymatgen.core.periodic_table.Element
 
     Returns
     -------
     dict of str, float
         Number of each species needed to fill the box with the given density.
     """
-    masses = {el.Z: el.atomic_mass for el in Element}
+    atomic_masses = atomic_masses or DEFAULT_ATOMIC_MASSES
 
-    na = 6.02214076e23  # Avogadro's number
     volume = (side_length * 1e-7) ** 3  # Convert from nm3 to cm^3
     total_mass = volume * density  # grams
 
     # Calculate molecular weights
     mol_weights = []
     for smile in species:
         mol = tk.Molecule.from_smiles(smile, allow_undefined_stereo=True)
-        mol_weights.append(sum(masses[atom.atomic_number] for atom in mol.atoms))
+        mol_weights.append(sum(atomic_masses[atom.atomic_number] for atom in mol.atoms))
     mean_mw = np.mean(mol_weights)
-    n_mol = (total_mass / mean_mw) * na
+    n_mol = (total_mass / mean_mw) * Avogadro
 
     # Calculate the number of moles needed for each species
     mol_ratio = np.array(list(species.values())) / np.array(mol_weights)

src/atomate2/vasp/flows/matpes.py

@@ -58,16 +58,7 @@ class MatPesStaticFlowMaker(Maker):
             copy_vasp_kwargs={"additional_vasp_files": ("WAVECAR",)}
         )
     )
-    # optional 3rd PBE+U static in case structure contains elements with +U corrections
-    static3: Maker | None = field(
-        default_factory=lambda: MatPesGGAStaticMaker(
-            name="MatPES GGA+U static",
-            input_set_generator=MatPESStaticSet(
-                user_incar_settings={"LDAU:": True},  # enable +U corrections
-            ),
-            copy_vasp_kwargs={"additional_vasp_files": ("WAVECAR",)},
-        )
-    )
+    static3: Maker | None = None
 
     def __post_init__(self) -> None:
         """Validate flow."""

tests/conftest.py

@@ -7,7 +7,6 @@
 from unittest import mock
 
 import pytest
-from fireworks import LaunchPad
 from jobflow import JobStore
 from jobflow.settings import JobflowSettings
 from maggma.stores import MemoryStore
@@ -69,6 +68,12 @@ def debug_mode():
 
 @pytest.fixture(scope="session")
 def lpad(database, debug_mode):
+    try:
+        from fireworks import LaunchPad
+    except ImportError as exc:
+        raise ImportError(
+            "Please pip install fireworks to use this test fixture."
+        ) from exc
     lpad = LaunchPad(name=database)
     lpad.reset("", require_password=False)
     yield lpad

tests/openff_md/test_utils.py

@@ -26,6 +26,17 @@
     mol_graph_to_openff_mol,
 )
 
+# Ensure consistency of test results by fixing atomic masses (no updates)
+LEGACY_ATOMIC_MASSES = {
+    1: 1.00794,
+    3: 6.941,
+    6: 12.0107,
+    7: 14.0067,
+    8: 15.9994,
+    9: 18.9984032,
+    15: 30.973762,
+}
+
 
 def test_molgraph_to_openff_pf6(mol_files):
     """transform a water MoleculeGraph to a OpenFF water molecule"""
@@ -229,10 +240,14 @@ def test_calculate_elyte_composition():
     solvent_densities = {"O": 1.0, "CCO": 0.8}
 
     comp_dict = calculate_elyte_composition(
-        vol_ratio, salts, solvent_densities, "volume"
+        vol_ratio,
+        salts,
+        solvent_densities,
+        "volume",
+        atomic_masses=LEGACY_ATOMIC_MASSES,
     )
-    counts = counts_from_masses(comp_dict, 100)
-    assert sum(counts.values()) == 100
+    counts = counts_from_masses(comp_dict, 100, atomic_masses=LEGACY_ATOMIC_MASSES)
+    assert counts == {"CCO": 23, "F[P-](F)(F)(F)(F)F": 3, "O": 72, "[Li+]": 3}
 
     mol_ratio = {
         "[Li+]": 0.00616,
@@ -242,16 +257,21 @@ def test_calculate_elyte_composition():
         "CC#N": 0.130,
         "FC1COC(=O)O1": 0.028,
     }
-    counts2 = counts_from_masses(mol_ratio, 1000)
+    counts2 = counts_from_masses(mol_ratio, 1000, atomic_masses=LEGACY_ATOMIC_MASSES)
     assert np.allclose(sum(counts2.values()), 1000, atol=5)
 
 
 def test_counts_calculators():
     mass_fractions = {"O": 0.5, "CCO": 0.5}
 
-    counts_size = counts_from_box_size(mass_fractions, 3)
-    counts_number = counts_from_masses(mass_fractions, 324)
+    counts_size = counts_from_box_size(
+        mass_fractions, 3, atomic_masses=LEGACY_ATOMIC_MASSES
+    )
+
+    assert sum(counts_size.values()) == 406
 
-    assert 200 < sum(counts_size.values()) < 500
+    counts_number = counts_from_masses(
+        mass_fractions, 406, atomic_masses=LEGACY_ATOMIC_MASSES
+    )
 
     assert counts_size == counts_number

tests/vasp/flows/test_matpes.py

@@ -6,6 +6,7 @@
 from pymatgen.core import Structure
 
 from atomate2.vasp.flows.matpes import MatPesStaticFlowMaker
+from atomate2.vasp.jobs.matpes import MatPesGGAStaticMaker
 
 
 def test_matpes_static_flow_maker(mock_vasp, clean_dir, vasp_test_dir):
@@ -59,7 +60,9 @@ def test_matpes_static_flow_maker(mock_vasp, clean_dir, vasp_test_dir):
     assert flow[0].name == "MatPES GGA static"
 
     # Test setting two makers to None
-    flow_maker = MatPesStaticFlowMaker(static1=None, static2=None)
+    flow_maker = MatPesStaticFlowMaker(
+        static1=None, static2=None, static3=MatPesGGAStaticMaker()
+    )
     flow = flow_maker.make(si_struct)
     assert len(flow) == 0
 
@@ -68,7 +71,9 @@ def test_matpes_static_flow_maker(mock_vasp, clean_dir, vasp_test_dir):
         MatPesStaticFlowMaker(static1=None, static2=None, static3=None)
 
     # Test no static3 if structure requires no Hubbard U corrections
-    flow_maker = MatPesStaticFlowMaker(static1=None, static2=None)
+    flow_maker = MatPesStaticFlowMaker(
+        static1=None, static2=None, static3=MatPesGGAStaticMaker()
+    )
     flow = flow_maker.make(si_struct)
     assert len(flow) == 0
     assert flow.output == {}