NPT support (#1972)

add NPT support + log T, P, KE at each frame as well.

Author: mshuaibii

src/fairchem/core/components/calculate/__init__.py

@@ -25,6 +25,7 @@
 from ._single.relaxation_runner import RelaxationRunner
 from ._single.singlepoint_runner import SinglePointRunner
 from .simulation_tools.thermostats import (
+    BerendsenNPT,
     BussiThermostat,
     LangevinThermostat,
     NoseHooverNVT,
@@ -37,6 +38,7 @@
     "AdsorbMLRunner",
     "AdsorptionRunner",
     "AdsorptionSinglePointRunner",
+    "BerendsenNPT",
     "BussiThermostat",
     "ElasticityRunner",
     "KappaRunner",

src/fairchem/core/components/calculate/simulation_tools/thermostats.py

@@ -16,6 +16,7 @@
 from ase.md.bussi import Bussi
 from ase.md.langevin import Langevin
 from ase.md.nose_hoover_chain import NoseHooverChainNVT
+from ase.md.nptberendsen import NPTBerendsen
 from ase.md.verlet import VelocityVerlet
 from monty.json import jsanitize
 
@@ -210,3 +211,33 @@ def restore_state(self, dyn: MolecularDynamics, state: dict[str, Any]) -> None:
                 rng["cached_gaussian"],
             )
         )
+
+
+@dataclass
+class BerendsenNPT(Thermostat):
+    """
+    Berendsen NPT thermostat/barostat for constant pressure simulations.
+    """
+
+    temperature_K: float
+    pressure_bar: float = 1.0
+    taut_fs: float = 5.0
+    taup_fs: float = 500.0
+    compressibility_bar: float = 5e-7
+
+    def build(self, atoms: Atoms, timestep_fs: float) -> MolecularDynamics:
+        return NPTBerendsen(
+            atoms=atoms,
+            timestep=timestep_fs * units.fs,
+            temperature_K=self.temperature_K,
+            pressure_au=self.pressure_bar * units.bar,
+            taut=self.taut_fs * units.fs,
+            taup=self.taup_fs * units.fs,
+            compressibility_au=self.compressibility_bar / units.bar,
+        )
+
+    def save_state(self, dyn: MolecularDynamics) -> dict[str, Any]:
+        return {"class_name": "BerendsenNPT"}
+
+    def restore_state(self, dyn: MolecularDynamics, state: dict[str, Any]) -> None:
+        pass

src/fairchem/core/components/calculate/simulation_tools/trajectory.py

@@ -11,6 +11,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
+import ase.units
 import numpy as np
 import pyarrow as pa
 import pyarrow.parquet as pq
@@ -42,6 +43,9 @@ class TrajectoryFrame:
     energy: float | None = None
     forces: np.ndarray | None = None  # (N, 3)
     stress: np.ndarray | None = None  # (6,) Voigt notation
+    temperature: float | None = None  # Kelvin
+    kinetic_energy: float | None = None  # eV
+    pressure: float | None = None  # bar
     sid: str | int | None = None
 
     def to_dict(self) -> dict:
@@ -67,6 +71,12 @@ def to_dict(self) -> dict:
             d["forces"] = self.forces.tolist()
         if self.stress is not None:
             d["stress"] = self.stress.tolist()
+        if self.temperature is not None:
+            d["temperature"] = self.temperature
+        if self.kinetic_energy is not None:
+            d["kinetic_energy"] = self.kinetic_energy
+        if self.pressure is not None:
+            d["pressure"] = self.pressure
         return d
 
     @classmethod
@@ -104,6 +114,22 @@ def from_atoms(
         except (PropertyNotImplementedError, RuntimeError):
             velocities = None
 
+        try:
+            temperature = atoms.get_temperature()
+        except Exception:
+            temperature = None
+
+        try:
+            kinetic_energy = atoms.get_kinetic_energy()
+        except Exception:
+            kinetic_energy = None
+
+        # Pressure from stress: P = -trace(stress)/3, converted to bar
+        if stress is not None:
+            pressure = -stress[:3].mean() / ase.units.bar
+        else:
+            pressure = None
+
         return cls(
             step=step,
             time=time,
@@ -115,6 +141,9 @@ def from_atoms(
             energy=energy,
             forces=forces,
             stress=stress,
+            temperature=temperature,
+            kinetic_energy=kinetic_energy,
+            pressure=pressure,
             sid=atoms.info.get("sid"),
         )
 

tests/core/components/test_md_runner.py

@@ -26,6 +26,7 @@
 from ase.md.verlet import VelocityVerlet
 
 from fairchem.core.components.calculate import (
+    BerendsenNPT,
     BussiThermostat,
     LangevinThermostat,
     MDRunner,
@@ -148,8 +149,15 @@ def test_md_correctness_vs_ase(self, cu_atoms, results_dir):
             NoseHooverNVT(temperature_K=300.0, tdamp_fs=25.0),
             BussiThermostat(temperature_K=300.0, taut_fs=25.0),
             LangevinThermostat(temperature_K=300.0, friction_per_fs=0.01),
+            BerendsenNPT(
+                temperature_K=300.0,
+                pressure_bar=1.0,
+                taut_fs=500.0,
+                taup_fs=1000.0,
+                compressibility_bar=1.0 / 140e9,
+            ),
         ],
-        ids=["VelocityVerlet", "NoseHoover", "Bussi", "Langevin"],
+        ids=["VelocityVerlet", "NoseHoover", "Bussi", "Langevin", "BerendsenNPT"],
     )
     def test_checkpoint_resume(self, cu_atoms, results_dir, thermostat):
         """
@@ -333,3 +341,38 @@ def test_stopfair_graceful_stop(self, cu_atoms, results_dir):
 
         traj_df = pd.read_parquet(results["trajectory_file"])
         assert list(traj_df["step"]) == [0, 10, 20]
+
+    def test_npt_cell_changes(self, cu_atoms, results_dir):
+        """
+        Verify that NPT simulation changes the cell volume.
+        """
+        md_results_dir = results_dir / "results"
+        md_results_dir.mkdir()
+
+        atoms = cu_atoms.copy()
+        initial_volume = atoms.get_volume()
+
+        # Use a large pressure to drive a noticeable volume change
+        thermostat = BerendsenNPT(
+            temperature_K=300.0,
+            pressure_bar=1e5,
+            taut_fs=100.0,
+            taup_fs=100.0,
+            compressibility_bar=1.0 / 140e9,
+        )
+
+        runner = MDRunner(
+            calculator=EMT(),
+            atoms=atoms,
+            thermostat=thermostat,
+            timestep_fs=1.0,
+            steps=200,
+            trajectory_interval=50,
+            log_interval=50,
+            trajectory_writer=partial(ParquetTrajectoryWriter, flush_interval=1000),
+        )
+        runner._job_config = _create_mock_job_config(str(md_results_dir))
+        runner.calculate(job_num=0, num_jobs=1)
+
+        final_volume = atoms.get_volume()
+        assert initial_volume != pytest.approx(final_volume, rel=1e-6)