Computing LLPR uncertainties

This tutorial demonstrates how to train a model with uncertainties using metatrain. It involves the computation of the uncertainties on ethanol molecules, using the last-layer prediction rigidity (LLPR) approximation. Both total and local (LPR) uncertainties are computed

The baseline model was trained using the following training options, where the training set consists of 100 structures from the QM9 dataset.

device: cpu
base_precision: 64
seed: 42

architecture:
  name: soap_bpnn
  training:
    batch_size: 10
    num_epochs: 10
    learning_rate: 0.01

# Section defining the parameters for system and target data
training_set:
  systems: qm9_reduced_100.xyz
  targets:
    energy:
      key: U0
      unit: hartree  # very important to run simulations

validation_set: 0.1
test_set: 0.0

Once a model is trained, you can add LLPR uncertainties to it by launching a training run with the “llpr” architecture, on the same data. In this case, the training options to add LLPR uncertainties are as follows:

device: cpu
base_precision: 64
seed: 42

architecture:
  name: llpr
  training:
    model_checkpoint: model.ckpt
    batch_size: 4

# Section defining the parameters for system and target data
training_set:
  systems: qm9_reduced_100.xyz
  targets:
    energy:
      key: U0
      unit: hartree  # very important to run simulations

validation_set: 0.1
test_set: 0.0

Adding LLPR uncertainties is very cheap compared to training a model, as it only involves one pass through the training data (equivalent to one epoch of training).

You can repeat the same training yourself with

import subprocess

import ase.io
import matplotlib.pyplot as plt
from ase.visualize.plot import plot_atoms
from matplotlib.colors import LogNorm
from metatomic.torch import ModelOutput
from metatomic_ase import MetatomicCalculator

We first train the baseline model without uncertainties and then the LLPR model.

# Here, we run training as a subprocess, in reality you would run this from the command
# line, e.g. ``mtt train options-model.yaml -o model.pt``.
subprocess.run(["mtt", "train", "options-model.yaml", "-o", "model.pt"], check=True)
subprocess.run(["mtt", "train", "options-llpr.yaml", "-o", "model-llpr.pt"], check=True)

CompletedProcess(args=['mtt', 'train', 'options-llpr.yaml', '-o', 'model-llpr.pt'], returncode=0)

A detailed step-by-step introduction on how to train a model is provided in the Basic Usage tutorial.

As an example, we will compute the energies and uncertainties of the LLPR model on a few ethanol structures.

Next we load 5 ethanol structures

structures = ase.io.read("ethanol_reduced_100.xyz", ":5")

and load the model as an ASE calculator

calc = MetatomicCalculator(
    "model-llpr.pt", extensions_directory="extensions/", device="cpu"
)

# the uncertainties are available throguh the ``run_model`` method of the calculator
predictions = calc.run_model(
    structures,
    {
        "energy": ModelOutput(sample_kind="system"),
        "energy_uncertainty": ModelOutput(sample_kind="system"),
    },
)

# print the energies and uncertainties

energies = predictions["energy"].block().values.squeeze().numpy()
uncertainties = predictions["energy_uncertainty"].block().values.squeeze().numpy()

print(energies)
print(uncertainties)

/home/runner/work/metatrain/metatrain/.tox/docs/lib/python3.14/site-packages/torch/jit/_serialization.py:170: FutureWarning: `torch.jit.load` is not supported in Python 3.14+ and may break. Please switch to `torch.export`.
  warnings.warn(
/home/runner/work/metatrain/metatrain/.tox/docs/lib/python3.14/site-packages/metatensor_torch/_module.py:240: UserWarning: module does not have '_mts_buffer_names'; falling back to processing all attributes. This is deprecated and will be removed in a future version, update your version of metatensor-learn and use `register_buffer` explicitly to remove this warning.
  _apply_metatensor(self)
/home/runner/work/metatrain/metatrain/.tox/docs/lib/python3.14/site-packages/metatensor_torch/_module.py:244: UserWarning: module does not have '_mts_buffer_names'; falling back to processing all attributes. This is deprecated and will be removed in a future version, update your version of metatensor-learn and use `register_buffer` explicitly to remove this warning.
  _apply_metatensor(module)
[-154.98422258 -154.98419238 -154.98523923 -154.98111695 -154.97973745]
[204.39436553 160.16873605 138.78221312 131.58916761 128.86355591]

We can also obtain per-atom uncertainties (local prediction rigidity, LPR). As an example, we will compute the uncertainties on an ethanol structure.

structure = structures[0]
predictions = calc.run_model(
    structure,
    {
        # here, we use sample_kind="atom" to request per-atom uncertainties
        "energy_uncertainty": ModelOutput(sample_kind="atom"),
    },
)
local_uncertainty = predictions["energy_uncertainty"].block().values.squeeze().numpy()
local_uncertainty = local_uncertainty * 1000.0  # convert from eV to meV

norm = LogNorm(vmin=min(local_uncertainty), vmax=max(local_uncertainty))
colormap = plt.get_cmap("viridis")
colors = colormap(norm(local_uncertainty))
ax = plot_atoms(structure, colors=colors, rotation="180x,0y,0z")
custom_ticks = [0.7, 0.8, 0.9, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5]
cbar = plt.colorbar(
    plt.cm.ScalarMappable(norm=norm, cmap=colormap),
    ax=ax,
    label="Local energy uncertainty (meV)",
    ticks=custom_ticks,
)
cbar.ax.set_yticklabels([f"{tick}" for tick in custom_ticks])
cbar.minorticks_off()
ax.set_xticks([])
ax.set_yticks([])
plt.show()