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-07-08 12:27:37,214 - INFO - Found model version in url_resolve
2026-07-08 12:27:37,215 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_0.pt?download=1
2026-07-08 12:27:37,215 - INFO - Model not found in cache, downloading it
2026-07-08 12:27:46,166 - INFO - Downloaded ShiftML30 and saved to /home/runner/.cache/shiftml/ShiftML30
2026-07-08 12:27:46,323 - INFO - Found model version in url_resolve
2026-07-08 12:27:46,323 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_1.pt?download=1
2026-07-08 12:27:46,324 - INFO - Model not found in cache, downloading it
2026-07-08 12:27:57,553 - INFO - Downloaded ShiftML31 and saved to /home/runner/.cache/shiftml/ShiftML31
2026-07-08 12:27:57,705 - INFO - Found model version in url_resolve
2026-07-08 12:27:57,705 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_2.pt?download=1
2026-07-08 12:27:57,705 - INFO - Model not found in cache, downloading it
2026-07-08 12:28:02,064 - INFO - Downloaded ShiftML32 and saved to /home/runner/.cache/shiftml/ShiftML32
2026-07-08 12:28:02,217 - INFO - Found model version in url_resolve
2026-07-08 12:28:02,217 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_3.pt?download=1
2026-07-08 12:28:02,218 - INFO - Model not found in cache, downloading it
2026-07-08 12:28:11,017 - INFO - Downloaded ShiftML33 and saved to /home/runner/.cache/shiftml/ShiftML33
2026-07-08 12:28:11,168 - INFO - Found model version in url_resolve
2026-07-08 12:28:11,169 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_4.pt?download=1
2026-07-08 12:28:11,169 - INFO - Model not found in cache, downloading it
2026-07-08 12:28:17,713 - INFO - Downloaded ShiftML34 and saved to /home/runner/.cache/shiftml/ShiftML34
2026-07-08 12:28:17,864 - INFO - Found model version in url_resolve
2026-07-08 12:28:17,864 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_5.pt?download=1
2026-07-08 12:28:17,864 - INFO - Model not found in cache, downloading it
2026-07-08 12:28:26,774 - INFO - Downloaded ShiftML35 and saved to /home/runner/.cache/shiftml/ShiftML35
2026-07-08 12:28:26,926 - INFO - Found model version in url_resolve
2026-07-08 12:28:26,926 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_6.pt?download=1
2026-07-08 12:28:26,926 - INFO - Model not found in cache, downloading it
2026-07-08 12:28:30,024 - INFO - Downloaded ShiftML36 and saved to /home/runner/.cache/shiftml/ShiftML36
2026-07-08 12:28:30,175 - INFO - Found model version in url_resolve
2026-07-08 12:28:30,176 - INFO - Resolving model version to model files at url: https://zenodo.org/records/15767390/files/model_7.pt?download=1
2026-07-08 12:28:30,176 - INFO - Model not found in cache, downloading it
2026-07-08 12:28:39,294 - 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 0x7f16d855b850>]

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 0x7f16d458b710>]

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"],
        },
    ),
)

×

×