Updating thermostat with logging and capability for temperature ramp (#9)

* Updating thermostat with logging and capability for temperature ramp

Signed-off-by: Fabian Thiemann <fabian.thiemann@ibm.com>

* Adjusting github action tests.yaml

Signed-off-by: Fabian Thiemann <fabian.thiemann@ibm.com>

* Updating unittests and minor cosmetics

Signed-off-by: Fabian Thiemann <fabian.thiemann@ibm.com>

* Updating temperature in forecast for logging

Signed-off-by: Fabian Thiemann <fabian.thiemann@ibm.com>

* Fixing cuda problem for notebook

Signed-off-by: Fabian Thiemann <fabian.thiemann@ibm.com>

* Add PR template

Signed-off-by: Fabian Thiemann <fabian.thiemann@ibm.com>

* Update changelog

Signed-off-by: Fabian Thiemann <fabian.thiemann@ibm.com>

---------

Signed-off-by: Fabian Thiemann <fabian.thiemann@ibm.com>

Author: flt17

.github/workflows/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,14 @@
+**Summary**
+<!-- Briefly describe the purpose of this PR and what it changes. -->
+
+**Key Features**
+<!-- List the main changes or additions introduced by this PR. -->
+
+**Dependencies**
+<!-- List any new dependencies added or changes to existing requirements. -->
+
+**Tests**
+<!-- Describe testing performed and any tests added. -->
+
+**Related Issues / Branches**
+<!-- Link any related issues or pull requests, e.g., “Closes #123”. -->

.github/workflows/tests.yaml

@@ -61,6 +61,7 @@ jobs:
           source venv/bin/activate
           cd $GITHUB_WORKSPACE
           coverage run --source=trajcast -m unittest discover  -s  tests -p "test_*.py"
+          exit $?
           coverage xml
 
       - name: Generate coverage report

CHANGELOG.md

@@ -6,13 +6,19 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+
+## [1.1.0] - 2025-09-23
+
 ### Added
+- Added support for temperature ramps in the thermostat to enable controlled cooling/heating in simulations.
+- Improved logging for forecast runs to provide more informative output.
 - Added .github/workflows/tests.yaml to handle unittesting on github.
 - Added files for unittests which were missing in [1.0.0].
 - Added this changelog.
 
 ### Fixed
 - Fixed small issue in trajcast/data/wrappers/_lammps.py to not save the timestep which results in unittests failing with ase versions >= 3.25.0.
+- Fixed a CUDA error in generating inertia tensor reported by a user.
 
 ### Removed
 - Removed travis.yml as unittest will be handled by .github/workflows/tests.yaml.

examples/inference/forecasting.ipynb

@@ -463,24 +463,23 @@ def test_returns_thermostat_is_initialised_correctly(self):
             timestep=5.0,
             damping=1000.0,
             temperature_handler=temperature,
+            n_steps=50,
         )
 
-        self.assertIsInstance(csvr.e_kin_target, torch.Tensor)
-        self.assertEqual(csvr.e_kin_target.shape, torch.Size([1]))
-
         self.assertIsInstance(csvr.temp.conv_fac, torch.Tensor)
         self.assertEqual(csvr.temp.conv_fac.shape, torch.Size([1]))
 
         self.assertTrue(csvr.temp._n_dofs == 90)
 
     def test_returns_sum_noises_are_sampled_correclty_when_ndof_odd(self):
-
+        n_steps = 500000
         temperature = Temperature(units="real", n_atoms=162, n_extra_dofs=3)
         csvr = CSVRThermostat(
             target_temp=300.0,
             timestep=20.0,
             damping=2000,
             temperature_handler=temperature,
+            n_steps=n_steps,
         )
 
         gamm = []
@@ -499,18 +498,20 @@ def test_returns_sum_noises_are_sampled_correclty_when_ndof_odd(self):
         self.assertTrue(torch.isclose(gamm.std(), gauss.std(), atol=1e-1))
 
     def test_returns_sum_noises_are_sampled_correclty_when_ndof_even(self):
+        n_steps = 500000
         temperature = Temperature(units="real", n_atoms=22, n_extra_dofs=6)
         csvr = CSVRThermostat(
             target_temp=150.0,
             timestep=5.0,
             damping=500,
             temperature_handler=temperature,
+            n_steps=n_steps,
         )
 
         gamm = []
         gauss = []
         n_dofs = temperature._n_dofs
-        for i in range(500000):
+        for i in range(n_steps):
             # gaussian direct
             rns = torch.randn(n_dofs - 1)
             gauss.append(rns.pow(2).sum())
@@ -540,18 +541,18 @@ def test_returns_forward_produces_rescaled_velocities(self):
         )
 
         graph = f.start_graph
-
+        n_steps = 500000
         temperature = Temperature(units="real", n_atoms=graph.num_nodes, n_extra_dofs=0)
         csvr = CSVRThermostat(
             target_temp=150.0,
             timestep=5.0,
             damping=50.0,
             temperature_handler=temperature,
+            n_steps=n_steps,
         )
