density updates

CHANGELOG.md

@@ -18,3 +18,4 @@
 - add `pbc: bool` to `Packmol` to fix PBC not supported by PACKMOL
 - add `num_ramp_oscillations: float` to `BoxOscillatingRampModifier`
 - remove `train_log_file` from apax Node
+- add `density` to `RescaleBoxModifier`

ipsuite/analysis/__init__.py

@@ -8,6 +8,7 @@
 )
 from ipsuite.analysis.bond_stretch import BondStretchAnalyses
 from ipsuite.analysis.ensemble import ModelEnsembleAnalysis
+from ipsuite.analysis.md import AnalyseDensity
 from ipsuite.analysis.model import (
     BoxHeatUp,
     BoxScale,
@@ -54,4 +55,5 @@
     "StressHistogram",
     "ForcesUncertaintyHistogram",
     "EnergyUncertaintyHistogram",
+    "AnalyseDensity",
 ]

ipsuite/analysis/md.py

@@ -0,0 +1,38 @@
+import pathlib
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import zntrack
+
+from ipsuite import base
+from ipsuite.utils.ase_sim import get_density_from_atoms
+
+
+class AnalyseDensity(base.AnalyseAtoms):
+    window: int = zntrack.params(1000)
+    start: int = zntrack.params(0)
+    end: int = zntrack.params(None)
+
+    density: dict = zntrack.metrics()
+    results: pd.DataFrame = zntrack.plots()
+
+    figure: pathlib.Path = zntrack.plots_path(zntrack.nwd / "density.png")
+
+    def run(self):
+        densities = [get_density_from_atoms(x) for x in self.data]
+
+        fig, ax = plt.subplots()
+        ax.plot(densities)
+        ax.plot(np.convolve(densities, np.ones(self.window) / self.window, mode="valid"))
+        ax.set_ylabel(r"Density $\rho$ / kg $\cdot$ m$^{-3}$")
+        ax.set_xlabel("Step")
+        fig.tight_layout()
+        fig.savefig(self.figure)
+
+        self.density = {
+            "density": np.mean(densities[self.start : self.end]),
+            "std": np.std(densities[self.start : self.end]),
+        }
+
+        self.results = pd.DataFrame(densities, columns=["density"])

ipsuite/calculators/ase_md.py

@@ -18,17 +18,28 @@
 from tqdm import trange
 
 from ipsuite import base
-from ipsuite.utils.ase_sim import freeze_copy_atoms, get_energy
+from ipsuite.utils.ase_sim import freeze_copy_atoms, get_box_from_density, get_energy
 
 log = logging.getLogger(__name__)
 
 
 class RescaleBoxModifier(base.IPSNode):
-    cell: int = zntrack.zn.params()
+    cell: int = zntrack.params(None)
+    density: float = zntrack.params(None)
     _initial_cell = None
 
+    # def _post_init_(self):
+    #     super()._post_init_()
+    #     if self.density is not None and self.cell is not None:
+    #         raise ValueError("Only one of density or cell can be given.")
+    #     if self.density is None and self.cell is None:
+    #         raise ValueError("Either density or cell has to be given.")
+    # Currently not possible due to a ZnTrack bug
+
     def modify(self, thermostat, step, total_steps):
         # we use the thermostat, so we can also modify e.g. temperature
