Fix NPH and NVT_NH (#429)

* Fix NPH and NVT_NH logging

* Add stats file test

Author: k-harris27

janus_core/calculations/md.py

@@ -1250,8 +1250,7 @@ def _set_param_prefix(self, file_prefix: PathLike | None = None) -> str:
         if file_prefix is not None:
             return ""
 
-        pressure = f"-p{self.pressure}" if not isinstance(self, NVT_NH) else ""
-        return f"{super()._set_param_prefix(file_prefix)}{pressure}"
+        return f"{super()._set_param_prefix(file_prefix)}-p{self.pressure}"
 
     def get_stats(self) -> dict[str, float]:
         """
@@ -1262,7 +1261,7 @@ def get_stats(self) -> dict[str, float]:
         dict[str, float]
             Thermodynamical statistics to be written out.
         """
-        stats = MolecularDynamics.get_stats(self)
+        stats = super().get_stats()
         stats |= {"Target_P": self.pressure, "Target_T": self.temp}
         return stats
 
@@ -1434,7 +1433,7 @@ def __init__(
         )
 
 
-class NVT_NH(NPT):  # noqa: N801 (invalid-class-name)
+class NVT_NH(MolecularDynamics):  # noqa: N801 (invalid-class-name)
     """
     Configure NVT Nosé-Hoover simulation.
 
@@ -1446,6 +1445,9 @@ class NVT_NH(NPT):  # noqa: N801 (invalid-class-name)
         Thermostat time, in fs. Default is 50.0.
     ensemble
         Name for thermodynamic ensemble. Default is "nvt-nh".
+    file_prefix
+        Prefix for output filenames. Default is inferred from structure, ensemble,
+        temperature, and pressure.
     ensemble_kwargs
         Keyword arguments to pass to ensemble initialization. Default is {}.
     **kwargs
@@ -1457,6 +1459,7 @@ def __init__(
         *args,
         thermostat_time: float = 50.0,
         ensemble: Ensembles = "nvt-nh",
+        file_prefix: PathLike | None = None,
         ensemble_kwargs: dict[str, Any] | None = None,
         **kwargs,
     ) -> None:
@@ -1471,19 +1474,26 @@ def __init__(
             Thermostat time, in fs. Default is 50.0.
         ensemble
             Name for thermodynamic ensemble. Default is "nvt-nh".
+        file_prefix
+            Prefix for output filenames. Default is inferred from structure, ensemble,
+            temperature, and pressure.
         ensemble_kwargs
             Keyword arguments to pass to ensemble initialization. Default is {}.
         **kwargs
             Additional keyword arguments.
         """
+        super().__init__(*args, ensemble=ensemble, file_prefix=file_prefix, **kwargs)
+
         (ensemble_kwargs,) = none_to_dict(ensemble_kwargs)
-        super().__init__(
-            *args,
-            ensemble=ensemble,
-            thermostat_time=thermostat_time,
-            barostat_time=None,
-            ensemble_kwargs=ensemble_kwargs,
-            **kwargs,
+        self.ttime = thermostat_time * units.fs
+
+        self.dyn = ASE_NPT(
+            self.struct,
+            timestep=self.timestep,
+            temperature_K=self.temp,
+            ttime=self.ttime,
+            append_trajectory=self.traj_append,
+            **ensemble_kwargs,
         )
 
     def get_stats(self) -> dict[str, float]:
@@ -1495,7 +1505,7 @@ def get_stats(self) -> dict[str, float]:
         dict[str, float]
             Thermodynamical statistics to be written out.
         """
-        stats = MolecularDynamics.get_stats(self)
+        stats = super().get_stats()
         stats |= {"Target_T": self.temp}
         return stats
 
@@ -1585,16 +1595,16 @@ def __init__(
         )
 
 
-class NPH(NPT):
+class NPH(MolecularDynamics):
     """
     Configure NPH simulation.
 
     Parameters
     ----------
     *args
         Additional arguments.
-    thermostat_time
-        Thermostat time, in fs. Default is 50.0.
+    barostat_time
+        Barostat time, in fs. Default is 75.0.
     bulk_modulus
         Bulk modulus, in GPa. Default is 2.0.
     pressure
@@ -1613,7 +1623,7 @@ class NPH(NPT):
     def __init__(
         self,
         *args,
-        thermostat_time: float = 50.0,
+        barostat_time: float = 75.0,
         bulk_modulus: float = 2.0,
         pressure: float = 0.0,
         ensemble: Ensembles = "nph",
@@ -1628,8 +1638,8 @@ def __init__(
         ----------
         *args
             Additional arguments.
-        thermostat_time
-            Thermostat time, in fs. Default is 50.0.
+        barostat_time
+            Barostat time, in fs. Default is 75.0.
         bulk_modulus
             Bulk modulus, in GPa. Default is 2.0.
         pressure
@@ -1644,19 +1654,87 @@ def __init__(
         **kwargs
             Additional keyword arguments.
         """
+        self.pressure = pressure
+        super().__init__(*args, ensemble=ensemble, file_prefix=file_prefix, **kwargs)
+
         (ensemble_kwargs,) = none_to_dict(ensemble_kwargs)
-        super().__init__(
-            *args,
-            thermostat_time=thermostat_time,
-            barostat_time=None,
-            bulk_modulus=bulk_modulus,
-            pressure=pressure,
-            ensemble=ensemble,
-            file_prefix=file_prefix,
-            ensemble_kwargs=ensemble_kwargs,
-            **kwargs,
+
+        pfactor = barostat_time**2 * bulk_modulus
+        if self.logger:
+            self.logger.info("NPT pfactor=%s GPa fs^2", pfactor)
+
+        # convert the pfactor to ASE internal units
+        pfactor *= units.fs**2 * units.GPa
+
+        self.dyn = ASE_NPT(
+            self.struct,
+            timestep=self.timestep,
+            temperature_K=self.temp,
+            ttime=None,
+            pfactor=pfactor,
+            append_trajectory=self.traj_append,
+            externalstress=self.pressure * units.GPa,
+            **ensemble_kwargs,
         )
 
+    def _set_param_prefix(self, file_prefix: PathLike | None = None) -> str:
+        """
+        Set ensemble parameters for output files.
+
+        Parameters
+        ----------
+        file_prefix
+            Prefix for output filenames on class init. If not None, param_prefix = "".
+
+        Returns
+        -------
+        str
+           Formatted ensemble parameters, including pressure and temperature(s).
+        """
+        if file_prefix is not None:
+            return ""
+
+        return f"{super()._set_param_prefix(file_prefix)}-p{self.pressure}"
+
+    def get_stats(self) -> dict[str, float]:
+        """
+        Get thermodynamical statistics to be written to file.
+
+        Returns
+        -------
+        dict[str, float]
+            Thermodynamical statistics to be written out.
+        """
+        stats = super().get_stats()
+        stats |= {"Target_P": self.pressure}
+        return stats
+
+    @property
+    def unit_info(self) -> dict[str, str]:
+        """
+        Get units of returned statistics.
+
+        Returns
+        -------
+        dict[str, str]
+            Units attached to statistical properties.
+        """
+        return super().unit_info | {
+            "Target_P": JANUS_UNITS["pressure"],
+        }
+
+    @property
+    def default_formats(self) -> dict[str, str]:
+        """
+        Default format of returned statistics.
+
+        Returns
+        -------
+        dict[str, str]
+            Default formats attached to statistical properties.
+        """
+        return super().default_formats | {"Target_P": ".5f"}
+
 
 class NPT_MTK(MolecularDynamics):  # noqa: N801 (invalid-class-name)
     """

janus_core/cli/md.py

@@ -42,7 +42,7 @@ def md(
         Option(
             help=(
                 """
-                Thermostat time for NPT, NPT-MTK, NVT Nosé-Hoover, or NPH simulation,
+                Thermostat time for NPT, NPT-MTK or NVT Nosé-Hoover simulation,
                 in fs. Default is 50 fs for NPT and NVT Nosé-Hoover, or 100 fs for
                 NPT-MTK.
                 """
@@ -235,8 +235,8 @@ def md(
     temp
         Temperature, in K. Default is 300.
     thermostat_time
-        Thermostat time for NPT, NPT-MTK, NVT Nosé-Hoover or NPH simulation,
-        in fs. Default is 50 fs for NPT, NVT Nosé-Hoover and NPH, or 100 fs for NPT-MTK.
+        Thermostat time for NPT, NPT-MTK or NVT Nosé-Hoover simulation,
+        in fs. Default is 50 fs for NPT and NVT Nosé-Hoover, or 100 fs for NPT-MTK.
     barostat_time
         Barostat time for NPT, NPT-MTK or NPH simulation, in fs.
         Default is 75 fs for NPT and NPH, or 1000 fs for NPT-MTK.
@@ -487,7 +487,7 @@ def md(
     elif ensemble == "nph":
         for key in (
             "friction",
-            "barostat_time",
+            "thermostat_time",
             "taut",
             "thermostat_chain",
             "barostat_chain",

tests/test_md.py

@@ -667,6 +667,33 @@ def test_atoms_struct(tmp_path):
         assert len(lines) == 6
 
 
+@pytest.mark.parametrize("ensemble, tag", test_data)
+def test_stats(tmp_path, ensemble, tag):
+    """Test stats file has correct structure and entries for all ensembles."""
+    file_prefix = tmp_path / tag / "NaCl"
+    single_point = SinglePoint(
+        struct_path=DATA_PATH / "NaCl.cif",
+        arch="mace",
+        calc_kwargs={"model": MODEL_PATH},
+    )
+    md = ensemble(
+        struct=single_point.struct,
+        steps=2,
+        stats_every=1,
+        file_prefix=file_prefix,
+    )
+    md.run()
+
+    stat_data = Stats(md.stats_file)
+
+    etot_index = stat_data.labels.index("ETot/N")
+
+    assert stat_data.columns == len(stat_data.labels)
+    assert stat_data.columns == len(stat_data.units)
+    assert stat_data.columns >= 16
+    assert stat_data.units[etot_index] == "eV"
+
+
 def test_heating(tmp_path):
     """Test heating with no MD."""
     file_prefix = tmp_path / "NaCl-heating"