Computing NMR shielding tensors using ShiftML

Authors:

Michele Ceriotti @ceriottm

This example shows how to compute NMR shielding tensors using a point-edge transformer model trained on the ShiftML dataset.

import os
import zipfile

import chemiscope
import matplotlib.pyplot as plt
import numpy as np

from ase.io import read
from atomistic_cookbook_utils import download_with_retry
from shiftml.ase import ShiftML

Create a ShiftML calculator and fetch a dataset

calculator = ShiftML("ShiftML3")

filename = "ShiftML_poly.zip"
download_with_retry(
    "https://archive.materialscloud.org/records/j2fka-sda13/files/ShiftML_poly.zip",
    filename,
)


with zipfile.ZipFile(filename, "r") as zip_ref:
    for file in ["ShiftML_poly/Cocaine/cocaine_QuantumEspresso.xyz"]:
        target = os.path.basename(file)
        with zip_ref.open(file) as source, open(target, "wb") as dest:
            dest.write(source.read())

frames_cocaine = read("cocaine_QuantumEspresso.xyz", index=":16")
reference = [frame.arrays["CS"] for frame in frames_cocaine]

2026-05-19 13:46:51,917 - INFO - Found model version in url_resolve
2026-05-19 13:46:51,918 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_0.pt?download=1
2026-05-19 13:46:51,918 - INFO - Model not found in cache, downloading it
2026-05-19 13:47:04,893 - INFO - Downloaded ShiftML30 and saved to /home/runner/.cache/shiftml/ShiftML30
2026-05-19 13:47:05,058 - INFO - Found model version in url_resolve
2026-05-19 13:47:05,058 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_1.pt?download=1
2026-05-19 13:47:05,058 - INFO - Model not found in cache, downloading it
2026-05-19 13:47:17,724 - INFO - Downloaded ShiftML31 and saved to /home/runner/.cache/shiftml/ShiftML31
2026-05-19 13:47:17,883 - INFO - Found model version in url_resolve
2026-05-19 13:47:17,883 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_2.pt?download=1
2026-05-19 13:47:17,883 - INFO - Model not found in cache, downloading it
2026-05-19 13:47:30,724 - INFO - Downloaded ShiftML32 and saved to /home/runner/.cache/shiftml/ShiftML32
2026-05-19 13:47:30,882 - INFO - Found model version in url_resolve
2026-05-19 13:47:30,882 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_3.pt?download=1
2026-05-19 13:47:30,882 - INFO - Model not found in cache, downloading it
2026-05-19 13:47:43,372 - INFO - Downloaded ShiftML33 and saved to /home/runner/.cache/shiftml/ShiftML33
2026-05-19 13:47:43,530 - INFO - Found model version in url_resolve
2026-05-19 13:47:43,530 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_4.pt?download=1
2026-05-19 13:47:43,531 - INFO - Model not found in cache, downloading it
2026-05-19 13:47:58,835 - INFO - Downloaded ShiftML34 and saved to /home/runner/.cache/shiftml/ShiftML34
2026-05-19 13:47:58,993 - INFO - Found model version in url_resolve
2026-05-19 13:47:58,993 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_5.pt?download=1
2026-05-19 13:47:58,993 - INFO - Model not found in cache, downloading it
2026-05-19 13:48:12,442 - INFO - Downloaded ShiftML35 and saved to /home/runner/.cache/shiftml/ShiftML35
2026-05-19 13:48:12,600 - INFO - Found model version in url_resolve
2026-05-19 13:48:12,600 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_6.pt?download=1
2026-05-19 13:48:12,601 - INFO - Model not found in cache, downloading it
2026-05-19 13:48:20,125 - INFO - Downloaded ShiftML36 and saved to /home/runner/.cache/shiftml/ShiftML36
2026-05-19 13:48:20,284 - INFO - Found model version in url_resolve
2026-05-19 13:48:20,284 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_7.pt?download=1
2026-05-19 13:48:20,284 - INFO - Model not found in cache, downloading it
2026-05-19 13:48:35,758 - INFO - Downloaded ShiftML37 and saved to /home/runner/.cache/shiftml/ShiftML37

Predicts isotropic chemical shielding tensors, including uncertainty

predicted = [calculator.get_cs_iso_ensemble(frame) for frame in frames_cocaine]

Make a plot for all the H shielding values

h_shieldings = np.hstack(
    [
        [
            r[f.symbols == "H"],
            p[f.symbols == "H"].mean(axis=1),
            p[f.symbols == "H"].std(axis=1),
        ]
        for r, p, f in zip(reference, predicted, frames_cocaine)
    ]
)

and of the uncertainty

fig, ax = plt.subplots(figsize=(4, 3))

ax.plot(h_shieldings[0], h_shieldings[1], "o", markersize=2, alpha=0.5)

[<matplotlib.lines.Line2D object at 0x7f3ebf7a1290>]

fig, ax = plt.subplots(figsize=(4, 3))
ax.loglog(
    h_shieldings[2],
    np.abs(h_shieldings[0] - h_shieldings[1]),
    "o",
    markersize=2,
    alpha=0.5,
)
ax.plot([0.2, 0.8], [0.2, 0.8], "k--", lw=0.5)

[<matplotlib.lines.Line2D object at 0x7f3ebace91d0>]

Anisotropic shielding tensors

tensors = [calculator.get_cs_tensor(frame) for frame in frames_cocaine]

h_tensors = np.array(
    [
        (t[iat] if f.symbols[iat] == "H" else np.zeros((3, 3)))
        for t, f in zip(tensors, frames_cocaine)
        for iat in range(len(f))
    ]
)

chemiscope.show(
    frames_cocaine,
    shapes={
        "cs_ellipsoid": {
            "kind": "ellipsoid",
            "parameters": {
                "global": {},
                "atom": [
                    chemiscope.ellipsoid_from_tensor(cs_h * 0.05) for cs_h in h_tensors
                ],
            },
        }
    },
    mode="structure",
    settings=chemiscope.quick_settings(
        periodic=True,
        structure_settings={
            "shape": ["cs_ellipsoid"],
        },
    ),
)

×

×