fix heating fixes #133 (#134)

* Fix temperatures during heating

* Fix cooling ramp

* Save structure when 0K in heating ramp

---------

Co-authored-by: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>

janus_core/calculations/md.py

@@ -2,6 +2,7 @@
 """Run molecular dynamics simulations."""
 
 import datetime
+from math import isclose
 from pathlib import Path
 import random
 from typing import Any, Optional
@@ -292,16 +293,25 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-locals,too-many-sta
 
         # Warn if mix of None and not None
         self.ramp_temp = (
-            self.temp_start and self.temp_end and self.temp_step and self.temp_time
+            self.temp_start is not None
+            and self.temp_end is not None
+            and self.temp_step
+            and self.temp_time
         )
         if (
-            self.temp_start or self.temp_end or self.temp_step or self.temp_time
+            self.temp_start is not None
+            or self.temp_end is not None
+            or self.temp_step
+            or self.temp_time
         ) and not self.ramp_temp:
             warn(
                 "`temp_start`, `temp_end` and `temp_step` must all be specified for "
                 "heating to run"
             )
 
+        if self.ramp_temp and (self.temp_start < 0 or self.temp_end < 0):
+            raise ValueError("Start and end temperatures must be positive")
+
         self.minimize_kwargs = minimize_kwargs if minimize_kwargs else {}
         self.restart_files = []
         self.dyn = None
@@ -335,12 +345,16 @@ def _set_velocity_distribution(self) -> None:
         Sets Maxwell-Boltzmann velocity distribution, as well as removing
         centre-of-mass momentum, and (optionally) total angular momentum.
         """
-        MaxwellBoltzmannDistribution(self.struct, temperature_K=self.temp)
-        Stationary(self.struct)
+        atoms = self.struct
+        if self.dyn.nsteps >= 0:
+            atoms = self.dyn.atoms
+
+        MaxwellBoltzmannDistribution(atoms, temperature_K=self.temp)
+        Stationary(atoms)
         if self.logger:
             self.logger.info("Velocities reset at step %s", self.dyn.nsteps)
         if self.remove_rot:
-            ZeroRotation(self.struct)
+            ZeroRotation(atoms)
             if self.logger:
                 self.logger.info("Rotation reset at step %s", self.dyn.nsteps)
 
@@ -373,7 +387,7 @@ def _parameter_prefix(self) -> str:
         """
 
         temperature_prefix = ""
-        if self.temp_start and self.temp_end:
+        if self.temp_start is not None and self.temp_end is not None:
             temperature_prefix = f"-T{self.temp_start}-T{self.temp_end}"
 
         if self.steps > 0:
