ASE Calculator stress

mace/calculators/mace.py

@@ -15,7 +15,7 @@
 class MACECalculator(Calculator):
     """MACE ASE Calculator"""
 
-    implemented_properties = ["energy", "forces"]
+    implemented_properties = ["energy", "free_energy", "forces", "stress"]
 
     def __init__(
         self,
@@ -41,7 +41,7 @@ def __init__(
         torch_tools.set_default_dtype(default_dtype)
 
     # pylint: disable=dangerous-default-value
-    def calculate(self, atoms=None, properties=["energy"], system_changes=all_changes):
+    def calculate(self, atoms=None, properties=None, system_changes=all_changes):
         """
         Calculate properties.
         :param atoms: ase.Atoms object
@@ -67,17 +67,28 @@ def calculate(self, atoms=None, properties=["energy"], system_changes=all_change
         batch = next(iter(data_loader)).to(self.device)
 
         # predict + extract data
-        out = self.model(batch)
-        forces = out["forces"].detach().cpu().numpy()
+        out = self.model(batch, compute_stress=True)
         energy = out["energy"].detach().cpu().item()
+        forces = out["forces"].detach().cpu().numpy()
 
         # store results
+        E = energy * self.energy_units_to_eV
         self.results = {
-            "energy": energy * self.energy_units_to_eV,
+            "energy": E,
+            "free_energy": E,
             # force has units eng / len:
             "forces": forces * (self.energy_units_to_eV / self.length_units_to_A),
         }
 
+        # even though compute_stress is True, stress can be none if pbc is False
+        # not sure if correct ASE thing is to have no dict key, or dict key with value None
+        if out["stress"] is not None:
+            stress = out["stress"].detach().cpu().numpy()
+            # stress has units eng / len^3:
+            self.results["stress"] = (
+                stress * (self.energy_units_to_eV / self.length_units_to_A**3)
+            )[0]
+
 
 class DipoleMACECalculator(Calculator):
     """MACE ASE Calculator for predicting dipoles"""
@@ -113,7 +124,7 @@ def __init__(
         torch_tools.set_default_dtype(default_dtype)
 
     # pylint: disable=dangerous-default-value
-    def calculate(self, atoms=None, properties=["dipole"], system_changes=all_changes):
+    def calculate(self, atoms=None, properties=None, system_changes=all_changes):
         """
         Calculate properties.
         :param atoms: ase.Atoms object
@@ -153,7 +164,9 @@ class EnergyDipoleMACECalculator(Calculator):
 
     implemented_properties = [
         "energy",
+        "free_energy",
         "forces",
+        "stress",
         "dipole",
     ]
 
@@ -186,7 +199,7 @@ def __init__(
         torch_tools.set_default_dtype(default_dtype)
 
     # pylint: disable=dangerous-default-value
-    def calculate(self, atoms=None, properties=["dipole"], system_changes=all_changes):
+    def calculate(self, atoms=None, properties=None, system_changes=all_changes):
         """
         Calculate properties.
         :param atoms: ase.Atoms object
@@ -212,15 +225,26 @@ def calculate(self, atoms=None, properties=["dipole"], system_changes=all_change
         batch = next(iter(data_loader)).to(self.device)
 
         # predict + extract data
-        out = self.model(batch)
-        forces = out["forces"].detach().cpu().numpy()
+        out = self.model(batch, compute_stress=True)
         energy = out["energy"].detach().cpu().item()
+        forces = out["forces"].detach().cpu().numpy()
         dipole = out["dipole"].detach().cpu().numpy()
 
         # store results
+        E = energy * self.energy_units_to_eV
         self.results = {
-            "energy": energy * self.energy_units_to_eV,
+            "energy": E,
+            "free_energy": E,
             # force has units eng / len:
             "forces": forces * (self.energy_units_to_eV / self.length_units_to_A),
+            # stress has units eng / len:
             "dipole": dipole,
         }
+
+        # even though compute_stress is True, stress can be none if pbc is False
+        # not sure if correct ASE thing is to have no dict key, or dict key with value None
+        if out["stress"] is not None:
+            stress = out["stress"].detach().cpu().numpy()
+            self.results["stress"] = (
+                stress * (self.energy_units_to_eV / self.length_units_to_A**3)
+            )[0]

mace/modules/models.py

@@ -153,7 +153,7 @@ def forward(
         # Setup
         data.positions.requires_grad = True
         displacement = None
-        if compute_virials:
+        if compute_virials or compute_stress:
             data.positions, data.shifts, displacement = get_symmetric_displacement(
                 positions=data.positions,
                 unit_shifts=data.unit_shifts,
@@ -245,7 +245,7 @@ def forward(
         # Setup
         data.positions.requires_grad = True
         displacement = None
-        if compute_virials:
+        if compute_virials or compute_stress:
             data.positions, data.shifts, displacement = get_symmetric_displacement(
                 positions=data.positions,
                 unit_shifts=data.unit_shifts,

mace/modules/utils.py

@@ -57,27 +57,21 @@ def compute_forces_virials(
         allow_unused=True,
     )
     stress = None
-    if compute_stress:
+    if compute_stress and virials is not None:
         cell = cell.view(-1, 3, 3)
         volume = torch.einsum(
             "zi,zi->z",
             cell[:, 0, :],
             torch.cross(cell[:, 1, :], cell[:, 2, :], dim=1),
         ).unsqueeze(-1)
         stress = virials / volume.view(-1, 1, 1)
-    if forces is None and virials is None:
+    if forces is None:
         logging.warning("Gradient is None, padded with zeros")
-        return (
-            torch.zeros_like(positions),
-            torch.zeros_like(positions).expand(1, 1, 3),
-            None,
-        )
-    if forces is not None and virials is None:
+        forces = torch.zeros_like(positions)
+    if virials is None:
         logging.warning("Virial is None, padded with zeros")
-        return -1 * forces, torch.zeros_like(positions).expand(1, 1, 3), None
-    if forces is None and virials is not None:
-        logging.warning("Virial is None, padded with zeros")
-        return torch.zeros_like(positions), -1 * virials, None
+        virials = torch.zeros((1, 3, 3))
+
     return -1 * forces, -1 * virials, stress
 
 
@@ -130,7 +124,8 @@ def get_outputs(
     compute_virials: bool = True,
     compute_stress: bool = True,
 ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
-    if compute_force and compute_virials:
+    if compute_virials or compute_stress:
+        # forces come for free
         forces, virials, stress = compute_forces_virials(
             energy=energy,
             positions=positions,
@@ -139,13 +134,12 @@ def get_outputs(
             compute_stress=compute_stress,
             training=training,
         )
-    elif compute_force and not compute_stress:
+    elif compute_force:
         forces, virials, stress = (
             compute_forces(energy=energy, positions=positions, training=training),
             None,
             None,
         )
-        stress = None
     else:
         forces, virials, stress = (None, None, None)
     return forces, virials, stress

tests/test_calculator.py

@@ -0,0 +1,127 @@
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+import ase.io
+import numpy as np
+from ase.atoms import Atoms
+from ase.calculators.test import gradient_test
+from ase.constraints import ExpCellFilter
+
+from mace.calculators.mace import MACECalculator
+
+pytest_mace_dir = Path(__file__).parent.parent
+run_train = Path(__file__).parent.parent / "scripts" / "run_train.py"
+
+
+@pytest.fixture(scope="module", name="fitting_configs")
+def fitting_configs_fixture():
+    water = Atoms(
+        numbers=[8, 1, 1],
+        positions=[[0, -2.0, 0], [1, 0, 0], [0, 1, 0]],
+        cell=[4] * 3,
+        pbc=[True] * 3,
+    )
+    fit_configs = [
+        Atoms(numbers=[8], positions=[[0, 0, 0]], cell=[6] * 3),
+        Atoms(numbers=[1], positions=[[0, 0, 0]], cell=[6] * 3),
+    ]
+    fit_configs[0].info["REF_energy"] = 0.0
+    fit_configs[0].info["config_type"] = "IsolatedAtom"
+    fit_configs[1].info["REF_energy"] = 0.0
+    fit_configs[1].info["config_type"] = "IsolatedAtom"
+
+    np.random.seed(5)
+    for _ in range(20):
+        c = water.copy()
+        c.positions += np.random.normal(0.1, size=c.positions.shape)
+        c.info["REF_energy"] = np.random.normal(0.1)
+        c.new_array("REF_forces", np.random.normal(0.1, size=c.positions.shape))
+        c.info["REF_stress"] = np.random.normal(0.1, size=6)
+        fit_configs.append(c)
+
+    return fit_configs
+
+
+@pytest.fixture(scope="module", name="trained_model")
+def trained_model_fixture(tmp_path_factory, fitting_configs):
+    _mace_params = {
+        "name": "MACE",
+        "valid_fraction": 0.05,
+        "energy_weight": 1.0,
+        "forces_weight": 10.0,
+        "stress_weight": 1.0,
+        "model": "MACE",
+        "hidden_irreps": "128x0e",
+        "r_max": 3.5,
+        "batch_size": 5,
+        "max_num_epochs": 10,
+        "swa": None,
+        "start_swa": 5,
+        "ema": None,
+        "ema_decay": 0.99,
+        "amsgrad": None,
+        "restart_latest": None,
+        "device": "cpu",
+        "seed": 5,
+        "loss": "stress",
+        "energy_key": "REF_energy",
+        "forces_key": "REF_forces",
+        "stress_key": "REF_stress",
+    }
+
+    tmp_path = tmp_path_factory.mktemp("run_")
+
+    ase.io.write(tmp_path / "fit.xyz", fitting_configs)
+
+    mace_params = _mace_params.copy()
+    mace_params["checkpoints_dir"] = str(tmp_path)
+    mace_params["model_dir"] = str(tmp_path)
+    mace_params["train_file"] = tmp_path / "fit.xyz"
+
+    # make sure run_train.py is using the mace that is currently being tested
+    run_env = os.environ
+    run_env["PYTHONPATH"] = str(pytest_mace_dir) + ":" + os.environ["PYTHONPATH"]
+
+    cmd = (
+        sys.executable
+        + " "
+        + str(run_train)
+        + " "
+        + " ".join(
+            [
+                (f"--{k}={v}" if v is not None else f"--{k}")
+                for k, v in mace_params.items()
+            ]
+        )
+    )
+
+    p = subprocess.run(cmd.split(), env=run_env, check=True)
+
+    assert p.returncode == 0
+
+    return MACECalculator(tmp_path / "MACE.model", device="cpu")
+
+
+def test_calculator_forces(fitting_configs, trained_model):
+    at = fitting_configs[2].copy()
+    at.calc = trained_model
+
+    # test just forces
+    grads = gradient_test(at)
+
+    assert np.allclose(grads[0], grads[1])
+
+
+def test_calculator_stress(fitting_configs, trained_model):
+    at = fitting_configs[2].copy()
+    at.calc = trained_model
+
+    # test forces and stress
+    at_wrapped = ExpCellFilter(at)
+    grads = gradient_test(at_wrapped)
+
+    assert np.allclose(grads[0], grads[1])

tests/test_run_train.py

@@ -204,6 +204,7 @@ def test_run_train_missing_data(tmp_path, fitting_configs):
     ]
     assert np.allclose(Es, ref_Es)
 
+
 def test_run_train_no_stress(tmp_path, fitting_configs):
     del fitting_configs[5].info["REF_energy"]
     del fitting_configs[6].arrays["REF_forces"]
@@ -270,6 +271,6 @@ def test_run_train_no_stress(tmp_path, fitting_configs):
         -0.06608313576078294,
         -0.36358220540264646,
         -0.12097397940768086,
-        0.002021055463491156
-        ]
+        0.002021055463491156,
+    ]
     assert np.allclose(Es, ref_Es)