Fix info

janus_core/calculations/md.py

@@ -509,12 +509,21 @@ def __init__(
 
         self._parse_correlations()
 
-    def _set_time_step(self):
-        """Set time in fs and current dynamics step to info."""
+    def _set_info(self):
+        """Set time in fs, current dynamics step, and density to info."""
         time = (self.offset * self.timestep + self.dyn.get_time()) / units.fs
         step = self.offset + self.dyn.nsteps
         self.dyn.atoms.info["time_fs"] = time
         self.dyn.atoms.info["step"] = step
+        try:
+            density = (
+                np.sum(self.dyn.atoms.get_masses())
+                / self.dyn.atoms.get_volume()
+                * DENS_FACT
+            )
+            self.dyn.atoms.info["density"] = density
+        except ValueError:
+            self.dyn.atoms.info["density"] = 0.0
 
     def _prepare_restart(self) -> None:
         """Prepare restart files, structure and offset."""
@@ -726,19 +735,13 @@ def get_stats(self) -> dict[str, float]:
         e_kin = self.dyn.atoms.get_kinetic_energy() / self.n_atoms
         current_temp = e_kin / (1.5 * units.kB)
 
-        self._set_time_step()
+        self._set_info()
 
         time_now = datetime.datetime.now()
         real_time = time_now - self.dyn.atoms.info["real_time"]
         self.dyn.atoms.info["real_time"] = time_now
 
         try:
-            density = (
-                np.sum(self.dyn.atoms.get_masses())
-                / self.dyn.atoms.get_volume()
-                * DENS_FACT
-            )
-            self.dyn.atoms.info["density"] = density
             volume = self.dyn.atoms.get_volume()
             pressure = (
                 -np.trace(
@@ -754,7 +757,6 @@ def get_stats(self) -> dict[str, float]:
         except ValueError:
             volume = 0.0
             pressure = 0.0
-            density = 0.0
             pressure_tensor = np.zeros(6)
 
         return {
@@ -765,7 +767,7 @@ def get_stats(self) -> dict[str, float]:
             "EKin/N": e_kin,
             "T": current_temp,
             "ETot/N": e_pot + e_kin,
-            "Density": density,
+            "Density": self.dyn.atoms.info["density"],
             "Volume": volume,
             "P": pressure,
             "Pxx": pressure_tensor[0],
@@ -874,7 +876,7 @@ def _write_traj(self) -> None:
                 self.dyn.nsteps > self.traj_start + self.traj_start % self.traj_every
             )
 
-            self._set_time_step()
+            self._set_info()
             write_kwargs = self.write_kwargs
             write_kwargs["filename"] = self.traj_file
             write_kwargs["append"] = append
@@ -895,7 +897,7 @@ def _write_final_state(self) -> None:
         # Append if final file has been created
         append = self.created_final_file
 
-        self._set_time_step()
+        self._set_info()
         write_kwargs = self.write_kwargs
         write_kwargs["filename"] = self.final_file
         write_kwargs["append"] = append
@@ -998,7 +1000,7 @@ def _write_restart(self) -> None:
         if step > 0:
             write_kwargs = self.write_kwargs
             write_kwargs["filename"] = self._restart_file
-            self._set_time_step()
+            self._set_info()
 
             output_structs(
                 images=self.struct,

tests/test_md.py

@@ -1058,3 +1058,30 @@ def test_auto_restart_restart_stem(tmp_path):
 
     final_traj = read(traj_path, index=":")
     assert len(final_traj) == 9
+
+
+def test_set_info(tmp_path):
+    """Test info is set at correct frequency."""
+    file_prefix = tmp_path / "npt"
+    traj_path = tmp_path / "npt-traj.extxyz"
+
+    single_point = SinglePoint(
+        struct_path=DATA_PATH / "NaCl.cif",
+        arch="mace",
+        calc_kwargs={"model": MODEL_PATH},
+    )
+
+    npt = NPT(
+        struct=single_point.struct,
+        steps=10,
+        temp=1000,
+        stats_every=7,
+        file_prefix=file_prefix,
+        seed=2024,
+        traj_every=10,
+    )
+
+    npt.run()
+    final_struct = read(traj_path, index="-1")
+    assert npt.struct.info["density"] == pytest.approx(2.120952627887493)
+    assert final_struct.info["density"] == pytest.approx(2.120952627887493)