Atomic tensor models¶

FYI, you can open this documentation as a Google Colab notebook to follow along interactively

graph-pes provides models that target atomic “tensorial” properties, ranging from atomic energies and charges, dipoles, NMR anisotropic parameters, to higher rank tensors.

Available models¶

The currently available models extend the MACE and NequIP architectures:

Both NequIP- and MACE-based models implement two learning approaches: direct and tensor_product. To learn tensor components with nonstandard spherical harmonics, such as 0o; 1e; 2o; 3e;..., using a MACE-based model, we recommend the tensor_product approach.

Data preparation¶

For the remainder of this notebook, we will reconstruct lightweight TensorNequIP NMR models targeting the magnetic shielding tensor (MS) used for amorphous silica in this paper.

First we download the training data:

[1]:

!wget https://github.com/cbenmahm/anistropic-nmr-parameters-data/raw/refs/heads/main/data/train_test/train.xyz

--2025-11-27 12:51:12--  https://github.com/cbenmahm/anistropic-nmr-parameters-data/raw/refs/heads/main/data/train_test/train.xyz
Resolving github.com (github.com)... 20.26.156.215
Connecting to github.com (github.com)|20.26.156.215|:443... connected.
HTTP request sent, awaiting response... 302 Found
Location: https://media.githubusercontent.com/media/cbenmahm/anistropic-nmr-parameters-data/refs/heads/main/data/train_test/train.xyz [following]
--2025-11-27 12:51:12--  https://media.githubusercontent.com/media/cbenmahm/anistropic-nmr-parameters-data/refs/heads/main/data/train_test/train.xyz
Resolving media.githubusercontent.com (media.githubusercontent.com)... 185.199.109.133, 185.199.111.133, 185.199.110.133, ...
Connecting to media.githubusercontent.com (media.githubusercontent.com)|185.199.109.133|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 52559049 (50M) [application/octet-stream]
Saving to: ‘train.xyz.1’

train.xyz.1         100%[===================>]  50.12M  19.1MB/s    in 2.6s

2025-11-27 12:51:15 (19.1 MB/s) - ‘train.xyz.1’ saved [52559049/52559049]

[2]:

import load_atoms

structures = load_atoms.load_dataset("./train.xyz")

[3]:

import graph_pes  # noqa: F401

A generic rank 2 Cartesian tensor \(T\) (of size 3 \(\times\) 3), e.g. magnetic shielding tensor, can be decomposed into symmetric and antisymmetric parts:

\[T_{\mathrm{symm}} = \dfrac{1}{2}(T+T^\mathrm{T}), \quad T_{\mathrm{antisymm}} = \dfrac{1}{2}(T-T^\mathrm{T})\]

These Cartesian coordinates have their spherical conterparts, that are usually expressed in terms of spherical harmonics (SHs). In this case we obtain the following:

a scalar part \(\ell=0\): the trace, a rotationally invariant scalar (dimension=1).
a pseudovector part \(\ell=1\): dual to the antisymmetric part (dimension=3)
a quadropole part \(\ell=2\): the traceless symmetric part (dimension=5)

In equivariant models, such as MACE and NequIP, we use features that transform irreducibly under rotations. e3nn provides tools to compute exactly these irreducible components from Cartesian tensors, allowing us to use them directly as equivariant features or targets.

The symmetry-aware conversion between the Cartesian components and the corresponding irreducible spherical representation is handled by the class e3nn.io.CartesianTensor.

[4]:

import torch
from e3nn import io

cartesian_symm = io.CartesianTensor("ij=ji")  # symmetry constarint
cartesian_antisymm = io.CartesianTensor("ij=-ji")  # antisymmetry constraint

We use the two CartesianTensor objects to transform the magnetic shielding tensors from Cartesian coordiantes to their irreducible spherical representation, which is the natural way to describe these properties:

cartesian_symm extracts the symmetric part of the tensor, which consists of the scalar part and the quadropole part. In terms of e3nn’s Irreps notations, we obtain the 0e (scalar) and 2e (quadropole) terms.
cartesian_antisymm: extracts the pseudovector part of the tensor. For a rank 2 tensor, this corresonds a pseudovector. In terms of e3nn’s Irreps notations, we obtain the 1e (pseudovector) term.