@@ -601,13 +615,26 @@ def _run_dynamics(self) -> None:
         if self.ramp_temp:
             heating_steps = int(self.temp_time // self.timestep)
 
-            n_temps = int(1 + (self.temp_end - self.temp_start) // self.temp_step)
-            temps = [self.temp_start + i * self.temp_step for i in range(n_temps)]
+            # Always include start temperature in ramp, and include end temperature
+            # if separated by an integer number of temperature steps
+            n_temps = int(1 + abs(self.temp_end - self.temp_start) // self.temp_step)
+
+            # Add or subtract temperatures
+            ramp_sign = 1 if (self.temp_end - self.temp_start) > 0 else -1
+            temps = [
+                self.temp_start + ramp_sign * i * self.temp_step for i in range(n_temps)
+            ]
 
             if self.logger:
                 self.logger.info("Beginning temperature ramp at %sK", temps[0])
             for temp in temps:
                 self.temp = temp
+                self._set_velocity_distribution()
+                if isclose(temp, 0.0):
+                    self._write_final_state()
+                    continue
+                if not isinstance(self, NVE):
+                    self.dyn.set_temperature(temperature_K=self.temp)
                 self.dyn.run(heating_steps)
                 self._write_final_state()
             if self.logger:
@@ -618,6 +645,10 @@ def _run_dynamics(self) -> None:
             if self.logger:
                 self.logger.info("Starting molecular dynamics simulation")
             self.temp = md_temp
+            if self.ramp_temp:
+                self._set_velocity_distribution()
+                if not isinstance(self, NVE):
+                    self.dyn.set_temperature(temperature_K=self.temp)
             self.dyn.run(self.steps)
             self._write_final_state()
             if self.logger:
@@ -697,7 +728,6 @@ def __init__(
             pfactor = (barostat_time * units.fs) ** 2 * scaled_bulk_modulus
         else:
             pfactor = None
-
         self.dyn = ASE_NPT(
             self.struct,
             timestep=self.timestep,

tests/test_md.py

@@ -549,7 +549,10 @@ def test_atoms_struct(tmp_path):
 
 def test_heating(tmp_path):
     """Test heating with no MD."""
+    # pylint: disable=invalid-name
     file_prefix = tmp_path / "NaCl-heating"
+    final_T0_file = tmp_path / "NaCl-heating-T0.0-final.xyz"
+    final_T20_file = tmp_path / "NaCl-heating-T20.0-final.xyz"
     log_file = tmp_path / "nvt.log"
 
     single_point = SinglePoint(
@@ -563,10 +566,10 @@ def test_heating(tmp_path):
         temp=300.0,
         steps=0,
         traj_every=10,
-        stats_every=1,
+        stats_every=10,
         file_prefix=file_prefix,
-        temp_start=10,
-        temp_end=40,
+        temp_start=0.0,
+        temp_end=20.0,
         temp_step=20,
         temp_time=0.5,
         log_kwargs={"filename": log_file},
@@ -575,11 +578,13 @@ def test_heating(tmp_path):
     assert_log_contains(
         log_file,
         includes=[
-            "Beginning temperature ramp at 10K",
-            "Temperature ramp complete at 30K",
+            "Beginning temperature ramp at 0.0K",
+            "Temperature ramp complete at 20.0K",
         ],
         excludes=["Starting molecular dynamics simulation"],
     )
+    assert final_T0_file.exists()
+    assert final_T20_file.exists()
 
 
 def test_noramp_heating(tmp_path):
@@ -757,3 +762,81 @@ def test_heating_md_files():
         final_10_path.unlink(missing_ok=True)
         final_20_path.unlink(missing_ok=True)
         final_path.unlink(missing_ok=True)
+
+
+def test_ramp_negative(tmp_path):
+    """Test ValueError is thrown for negative values in temperature ramp."""
+    file_prefix = tmp_path / "NaCl-heating"
+
+    single_point = SinglePoint(
+        struct_path=DATA_PATH / "NaCl.cif",
+        architecture="mace",
+        calc_kwargs={"model": MODEL_PATH},
+    )
+
+    with pytest.raises(ValueError):
+        NVT(
+            struct=single_point.struct,
+            file_prefix=file_prefix,
+            temp_start=-5,
+            temp_end=10,
+            temp_step=20,
+            temp_time=10,
+        )
+    with pytest.raises(ValueError):
+        NVT(
+            struct=single_point.struct,
+            file_prefix=file_prefix,
+            temp_start=5,
+            temp_end=-5,
+            temp_step=20,
+            temp_time=10,
+        )
+
+
+def test_cooling(tmp_path):
+    """Test cooling with no MD."""
+    # pylint: disable=invalid-name
+    file_prefix = tmp_path / "NaCl-cooling"
+    final_T0_file = tmp_path / "NaCl-cooling-T0.0-final.xyz"
+    final_T10_file = tmp_path / "NaCl-cooling-T10.0-final.xyz"
+    final_T20_file = tmp_path / "NaCl-cooling-T20.0-final.xyz"
+    stats_cooling_path = tmp_path / "NaCl-cooling-stats.dat"
+    log_file = tmp_path / "nvt.log"
+
+    single_point = SinglePoint(
+        struct_path=DATA_PATH / "NaCl.cif",
+        architecture="mace",
+        calc_kwargs={"model": MODEL_PATH},
+    )
+
+    nvt = NVT(
+        struct=single_point.struct,
+        temp=300.0,
+        steps=0,
+        traj_every=10,
+        stats_every=1,
+        file_prefix=file_prefix,
+        temp_start=20.0,
+        temp_end=0.0,
+        temp_step=10,
+        temp_time=0.5,
+        log_kwargs={"filename": log_file},
+    )
+    nvt.run()
+    assert_log_contains(
+        log_file,
+        includes=[
+            "Beginning temperature ramp at 20.0K",
+            "Temperature ramp complete at 0.0K",
+        ],
+        excludes=["Starting molecular dynamics simulation"],
+    )
+    assert final_T0_file.exists()
+    assert final_T10_file.exists()
+    assert final_T20_file.exists()
+
+    stats = Stats(source=stats_cooling_path)
+    assert stats.rows == 2
+    assert stats.data[0, 16] == 20.0
+    assert stats.data[1, 16] == 10.0