Add RescaleBoxModifier and TemperatureRampModifier

ipsuite/calculators/__init__.pyi

@@ -1,8 +1,14 @@
 from .apax_jax_md import ApaxJaxMD
 from .ase_geoopt import ASEGeoOpt
-from .ase_md import ASEMD, FixedSphereASEMD, LangevinThermostat
+from .ase_md import (
+    ASEMD,
+    FixedSphereASEMD,
+    LangevinThermostat,
+    RescaleBoxModifier,
+    TemperatureRampModifier,
+)
+from .ase_standard import EMTSinglePoint, LJSinglePoint
 from .cp2k import CP2KSinglePoint, CP2KYaml
-from .lj import LJSinglePoint
 from .xtb import xTBSinglePoint
 
 __all__ = [
@@ -15,4 +21,7 @@ __all__ = [
     "LJSinglePoint",
     "LangevinThermostat",
     "ApaxJaxMD",
+    "RescaleBoxModifier",
+    "EMTSinglePoint",
+    "TemperatureRampModifier",
 ]

ipsuite/calculators/ase_md.py

@@ -20,6 +20,40 @@
 log = logging.getLogger(__name__)
 
 
+class RescaleBoxModifier(base.IPSNode):
+    cell: int = zntrack.zn.params()
+    _initial_cell = None
+
+    def modify(self, thermostat, step, total_steps):
+        # we use the thermostat, so we can also modify e.g. temperature
+        if isinstance(self.cell, int):
+            self.cell = np.array(
+                [[self.cell, 0, 0], [0, self.cell, 0], [0, 0, self.cell]]
+            )
+        elif isinstance(self.cell, list):
+            self.cell = np.array(
+                [[self.cell[0], 0, 0], [0, self.cell[1], 0], [0, 0, self.cell[2]]]
+            )
+
+        if self._initial_cell is None:
+            self._initial_cell = thermostat.atoms.get_cell()
+        percentage = step / (total_steps - 1)
+        new_cell = (1 - percentage) * self._initial_cell + percentage * self.cell
+        thermostat.atoms.set_cell(new_cell, scale_atoms=True)
+
+
+class TemperatureRampModifier(base.IPSNode):
+    temperature: float = zntrack.zn.params()
+
+    def modify(self, thermostat, step, total_steps):
+        # we use the thermostat, so we can also modify e.g. temperature
+        percentage = step / (total_steps - 1)
+        new_temperature = (
+            1 - percentage
+        ) * thermostat.temp + percentage * self.temperature
+        thermostat.set_temperature(temperature_K=new_temperature)
+
+
 class LangevinThermostat(base.IPSNode):
     """Initialize the langevin thermostat
 
@@ -89,6 +123,7 @@ class ASEMD(base.ProcessSingleAtom):
     model = zntrack.zn.deps()
     model_outs = zntrack.dvc.outs(zntrack.nwd / "model/")
     checker_list: list = zntrack.zn.nodes(None)
+    modifier: list = zntrack.zn.nodes(None)
     thermostat = zntrack.zn.nodes()
 
     steps: int = zntrack.zn.params()
@@ -121,6 +156,8 @@ def run(self):  # noqa: C901
         """Run the simulation."""
         if self.checker_list is None:
             self.checker_list = []
+        if self.modifier is None:
+            self.modifier = []
 
         self.model_outs.mkdir(parents=True, exist_ok=True)
         (self.model_outs / "outs.txt").write_text("Lorem Ipsum")
@@ -168,6 +205,8 @@ def run(self):  # noqa: C901
             for idx in range(self.steps):
                 desc = []
                 stop = []
+                for modifier in self.modifier:
+                    modifier.modify(thermostat, step=idx, total_steps=self.steps)
                 thermostat.run(self.sampling_rate)
                 _, energy = get_energy(atoms)
                 metrics_dict["energy"].append(energy)

ipsuite/calculators/lj.py → ipsuite/calculators/ase_standard.py

@@ -1,4 +1,5 @@
 import tqdm
+from ase.calculators.emt import EMT
 from ase.calculators.lj import LennardJones
 
 from ipsuite import base
@@ -21,7 +22,28 @@ def run(self):
             atom.get_potential_energy()
             atom.get_stress()
 
-    def get_calculator(self):
+    def get_calculator(self, **kwargs):
         """Get an LJ ase calculator."""
 
         return LennardJones()
+
+
+class EMTSinglePoint(base.ProcessAtoms):
+    """This is a testing Node!
+    It uses ASE'S EMT calculator with default arguments.
+    The calculator accept all elements and implements energy, forces,
+    making it very useful for creating dummy data.
+    """
+
+    def run(self):
+        self.atoms = self.get_data()
+
+        calculator = self.get_calculator()
+
+        for atom in tqdm.tqdm(self.atoms):
+            atom.calc = calculator
+            atom.get_potential_energy()
+
+    def get_calculator(self, **kwargs):
+        """Get an LJ ase calculator."""
+        return EMT()

ipsuite/nodes.py

@@ -63,6 +63,7 @@ class _Nodes:
     AnalyseSingleForceSensitivity = "ipsuite.analysis.AnalyseSingleForceSensitivity"
     ForceDecomposition = "ipsuite.analysis.ForceDecomposition"
     ThresholdCheck = "ipsuite.analysis.ThresholdCheck"
+    TemperatureCheck = "ipsuite.analysis.TemperatureCheck"
 
     # calculators
     CP2KSinglePoint = "ipsuite.calculators.CP2KSinglePoint"
@@ -72,8 +73,11 @@ class _Nodes:
     FixedSphereASEMD = "ipsuite.calculators.FixedSphereASEMD"
     xTBSinglePoint = "ipsuite.calculators.xTBSinglePoint"
     LJSinglePoint = "ipsuite.calculators.LJSinglePoint"
+    EMTSinglePoint = "ipsuite.calculators.EMTSinglePoint"
 
     LangevinThermostat = "ipsuite.calculators.LangevinThermostat"
+    RescaleBoxModifier = "ipsuite.calculators.RescaleBoxModifier"
+    TemperatureRampModifier = "ipsuite.calculators.TemperatureRampModifier"
 
     # Geometry
     BarycenterMapping = "ipsuite.geometry.BarycenterMapping"

tests/conftest.py

@@ -18,6 +18,7 @@
 import numpy as np
 import pytest
 from ase import Atoms
+from ase.lattice.cubic import FaceCenteredCubic
 
 import ipsuite as ips
 
@@ -50,6 +51,18 @@ def atoms_list() -> typing.List[ase.Atoms]:
     return atoms
 
 
+@pytest.fixture
+def cu_box() -> typing.List[ase.Atoms]:
+    return [
+        FaceCenteredCubic(
+            directions=[[1, 0, 0], [0, 1, 0], [0, 0, 1]],
+            symbol="Cu",
+            size=(2, 2, 2),
+            pbc=True,
+        )
+    ]
+
+
 @pytest.fixture()
 def traj_file(tmp_path_factory, atoms_list) -> str:
     """Save an extxyz trajectory file based on atoms_list."""

tests/integration/calculators/test_ase_md.py

@@ -1,23 +1,29 @@
+import ase.io
+
 import ipsuite as ips
 
 
-def test_ase_md(proj_path, traj_file):
+def test_ase_md(proj_path, cu_box):
+    ase.io.write("cu_box.xyz", cu_box)
     checker = ips.analysis.TemperatureCheck()
     thermostat = ips.calculators.LangevinThermostat(
         time_step=1,
         temperature=1,
         friction=1,
     )
+    rescale_box = ips.calculators.RescaleBoxModifier(cell=10)
+    temperature_ramp = ips.calculators.TemperatureRampModifier(temperature=300)
     with ips.Project() as project:
-        data = ips.AddData(file=traj_file)
-        model = ips.models.GAP(data=data.atoms)
+        data = ips.AddData(file="cu_box.xyz")
+        model = ips.calculators.EMTSinglePoint(data=data.atoms)
         md = ips.calculators.ASEMD(
             data=data.atoms,
             model=model,
             checker_list=[checker],
+            modifier=[rescale_box, temperature_ramp],
             thermostat=thermostat,
             init_temperature=1.0,
-            steps=100,
+            steps=30,
             sampling_rate=1,
             dump_rate=33,
         )
@@ -26,4 +32,8 @@ def test_ase_md(proj_path, traj_file):
 
     md.load()
 
-    assert len(md.atoms) == 100
+    assert len(md.atoms) == 30
+    assert md.atoms[0].cell[0, 0] == 7.22
+    assert md.atoms[1].cell[0, 0] > 7.22
+    assert md.atoms[1].cell[0, 0] < 10
+    assert md.atoms[-1].cell[0, 0] == 10