[5]:

for frm in structures:
    # extract the magnetic shielding tensor from the ase.Atoms object
    torch_ms = torch.from_numpy(frm.arrays["ms"].reshape(-1, 3, 3))

    # extract its symmetric and antisummetric components
    symm = cartesian_symm.from_cartesian(torch_ms)
    anti = cartesian_antisymm.from_cartesian(torch_ms)

    # rearrange terms to follow the $ell$ order: 0, 1, 2
    ms = torch.cat((symm[..., :1], anti, symm[..., 1:]), dim=-1)
    ms = ms.numpy()
    frm.arrays["tensor"] = ms

[6]:

import ase.io

train, val, test = structures.random_split([0.8, 0.1, 0.1])

ase.io.write("train-nmr.xyz", train)
ase.io.write("val-nmr.xyz", val)
ase.io.write("test-nmr.xyz", test)

Configuration file¶

Now that we’ve saved our labelled structures to suitable files, we’re ready to train a model.

To do this, we have specified the following in the tensornequip-direct-sio2.yaml file:

the model architecture to instantiate and train, here TensorNequIP. Note that we also include a FixedTensorOffset component to account for the fact that the amophous silica labels have an arbitrary offset.
the data to train on, here a random split of the amorphous silica dataset we just downloaded
the loss function to use, here a per-atom RMSE
and various other training hyperparameters (e.g. the learning rate, batch size, etc.)

You can download this config file for the direct approach using wget:

[7]:

%%bash

if [ ! -f tensornequip-direct-sio2.yaml ]; then
    wget https://raw.githubusercontent.com/jla-gardner/graph-pes/refs/heads/main/docs/source/quickstart/tensornequip-direct-sio2.yaml -O tensornequip-direct-sio2.yaml
fi

or this one for the tensor_product approach:

[8]:

%%bash

if [ ! -f tensornequip-tp-sio2.yaml ]; then
    wget https://raw.githubusercontent.com/jla-gardner/graph-pes/refs/heads/main/docs/source/quickstart/tensornequip-direct-sio2.yaml -O tensornequip-direct-sio2.yaml
fi

Training¶

The models are trained in the same way as the usual GraphPES models using the graph-pes-train command.

[9]:

%%bash

graph-pes-train tensornequip-direct-sio2.yaml \
    general/run_id=train-nequip-tensor

[graph-pes INFO]: Started `graph-pes-train` at 2025-11-27 12:55:54.216
[graph-pes INFO]: Successfully parsed config.
[graph-pes WARNING]: Specified run ID "train-nequip-tensor" already exists. Using train-nequip-tensor-1 instead.
[graph-pes INFO]: Logging to graph-pes-results/train-nequip-tensor-1/rank-0.log
[graph-pes INFO]: ID for this training run: train-nequip-tensor-1
[graph-pes INFO]:
Output for this training run can be found at:
    └─ graph-pes-results/train-nequip-tensor-1
        ├─ rank-0.log         # find a verbose log here
        ├─ model.pt           # the best model (according to valid/loss/total)
        ├─ lammps_model.pt    # the best model deployed to LAMMPS
        ├─ train-config.yaml  # the complete config used for this run
        └─ summary.yaml       # the summary of the training run

GPU available: True (mps), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs

[graph-pes INFO]: Preparing data
[graph-pes INFO]: Caching neighbour lists for 50 structures with cutoff 5.5329999923706055, property mapping None and torch dtype torch.float32
[graph-pes INFO]: Caching neighbour lists for 85 structures with cutoff 5.5329999923706055, property mapping None and torch dtype torch.float32
[graph-pes INFO]: Setting up datasets
[graph-pes INFO]:
Number of learnable params:
    offset (LearnableTensorOffset): 18
    many_body (TensorNequIP)      : 380,178

[graph-pes INFO]: Sanity checking the model...
[graph-pes INFO]: Starting fit...

