{ "cells": [ { "cell_type": "markdown", "id": "978086fe", "metadata": {}, "source": [ "# Atomic tensor models" ] }, { "cell_type": "markdown", "id": "780fae38", "metadata": {}, "source": [ "> **FYI**, you can open this documentation as a [Google Colab notebook](https://colab.research.google.com/github/jla-gardner/graph-pes/blob/main/docs/source/quickstart/atomic-tensor-example.ipynb) to follow along interactively\n", "\n", "``graph-pes`` provides models that target atomic \"tensorial\" properties, ranging from atomic energies and charges, dipoles, NMR anisotropic parameters, to higher rank tensors." ] }, { "cell_type": "markdown", "id": "b2eaef28", "metadata": {}, "source": [ "## Available models" ] }, { "cell_type": "markdown", "id": "d6fdbfb3", "metadata": {}, "source": [ "The currently available models extend the `MACE` and `NequIP` architectures:\n", "1. [TensorMACE](https://jla-gardner.github.io/graph-pes/models/many-body/tensormace.html)\n", "\n", "1. [ZEmbeddingTensorMACE](https://jla-gardner.github.io/graph-pes/models/many-body/tensormace.html#graph_pes.models.ZEmbeddingTensorMACE)\n", "\n", "1. [TensorNequIP](https://jla-gardner.github.io/graph-pes/models/many-body/tensornequip.html)\n", "\n", "1. [ZEmbeddingTensorNequIP](https://jla-gardner.github.io/graph-pes/models/many-body/tensornequip.html#graph_pes.models.ZEmbeddingTensorNequIP)\n", "\n", "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.\n" ] }, { "cell_type": "markdown", "id": "b1e67968", "metadata": { "vscode": { "languageId": "plaintext" } }, "source": [ "## Data preparation" ] }, { "cell_type": "markdown", "id": "e015d3d1", "metadata": {}, "source": [ "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](https://doi.org/10.1063/5.0274240)." ] }, { "cell_type": "markdown", "id": "d7b32a41", "metadata": { "vscode": { "languageId": "plaintext" } }, "source": [ "First we download the training data:" ] }, { "cell_type": "code", "execution_count": 1, "id": "a1502c42", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "--2025-11-27 12:51:12-- https://github.com/cbenmahm/anistropic-nmr-parameters-data/raw/refs/heads/main/data/train_test/train.xyz\n", "Resolving github.com (github.com)... 20.26.156.215\n", "Connecting to github.com (github.com)|20.26.156.215|:443... connected.\n", "HTTP request sent, awaiting response... 302 Found\n", "Location: https://media.githubusercontent.com/media/cbenmahm/anistropic-nmr-parameters-data/refs/heads/main/data/train_test/train.xyz [following]\n", "--2025-11-27 12:51:12-- https://media.githubusercontent.com/media/cbenmahm/anistropic-nmr-parameters-data/refs/heads/main/data/train_test/train.xyz\n", "Resolving media.githubusercontent.com (media.githubusercontent.com)... 185.199.109.133, 185.199.111.133, 185.199.110.133, ...\n", "Connecting to media.githubusercontent.com (media.githubusercontent.com)|185.199.109.133|:443... connected.\n", "HTTP request sent, awaiting response... 200 OK\n", "Length: 52559049 (50M) [application/octet-stream]\n", "Saving to: ‘train.xyz.1’\n", "\n", "train.xyz.1 100%[===================>] 50.12M 19.1MB/s in 2.6s \n", "\n", "2025-11-27 12:51:15 (19.1 MB/s) - ‘train.xyz.1’ saved [52559049/52559049]\n", "\n" ] } ], "source": [ "!wget https://github.com/cbenmahm/anistropic-nmr-parameters-data/raw/refs/heads/main/data/train_test/train.xyz" ] }, { "cell_type": "code", "execution_count": 2, "id": "9aba4865", "metadata": {}, "outputs": [], "source": [ "import load_atoms\n", "\n", "structures = load_atoms.load_dataset(\"./train.xyz\")" ] }, { "cell_type": "code", "execution_count": 3, "id": "8aa6dc78", "metadata": {}, "outputs": [], "source": [ "import graph_pes # noqa: F401" ] }, { "cell_type": "markdown", "id": "a1b9e1b8", "metadata": {}, "source": [ "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:\n", "\n", "$$T_{\\mathrm{symm}} = \\dfrac{1}{2}(T+T^\\mathrm{T}), \\quad T_{\\mathrm{antisymm}} = \\dfrac{1}{2}(T-T^\\mathrm{T})$$\n", "\n", "These Cartesian coordinates have their spherical conterparts, that are usually expressed in terms of spherical harmonics (SHs). In this case we obtain the following:\n", "\n", "- a scalar part $\\ell=0$: the trace, a rotationally invariant scalar (dimension=1).\n", "- a pseudovector part $\\ell=1$: dual to the antisymmetric part (dimension=3)\n", "- a quadropole part $\\ell=2$: the traceless symmetric part (dimension=5)\n", "\n", "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.\n", "\n", "The symmetry-aware conversion between the Cartesian components and the corresponding irreducible spherical representation is handled by the class `e3nn.io.CartesianTensor`. " ] }, { "cell_type": "code", "execution_count": 4, "id": "6cfa468b", "metadata": {}, "outputs": [], "source": [ "import torch\n", "from e3nn import io\n", "\n", "cartesian_symm = io.CartesianTensor(\"ij=ji\") # symmetry constarint\n", "cartesian_antisymm = io.CartesianTensor(\"ij=-ji\") # antisymmetry constraint" ] }, { "cell_type": "markdown", "id": "1ef43397", "metadata": {}, "source": [ "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:\n", "\n", "* `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.\n", "\n", "* `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." ] }, { "cell_type": "code", "execution_count": 5, "id": "96d16963", "metadata": {}, "outputs": [], "source": [ "for frm in structures:\n", " # extract the magnetic shielding tensor from the ase.Atoms object\n", " torch_ms = torch.from_numpy(frm.arrays[\"ms\"].reshape(-1, 3, 3))\n", "\n", " # extract its symmetric and antisummetric components\n", " symm = cartesian_symm.from_cartesian(torch_ms)\n", " anti = cartesian_antisymm.from_cartesian(torch_ms)\n", "\n", " # rearrange terms to follow the $ell$ order: 0, 1, 2\n", " ms = torch.cat((symm[..., :1], anti, symm[..., 1:]), dim=-1)\n", " ms = ms.numpy()\n", " frm.arrays[\"tensor\"] = ms" ] }, { "cell_type": "code", "execution_count": 6, "id": "72ce4a80", "metadata": {}, "outputs": [], "source": [ "import ase.io\n", "\n", "train, val, test = structures.random_split([0.8, 0.1, 0.1])\n", "\n", "ase.io.write(\"train-nmr.xyz\", train)\n", "ase.io.write(\"val-nmr.xyz\", val)\n", "ase.io.write(\"test-nmr.xyz\", test)" ] }, { "cell_type": "markdown", "id": "c5677b4f", "metadata": {}, "source": [ "## Configuration file" ] }, { "cell_type": "markdown", "id": "8026fb8c", "metadata": {}, "source": [ "Now that we've saved our labelled structures to suitable files, we're ready to train a model.\n", "\n", "To do this, we have specified the following in the ``tensornequip-direct-sio2.yaml`` file:\n", "\n", "* the model architecture to instantiate and train, here [TensorNequIP](https://jla-gardner.github.io/graph-pes/models/many-body/tensornequip.html). Note that we also include a [FixedTensorOffset](https://jla-gardner.github.io/graph-pes/models/offsets.html#graph_pes.models.FixedTensorOffset) component to account for the fact that the amophous silica labels have an arbitrary offset.\n", "* the data to train on, here a random split of the [amorphous silica](https://github.com/cbenmahm/anistropic-nmr-parameters-data/raw/refs/heads/main/data/train_test/train.xyz) dataset we just downloaded\n", "* the loss function to use, here a per-atom RMSE\n", "* and various other training hyperparameters (e.g. the learning rate, batch size, etc.)" ] }, { "cell_type": "raw", "id": "097ccc94", "metadata": { "raw_mimetype": "text/restructuredtext", "vscode": { "languageId": "raw" } }, "source": [ ".. dropdown:: ``tensornequip-direct-sio2.yaml``\n", "\n", " .. literalinclude:: tensornequip-direct-sio2.yaml\n", " :language: yaml" ] }, { "cell_type": "markdown", "id": "1d9f85a6", "metadata": {}, "source": [ "You can download [this config file](https://raw.githubusercontent.com/jla-gardner/graph-pes/refs/heads/main/docs/source/quickstart/tensornequip-direct-sio2.yaml) for the direct approach using wget:" ] }, { "cell_type": "code", "execution_count": 7, "id": "41736019", "metadata": {}, "outputs": [], "source": [ "%%bash\n", "\n", "if [ ! -f tensornequip-direct-sio2.yaml ]; then\n", " wget https://raw.githubusercontent.com/jla-gardner/graph-pes/refs/heads/main/docs/source/quickstart/tensornequip-direct-sio2.yaml -O tensornequip-direct-sio2.yaml\n", "fi" ] }, { "cell_type": "markdown", "id": "0ca8a6dd", "metadata": { "vscode": { "languageId": "plaintext" } }, "source": [ "or this one for the tensor_product approach:" ] }, { "cell_type": "code", "execution_count": 8, "id": "c9696b95", "metadata": {}, "outputs": [], "source": [ "%%bash\n", "\n", "if [ ! -f tensornequip-tp-sio2.yaml ]; then\n", " wget https://raw.githubusercontent.com/jla-gardner/graph-pes/refs/heads/main/docs/source/quickstart/tensornequip-direct-sio2.yaml -O tensornequip-direct-sio2.yaml\n", "fi" ] }, { "cell_type": "markdown", "id": "3586d344", "metadata": { "vscode": { "languageId": "plaintext" } }, "source": [ "## Training" ] }, { "cell_type": "markdown", "id": "4e1e71b6", "metadata": { "vscode": { "languageId": "plaintext" } }, "source": [ "The models are trained in the same way as the usual ``GraphPES`` models using the [graph-pes-train](https://jla-gardner.github.io/graph-pes/cli/graph-pes-train/root.html) command." ] }, { "cell_type": "code", "execution_count": 9, "id": "b5c9b238", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[graph-pes INFO]: Started `graph-pes-train` at 2025-11-27 12:55:54.216\n", "[graph-pes INFO]: Successfully parsed config.\n", "[graph-pes WARNING]: Specified run ID \"train-nequip-tensor\" already exists. Using train-nequip-tensor-1 instead.\n", "[graph-pes INFO]: Logging to graph-pes-results/train-nequip-tensor-1/rank-0.log\n", "[graph-pes INFO]: ID for this training run: train-nequip-tensor-1\n", "[graph-pes INFO]: \n", "Output for this training run can be found at:\n", " └─ graph-pes-results/train-nequip-tensor-1\n", " ├─ rank-0.log # find a verbose log here\n", " ├─ model.pt # the best model (according to valid/loss/total)\n", " ├─ lammps_model.pt # the best model deployed to LAMMPS\n", " ├─ train-config.yaml # the complete config used for this run\n", " └─ summary.yaml # the summary of the training run\n", "\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "GPU available: True (mps), used: True\n", "TPU available: False, using: 0 TPU cores\n", "HPU available: False, using: 0 HPUs\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "[graph-pes INFO]: Preparing data\n", "[graph-pes INFO]: Caching neighbour lists for 50 structures with cutoff 5.5329999923706055, property mapping None and torch dtype torch.float32\n", "[graph-pes INFO]: Caching neighbour lists for 85 structures with cutoff 5.5329999923706055, property mapping None and torch dtype torch.float32\n", "[graph-pes INFO]: Setting up datasets\n", "[graph-pes INFO]: \n", "Number of learnable params:\n", " offset (LearnableTensorOffset): 18\n", " many_body (TensorNequIP) : 380,178\n", "\n", "[graph-pes INFO]: Sanity checking the model...\n", "[graph-pes INFO]: Starting fit...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/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.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ " timer/its_per_s timer/its_per_s\n", " epoch time train valid\n", " 1 174.9 0.10544 0.27718\n", " 2 348.3 0.09474 0.29139\n", " 3 521.3 0.09121 0.29418\n", " 4 697.5 0.08887 0.28558\n", " 5 878.9 0.08652 0.28721\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "`Trainer.fit` stopped: `max_epochs=5` reached.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "[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\"\n", "[graph-pes INFO]: Training complete.\n", "[graph-pes INFO]: Testing best model...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "💡 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.\n", "GPU available: True (mps), used: True\n", "TPU available: False, using: 0 TPU cores\n", "HPU available: False, using: 0 HPUs\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Testing DataLoader 2: 100%|██████████| 17/17 [01:08<00:00, 0.25it/s] | 0/17 [00:00