Fix CSVR temperature ramp (#414)

* Fix CSVR temperature ramp

* Replace warnings with errors

* Add temperature units

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

* Fix MTK _set_target_temperature

* Add heating tests for all ensembles

* Rework NPT_MTK test skip

* Fix tests

* Fix tests

---------

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

Author: k-harris27

janus_core/calculations/md.py

@@ -423,6 +423,10 @@ def __init__(
                 "heating to run",
                 stacklevel=2,
             )
+        if self.ramp_temp and self.ensemble in ("nve", "nph"):
+            raise ValueError(
+                "Temperature ramp requested for ensemble with no thermostat."
+            )
 
         # Check validate start and end temperatures
         if self.ramp_temp and (self.temp_start < 0 or self.temp_end < 0):
@@ -1021,6 +1025,28 @@ def _write_restart(self) -> None:
                 self.restart_files.append(self._restart_file)
                 self._rotate_restart_files()
 
+    def _set_target_temperature(self, temperature: float):
+        """
+        Set the target temperature of the thermostat.
+
+        Parameters
+        ----------
+        temperature
+            New target temperature, in K.
+        """
+        if hasattr(self.dyn, "set_temperature"):
+            self.dyn.set_temperature(temperature_K=temperature)
+        elif isinstance(self, NVT_CSVR):
+            self.dyn.temp = temperature * units.kB
+            self.dyn.target_kinetic_energy = 0.5 * self.dyn.temp * self.dyn.ndof
+        elif isinstance(self, NPT_MTK):
+            kt = temperature * units.kB
+            self.dyn._kT = kt
+            self.dyn._thermostat._kT = kt
+            self.dyn._barostat._kT = kt
+        else:
+            raise ValueError("Temperature set for ensemble with no thermostat.")
+
     def run(self) -> None:
         """Run molecular dynamics simulation and/or temperature ramp."""
         unit_keys = (
@@ -1105,8 +1131,7 @@ def _run_dynamics(self) -> None:
                     self._write_final_state()
                     self.created_final_file = True
                     continue
-                if not isinstance(self, NVE):
-                    self.dyn.set_temperature(temperature_K=self.temp)
+                self._set_target_temperature(temp)
                 self.dyn.run(heating_steps)
                 self._write_final_state()
                 self.created_final_file = True
@@ -1126,8 +1151,7 @@ def _run_dynamics(self) -> None:
             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._set_target_temperature(self.temp)
             self.dyn.run(self.steps - self.offset)
             self._write_final_state()
             self.created_final_file = True

tests/test_md.py

@@ -7,6 +7,7 @@
 from ase import Atoms
 from ase.io import read
 import ase.md.nose_hoover_chain
+from ase.md.npt import NPT as ASE_NPT
 import numpy as np
 import pytest
 
@@ -21,6 +22,7 @@
 
     MTK_IMPORT_FAILED = False
 else:
+    NPT_MTK = None
     MTK_IMPORT_FAILED = True
 
 DATA_PATH = Path(__file__).parent / "data"
@@ -32,9 +34,29 @@
     (NPT, "npt"),
     (NVT_NH, "nvt-nh"),
     (NPH, "nph"),
+    (NVT_CSVR, "nvt-csvr"),
+    pytest.param(
+        NPT_MTK,
+        "npt-mtk",
+        marks=pytest.mark.skipif(
+            MTK_IMPORT_FAILED, reason="Requires updated version of ASE"
+        ),
+    ),
 ]
-if not MTK_IMPORT_FAILED:
-    test_data.append((NPT_MTK, "npt-mtk"))
+
+ensembles_without_thermostat = (NVE, NPH)
+ensembles_with_thermostat = (
+    NVT,
+    NPT,
+    NVT_NH,
+    NVT_CSVR,
+    pytest.param(
+        NPT_MTK,
+        marks=pytest.mark.skipif(
+            MTK_IMPORT_FAILED, reason="Requires updated version of ASE"
+        ),
+    ),
+)
 
 
 @pytest.mark.parametrize("ensemble, expected", test_data)
@@ -694,19 +716,20 @@ def test_stats(tmp_path, ensemble, tag):
     assert stat_data.units[etot_index] == "eV"
 
 
-def test_heating(tmp_path):
+@pytest.mark.parametrize("ensemble", ensembles_with_thermostat)
+def test_heating(tmp_path, ensemble):
     """Test heating with no MD."""
     file_prefix = tmp_path / "NaCl-heating"
     final_file = tmp_path / "NaCl-heating-final.extxyz"
-    log_file = tmp_path / "nvt.log"
+    log_file = tmp_path / "NaCl.log"
 
     single_point = SinglePoint(
         struct_path=DATA_PATH / "NaCl.cif",
         arch="mace",
         calc_kwargs={"model": MODEL_PATH},
     )
 
-    nvt = NVT(
+    md = ensemble(
         struct=single_point.struct,
         temp=300.0,
         steps=0,
@@ -719,7 +742,7 @@ def test_heating(tmp_path):
         temp_time=0.5,
         log_kwargs={"filename": log_file},
     )
-    nvt.run()
+    md.run()
     assert_log_contains(
         log_file,
         includes=[
@@ -728,10 +751,42 @@ def test_heating(tmp_path):
         ],
         excludes=["Starting molecular dynamics simulation"],
     )
+
     assert final_file.exists()
 
 
-def test_noramp_heating(tmp_path):
+@pytest.mark.parametrize("ensemble", ensembles_without_thermostat)
+def test_no_thermostat_heating(tmp_path, ensemble):
+    """Test that temperature ramp with no thermostat throws an error."""
+    file_prefix = tmp_path / "NaCl-heating"
+    final_file = tmp_path / "NaCl-heating-final.extxyz"
+    log_file = tmp_path / "NaCl.log"
+
+    single_point = SinglePoint(
+        struct_path=DATA_PATH / "NaCl.cif",
+        arch="mace",
+        calc_kwargs={"model": MODEL_PATH},
+    )
+    with pytest.raises(ValueError, match="no thermostat"):
+        md = ensemble(
+            struct=single_point.struct,
+            temp=300.0,
+            steps=0,
+            traj_every=10,
+            stats_every=10,
+            file_prefix=file_prefix,
+            temp_start=0.0,
+            temp_end=20.0,
+            temp_step=20,
+            temp_time=0.5,
+            log_kwargs={"filename": log_file},
+        )
+        md.run()
+    assert not final_file.exists()
+
+
+@pytest.mark.parametrize("ensemble", ensembles_with_thermostat)
+def test_noramp_heating(tmp_path, ensemble):
     """Test ValueError is thrown for invalid temperature ramp."""
     file_prefix = tmp_path / "NaCl-heating"
 
@@ -742,7 +797,7 @@ def test_noramp_heating(tmp_path):
     )
 
     with pytest.raises(ValueError):
-        NVT(
+        ensemble(
             struct=single_point.struct,
             file_prefix=file_prefix,
             temp_start=10,
@@ -751,18 +806,19 @@ def test_noramp_heating(tmp_path):
         )
 
 
-def test_heating_md(tmp_path):
+@pytest.mark.parametrize("ensemble", ensembles_with_thermostat)
+def test_heating_md(tmp_path, ensemble):
     """Test heating followed by MD."""
     file_prefix = tmp_path / "NaCl-heating"
     stats_path = tmp_path / "NaCl-heating-stats.dat"
-    log_file = tmp_path / "nvt.log"
+    log_file = tmp_path / "NaCl.log"
 
     single_point = SinglePoint(
         struct_path=DATA_PATH / "NaCl.cif",
         arch="mace",
         calc_kwargs={"model": MODEL_PATH},
     )
-    nvt = NVT(
+    md = ensemble(
         struct=single_point.struct,
         temp=25.0,
         steps=5,
@@ -775,7 +831,7 @@ def test_heating_md(tmp_path):
         temp_time=2,
         log_kwargs={"filename": log_file},
     )
-    nvt.run()
+    md.run()
     assert_log_contains(
         log_file,
         includes=[
@@ -786,15 +842,23 @@ def test_heating_md(tmp_path):
         ],
     )
     stat_data = Stats(stats_path)
-    assert stat_data.rows == 5
-    assert stat_data.columns == 17
-    assert stat_data.data[0, 16] == pytest.approx(10.0)
-    assert stat_data.data[2, 16] == pytest.approx(20.0)
-    assert stat_data.data[4, 16] == pytest.approx(25.0)
+    target_t_col = stat_data.labels.index("Target_T")
+
+    is_ase_npt = isinstance(md.dyn, ASE_NPT)
+
+    # ASE_NPT skips first row of output - ASE merge request submitted to fix:
+    # https://gitlab.com/ase/ase/-/merge_requests/3598
+    assert stat_data.rows == 4 if is_ase_npt else 5
+    assert stat_data.data[0, target_t_col] == pytest.approx(10.0)
+    if is_ase_npt:
+        assert stat_data.data[1, target_t_col] == pytest.approx(20.0)
+        assert stat_data.data[3, target_t_col] == pytest.approx(25.0)
+    else:
+        assert stat_data.data[2, target_t_col] == pytest.approx(20.0)
+        assert stat_data.data[4, target_t_col] == pytest.approx(25.0)
     assert stat_data.labels[0] == "# Step"
     assert stat_data.units[0] == ""
-    assert stat_data.units[16] == "K"
-    assert stat_data.labels[16] == "Target_T"
+    assert stat_data.units[target_t_col] == "K"
 
 
 def test_heating_files():

tests/test_md_cli.py

@@ -32,7 +32,20 @@ def test_md_help():
     assert "Usage: janus md [OPTIONS]" in strip_ansi_codes(result.stdout)
 
 
-test_data = [("nvt"), ("nve"), ("npt"), ("nvt-nh"), ("nph"), ("nvt-csvr"), ("npt-mtk")]
+test_data = [
+    ("nvt"),
+    ("nve"),
+    ("npt"),
+    ("nvt-nh"),
+    ("nph"),
+    ("nvt-csvr"),
+    pytest.param(
+        "npt-mtk",
+        marks=pytest.mark.skipif(
+            MTK_IMPORT_FAILED, reason="Requires updated version of ASE"
+        ),
+    ),
+]
 
 
 @pytest.mark.parametrize("ensemble", test_data)
@@ -49,9 +62,6 @@ def test_md(ensemble):
         "npt-mtk": "NaCl-npt-mtk-T300.0-p0.0-",
     }
 
-    if ensemble == "npt-mtk" and MTK_IMPORT_FAILED:
-        pytest.skip(reason="Requires updated version of ASE")
-
     final_path = Path(f"{file_prefix[ensemble]}final.extxyz").absolute()
     restart_path = Path(f"{file_prefix[ensemble]}res-2.extxyz").absolute()
     stats_path = Path(f"{file_prefix[ensemble]}stats.dat").absolute()
@@ -342,7 +352,8 @@ def test_config(tmp_path):
     assert_log_contains(log_path, includes=["hydrostatic_strain: True"])
 
 
-def test_heating(tmp_path):
+@pytest.mark.parametrize("ensemble", test_data)
+def test_heating(tmp_path, ensemble):
     """Test heating before MD."""
     file_prefix = tmp_path / "nvt-T300"
 
@@ -351,7 +362,7 @@ def test_heating(tmp_path):
         [
             "md",
             "--ensemble",
-            "nvt",
+            ensemble,
             "--struct",
             DATA_PATH / "NaCl.cif",
             "--file-prefix",
@@ -370,7 +381,10 @@ def test_heating(tmp_path):
             0.05,
         ],
     )
-    assert result.exit_code == 0
+    if ensemble in ("nve", "nph"):
+        assert result.exit_code != 0
+    else:
+        assert result.exit_code == 0
 
 
 def test_invalid_config():