/Users/work/source/miniconda3/envs/uv/lib/python3.11/site-packages/pytorch_lightning/loops/fit_loop.py:310: The number of training batches (10) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.

                      timer/its_per_s   timer/its_per_s
   epoch       time             train             valid
       1      174.9           0.10544           0.27718
       2      348.3           0.09474           0.29139
       3      521.3           0.09121           0.29418
       4      697.5           0.08887           0.28558
       5      878.9           0.08652           0.28721

`Trainer.fit` stopped: `max_epochs=5` reached.

[graph-pes INFO]: Loading best weights from "/Users/work/source/graph-pes/docs/source/quickstart/graph-pes-results/train-nequip-tensor-1/checkpoints/best.ckpt"
[graph-pes INFO]: Training complete.
[graph-pes INFO]: Testing best model...

💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.
GPU available: True (mps), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs

Testing DataLoader 2: 100%|██████████| 17/17 [01:08<00:00,  0.25it/s]        | 0/17 [00:00<?, ?it/s]
[graph-pes INFO]: Testing complete.
[graph-pes INFO]: Awaiting final Lightning and W&B shutdown...

Model analysis¶

First, we load the model

[ ]:

import torch

from graph_pes.models import load_model

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
best_model = (
    load_model(
        "graph-pes-results/train-nequip-tensor/model.pt"
    )  # load the model
    .to(device)  # move to GPU if available
    .eval()  # set to evaluation mode
)

[ ]:

from graph_pes.utils.calculator import GraphPESCalculator

calculator = GraphPESCalculator(best_model)

Then, we need to tranform the predictions back to Cartesian coordinates

[ ]:

for frm in test:
    calculator.calculate(frm, properties=["tensor"])
    tensor = calculator.results["tensor"]

    tensor = torch.from_numpy(tensor)

    symm = torch.cat((tensor[..., :1], tensor[..., 4:]), dim=-1)
    tensor_symm = cartesian_symm.to_cartesian(symm)

    tensor_antisymm = cartesian_antisymm.to_cartesian(tensor[..., 1:4])

    tensor = tensor_symm + tensor_antisymm
    tensor = tensor.cpu().numpy()
    frm.arrays["ms_ML"] = tensor

and now you can use libraries like soprano to exctract NMR tensor properties or MRSimulator to simulate spectra!

In what follows, we compare the magnetic shielding isotropy obtained from DFT and ML calculations. We also examine one quantity that characterizes the anisotropy of the shielding tensor: the asymmetry parameter \(\zeta\).

To extract all of these values, we will use the soprano package.

[ ]:

!pip install soprano

import numpy as np
from soprano.properties import nmr

[ ]:

# soprano expects the tensors to have the shape 3x3
for frm in test:
    frm.arrays["ms"] = frm.arrays["ms"].reshape(-1, 3, 3)
    frm.arrays["ms_ML"] = frm.arrays["ms_ML"].reshape(-1, 3, 3)

[ ]:

shielding_dft = np.concatenate(
    [nmr.MSIsotropy.get(frame, tag="ms") for frame in test]
)

shielding_ml = np.concatenate(
    [nmr.MSIsotropy.get(frame, tag="ms_ML") for frame in test]
)

[ ]:

import matplotlib.pyplot as plt

%config InlineBackend.figure_format = 'retina'

[ ]:

plt.scatter(shielding_dft, shielding_ml, color="crimson", s=4)

plt.axline((shielding_dft[0], shielding_dft[0]), axline=1, lw=1, color="k")
plt.xlabel("DFT isotropic shielding (ppm)")
plt.ylabel("ML isotropic shielding (ppm)")

[ ]:

asymmetry_dft = np.concatenate(
    [nmr.MSAsymmetry.get(frame, tag="ms") for frame in test]
)

asymmetry_ml = np.concatenate(
    [nmr.MSAsymmetry.get(frame, tag="ms_ML") for frame in test]
)

[ ]:

plt.scatter(asymmetry_dft, asymmetry_ml, color="crimson", s=4)

plt.axline((asymmetry_dft[0], asymmetry_dft[0]), axline=1, lw=1, color="k")
plt.xlabel("DFT asymmetry")
plt.ylabel("ML asymmetry")