+        if self.cell is None:
+            self.cell = get_box_from_density([[thermostat.atoms]], [1], self.density)
         if isinstance(self.cell, int):
             self.cell = np.array(
                 [[self.cell, 0, 0], [0, self.cell, 0], [0, 0, self.cell]]

ipsuite/configuration_generation/packmol.py

@@ -11,6 +11,7 @@
 from ase.visualize import view
 
 from ipsuite import base, fields
+from ipsuite.utils.ase_sim import get_box_from_density
 
 log = logging.getLogger(__name__)
 
@@ -99,17 +100,7 @@ def run(self):
 
     def _get_box_from_molar_volume(self):
         """Get the box size from the molar volume"""
-        molar_mass = [
-            sum(atoms[0].get_masses()) * count
-            for atoms, count in zip(self.data, self.count)
-        ]
-        molar_mass = sum(molar_mass)  #  g / mol
-        molar_volume = molar_mass / self.density / 1000  # m^3 / mol
-
-        # convert to particles / A^3
-        volume = molar_volume * (ase.units.m**3) / ase.units.mol
-
-        self.box = [volume ** (1 / 3) for _ in range(3)]
+        self.box = get_box_from_density(self.data, self.count, self.density)
         log.info(f"estimated box size: {self.box}")
 
     def view(self) -> view:

ipsuite/nodes.py

@@ -75,6 +75,7 @@ class _Nodes:
     ThresholdCheck = "ipsuite.analysis.ThresholdCheck"
     TemperatureCheck = "ipsuite.analysis.TemperatureCheck"
     FixedSphereConstraint = "ipsuite.calculators.FixedSphereConstraint"
+    AnalyseDensity = "ipsuite.analysis.AnalyseDensity"
 
     # calculators
     CP2KSinglePoint = "ipsuite.calculators.CP2KSinglePoint"

ipsuite/utils/ase_sim.py

@@ -50,3 +50,36 @@ def freeze_copy_atoms(atoms) -> ase.Atoms:
 @freeze_copy_atoms.register
 def _(atoms: list) -> list[ase.Atoms]:
     return [freeze_copy_atoms(x) for x in atoms]
+
+
+def get_box_from_density(
+    data: list[list[ase.Atoms]], count: list[int], density: float
+) -> list[float]:
+    """Get the box size from the molar volume.
+
+    Attributes
+    ----------
+    data: list[list[ase.Atoms]]
+        List of list of atoms objects. The last atoms object is used to compute the
+        molar volume.
+    count: list[int]
+        Number of molecules for each entry in data.
+    density: float
+        Density of the system in kg/m^3.
+    """
+    molar_mass = [sum(atoms[0].get_masses()) * count for atoms, count in zip(data, count)]
+    molar_mass = sum(molar_mass)  #  g / mol
+    molar_volume = molar_mass / density / 1000  # m^3 / mol
+
+    # convert to particles / A^3
+    volume = molar_volume * (ase.units.m**3) / ase.units.mol
+
+    box = [volume ** (1 / 3) for _ in range(3)]
+    return box
+
+
+def get_density_from_atoms(atoms: ase.Atoms) -> float:
+    """Compute the density of the atoms in kg/m3."""
+    volume = atoms.get_volume()
+    molar_volume = volume / (units.m**3 / units.mol)
+    return atoms.get_masses().sum() / 1000 / molar_volume

tests/integration/calculators/test_ase_md.py

@@ -1,8 +1,10 @@
 import ase.io
 import numpy as np
+import numpy.testing as npt
 from ase import units
 
 import ipsuite as ips
+from ipsuite.utils.ase_sim import get_density_from_atoms
 
 
 def test_ase_md(proj_path, cu_box):
@@ -40,6 +42,37 @@ def test_ase_md(proj_path, cu_box):
     assert md.atoms[-1].cell[0, 0] == 10
 
 
+def test_ase_md_target_density(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(density=1000)
+
+    with ips.Project() as project:
+        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],
+            thermostat=thermostat,
+            steps=30,
+            sampling_rate=1,
+            dump_rate=33,
+        )
+
+    project.run()
+
+    md.load()
+    npt.assert_almost_equal(get_density_from_atoms(md.atoms[0]), 8971.719659196913)
+    npt.assert_almost_equal(get_density_from_atoms(md.atoms[-1]), 1000)
+
+
 def test_ase_md_box_ramp(proj_path, cu_box):
     ase.io.write("cu_box.xyz", cu_box)
     thermostat = ips.calculators.LangevinThermostat(

tests/unit_tests/test_utils_sim.py

@@ -0,0 +1,11 @@
+import ase.build
+import numpy.testing as npt
+
+from ipsuite.utils.ase_sim import get_box_from_density
+
+
+def test_get_box():
+    water = [ase.build.molecule("H2O")]
+
+    box = get_box_from_density([water], [10], 997)
+    npt.assert_almost_equal(box, [6.6946735, 6.6946735, 6.6946735])