deploy: 96f2c74

metatensor · Jul 28, 2024 · 05dfea2 · 05dfea2
commit 05dfea2
Show file tree

Hide file tree

Showing 205 changed files with 155,240 additions and 0 deletions.
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,6 @@
+/.tox/
+/build/
+/dist/
+*.egg-info
+__pycache__/
+
diff --git a/.nojekyll b/.nojekyll
diff --git a/index.html b/index.html
@@ -0,0 +1,11 @@
+<!DOCTYPE html>
+<html>
+
+<head>
+  <meta charset="utf-8" />
+  <meta http-equiv="X-UA-Compatible" content="IE=edge,chrome=1" />
+  <meta http-equiv="refresh" content="0;URL=latest/index.html" />
+</head>
+
+<body></body>
+</html>
diff --git a/latest/.buildinfo b/latest/.buildinfo
@@ -0,0 +1,4 @@
+# Sphinx build info version 1
+# This file hashes the configuration used when building these files. When it is not found, a full rebuild will be done.
+config: de62692a86e8789db1a2d38f12e3fc2e
+tags: 645f666f9bcd5a90fca523b33c5a78b7
diff --git a/latest/_downloads/1403485000e5389f5f360017c3c8ecb2/llpr.ipynb b/latest/_downloads/1403485000e5389f5f360017c3c8ecb2/llpr.ipynb
@@ -0,0 +1,151 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n# Computing LLPR uncertainties\n\nThis tutorial demonstrates how to use an already trained and exported model\nfrom Python. It involves the computation of the local prediction rigidity\n([LPR](LPR_)) for every atom of a single ethanol molecule, using the\nlast-layer prediction rigidity ([LLPR](LLPR_)) approximation.\n\n\nThe model was trained using the following training options.\n\n.. literalinclude:: options.yaml\n   :language: yaml\n\nYou can train the same model yourself with\n\n.. literalinclude:: train.sh\n   :language: bash\n\nA detailed step-by-step introduction on how to train a model is provided in\nthe `label_basic_usage` tutorial.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "import torch\nfrom metatensor.torch.atomistic import load_atomistic_model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Exported models can be loaded using the `load_atomistic_model` function from the\nmetatensor.torch.atomistic` module. The function requires the path to the exported\nmodel and, for many models, also the path to the respective extensions directory.\nBoth are produced during the training process.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "model = load_atomistic_model(\"model.pt\", extensions_directory=\"extensions/\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "In metatrain, a Dataset is composed of a list of systems and a dictionary of targets.\nThe following lines illustrate how to read systems and targets from xyz files, and\nhow to create a Dataset object from them.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "from metatrain.utils.data import Dataset, read_systems, read_targets  # noqa: E402\nfrom metatrain.utils.neighbor_lists import get_system_with_neighbor_lists  # noqa: E402\n\n\nqm9_systems = read_systems(\"qm9_reduced_100.xyz\")\n\ntarget_config = {\n    \"energy\": {\n        \"quantity\": \"energy\",\n        \"read_from\": \"ethanol_reduced_100.xyz\",\n        \"reader\": \"ase\",\n        \"key\": \"energy\",\n        \"unit\": \"kcal/mol\",\n        \"forces\": False,\n        \"stress\": False,\n        \"virial\": False,\n    },\n}\ntargets, _ = read_targets(target_config)\n\nrequested_neighbor_lists = model.requested_neighbor_lists()\nqm9_systems = [\n    get_system_with_neighbor_lists(system, requested_neighbor_lists)\n    for system in qm9_systems\n]\ndataset = Dataset({\"system\": qm9_systems, **targets})\n\n# We also load a single ethanol molecule on which we will compute properties.\n# This system is loaded without targets, as we are only interested in the LPR\n# values.\nethanol_system = read_systems(\"ethanol_reduced_100.xyz\")[0]\nethanol_system = get_system_with_neighbor_lists(\n    ethanol_system, requested_neighbor_lists\n)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "The dataset is fully compatible with torch. For example, be used to create\na DataLoader object.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "from metatrain.utils.data import collate_fn  # noqa: E402\n\n\ndataloader = torch.utils.data.DataLoader(\n    dataset,\n    batch_size=10,\n    shuffle=False,\n    collate_fn=collate_fn,\n)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "We now wrap the model in a LLPRUncertaintyModel object, which will allows us\nto compute prediction rigidity metrics, which are useful for uncertainty\nquantification and model introspection.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "from metatensor.torch.atomistic import (  # noqa: E402\n    MetatensorAtomisticModel,\n    ModelMetadata,\n)\n\nfrom metatrain.utils.llpr import LLPRUncertaintyModel  # noqa: E402\n\n\nllpr_model = LLPRUncertaintyModel(model)\nllpr_model.compute_covariance(dataloader)\nllpr_model.compute_inverse_covariance(regularizer=1e-4)\n\n# calibrate on the same dataset for simplicity. In reality, a separate\n# calibration/validation dataset should be used.\nllpr_model.calibrate(dataloader)\n\nexported_model = MetatensorAtomisticModel(\n    llpr_model.eval(),\n    ModelMetadata(),\n    llpr_model.capabilities,\n)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "We can now use the model to compute the LPR for every atom in the ethanol molecule.\nTo do so, we create a ModelEvaluationOptions object, which is used to request\nspecific outputs from the model. In this case, we request the uncertainty in the\natomic energy predictions.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "from metatensor.torch.atomistic import ModelEvaluationOptions, ModelOutput  # noqa: E402\n\n\nevaluation_options = ModelEvaluationOptions(\n    length_unit=\"angstrom\",\n    outputs={\n        # request the uncertainty in the atomic energy predictions\n        \"mtt::aux::energy_uncertainty\": ModelOutput(per_atom=True),\n        # `per_atom=False` would return the total uncertainty for the system,\n        # or (the inverse of) the TPR (total prediction rigidity)\n        # you also can request other outputs from the model here, for example:\n        # \"energy\": ModelOutput(per_atom=True),\n        # \"mtt::aux::last_layer_features\": ModelOutput(per_atom=True),\n    },\n    selected_atoms=None,\n)\n\noutputs = exported_model([ethanol_system], evaluation_options, check_consistency=False)\nlpr = outputs[\"mtt::aux::energy_uncertainty\"].block().values.detach().cpu().numpy()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "We can now visualize the LPR values using the `plot_atoms` function from\n``ase.visualize.plot``.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "import ase.io  # noqa: E402\nimport matplotlib.pyplot as plt  # noqa: E402\nfrom ase.visualize.plot import plot_atoms  # noqa: E402\nfrom matplotlib.colors import LogNorm  # noqa: E402\n\n\nstructure = ase.io.read(\"ethanol_reduced_100.xyz\")\nnorm = LogNorm(vmin=min(lpr), vmax=max(lpr))\ncolormap = plt.get_cmap(\"viridis\")\ncolors = colormap(norm(lpr))\nax = plot_atoms(structure, colors=colors, rotation=\"180x,0y,0z\")\ncustom_ticks = [1e10, 2e10, 5e10, 1e11, 2e11]\ncbar = plt.colorbar(\n    plt.cm.ScalarMappable(norm=norm, cmap=colormap),\n    ax=ax,\n    label=\"LPR\",\n    ticks=custom_ticks,\n)\ncbar.ax.set_yticklabels([f\"{tick:.0e}\" for tick in custom_ticks])\ncbar.minorticks_off()\nplt.show()"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.12.4"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
diff --git a/latest/_downloads/2ab49b443a2890c2d9b24a2fcebdf275/llpr.py b/latest/_downloads/2ab49b443a2890c2d9b24a2fcebdf275/llpr.py
@@ -0,0 +1,181 @@
+"""
+Computing LLPR uncertainties
+============================
+
+This tutorial demonstrates how to use an already trained and exported model
+from Python. It involves the computation of the local prediction rigidity
+(`LPR <LPR_>`_) for every atom of a single ethanol molecule, using the
+last-layer prediction rigidity (`LLPR <LLPR_>`_) approximation.
+
+.. _LPR: https://pubs.acs.org/doi/10.1021/acs.jctc.3c00704
+.. _LLPR: https://arxiv.org/html/2403.02251v1
+
+The model was trained using the following training options.
+
+.. literalinclude:: options.yaml
+   :language: yaml
+
+You can train the same model 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.
+"""
+
+# %%
+#
+
+import torch
+from metatensor.torch.atomistic import load_atomistic_model
+
+
+# %%
+#
+# Exported models can be loaded using the `load_atomistic_model` function from the
+# metatensor.torch.atomistic` module. The function requires the path to the exported
+# model and, for many models, also the path to the respective extensions directory.
+# Both are produced during the training process.
+
+
+model = load_atomistic_model("model.pt", extensions_directory="extensions/")
+
+# %%
+#
+# In metatrain, a Dataset is composed of a list of systems and a dictionary of targets.
+# The following lines illustrate how to read systems and targets from xyz files, and
+# 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
+
+
+qm9_systems = read_systems("qm9_reduced_100.xyz")
+
+target_config = {
+    "energy": {
+        "quantity": "energy",
+        "read_from": "ethanol_reduced_100.xyz",
+        "reader": "ase",
+        "key": "energy",
+        "unit": "kcal/mol",
+        "forces": False,
+        "stress": False,
+        "virial": False,
+    },
+}
+targets, _ = read_targets(target_config)
+
+requested_neighbor_lists = model.requested_neighbor_lists()
+qm9_systems = [
+    get_system_with_neighbor_lists(system, requested_neighbor_lists)
+    for system in qm9_systems
+]
+dataset = Dataset({"system": qm9_systems, **targets})
+
+# We also load a single ethanol molecule on which we will compute properties.
+# This system is loaded without targets, as we are only interested in the LPR
+# values.
+ethanol_system = read_systems("ethanol_reduced_100.xyz")[0]
+ethanol_system = get_system_with_neighbor_lists(
+    ethanol_system, requested_neighbor_lists
+)
+
+# %%
+#
+# The dataset is fully compatible with torch. For example, be used to create
+# a DataLoader object.
+
+from metatrain.utils.data import collate_fn  # noqa: E402
+
+
+dataloader = torch.utils.data.DataLoader(
+    dataset,
+    batch_size=10,
+    shuffle=False,
+    collate_fn=collate_fn,
+)
+
+
+# %%
+#
+# We now wrap the model in a LLPRUncertaintyModel object, which will allows us
+# to compute prediction rigidity metrics, which are useful for uncertainty
+# quantification and model introspection.
+
+from metatensor.torch.atomistic import (  # noqa: E402
+    MetatensorAtomisticModel,
+    ModelMetadata,
+)
+
+from metatrain.utils.llpr import LLPRUncertaintyModel  # noqa: E402
+
+
+llpr_model = LLPRUncertaintyModel(model)
+llpr_model.compute_covariance(dataloader)
+llpr_model.compute_inverse_covariance(regularizer=1e-4)
+
+# calibrate on the same dataset for simplicity. In reality, a separate
+# calibration/validation dataset should be used.
+llpr_model.calibrate(dataloader)
+
+exported_model = MetatensorAtomisticModel(
+    llpr_model.eval(),
+    ModelMetadata(),
+    llpr_model.capabilities,
+)
+
+# %%
+#
+# We can now use the model to compute the LPR for every atom in the ethanol molecule.
+# To do so, we create a ModelEvaluationOptions object, which is used to request
+# specific outputs from the model. In this case, we request the uncertainty in the
+# atomic energy predictions.
+
+from metatensor.torch.atomistic import ModelEvaluationOptions, ModelOutput  # noqa: E402
+
+
+evaluation_options = ModelEvaluationOptions(
+    length_unit="angstrom",
+    outputs={
+        # request the uncertainty in the atomic energy predictions
+        "mtt::aux::energy_uncertainty": ModelOutput(per_atom=True),
+        # `per_atom=False` would return the total uncertainty for the system,
+        # or (the inverse of) the TPR (total prediction rigidity)
+        # you also can request other outputs from the model here, for example:
+        # "energy": ModelOutput(per_atom=True),
+        # "mtt::aux::last_layer_features": ModelOutput(per_atom=True),
+    },
+    selected_atoms=None,
+)
+
+outputs = exported_model([ethanol_system], evaluation_options, check_consistency=False)
+lpr = outputs["mtt::aux::energy_uncertainty"].block().values.detach().cpu().numpy()
+
+# %%
+#
+# We can now visualize the LPR values using the `plot_atoms` function from
+# ``ase.visualize.plot``.
+
+import ase.io  # noqa: E402
+import matplotlib.pyplot as plt  # noqa: E402
+from ase.visualize.plot import plot_atoms  # noqa: E402
+from matplotlib.colors import LogNorm  # noqa: E402
+
+
+structure = ase.io.read("ethanol_reduced_100.xyz")
+norm = LogNorm(vmin=min(lpr), vmax=max(lpr))
+colormap = plt.get_cmap("viridis")
+colors = colormap(norm(lpr))
+ax = plot_atoms(structure, colors=colors, rotation="180x,0y,0z")
+custom_ticks = [1e10, 2e10, 5e10, 1e11, 2e11]
+cbar = plt.colorbar(
+    plt.cm.ScalarMappable(norm=norm, cmap=colormap),
+    ax=ax,
+    label="LPR",
+    ticks=custom_ticks,
+)
+cbar.ax.set_yticklabels([f"{tick:.0e}" for tick in custom_ticks])
+cbar.minorticks_off()
+plt.show()
diff --git a/latest/_downloads/3cc6812a4940272d54fcd07ee3d9c396/llpr.zip b/latest/_downloads/3cc6812a4940272d54fcd07ee3d9c396/llpr.zip