A notebook for non-conservative MTS dynamics with PET-MAD

README.rst

@@ -1,10 +1,10 @@
-The Atomistic Cookbook repository
-=================================
+Source code for the Atomistic Cookbook
+======================================
 
 .. image:: ./docs/src/_static/cookbook-icon.svg
    :alt: A cookbook with a cover showing a water molecule and mathematical symbols
    :align: center
-   :width: 50%
+   :width: 30%
 
 
 This repository contains the source code for a collection of software recipes,

docs/src/software/chemiscope.sec

@@ -21,3 +21,4 @@ repository <https://github.com/lab-cosmo/chemiscope>`_.
 - examples/water-model/water-model
 - examples/polarizability/polarizability
 - examples/pet-mad/pet-mad
+- examples/pet-mad-nc/pet-mad-nc  

docs/src/software/i-pi.sec

@@ -15,3 +15,4 @@ it on the `ipi-code website <http://ipi-code.org>`_, the `documentation pages
 - examples/pi-mts-rpc/mts-rpc
 - examples/water-model/water-model
 - examples/pet-mad/pet-mad
+- examples/pet-mad-nc/pet-mad-nc

docs/src/software/metatensor.sec

@@ -12,3 +12,4 @@ training and evaluating ML models.
 - examples/water-model/water-model
 - examples/polarizability/polarizability
 - examples/pet-mad/pet-mad
+- examples/pet-mad-nc/pet-mad-nc

examples/pet-mad-nc/README.rst

@@ -0,0 +1,5 @@
+Non-conservative MTS MD 
+=======================
+
+An example of running hybrid conservative/non-conservative MD
+using PET-MAD and i-PI.

examples/pet-mad-nc/data/input-nc-nve-mts.xml

@@ -0,0 +1,45 @@
+<simulation verbosity="high">
+<output prefix="nve-nc-mts">
+    <properties stride="1" filename="out">
+		[step, time{picosecond}, conserved{electronvolt}, temperature{kelvin},
+        kinetic_md{electronvolt}, potential{electronvolt}, 
+        pot_component(0){electronvolt}, pot_component(1){electronvolt}
+		]
+    </properties>
+    <trajectory filename="pos" stride="2" format="ase"> positions </trajectory>
+    <trajectory filename="forces_c" stride="2" format="ase"> forces_component(0) </trajectory>
+    <trajectory filename="forces_nc" stride="2" format="ase"> forces_component(1) </trajectory>
+    <checkpoint stride="1000"/>
+</output>
+<total_steps>50</total_steps>
+<prng><seed>12345</seed></prng>
+
+<ffdirect  name='cons' pbc="false">
+    <pes>metatensor</pes>
+    <parameters>{template:data/bmimcl.xyz,model:pet-mad-latest.pt,device:cpu,non_conservative:False} </parameters>
+</ffdirect>
+<ffdirect  name='nocons' pbc="false">
+    <pes>metatensor</pes>
+    <parameters>{template:data/bmimcl.xyz,model:pet-mad-latest.pt,device:cpu,non_conservative:True} </parameters>
+</ffdirect>
+<system>
+    <initialize nbeads="1">
+        <file mode="ase"> data/bmimcl.xyz </file>
+        <velocities mode="thermal" units="kelvin"> 400.0 </velocities>
+    </initialize>
+    <forces>
+        <force forcefield="cons">
+            <mts_weights>[1,0]</mts_weights>
+        </force>
+        <force forcefield="nocons">
+            <mts_weights>[-1,1]</mts_weights>
+        </force>
+    </forces>
+    <motion mode="dynamics">
+    	<dynamics mode="nve">
+        	<timestep units="femtosecond"> 4 </timestep>
+            <nmts>[1,4]</nmts>
+        </dynamics>
+    </motion>
+</system>
+</simulation>

examples/pet-mad-nc/data/input-nc-nve.xml

@@ -0,0 +1,36 @@
+<simulation verbosity="high">
+<output prefix="nve-nc">
+    <properties stride="4" filename="out">
+		[step, time{picosecond}, conserved{electronvolt}, temperature{kelvin},
+        kinetic_md{electronvolt}, potential{electronvolt}, 
+        pot_component(0){electronvolt}
+		]
+    </properties>
+    <trajectory filename="pos" stride="8" format="ase"> positions </trajectory>
+    <trajectory filename="forces_nc" stride="8" format="ase"> forces_component(0) </trajectory>
+    <checkpoint stride="1000"/>
+</output>
+<total_steps>200</total_steps>
+<prng><seed>12345</seed></prng>
+<ffdirect  name='nocons' pbc="false">
+    <pes>metatensor</pes>
+    <parameters>{template:data/bmimcl.xyz,model:pet-mad-latest.pt,device:cpu,non_conservative:True} </parameters>
+</ffdirect>
+<system>
+    <initialize nbeads="1">
+        <file mode="ase"> data/bmimcl.xyz </file>
+        <velocities mode="thermal" units="kelvin"> 400.0 </velocities>
+    </initialize>
+    <forces>
+        <force forcefield="nocons">
+        </force>
+    </forces>
+
+    <motion mode="dynamics">
+    	<dynamics mode="nve">
+        	<timestep units="femtosecond"> 1 </timestep>
+        </dynamics>
+    </motion>
+
+</system>
+</simulation>

examples/pet-mad-nc/data/input-nve.xml

@@ -0,0 +1,36 @@
+<simulation verbosity="high">
+<output prefix="nve-c">
+    <properties stride="4" filename="out">
+		[step, time{picosecond}, conserved{electronvolt}, temperature{kelvin},
+        kinetic_md{electronvolt}, potential{electronvolt}, 
+        pot_component(0){electronvolt}
+		]
+    </properties>
+    <trajectory filename="pos" stride="8" format="ase"> positions </trajectory>
+    <trajectory filename="forces_c" stride="8" format="ase"> forces_component(0) </trajectory>
+    <checkpoint stride="1000"/>
+</output>
+<total_steps>32</total_steps>
+<prng><seed>12345</seed></prng>
+<ffdirect  name='cons' pbc="false">
+    <pes>metatensor</pes>
+    <parameters>{template:data/bmimcl.xyz,model:pet-mad-latest.pt,device:cpu,non_conservative:False} </parameters>
+</ffdirect>
+<system>
+    <initialize nbeads="1">
+        <file mode="ase"> data/bmimcl.xyz </file>
+        <velocities mode="thermal" units="kelvin"> 400.0 </velocities>
+    </initialize>
+    <forces>
+        <force forcefield="cons">
+        </force>
+    </forces>
+
+    <motion mode="dynamics">
+    	<dynamics mode="nve">
+        	<timestep units="femtosecond"> 1 </timestep>
+        </dynamics>
+    </motion>
+
+</system>
+</simulation>

examples/pet-mad-nc/environment.yml

@@ -0,0 +1,13 @@
+channels:
+  - metatensor
+  - conda-forge
+dependencies:
+  - python=3.12
+  - pip
+  - pip:
+    - ase>=3.23
+    - pet-mad>=1.1.0rc4,<2.0  
+    - ipi>=3.1.3,<4.0
+    - chemiscope>=0.8.5,<0.9
+    - matplotlib>=3.10,<4.0
+    - vesin[torch]>=0.3.2,<0.4

examples/pet-mad-nc/pet-mad-nc.py

@@ -0,0 +1,228 @@
+"""
+Using non-conservative forces with PET-MAD
+===========================================
+
+:Authors: Michele Ceriotti `@ceriottm <https://github.com/ceriottm>`_,
+          Filippo Bigi `@frostedoyster <https://github.com/frostedoyster>`_
+
+PET-MAD is introduced, and benchmarked for several challenging modeling tasks,
+in `this preprint <https://arxiv.org/abs/2503.14118>`_. To get a first taste
+of PET-MAD, for basic tasks such as geometry optimization and conservative
+MD, see also
+`this recipe <https://atomistic-cookbook.org/examples/pet-mad/pet-mad.html>`_.
+"""
+
+# %%
+#
+# If you don't want to use the conda environment for this recipe,
+# you can get all dependencies installing
+# the `PET-MAD package <https://github.com/lab-cosmo/pet-mad>`_:
+#
+# .. code-block:: bash
+#
+#     pip install git+https://github.com/lab-cosmo/pet-mad.git
+#
+
+import ase.io
+
+# i-PI scripting utilities
+import chemiscope
+import matplotlib.pyplot as plt
+import metatensor.torch.atomistic as mta
+
+# import numpy as np
+# import requests
+from ipi.utils.parsing import read_output, read_trajectory
+from ipi.utils.scripting import InteractiveSimulation
+
+# pet-mad ASE calculator
+from pet_mad.calculator import PETMADCalculator
+
+
+# from copy import copy, deepcopy
+
+
+if hasattr(__import__("builtins"), "get_ipython"):
+    get_ipython().run_line_magic("matplotlib", "inline")  # noqa: F821
+
+
+# %%
+# Fetch PET-MAD and export the model
+# ----------------------------------
+
+calculator = PETMADCalculator(version="latest", device="cpu")
+
+# %%
+#
+# The model can also be exported in a format that can be used with
+# external MD engines. This is done by saving the model to a file,
+# which includes the model architecture and weights.
+
+calculator.model.save("pet-mad-latest.pt")
+
+
+# %%
+#
+# Non-conservative forces
+# -----------------------
+#
+# Interatomic potentials are typically used to compute the forces acting
+# on the atoms by differentiating with respect to atomic positions, i.e. if
+#
+# .. math ::
+#
+#    V(\mathbf{r}_1, \mathbf{r}_2, \ldots \mathbf{r}_n)
+#
+# is the potential for an atomic configuration then the force acting on
+# atom :math:`i` is
+#
+# .. math ::
+#
+#    \mathbf{f}_i = \partial V/\partial \mathbf{r}_i
+#
+# Even though the early ML-based interatomic potentials followed this route,
+# .... blah blah and intro to be filled. Reference to our paper.
+
+structure = ase.io.read("data/bmimcl.xyz")
+
+structure.calc = calculator
+energy_c = structure.get_potential_energy()
+forces_c = structure.get_forces()
+
+calculator_nc = mta.ase_calculator.MetatensorCalculator(
+    calculator.model, device="cpu", non_conservative=True
+)
+
+structure.calc = calculator_nc
+energy_nc = structure.get_potential_energy()
+# forces_nc = structure.get_forces()
+
+# %%
+
+print(f"Energy:\n  Conservative: {energy_c:.8}\n  Non-conserv.: {energy_nc:.8}")
+
+# %%
+#
+# Constant-energy molecular dynamics without energy conservation
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+#
+# Some more blah blah. Explain the system (melting of Ni)
+
+# %%
+# We use an `XML` input file that instructs ``i-PI`` to perform a constant-
+# energy dynamics. The ``non_conservative:True`` option requires the model
+# to predict forces using a dedicated direct head, rather than backpropagation.
+
+with open("data/input-nc-nve.xml", "r") as file:
+    input_nve = file.read()
+print(input_nve)
+
+# %%
+# This input file uses a ``<ffdirect>`` block to run the metatensor PES
+# as a library -- it is also possible to use a socket interface that is useful
+# for parallelizing over multiple evaluators, see e.g.
+# `this recipe
+# <https://atomistic-cookbook.org/examples/path-integrals/path-integrals.html>`_.
+# ``i-PI`` can be run as a stand-alone command
+#
+# .. code-block:: bash
+#
+#    i-pi data/input-nc-nve.xml > log
+#
+# but here we use the Python API to integrate it in the notebook.
+
+sim = InteractiveSimulation(input_nve)
+sim.run(200)
+
+# %%
+#
+# The simulation generates output files that can be parsed and visualized from Python.
+
+data, info = read_output("nve-nc.out")
+trj = read_trajectory("nve-nc.pos_0.extxyz")
+
+fig, ax = plt.subplots(1, 1, figsize=(4, 3), constrained_layout=True)
+
+ax.set_facecolor("white")
+ax.plot(data["time"], data["potential"], "b-", label="potential")
+ax.plot(data["time"], data["conserved"] - 20, "k-", label="conserved")
+ax.set_xlabel("t / ps")
+ax.set_ylabel("energy / ev")
+ax.legend()
+
+# %%
+#
+# The trajectory (which is started from oxygen molecules placed on top of the surface)
+# shows quick relaxation to an oxide layer. If you look carefully, you'll also see that
+# Mg and Si atoms tend to cluster together, and accumulate at the surface.
+
+chemiscope.show(
+    frames=trj,
+    properties={
+        "time": data["time"][::2],
+        "potential": data["potential"][::2],
+        "temperature": data["temperature"][::2],
+    },
+    mode="default",
+    settings=chemiscope.quick_settings(
+        map_settings={
+            "x": {"property": "time", "scale": "linear"},
+            "y": {"property": "potential", "scale": "linear"},
+        },
+        structure_settings={
+            "unitCell": True,
+        },
+        trajectory=True,
+    ),
+)
+
+
+# %%
+#
+# Energy conservation at low-cost with multiple time stepping
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+#
+# Some more blah blah.
+
+# %%
+# We use an `XML` input file that instructs ``i-PI`` to perform a constant-
+# energy dynamics. The ``non_conservative:True`` option requires the model
+# to predict forces using a dedicated direct head, rather than backpropagation.
+
+with open("data/input-nc-nve-mts.xml", "r") as file:
+    input_nve_mts = file.read()
+print(input_nve_mts)
+
+# %%
+# This input file uses a ``<ffdirect>`` block to run the metatensor PES
+# as a library -- it is also possible to use a socket interface that is useful
+# for parallelizing over multiple evaluators, see e.g.
+# `this recipe
+# <https://atomistic-cookbook.org/examples/path-integrals/path-integrals.html>`_.
+# ``i-PI`` can be run as a stand-alone command
+#
+# .. code-block:: bash
+#
+#    i-pi data/input-nc-nve-mts.xml > log
+#
+# but here we use the Python API to integrate it in the notebook.
+
+sim = InteractiveSimulation(input_nve_mts)
+sim.run(50)
+
+# %%
+
+data_mts, info = read_output("nve-nc-mts.out")
+trj_mts = read_trajectory("nve-nc-mts.pos_0.extxyz")
+
+fig, ax = plt.subplots(1, 1, figsize=(4, 3), constrained_layout=True)
+
+ax.set_facecolor("white")
+ax.plot(data["time"], data["potential"], "b-", label="potential")
+ax.plot(data["time"], data["conserved"] - 20, "k-", label="conserved")
+ax.plot(data_mts["time"], data_mts["conserved"] - 20, "r-", label="conserved")
+ax.plot(data_mts["time"], data_mts["potential"] - 20, "r-", label="conserved")
+ax.set_xlabel("t / ps")
+ax.set_ylabel("energy / ev")
+ax.legend()
+# %%