Fix temperature saved during MD (#469)

* Rename stats units property

* Fix writing current MD temperature

* Test stats and traj data is consistent

Author: ElliottKasoar

janus_core/calculations/base.py

@@ -26,6 +26,7 @@
     "time": "fs",
     "real_time": "s",
     "temperature": "K",
+    "target_temperature": "K",
     "pressure": "GPa",
     "momenta": "(eV*u)^0.5",
     "density": "g/cm^3",

janus_core/calculations/md.py

@@ -517,7 +517,7 @@ def __init__(
         self._parse_correlations()
 
     def _set_info(self) -> None:
-        """Set time in fs, current dynamics step, and density to info."""
+        """Set time in fs, current dynamics step, density, and temperature to info."""
         time = (self.offset * self.timestep + self.dyn.get_time()) / units.fs
         step = self.offset + self.dyn.nsteps
         self.dyn.atoms.info["time"] = time
@@ -532,6 +532,13 @@ def _set_info(self) -> None:
         except ValueError:
             self.dyn.atoms.info["density"] = 0.0
 
+        e_kin = self.dyn.atoms.get_kinetic_energy() / self.n_atoms
+        current_temp = e_kin / (1.5 * units.kB)
+        self.struct.info["temperature"] = current_temp
+
+        if hasattr(self.dyn, "set_temperature") or isinstance(self, NVT_CSVR | NPT_MTK):
+            self.struct.info["target_temperature"] = self.temp
+
     def _prepare_restart(self) -> None:
         """Prepare restart files, structure and offset."""
         # Check offset can be read from steps
@@ -727,6 +734,28 @@ def _write_correlations(self) -> None:
                     data[str(cor)] = {"value": value.tolist(), "lags": lags.tolist()}
                 yaml.dump(data, out_file, default_flow_style=None)
 
+    @property
+    def info_unit_keys(self) -> tuple[str]:
+        """
+        Get Atoms.info keys to save units for.
+
+        Returns
+        -------
+        tuple[str]
+            Keys of Atoms.info.
+        """
+        return (
+            "energy",
+            "forces",
+            "stress",
+            "time",
+            "real_time",
+            "temperature",
+            "pressure",
+            "density",
+            "momenta",
+        )
+
     def get_stats(self) -> dict[str, float]:
         """
         Get thermodynamical statistics to be written to file.
@@ -738,7 +767,6 @@ def get_stats(self) -> dict[str, float]:
         """
         e_pot = self.dyn.atoms.get_potential_energy() / self.n_atoms
         e_kin = self.dyn.atoms.get_kinetic_energy() / self.n_atoms
-        current_temp = e_kin / (1.5 * units.kB)
 
         self._set_info()
 
@@ -770,7 +798,7 @@ def get_stats(self) -> dict[str, float]:
             "Time": self.dyn.atoms.info["time"],
             "Epot/N": e_pot,
             "EKin/N": e_kin,
-            "T": current_temp,
+            "T": self.dyn.atoms.info["temperature"],
             "ETot/N": e_pot + e_kin,
             "Density": self.dyn.atoms.info["density"],
             "Volume": volume,
@@ -784,7 +812,7 @@ def get_stats(self) -> dict[str, float]:
         }
 
     @property
-    def unit_info(self) -> dict[str, str]:
+    def stats_units(self) -> dict[str, str]:
         """
         Get units of returned statistics.
 
@@ -847,8 +875,8 @@ def _write_header(self) -> None:
             write_table(
                 "ascii",
                 file=stats_file,
-                units=self.unit_info,
-                **{key: () for key in self.unit_info},
+                units=self.stats_units,
+                **{key: () for key in self.stats_units},
             )
 
     def _write_stats_file(self) -> None:
@@ -863,7 +891,7 @@ def _write_stats_file(self) -> None:
             write_table(
                 "ascii",
                 file=stats_file,
-                units=self.unit_info,
+                units=self.stats_units,
                 formats=self.default_formats,
                 print_header=False,
                 **stats,
@@ -895,10 +923,6 @@ def _write_traj(self) -> None:
 
     def _write_final_state(self) -> None:
         """Write the final system state."""
-        self.struct.info["temperature"] = self.temp
-        if isinstance(self, NPT) and not isinstance(self, NVT_NH):
-            self.struct.info["pressure"] = self.pressure
-
         # Append if final file has been created
         append = self.created_final_file
 
@@ -1046,18 +1070,7 @@ def _set_target_temperature(self, temperature: float):
 
     def run(self) -> None:
         """Run molecular dynamics simulation and/or temperature ramp."""
-        unit_keys = (
-            "energy",
-            "forces",
-            "stress",
-            "time",
-            "real_time",
-            "temperature",
-            "pressure",
-            "density",
-            "momenta",
-        )
-        self._set_info_units(unit_keys)
+        self._set_info_units(self.info_unit_keys)
 
         if not self.restart:
             if self.minimize:
@@ -1299,6 +1312,18 @@ def _set_param_prefix(self, file_prefix: PathLike | None = None) -> str:
 
         return f"{super()._set_param_prefix(file_prefix)}-p{self.pressure}"
 
+    @property
+    def info_unit_keys(self) -> tuple[str]:
+        """
+        Get Atoms.info keys to save units for.
+
+        Returns
+        -------
+        tuple[str]
+            Keys of Atoms.info.
+        """
+        return super().info_unit_keys + ("target_temperature",)
+
     def get_stats(self) -> dict[str, float]:
         """
         Get thermodynamical statistics to be written to file.
@@ -1313,7 +1338,7 @@ def get_stats(self) -> dict[str, float]:
         return stats
 
     @property
-    def unit_info(self) -> dict[str, str]:
+    def stats_units(self) -> dict[str, str]:
         """
         Get units of returned statistics.
 
@@ -1322,7 +1347,7 @@ def unit_info(self) -> dict[str, str]:
         dict[str, str]
             Units attached to statistical properties.
         """
-        return super().unit_info | {
+        return super().stats_units | {
             "Target_P": JANUS_UNITS["pressure"],
             "Target_T": JANUS_UNITS["temperature"],
         }
@@ -1394,6 +1419,18 @@ def __init__(
             **ensemble_kwargs,
         )
 
+    @property
+    def info_unit_keys(self) -> tuple[str]:
+        """
+        Get Atoms.info keys to save units for.
+
+        Returns
+        -------
+        tuple[str]
+            Keys of Atoms.info.
+        """
+        return super().info_unit_keys + ("target_temperature",)
+
     def get_stats(self) -> dict[str, float]:
         """
         Get thermodynamical statistics to be written to file.
@@ -1408,7 +1445,7 @@ def get_stats(self) -> dict[str, float]:
         return stats
 
     @property
-    def unit_info(self) -> dict[str, str]:
+    def stats_units(self) -> dict[str, str]:
         """
         Get units of returned statistics.
 
@@ -1417,7 +1454,7 @@ def unit_info(self) -> dict[str, str]:
         dict[str, str]
             Units attached to statistical properties.
         """
-        return super().unit_info | {"Target_T": JANUS_UNITS["temperature"]}
+        return super().stats_units | {"Target_T": JANUS_UNITS["temperature"]}
 
     @property
     def default_formats(self) -> dict[str, str]:
@@ -1543,6 +1580,18 @@ def __init__(
             **ensemble_kwargs,
         )
 
+    @property
+    def info_unit_keys(self) -> tuple[str]:
+        """
+        Get Atoms.info keys to save units for.
+
+        Returns
+        -------
+        tuple[str]
+            Keys of Atoms.info.
+        """
+        return super().info_unit_keys + ("target_temperature",)
+
     def get_stats(self) -> dict[str, float]:
         """
         Get thermodynamical statistics to be written to file.
@@ -1557,7 +1606,7 @@ def get_stats(self) -> dict[str, float]:
         return stats
 
     @property
-    def unit_info(self) -> dict[str, str]:
+    def stats_units(self) -> dict[str, str]:
         """
         Get units of returned statistics.
 
@@ -1566,7 +1615,7 @@ def unit_info(self) -> dict[str, str]:
         dict[str, str]
             Units attached to statistical properties.
         """
-        return super().unit_info | {"Target_T": JANUS_UNITS["temperature"]}
+        return super().stats_units | {"Target_T": JANUS_UNITS["temperature"]}
 
     @property
     def default_formats(self) -> dict[str, str]:
@@ -1757,7 +1806,7 @@ def get_stats(self) -> dict[str, float]:
         return stats
 
     @property
-    def unit_info(self) -> dict[str, str]:
+    def stats_units(self) -> dict[str, str]:
         """
         Get units of returned statistics.
 
@@ -1766,7 +1815,7 @@ def unit_info(self) -> dict[str, str]:
         dict[str, str]
             Units attached to statistical properties.
         """
-        return super().unit_info | {
+        return super().stats_units | {
             "Target_P": JANUS_UNITS["pressure"],
         }
 
@@ -1906,6 +1955,18 @@ def _set_param_prefix(self, file_prefix: PathLike | None = None) -> str:
         pressure = f"-p{self.pressure}"
         return f"{super()._set_param_prefix(file_prefix)}{pressure}"
 
+    @property
+    def info_unit_keys(self) -> tuple[str]:
+        """
+        Get Atoms.info keys to save units for.
+
+        Returns
+        -------
+        tuple[str]
+            Keys of Atoms.info.
+        """
+        return super().info_unit_keys + ("target_temperature",)
+
     def get_stats(self) -> dict[str, float]:
         """
         Get thermodynamical statistics to be written to file.
@@ -1920,7 +1981,7 @@ def get_stats(self) -> dict[str, float]:
         return stats
 
     @property
-    def unit_info(self) -> dict[str, str]:
+    def stats_units(self) -> dict[str, str]:
         """
         Get units of returned statistics.
 
@@ -1929,7 +1990,7 @@ def unit_info(self) -> dict[str, str]:
         dict[str, str]
             Units attached to statistical properties.
         """
-        return super().unit_info | {
+        return super().stats_units | {
             "Target_P": JANUS_UNITS["pressure"],
             "Target_T": JANUS_UNITS["temperature"],
         }

tests/test_md_cli.py

@@ -13,6 +13,7 @@
 import yaml
 
 from janus_core.cli.janus import app
+from janus_core.helpers.stats import Stats
 from tests.utils import assert_log_contains, clear_log_handlers, strip_ansi_codes
 
 if hasattr(ase.md.nose_hoover_chain, "IsotropicMTKNPT"):
@@ -815,3 +816,64 @@ def test_no_carbon(tmp_path):
     with open(summary_path, encoding="utf8") as file:
         summary = yaml.safe_load(file)
     assert "emissions" not in summary
+
+
+@pytest.mark.parametrize("ensemble", ("nvt", "npt", "nvt-csvr"))
+@pytest.mark.parametrize("output_every", (1, 2))
+@pytest.mark.parametrize("heating", (True, False))
+def test_consistent_stats_traj(tmp_path, ensemble, output_every, heating):
+    """Test data saved to statistics is consistent with trajectory info."""
+    file_prefix = tmp_path / ensemble
+    stats_path = tmp_path / f"{ensemble}-stats.dat"
+    traj_path = tmp_path / f"{ensemble}-traj.extxyz"
+
+    inputs = [
+        "md",
+        "--ensemble",
+        "nvt",
+        "--struct",
+        DATA_PATH / "NaCl.cif",
+        "--no-tracker",
+        "--file-prefix",
+        file_prefix,
+        "--steps",
+        4,
+        "--stats-every",
+        output_every,
+        "--traj-every",
+        output_every,
+    ]
+
+    if heating:
+        inputs = inputs + [
+            "--temp-start",
+            10,
+            "--temp-end",
+            20,
+            "--temp-step",
+            10,
+            "--temp-time",
+            2,
+        ]
+
+    result = runner.invoke(app, inputs)
+    assert result.exit_code == 0
+
+    atoms = read(traj_path, index=":")
+    temps = [a.info["temperature"] for a in atoms]
+
+    data = Stats(stats_path)
+    temp_col = data.labels.index("T")
+    stats_temps = data.data[:, temp_col]
+
+    assert len(temps) == len(stats_temps)
+    assert temps == pytest.approx(stats_temps, rel=1e5)
+
+    if ensemble in ("nvt", "npt"):
+        target_temps = [a.info["target_temperature"] for a in atoms]
+
+        target_temp_col = data.labels.index("Target_T")
+        stats_target_temps = data.data[:, target_temp_col]
+
+        assert len(target_temps) == len(stats_target_temps)
+        assert target_temps == pytest.approx(stats_target_temps, rel=1e5)