From 486df7167d9ce2ca3dbc090314319bb0e7cebd46 Mon Sep 17 00:00:00 2001
From: ElliottKasoar <ElliottKasoar@users.noreply.github.com>
Date: Wed, 10 Apr 2024 16:51:25 +0100
Subject: [PATCH] Add heating CLI

---
 janus_core/cli.py    | 28 ++++++++++++++++++++++++++++
 janus_core/md.py     |  8 ++++----
 tests/test_md_cli.py | 37 +++++++++++++++++++++++++++++++++++++
 3 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/janus_core/cli.py b/janus_core/cli.py
index 3c395bf3..3e836c96 100644
--- a/janus_core/cli.py
+++ b/janus_core/cli.py
@@ -667,6 +667,20 @@ def md(
     traj_every: Annotated[
         int, typer.Option(help="Frequency of steps to save trajectory.")
     ] = 100,
+    temp_start: Annotated[
+        Optional[float],
+        typer.Option(help="Temperature to start heating, in K.  [default: None]"),
+    ] = None,
+    temp_end: Annotated[
+        Optional[float],
+        typer.Option(help="Maximum temperature for heating, in K.  [default: None]"),
+    ] = None,
+    temp_step: Annotated[
+        float, typer.Option(help="Size of temperature steps when heating, in K.")
+    ] = 10,
+    temp_time: Annotated[
+        float, typer.Option(help="Time between heating steps, in fs.")
+    ] = 10,
     log: LogPath = "md.log",
     seed: Annotated[
         Optional[int],
@@ -749,6 +763,16 @@ def md(
         Step to start saving trajectory. Default is 0.
     traj_every : int
         Frequency of steps to save trajectory. Default is 100.
+    temp_start : Optional[float]
+        Temperature to start heating, in K. Default is None, which disables
+        heating.
+    temp_end : Optional[float]
+        Maximum temperature for heating, in K. Default is None, which disables
+        heating.
+    temp_step : float
+        Size of temperature steps when heating, in K. Default is 10.
+    temp_time : float
+        Time between heating steps, in fs. Default is 10.
     log : Optional[Path]
         Path to write logs to. Default is "md.log".
     seed : Optional[int]
@@ -808,6 +832,10 @@ def md(
         "traj_append": traj_append,
         "traj_start": traj_start,
         "traj_every": traj_every,
+        "temp_start": temp_start,
+        "temp_end": temp_end,
+        "temp_step": temp_step,
+        "temp_time": temp_time,
         "log_kwargs": log_kwargs,
         "seed": seed,
     }
diff --git a/janus_core/md.py b/janus_core/md.py
index f1aab5ea..903b2538 100644
--- a/janus_core/md.py
+++ b/janus_core/md.py
@@ -88,7 +88,7 @@ class MolecularDynamics:  # pylint: disable=too-many-instance-attributes
     temp_start : Optional[float]
         Temperature to start heating, in K. Default is None, which disables heating.
     temp_end : Optional[float]
-        Temperature to finish heating, in K. Default is None, which disables heating.
+        Maximum temperature for heating, in K. Default is None, which disables heating.
     temp_step : float
         Size of temperature steps when heating, in K. Default is 10.
     temp_time : float
@@ -232,7 +232,7 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-locals,too-many-sta
             Temperature to start heating, in K. Default is None, which disables
             heating.
         temp_end : Optional[float]
-            Temperature to finish heating, in K. Default is None, which disables
+            Maximum temperature for heating, in K. Default is None, which disables
             heating.
         temp_step : float
             Size of temperature steps when heating, in K. Default is 10.
@@ -529,9 +529,9 @@ def _run_dynamics(self) -> None:
 
         # Run temperature ramp
         if self.temp_start and self.temp_end and self.temp_step:
-            heating_steps = self.temp_time // self.timestep
+            heating_steps = int(self.temp_time // self.timestep)
 
-            n_temps = 1 + (self.temp_end - self.temp_start) // self.temp_step
+            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)]
 
             if self.logger:
diff --git a/tests/test_md_cli.py b/tests/test_md_cli.py
index d40d39ec..479910dd 100644
--- a/tests/test_md_cli.py
+++ b/tests/test_md_cli.py
@@ -278,3 +278,40 @@ def test_config(tmp_path):
     assert md_summary["inputs"]["temp"] == 200
     # Check explicit option overwrites config
     assert md_summary["inputs"]["ensemble"] == "nve"
+
+
+def test_heating(tmp_path):
+    """Test heating before MD."""
+    file_prefix = tmp_path / "nvt-T300"
+    log_path = tmp_path / "test.log"
+    summary_path = tmp_path / "summary.yml"
+
+    result = runner.invoke(
+        app,
+        [
+            "md",
+            "--ensemble",
+            "nvt",
+            "--struct",
+            DATA_PATH / "NaCl.cif",
+            "--file-prefix",
+            file_prefix,
+            "--stats-every",
+            1,
+            "--steps",
+            5,
+            "--temp-start",
+            10,
+            "--temp-end",
+            20,
+            "--temp-step",
+            50,
+            "--temp-time",
+            0.05,
+            "--log",
+            log_path,
+            "--summary",
+            summary_path,
+        ],
+    )
+    assert result.exit_code == 0