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/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/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/pet-mad-nc

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

@@ -0,0 +1,45 @@
+<simulation verbosity="medium">
+<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="5" format="ase"> positions </trajectory>
+    <trajectory filename="forces_c" stride="5" format="ase"> forces_component(0) </trajectory>
+    <trajectory filename="forces_nc" stride="5" format="ase"> forces_component(1) </trajectory>
+    <checkpoint stride="1000"/>
+</output>
+<total_steps>100</total_steps>
+<prng><seed>12345</seed></prng>
+
+<ffdirect  name='cons' pbc="false">
+    <pes>metatensor</pes>
+    <parameters>{template:data/nickel.xyz,model:model_c.pt,device:cpu} </parameters>
+</ffdirect>
+<ffdirect  name='nocons' pbc="false">
+    <pes>metatensor</pes>
+    <parameters>{template:data/nickel.xyz,model:model_nc.pt,device:cpu,non_conservative:True} </parameters>
+</ffdirect>
+<system>
+    <initialize nbeads="1">
+        <file mode="ase"> data/nickel.xyz </file>
+        <velocities mode="thermal" units="kelvin"> 10000.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"> 16 </timestep>
+            <nmts>[1,8]</nmts>
+        </dynamics>
+    </motion>
+</system>
+</simulation>

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

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

examples/pet-mad/data/nickel.xyz

@@ -0,0 +1,34 @@
+32
+Lattice="7.199999999999999 0.0 0.0 0.0 7.199999999999999 0.0 0.0 0.0 7.199999999999999" Properties=species:S:1:pos:R:3 ipi_comment="Step:         100  Bead:       0 positions{ase}  cell{ase}" pbc="T T T"
+Al      -1.85390496      -0.95982674       1.97725575
+Al       2.86261250       3.15898022       0.57331454
+Al      -0.72359428      -2.62713427       3.64623439
+Al       4.29563893       1.35474805       1.94203711
+Al      -2.25383976      -0.51040542       6.71325518
+Al       1.03395315       1.04071714       4.46807275
+Al       0.93046776       1.41776842       6.93069925
+Al       3.21332631       3.13251253       3.29654208
+Ni       0.43653416       4.39201704      -1.17541302
+Ni       2.44317928       5.24507411       1.92192439
+Ni       6.13922636       4.31067580       0.27870937
+Ni       2.99738146       7.79472421      -2.94107677
+Ni       0.63257689       2.93422124       3.65847686
+Ni      -0.13903928       6.59208609       5.59856458
+Ni       4.44853393       2.13439743       5.58726078
+Ni       0.65202417       6.46319053       2.95105454
+Ni       3.07880245       1.28266141       0.24123238
+Ni       5.68508963       0.67317343       4.71454935
+Ni       6.68773449       1.01636173       2.31489816
+Ni       5.48738394       2.67065291      -3.56673154
+Ni       3.63945425      -2.52245666       6.19962471
+Ni       4.65160368      -1.38381038       4.32521624
+Ni      -0.29530481      -0.78118178       7.51117915
+Ni       4.61973821       4.24619744       1.92444695
+Ni       5.65027402       1.51590226       0.28430981
+Ni       2.03466520       8.19169094       2.14800988
+Ni       7.89866872       3.76495004       0.87094364
+Ni       2.67420027       6.89868302      -1.07012589
+Ni       1.83831272       4.94538488       4.39742971
+Ni       1.99506120       2.57229298       5.48248905
+Ni       5.48763500       3.91153783       5.57536938
+Ni       3.44996386       6.80436847       2.68806868

examples/pet-mad/inputs/nickel.xyz

@@ -0,0 +1,34 @@
+32
+Lattice="7.2 0.0 0.0 0.0 7.2 0.0 0.0 0.0 7.2" Properties=species:S:1:pos:R:3 pbc="T T T"
+Ni       0.04454187       0.15457874       0.12596108
+Ni       0.04370318       1.95000420       1.91184154
+Ni       1.83573509       0.12884331       1.91958398
+Ni       1.92303186       1.89265686       0.07165795
+Ni       0.19190924       0.04701160       3.61889659
+Ni       0.04116880       1.91517186       5.56317915
+Ni       1.91817295       0.03169102       5.58591245
+Ni       1.90862314       1.82490488       3.69025253
+Ni       0.16624021       3.78437918       0.07411722
+Ni       0.11502732       5.41155667       1.80895247
+Ni       1.88561487       3.68501992       1.91533902
+Ni       1.89092684       5.47334069       0.13616005
+Ni       0.13256096       3.78185764       3.77535577
+Ni       0.01409603       5.46866518       5.46642113
+Ni       1.97184883       3.60053345       5.44955058
+Ni       1.95063962       5.40970450       3.62949639
+Ni       3.61636855       0.00654232       0.19947714
+Ni       3.76729948       1.93207137       1.83696219
+Ni       5.47783206       0.09663835       1.87546754
+Ni       5.51789413       1.88532203       0.15452776
+Ni       3.66412964       0.13315923       3.68647837
+Ni       3.71818983       1.96085601       5.43543213
+Ni       5.54731704       0.16090601       5.59503045
+Ni       5.40527227       1.81865617       3.73525599
+Ni       3.77409855       3.77888268       0.04458552
+Ni       3.77575860       5.56995384       1.83599413
+Ni       5.49352058       3.71215650       1.92995287
+Ni       5.53270737       5.51022677       0.07537065
+Ni       3.60323764       3.73888965       3.76030754
+Ni       3.77150406       5.57225362       5.48239838
+Ni       5.48403719       3.68038819       5.43621244
+Ni       5.51535161       5.56333248       3.64169044

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

@@ -0,0 +1,184 @@
+"""
+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 subprocess
+
+# i-PI scripting utilities
+import chemiscope
+import matplotlib.pyplot as plt
+
+# import numpy as np
+# import requests
+from ipi.utils.parsing import read_output, read_trajectory
+from ipi.utils.scripting import InteractiveSimulation
+
+
+# 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
+# ----------------------------------
+
+# TODO
+
+subprocess.run("cp data/*pt .", shell=True, check=True)
+
+# %%
+#
+# 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.
+
+# %%
+#
+# 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(800)
+
+# %%
+#
+# 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"][::5],
+        "potential": data["potential"][::5],
+        "temperature": data["temperature"][::5],
+    },
+    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(100)
+
+# %%

examples/pet-mad/pet-mad.py

@@ -96,7 +96,7 @@
 if not os.path.exists(filename):
     url = (
         "https://huggingface.co/lab-cosmo/pet-mad/resolve/"
-        "main/benchmarks/mad-test-mad-settings-v1.0.xyz"
+        "main/benchmarks/mad-test-mad-settings.xyz"
     )
     response = requests.get(url)
     response.raise_for_status()