MagNET

MagNET is a family of neural networks for predicting NMR chemical shifts. This repository contains all models, datasets, and code to reproduce the data in the MagNET paper.

Contents

MagNET Models

All models use Equiformer-V2 and have approximately 10M weights. Separate weights are given for 1H and 13C prediction. Input structures can contain H, C, N, O, F, Cl, and S (inference should not be performed on structures with unsupported elements).

sigma Datasets

Installing MagNET

1. Install MagNET and its dependencies.

   Option A: Model Only via PyPI (won't work until publication)

   pip install torch==2.5.0 --index-url https://download.pytorch.org/whl/cpu
   pip install magnet-nmr
   pip install torch_scatter torch_cluster --find-links https://data.pyg.org/whl/torch-2.5.0+cpu.html
   

   Option B: Full Installation

   git clone https://github.com/ekwan/MagNET
   cd MagNET
   pip install torch==2.5.0 --index-url https://download.pytorch.org/whl/cpu
   pip install -r magnet/requirements.txt
   pip install --no-deps ./magnet
   

2. Fetch Model Weights

   pip install "huggingface_hub[cli]"
   hf download ekwan16/MagNET --local-dir . --include "model_checkpoints/*"
   

   The model weights will be downloaded into the working directory. No repository cloning is needed, as MagNET will automatically check the current directory for weights. Alternatively, you may pass the checkpoints_dir parameter to the inference methods.

3. Check the Installation

   python -c "import magnet"

   This should work with no errors.

Your First Prediction

Let's predict the 1H and 13C shifts of acetone in chloroform. predict_shifts takes an AIMNet2-optimized geometry and runs MagNET-Zero, MagNET-PCM, and the paper's scaling for you.

import numpy as np
import magnet

# acetone, (CH3)2C=O, on an AIMNet2-optimized geometry (Angstrom)
atomic_numbers = np.array([6, 6, 8, 6, 1, 1, 1, 1, 1, 1])   # supported elements: H C N O F Cl S
geometry = np.array([
    [ 1.2913, -0.5947, -0.0016],   # C  methyl
    [ 0.0029,  0.1931, -0.0010],   # C  carbonyl
    [-0.0174,  1.3994, -0.0003],   # O
    [-1.2743, -0.6189,  0.0004],   # C  methyl
    [-1.0822, -1.6899,  0.0014],   # H
    [-1.8597, -0.3513,  0.8791],   # H
    [-1.8605, -0.3530, -0.8783],   # H
    [ 1.3415, -1.2208,  0.8916],   # H
    [ 1.3170, -1.2669, -0.8614],   # H
    [ 2.1415,  0.0788, -0.0298],   # H
])

shifts = magnet.predict_shifts(atomic_numbers, geometry, solvent="chloroform")   # per-atom ppm

shifts has one value per atom (NaN where the atom is neither 1H nor 13C). Averaging acetone's chemically-equivalent atoms by index:

sites = {
    "carbonyl 13C": [1],
    "methyl 13C":   [0, 3],
    "methyl 1H":    [4, 5, 6, 7, 8, 9],
}
for name, indices in sites.items():
    print(f"{name:<13} {shifts[indices].mean():6.2f} ppm")

The predictions are close to experiment:

Notes:

• Passes and mirror averaging. n_passes (default 10) averages out equivariance error over the specified number of forward passes. If mirror_average (default True) is set, an additional n_passes are also performed on the mirror image of the input geometry (for a total of 2*n_passes). The passes are run as a single batch per head rather than one forward pass each, so raising n_passes costs far less than it used to; max_batch_graphs caps how many run at once and does not change the answer. mirror_average was called symmetrize before, and that name still works but warns: the option averages over the mirror image, not over symmetry-equivalent nuclei.
• Geometry. MagNET-Zero and MagNET-PCM require AIMNet2-optimized geometries. Do not use other geometries.
• Components. Pass return_components=True to also get the MagNET-Zero shielding, the MagNET-PCM correction, and the scaling coefficients behind each shift, as a dict.
• Other models. For raw shieldings and solvent corrections, see the API documentation.

Getting the datasets

The supporting datasets for this paper are large (42 GB) and are archived on Hugging Face. To fetch everything:

hf download ekwan16/MagNET --local-dir .

If you only want to download a single dataset, add --include "data/<name>/*" with one of the names below.

Reproducing Figures and Tables

Each figure and table in the paper is produced by a notebook under analysis/.

1. Use the Option B checkout.
2. Download the datasets.
3. Install the analysis dependencies: pip install -r requirements.txt
4. Run the analysis scripts: python reproduce.py

If you want to reproduce specific items: python reproduce.py fig3 s10

The notebooks are stored without outputs. Running a notebook outputs its figures into figures/ and its tables into documents/.

Running the Tests

pip install -r requirements.txt "pytest>=7"
pytest

The tests use small synthetic fixtures. If the large files are present, then more comprehensive tests will run. The magnet package tests need the Option B checkout with the inference stack installed (pip install -r magnet/requirements.txt); without PyTorch they are skipped.

API Documentation

You can render the API docs with pdoc:

pip install pdoc
python build_api_docs.py

Most people will only want to look at api_docs/magnet.html (the public API), though the other functions are documented as well.

How to Cite

"Chemical Shift Prediction Beyond the Electronic Structure Limit."

Adams, K.; Wagen, C.C.; Wolford, J.; Sak, M.H.; Saurí, J.; Feng, Z.; Bhadauria, A.S.; Bailey, M.A.; Downs, J.S.; Li, S.Z.; Liu, A.I.; Smidt, T.; Paton, R.S.; Liu, R.Y.; Coley, C.W.*; Kwan, E.E.*

submitted, July 2026.

License

• The original code, model weights, and datasets in this repository are released under the MIT License (see LICENSE).
• Third-party literature data redistributed here remains subject to its original publications' terms and should be cited accordingly: CP3 (data/cp3/, Smith and Goodman), NS372 (data/ns372/, Schattenberg and Kaupp), DFT8K (data/dft8k/, Guan and Paton), and DELTA50 (data/delta50/, Cohen et al.).