Merge branch 'main' into llpr

.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/conf.py

@@ -10,6 +10,7 @@
 # to include the documentation
 os.environ["METATENSOR_IMPORT_FOR_SPHINX"] = "1"
 os.environ["RASCALINE_IMPORT_FOR_SPHINX"] = "1"
+os.environ["PYTORCH_JIT"] = "0"
 
 import metatrain  # noqa: E402
 
@@ -53,9 +54,11 @@ def generate_examples():
     # METATENSOR_IMPORT_FOR_SPHINX=1). So instead we run it inside a small script, and
     # include the corresponding output later.
     del os.environ["METATENSOR_IMPORT_FOR_SPHINX"]
+    del os.environ["PYTORCH_JIT"]
     script = os.path.join(ROOT, "docs", "generate_examples", "generate-examples.py")
     subprocess.run([sys.executable, script])
     os.environ["METATENSOR_IMPORT_FOR_SPHINX"] = "1"
+    os.environ["PYTORCH_JIT"] = "0"
 
 
 def setup(app):

docs/src/dev-docs/utils/additive/composition.rst

@@ -0,0 +1,7 @@
+Composition model
+#################
+
+.. automodule:: metatrain.utils.additive.composition
+    :members:
+    :undoc-members:
+    :show-inheritance:

docs/src/dev-docs/utils/additive/index.rst

@@ -0,0 +1,12 @@
+Additive models
+===============
+
+API for handling additive models in ``metatrain``. These are models that
+can be added to one or more architectures.
+
+.. toctree::
+   :maxdepth: 1
+
+   remove_additive
+   composition
+   zbl

docs/src/dev-docs/utils/additive/remove_additive.rst

@@ -0,0 +1,7 @@
+Removing additive contributions
+###############################
+
+.. automodule:: metatrain.utils.additive.remove
+    :members:
+    :undoc-members:
+    :show-inheritance:

docs/src/dev-docs/utils/additive/zbl.rst

@@ -0,0 +1,7 @@
+ZBL short-range potential
+#########################
+
+.. automodule:: metatrain.utils.additive.zbl
+    :members:
+    :undoc-members:
+    :show-inheritance:

docs/src/dev-docs/utils/index.rst

@@ -6,9 +6,9 @@ This is the API for the ``utils`` module of ``metatrain``.
 .. toctree::
    :maxdepth: 1
 
+   additive/index
    data/index
    architectures
-   composition
    devices
    dtype
    errors
@@ -24,4 +24,5 @@ This is the API for the ``utils`` module of ``metatrain``.
    omegaconf
    output_gradient
    per_atom
+   transfer
    units

docs/src/dev-docs/utils/transfer.rst

@@ -0,0 +1,7 @@
+Data type and device transfers
+##############################
+
+.. automodule:: metatrain.utils.transfer
+    :members:
+    :undoc-members:
+    :show-inheritance:

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/programmatic/llpr/llpr.py

@@ -48,7 +48,10 @@
 # how to create a Dataset object from them.
 
 from metatrain.utils.data import Dataset, read_systems, read_targets  # noqa: E402
-from metatrain.utils.neighbor_lists import get_system_with_neighbor_lists  # noqa: E402
+from metatrain.utils.neighbor_lists import (  # noqa: E402
+    get_requested_neighbor_lists,
+    get_system_with_neighbor_lists,
+)
 
 
 qm9_systems = read_systems("qm9_reduced_100.xyz")
@@ -67,7 +70,7 @@
 }
 targets, _ = read_targets(target_config)
 
-requested_neighbor_lists = model.requested_neighbor_lists()
+requested_neighbor_lists = get_requested_neighbor_lists(model)
 qm9_systems = [
     get_system_with_neighbor_lists(system, requested_neighbor_lists)
     for system in qm9_systems

examples/zbl/README.rst

@@ -0,0 +1,2 @@
+Running molecular dynamics with ASE
+===================================

examples/zbl/dimers.py

@@ -0,0 +1,148 @@
+"""
+Training a model with ZBL corrections
+=====================================
+
+This tutorial demonstrates how to train a model with ZBL corrections.
+
+The training set for this example consists of 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
+
+As you can see, they are identical, except for the ``zbl`` key in the
+``model`` section.
+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()
+
+# %%
+#
+# It can be seen that all the dimer curves include a strong repulsion
+# at short distances, which is due to the ZBL contribution. Even the H-H dimer,
+# whose ZBL correction is very weak due to the small covalent radii of hydrogen,
+# would show a strong repulsion closer to the origin (here, we only plotted
+# starting from a distance of 0.5 Å). Let's zoom in on the H-H dimer to see
+# this effect more clearly.
+
+new_distances = np.linspace(0.1, 2.0, 200)
+
+structures = []
+for distance in new_distances:
+    atoms = ase.Atoms(
+        symbols=["H", "H"],
+        positions=[[0, 0, 0], [0, 0, distance]],
+    )
+    structures.append(atoms)
+
+for atoms in structures:
+    atoms.set_calculator(calc_zbl)
+with torch.jit.optimized_execution(False):
+    energies = [atoms.get_potential_energy() for atoms in structures]
+energies = np.array(energies) - energies[-1]
+plt.plot(new_distances, energies, label="H-H")
+plt.title("Dimer curve - H-H with ZBL")
+plt.xlabel("Distance (Å)")
+plt.ylabel("Energy (eV)")
+plt.legend()
+plt.tight_layout()
+plt.show()

examples/zbl/ethanol_reduced_100.xyz

@@ -0,0 +1 @@
+../ase/ethanol_reduced_100.xyz

examples/zbl/options_no_zbl.yaml

@@ -0,0 +1,21 @@
+seed: 42
+
+architecture:
+  name: experimental.soap_bpnn
+  model:
+    zbl: false
+  training:
+    num_epochs: 10
+
+# training set section
+training_set:
+  systems:
+    read_from: ethanol_reduced_100.xyz
+    length_unit: angstrom
+  targets:
+    energy:
+      key: "energy"
+      unit: "eV" # very important to run simulations
+
+validation_set: 0.1
+test_set: 0.0

examples/zbl/options_zbl.yaml

@@ -0,0 +1,21 @@
+seed: 42
+
+architecture:
+  name: experimental.soap_bpnn
+  model:
+    zbl: true
+  training:
+    num_epochs: 10
+
+# training set section
+training_set:
+  systems:
+    read_from: ethanol_reduced_100.xyz
+    length_unit: angstrom
+  targets:
+    energy:
+      key: "energy"
+      unit: "eV" # very important to run simulations
+
+validation_set: 0.1
+test_set: 0.0

examples/zbl/train.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+mtt train options_no_zbl.yaml -o model_no_zbl.pt
+mtt train options_zbl.yaml -o model_zbl.pt