fix atomic masses in openff + add optional kwarg for these

Author: Aaron Kaplan

pyproject.toml

@@ -121,7 +121,7 @@ strict-openff = [
     "monty==2025.3.3",
     "openmm-mdanalysis-reporter==0.1.0",
     "openmm==8.1.1",
-    "pymatgen==2024.11.13", # DO NOT CHANGE FROM 2024.11.13 - 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/openff/utils.py

@@ -11,6 +11,7 @@
 from emmet.core.openff import MoleculeSpec
 from pymatgen.core import Element, Molecule
 from pymatgen.io.openff import create_openff_mol
+from scipy.units import Avogadro
 
 if TYPE_CHECKING:
     import pathlib
@@ -145,6 +146,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 +160,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 +191,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 {el.Z: el.atomic_mass for el in Element}
     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 +212,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 +223,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 {el.Z: el.atomic_mass for el in Element}
 
     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 +248,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 +263,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 {el.Z: el.atomic_mass for el in Element}
 
-    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)

tests/openff_md/test_utils.py

@@ -26,6 +26,16 @@
     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,
+    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,9 +239,13 @@ 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)
+    counts = counts_from_masses(comp_dict, 100, atomic_masses=LEGACY_ATOMIC_MASSES)
     assert sum(counts.values()) == 100
 
     mol_ratio = {
@@ -242,15 +256,19 @@ 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
+    )
+    counts_number = counts_from_masses(
+        mass_fractions, 324, atomic_masses=LEGACY_ATOMIC_MASSES
+    )
 
     assert 200 < sum(counts_size.values()) < 500