-
         temps = []
-        for _ in range(500000):
-            graph = csvr(graph)
+        for step in range(n_steps):
+            graph = csvr(graph, step)
             T = temperature(graph)
             temps.append(T)
 
@@ -588,18 +589,22 @@ def test_returns_thermostatting_is_reproducible_via_seed(self):
                 },
             }
         )
-
+        n_steps = 500
         graph = f.start_graph
         csvr = CSVRThermostat(
-            target_temp=250.0, timestep=5.0, damping=500.0, temperature_handler=f.temp
+            target_temp=250.0,
+            timestep=5.0,
+            damping=500.0,
+            temperature_handler=f.temp,
+            n_steps=n_steps,
         )
 
         temps_s1 = []
         vel_s1 = []
-        for steps in range(500):
+        for steps in range(n_steps):
             T = f.temp(graph)
             vel_s1.append(graph[VELOCITIES_KEY][0])
-            graph = csvr(graph)
+            graph = csvr(graph, steps)
 
             temps_s1.append(T)
 
@@ -645,7 +650,7 @@ def test_returns_thermostatting_is_reproducible_via_seed(self):
         for steps in range(500):
             T = f.temp(graph)
             vel_s2.append(graph[VELOCITIES_KEY][0])
-            graph = csvr(graph)
+            graph = csvr(graph, steps)
 
             temps_s2.append(T)
 

tests/unit/model/test_forecast.py

@@ -1,7 +1,7 @@
 import os
 import warnings
 from typing import Dict, Optional
-
+from tqdm import tqdm
 import ase.io
 import torch
 import yaml
@@ -402,6 +402,7 @@ def __init__(self, protocol: Dict):
                 timestep=timestep,
                 damping=tau,
                 temperature_handler=self.temp,
+                n_steps=self.protocol.get(RUN_KEY),
             ).to(self.device)
 
         else:
@@ -435,6 +436,9 @@ def __init__(self, protocol: Dict):
                     f"Either set {VELOCITIES_KEY} to True or pass user requirements as dictionary."
                 )
 
