Add ZBL tutorial

.readthedocs.yml

@@ -15,6 +15,7 @@ build:
     pre_build:
       - set -e && cd examples/ase && bash train.sh
       - set -e && cd examples/programmatic/llpr && bash train.sh
+      - set -e && cd examples/zbl && bash train.sh
 
 # Build documentation in the docs/ directory with Sphinx
 sphinx:

docs/generate_examples/conf.py

@@ -13,8 +13,16 @@
 sphinx_gallery_conf = {
     "filename_pattern": "/*",
     "copyfile_regex": r".*\.(pt|sh|xyz|yaml)",
-    "examples_dirs": [os.path.join(ROOT, "examples", "ase"), os.path.join(ROOT, "examples", "programmatic", "llpr")],
-    "gallery_dirs": [os.path.join(ROOT, "docs", "src", "examples", "ase"), os.path.join(ROOT, "docs", "src", "examples", "programmatic", "llpr")],
+    "examples_dirs": [
+        os.path.join(ROOT, "examples", "ase"),
+        os.path.join(ROOT, "examples", "programmatic", "llpr"),
+        os.path.join(ROOT, "examples", "zbl")
+    ],
+    "gallery_dirs": [
+        os.path.join(ROOT, "docs", "src", "examples", "ase"),
+        os.path.join(ROOT, "docs", "src", "examples", "programmatic", "llpr"),
+        os.path.join(ROOT, "docs", "src", "examples", "zbl")
+    ],
     "min_reported_time": 5,
     "matplotlib_animations": True,
 }

docs/src/tutorials/index.rst

@@ -11,3 +11,4 @@ This sections includes some more advanced tutorials on the usage of the
 
    ../examples/ase/run_ase
    ../examples/programmatic/llpr/llpr
+   ../examples/zbl/dimers

examples/ase/run_ase.py

@@ -42,11 +42,6 @@
 
 # %%
 #
-# .. note::
-#    We have to import ``rascaline.torch`` even though it is not used explicitly in this
-#    tutorial. The SOAP-BPNN model contains compiled extensions and therefore the import
-#    is required.
-#
 # Setting up the simulation
 # -------------------------
 #

examples/zbl/dimers.py

@@ -0,0 +1,114 @@
+"""
+Training a model with ZBL corrections
+=====================================
+
+This tutorial demonstrates how to train a model with ZBL corrections.
+
+The models are trained on a
+subset of the ethanol moleculs from the `rMD17 dataset
+<https://iopscience.iop.org/article/10.1088/2632-2153/abba6f/meta>`_.
+
+The models are trained using the following training options, respectively:
+
+.. literalinclude:: options_no_zbl.yaml
+   :language: yaml
+
+.. literalinclude:: options_zbl.yaml
+    :language: yaml
+
+You can train the same models yourself with
+
+.. literalinclude:: train.sh
+   :language: bash
+
+A detailed step-by-step introduction on how to train a model is provided in
+the :ref:`label_basic_usage` tutorial.
+"""
+
+# %%
+#
+# First, we start by importing the necessary libraries, including the integration of ASE
+# calculators for metatensor atomistic models.
+
+import ase
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from metatensor.torch.atomistic.ase_calculator import MetatensorCalculator
+
+
+# %%
+#
+# Setting up the dimers
+# ---------------------
+#
+# We set up a series of dimers with different atom pairs and distances. We will
+# calculate the energies of these dimers using the models trained with and without ZBL
+# corrections.
+
+distances = np.linspace(0.5, 6.0, 200)
+pairs = {}
+for pair in [("H", "H"), ("H", "C"), ("C", "C"), ("C", "O"), ("O", "O"), ("H", "O")]:
+    structures = []
+    for distance in distances:
+        atoms = ase.Atoms(
+            symbols=[pair[0], pair[1]],
+            positions=[[0, 0, 0], [0, 0, distance]],
+        )
+        structures.append(atoms)
+    pairs[pair] = structures
+
+# %%
+#
+# We now load the two exported models, one with and one without ZBL corrections
+
+calc_no_zbl = MetatensorCalculator(
+    "model_no_zbl.pt", extensions_directory="extensions/"
+)
+calc_zbl = MetatensorCalculator("model_zbl.pt", extensions_directory="extensions/")
+
+
+# %%
+#
+# Calculate and plot energies without ZBL
+# ---------------------------------------
+#
+# We calculate the energies of the dimer curves for each pair of atoms and
+# plot the results, using the non-ZBL-corrected model.
+
+for pair, structures_for_pair in pairs.items():
+    energies = []
+    for atoms in structures_for_pair:
+        atoms.set_calculator(calc_no_zbl)
+        with torch.jit.optimized_execution(False):
+            energies.append(atoms.get_potential_energy())
+    energies = np.array(energies) - energies[-1]
+    plt.plot(distances, energies, label=f"{pair[0]}-{pair[1]}")
+plt.title("Dimer curves - no ZBL")
+plt.xlabel("Distance (Å)")
+plt.ylabel("Energy (eV)")
+plt.legend()
+plt.tight_layout()
+plt.show()
+
+# %%
+#
+# Calculate and plot energies from the ZBL-corrected model
+# --------------------------------------------------------
+#
+# We repeat the same procedure as above, but this time with the ZBL-corrected model.
+
+for pair, structures_for_pair in pairs.items():
+    energies = []
+    for atoms in structures_for_pair:
+        atoms.set_calculator(calc_zbl)
+        with torch.jit.optimized_execution(False):
+            energies.append(atoms.get_potential_energy())
+    energies = np.array(energies) - energies[-1]
+    plt.plot(distances, energies, label=f"{pair[0]}-{pair[1]}")
+plt.title("Dimer curves - with ZBL")
+plt.xlabel("Distance (Å)")
+plt.ylabel("Energy (eV)")
+plt.legend()
+plt.tight_layout()
+plt.show()