Add ASE tutorial

.gitignore

@@ -162,7 +162,12 @@ cython_debug/
 # Don't save model outputs
 *.pt
 !tests/resources/*.pt
+!examples/ase/*.pt
 
 # model output directories
 outputs/
 examples/*.xyz
+
+# sphinx gallery
+docs/src/tutorials/gallery
+*execution_times*

MANIFEST.in

@@ -3,6 +3,9 @@ graft src
 include LICENSE
 include README.md
 
+# TODO: remove once the hack is not needed anymore
+include scripts/setup.py
+
 prune docs
 prune examples
 prune tests

docs/requirements.txt

@@ -1,4 +1,7 @@
 furo
 sphinx >= 7
+sphinx-gallery
 sphinx-toggleprompt
 tomli
+
+ase  # for ASE tutorial

docs/src/conf.py

@@ -41,6 +41,7 @@
     "sphinx.ext.autodoc",
     "sphinx.ext.intersphinx",
     "sphinx_toggleprompt",
+    "sphinx_gallery.load_style",
 ]
 
 python_use_unqualified_type_names = True

docs/src/gallery.py

@@ -0,0 +1,68 @@
+import os
+import shutil
+import sys
+
+import sphinx_gallery.gen_gallery
+
+
+HERE = os.path.realpath(os.path.dirname(__file__))
+
+
+class AttrDict(dict):
+    def __init__(self):
+        super().__init__()
+        self.__dict__ = self
+
+
+class PseudoSphinxApp:
+    """
+    Class pretending to be a sphinx App, used to configure and run sphinx-gallery
+    from the command line, without having an actual sphinx project.
+    """
+
+    def __init__(self, example):
+        gallery_dir = os.path.join(
+            HERE, "tutorials", "gallery", os.path.basename(example)
+        )
+        if os.path.exists(gallery_dir):
+            shutil.rmtree(gallery_dir)
+
+        # the options here are the minimal set of options to get sphinx-gallery to run
+        # feel free to add more if sphinx-gallery uses more options in the future
+        self.config = AttrDict()
+        self.config.html_static_path = []
+        self.config.source_suffix = [".rst"]
+        self.config.default_role = ""
+        self.config.sphinx_gallery_conf = {
+            "filename_pattern": ".*",
+            "examples_dirs": os.path.join(HERE, "..", "..", example),
+            "gallery_dirs": gallery_dir,
+            "min_reported_time": 60,
+        }
+
+        self.builder = AttrDict()
+        self.builder.srcdir = os.path.join(HERE)
+        self.builder.outdir = ""
+        self.builder.name = os.path.basename(example)
+
+        self.extensions = []
+
+        self.builder.config = AttrDict()
+        self.builder.config.plot_gallery = "True"
+        self.builder.config.abort_on_example_error = True
+        self.builder.config.highlight_language = None
+
+    def add_css_file(self, path):
+        pass
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print(f"usage: {sys.argv[0]} <example/dir>")
+        sys.exit(1)
+
+    app = PseudoSphinxApp(example=sys.argv[1])
+
+    sphinx_gallery.gen_gallery.fill_gallery_conf_defaults(app, app.config)
+    sphinx_gallery.gen_gallery.update_gallery_conf_builder_inited(app)
+    sphinx_gallery.gen_gallery.generate_gallery_rst(app)

docs/src/getting-started/usage.rst

@@ -58,15 +58,31 @@ the current directory and type
 Evaluation
 ##########
 
-The sub-command to evaluate a already trained model is
+The sub-command to evaluate an already trained model is
 
 .. code-block:: bash
 
     metatensor-models eval
 
 .. literalinclude:: ../../../examples/usage.sh
     :language: bash
-    :lines: 9-
+    :lines: 9-25
+
+
+Exporting
+#########
+
+Exporting a model is very useful if you want to use it in other frameworks,
+especially in molecular dynamics simulations.
+The sub-command to export an already trained model is
+
+.. code-block:: bash
+
+    metatensor-models export
+
+.. literalinclude:: ../../../examples/usage.sh
+    :language: bash
+    :lines: 25-
 
 In the next tutorials we show how adjust the dataset section of ``options.yaml`` file
 to use it for your own datasets.

docs/src/index.rst

@@ -12,5 +12,6 @@ This is a collection of atomistic models interfaced with ``metatensor``.
    :hidden:
 
    getting-started/index
+   tutorials/index
    architectures/index
    dev-docs/index

docs/src/tutorials/index.rst

@@ -0,0 +1,10 @@
+Tutorials
+=========
+
+This sections includes some more advanced tutorials on the usage of the
+``metatensor-models`` package.
+
+.. toctree::
+   :maxdepth: 1
+
+   gallery/ase/run_ase

examples/ase/README.rst

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

examples/ase/run_ase.py

@@ -0,0 +1,91 @@
+"""
+Running molecular dynamics with ASE
+===================================
+
+This tutorial shows how to use an exported model to run an ASE simulation.
+"""
+
+# %%
+#
+# First, we import the necessary libraries:
+
+# NumPy
+import numpy as np
+
+# Tools to run the simulation
+import ase.md
+import ase.md.velocitydistribution
+import ase.units
+
+# Integration with ASE calculator for metatensor atomistic models
+from metatensor.torch.atomistic.ase_calculator import MetatensorCalculator
+
+# Plotting
+import matplotlib.pyplot as plt
+import ase.visualize.plot
+
+# The SOAP-BPNN model contains compiled extensions from rascaline.torch
+import rascaline.torch
+
+
+# %%
+#
+# Next, we initialize the simulation, including importing the model:
+
+# Initial positions (reading them from a file):
+atoms = ase.io.read("../../tests/resources/ethanol_reduced_100.xyz")
+
+# Initialize the velocities:
+ase.md.velocitydistribution.MaxwellBoltzmannDistribution(atoms, temperature_K=300)
+
+# Load the model and register it as the energy calculator for these ``atoms``:
+atoms.calc = MetatensorCalculator("exported-model.pt")
+
+# Define the integrator:
+integrator = ase.md.VelocityVerlet(
+    atoms,
+    timestep=1.0 * ase.units.fs
+)
+
+# Plot the initial configuration:
+ase.visualize.plot.plot_atoms(atoms)
+plt.show()
+
+# %%
+#
+# Run a short simulation:
+
+potential_energy = []
+kinetic_energy = []
+total_energy = []
+trajectory = []
+
+for step in range(100):
+    # run a single simulation step
+    integrator.run(1)
+
+    # collect data about the simulation
+    potential_energy.append(atoms.get_potential_energy())
+    kinetic_energy.append(atoms.get_kinetic_energy())
+    total_energy.append(atoms.get_total_energy())
+    trajectory.append(atoms.copy())
+
+# Convert to numpy arrays
+potential_energy = np.array(potential_energy)
+kinetic_energy = np.array(kinetic_energy)
+total_energy = np.array(total_energy)
+
+# Plot the final configuration:
+ase.visualize.plot.plot_atoms(trajectory[99])
+plt.show()
+
+# %%
+#
+# Plot the evolution of kinetic, potential, and total energy.
+# The total energy should approximately be conserved:
+
+plt.plot(range(100), potential_energy-potential_energy.mean(), label="potential energy")
+plt.plot(range(100), kinetic_energy-kinetic_energy.mean(), label="kinetic energy")
+plt.plot(range(100), total_energy-total_energy.mean(), label="total energy")
+plt.legend()
+plt.show()

examples/ase/train_model/options.yaml

@@ -0,0 +1,21 @@
+defaults:
+  - architecture: soap_bpnn  # architecture used to train the model
+  - _self_
+
+architecture:
+  training:
+    batch_size: 16
+    num_epochs: 1000
+    learning_rate: 0.003
+
+# Section defining the parameters for structure and target data
+training_set:
+  structures:
+    read_from: "../../../tests/resources/ethanol_reduced_100.xyz"
+  targets:
+    energy:
+      key: "energy"
+      unit: "kcal/mol"  # very important to run simulations
+
+validation_set: 0.1
+test_set: 0.0

examples/usage.sh

@@ -22,3 +22,16 @@ head -n 20 output.xyz
 # All command line flags of the eval sub-command can be listed via
 
 metatensor-models eval --help
+
+# However, before we export the model, we need to run the following command to
+# set up the environment for exporting
+
+python ../scripts/setup.py
+
+# The above script can be found in the `scripts` folder of the repository.
+
+# Finally, the `metatestor-models export`, i.e.,
+
+metatensor-models export model.pt
+
+# creates an `exported-model.pt` file that contains the exported model.

scripts/setup.py

@@ -0,0 +1,48 @@
+# Since torch.jit.save cannot handle Labels.single(), we need to replace it with
+# Labels(names=["_"], values=_dispatch.zeros_like(block.values, (1, 1)))
+# in metatensor-operations. This is a hacky way to do it.
+
+import os
+import metatensor.operations
+import metatensor.torch.atomistic
+
+file = os.path.join(
+    os.path.dirname(metatensor.operations.__file__),
+    "reduce_over_samples.py"
+)
+
+# Find the line that contains "Labels.single()"
+# and replace "Labels.single()" with 
+# "Labels(names=["_"], values=_dispatch.zeros_like(block.values, (1, 1)))"
+with open(file, "r") as f:
+    lines = f.readlines()
+    for i, line in enumerate(lines):
+        if "samples_label = Labels.single()" in line:
+            lines[i] = line.replace(
+                "samples_label = Labels.single()",
+                "samples_label = Labels(names=[\"_\"], values=_dispatch.zeros_like(block.values, (1, 1)))"
+            )
+            break
+
+with open(file, "w") as f:
+    f.writelines(lines)
+
+file = os.path.join(
+    os.path.dirname(metatensor.torch.atomistic.__file__),
+    "units.py"
+)
+
+# At the moment, unit handling is incorrect in metatensor.torch.atomistic
+# This will fix it:
+with open(file, "r") as f:
+    lines = f.readlines()
+    for i, line in enumerate(lines):
+        if "return self._conversions[to_unit] / self._conversions[from_unit]" in line:
+            lines[i] = line.replace(
+                "return self._conversions[to_unit] / self._conversions[from_unit]",
+                "return self._conversions[from_unit] / self._conversions[to_unit]"
+            )
+            break
+
+with open(file, "w") as f:
+    f.writelines(lines)

src/metatensor/models/cli/eval_model.py

@@ -57,8 +57,7 @@ def eval_model(model: str, structures: str, output: str = "output.xyz") -> None:
     loaded_model = load_model(model)
     structure_list = read_structures(structures)
 
-    # since the second argument is missing,
-    # this calculates all the available properties:
-    predictions = loaded_model(structure_list)
+    # this calculates all the properties that the model is capable of predicting:
+    predictions = loaded_model(structure_list, loaded_model.capabilities.outputs)
 
     write_predictions(output, predictions, structure_list)