+            if isinstance(temperature, list):
+                temperature = temperature[0]
+
             vel_init = init_velocity(
                 target_temperature=temperature,
                 graph=self.start_graph,
@@ -449,13 +453,30 @@ def __init__(self, protocol: Dict):
         # for writing to file
         # if no write frequency is given
         write_settings = self.protocol.get(WRITE_TRAJECTORY_KEY)
-        self.write_freq = self.protocol.get(RUN_KEY) + 1
+        self.write_freq_xyz = self.protocol.get(RUN_KEY) + 1
+        self.write_freq_temp = self.protocol.get(RUN_KEY) + 1
+
         if write_settings:
-            self.write_freq = write_settings.get("every", 1)
-            if not isinstance(self.write_freq, int):
-                raise TypeError("Write frequency must be specified as integer")
+            write_freq = write_settings.get("every", 1)
+
+            if isinstance(write_freq, int):
+                self.write_freq_xyz = write_freq
+                self.write_freq_temp = write_freq
+
+            elif isinstance(write_freq, Dict):
+                self.write_freq_xyz = write_freq.get("xyz", 1)
+                self.write_freq_temp = write_freq.get("temp", self.write_freq_xyz)
+
+            else:
+                raise TypeError("Write frequency must be specified as integer or dict.")
+            self.save_vels = write_settings.get("save_velocities", True)
             self.filename = write_settings[FILENAME_KEY]
             self.fileformat = write_settings.get("format", "extxyz")
+            logdir = os.path.dirname(self.filename)
+            self.logfile = os.path.join(
+                logdir, os.path.basename(self.filename).split(".")[0] + ".log"
+            )
+            os.makedirs(logdir, exist_ok=True)
             # write initial frame to file
             self._write_frame_to_file(
                 frame=self.start_graph,
@@ -480,21 +501,31 @@ def generate_trajectory(self):
         n_steps = self.protocol.get(RUN_KEY)
         # initialise frame
         frame = self.start_graph
+        temp = self.temp(frame)
+
+        with open(self.logfile, "w") as f:
+            f.write("step,Temperature\n")
+            f.write(f"0,{temp:3.3f}\n")
 
         # loop over all steps
         with torch.no_grad():
-            for step in range(1, n_steps + 1):
+            for step in tqdm(range(1, n_steps + 1)):
                 # make a step
                 frame = self._make_timestep(frame, step)
+                temp = self.temp(frame)
 
                 # check for writer
-                if step % self.write_freq == 0:
+                if step % self.write_freq_xyz == 0:
                     self._write_frame_to_file(
                         frame=frame,
                         step=step,
                         append=True,
                     )
 
+                if step % self.write_freq_temp == 0:
+                    with open(self.logfile, "a") as f:
+                        f.write(f"{step},{temp:3.3f}\n")
+
     def _write_frame_to_file(
         self,
         frame: AtomicGraph,
@@ -504,6 +535,9 @@ def _write_frame_to_file(
         # get ase.Atoms object
         ase_atoms = frame.ASEAtomsObject
 
+        if not self.save_vels:
+            ase_atoms.arrays.pop("velocities")
+
         # add timestamp
         ase_atoms.info[FRAME_KEY] = step
 
@@ -533,7 +567,7 @@ def _make_timestep(self, frame: AtomicGraph, step: int) -> AtomicGraph:
 
         # thermostatting
         if self.nvt:
-            frame = self.thermo(frame)
+            frame = self.thermo(frame, step)
 
         # manipulate momentum if required
         if self.momentum and step % self.momentum.adjust_freq == 0:

trajcast/model/forecast.py

@@ -45,10 +45,11 @@ def __init__(self, units: str, n_atoms: int, n_extra_dofs: int) -> None:
 
         self.register_buffer("kB", torch.tensor(kB))
 
-    def to_kinetic_energy(self, temperature: float) -> torch.Tensor:
+    def to_kinetic_energy(self, temperature: torch.Tensor) -> torch.Tensor:
         return torch.tensor(
             [0.5 * self._n_dofs * temperature * self.kB],
             dtype=torch.get_default_dtype(),
+            device=temperature.device,
         )
 
     def from_kinetic_energy(self, kinetic_energy: torch.Tensor) -> float:

trajcast/model/forecast_tools/_temperature.py

@@ -1,5 +1,5 @@
 import torch
-
+from typing import Union, List
 from trajcast.data._keys import ATOMIC_MASSES_KEY, VELOCITIES_KEY
 from trajcast.data.atomic_graph import AtomicGraph
 from trajcast.model.forecast_tools import Temperature
@@ -15,10 +15,11 @@ class CSVRThermostat(torch.nn.Module):
 
     def __init__(
         self,
-        target_temp: float,
+        target_temp: Union[float, int, List[float], List[int]],
         timestep: float,
         damping: float,
         temperature_handler: Temperature,
+        n_steps: int,
     ):
         super().__init__()
 
@@ -34,9 +35,23 @@ def __init__(
                 torch.tensor([-timestep / damping], dtype=torch.get_default_dtype())
             ),
         )
+        if isinstance(target_temp, float) or isinstance(target_temp, int):
+            self.register_buffer("start_temp", torch.tensor(float(target_temp)))
+            self.temp_style = "CONST"
 
-        # target kinetic energy, this is in eV
-        self.register_buffer("e_kin_target", self.temp.to_kinetic_energy(target_temp))
+        else:
+            if len(target_temp) != 2:
+                raise TypeError(
+                    "Temperature should be float or list with two floats (start and stop temperature)."
+                )
+
+            start_temp = float(target_temp[0])
+            stop_temp = float(target_temp[1])
+            self.temp_style = "RAMP" if start_temp != stop_temp else "CONST"
+            self.register_buffer(
+                "rate", torch.tensor((stop_temp - start_temp) / n_steps)
+            )
+            self.register_buffer("start_temp", torch.tensor(start_temp))
 
         # for sampling noise we initialise a gamma distribution
         if (self.n_dofs - 1) % 2 == 0:
@@ -48,7 +63,7 @@ def __init__(
                 concentration=(self.n_dofs - 2) / 2, rate=1.0
             )
 
-    def forward(self, data: AtomicGraph) -> AtomicGraph:
+    def forward(self, data: AtomicGraph, step: int) -> AtomicGraph:
         masses = data[ATOMIC_MASSES_KEY]
         velocities = data[VELOCITIES_KEY]
 
@@ -58,19 +73,28 @@ def forward(self, data: AtomicGraph) -> AtomicGraph:
             * self.temp.conv_fac
         )
         # compute rescale_factor
-        alpha = self._get_rescale_factor(e_kin)
+        alpha = self._get_rescale_factor(e_kin, step)
 
         # rescale velocities accordingly
         data[VELOCITIES_KEY] *= alpha
         return data
 
-    def _get_rescale_factor(self, e_kin_current: torch.Tensor) -> torch.Tensor:
+    def _get_rescale_factor(
+        self, e_kin_current: torch.Tensor, step: int
+    ) -> torch.Tensor:
+
+        if self.temp_style == "CONST":
+            temp_target = self.start_temp
+
+        else:
+            temp_target = self.start_temp + self.rate * step
+
+        # target kinetic energy, this is in eV
+        e_kin_target = self.temp.to_kinetic_energy(temp_target)
+
         # compute constant c2 with c1 and kinetic energies
         c2 = (
-            (torch.tensor(1.0) - self.c1)
-            * self.e_kin_target
-            / e_kin_current
-            / self.n_dofs
+            (torch.tensor(1.0) - self.c1) * e_kin_target / e_kin_current / self.n_dofs
         ).to(self.device)
         # draw random number from Gaussian distribution with unitary variance (R1)
         r1 = torch.randn(1, device=self.device)