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

item	contents
magnet/	the importable model package
data/	the sigma datasets and other auxiliary datasets
analysis/	scripts for reproducing figures and tables
model_checkpoints/	weights for the four MagNET models

MagNET Models

Model	Description	Details
MagNET	foundation model	for near-equilibrium geometries trained on PBE0/pcSseg-1/gas shieldings
MagNET-Zero	high-quality gas-phase solute shieldings	use AIMNet2-optimized geometries WP04/pcSseg-2 (1H shieldings) ωB97X-D/pcSseg-2 (13C shieldings)
MagNET-PCM	implicit solvent corrections	shielding(PCM) - shielding(gas) use AIMNet2-optimized geometries computed at B3LYP-D3(BJ)/pcSseg-2/chloroform
MagNET-x	explicit solvent corrections	shielding(solute+solvent) - shielding(solute) supports chloroform, benzene, methanol, and water use classical MD geometries trained at PBE0/pcSseg-1

All models use Equiformer-V2 and have approximately 10M weights. Separate weights are given for ¹H and ¹³C 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

Dataset	Details
sigma-shake	4.4M solutes from GDB-13/17 with functional group augmentation stationary and quasiclassically perturbed structures at B3LYP-D3(BJ)/6-31G* PBE0/pcSseg-1/gas shieldings
sigma-fresh	10K representative natural product, drug-like, sugar, and peptide solutes dissolved in benzene, chloroform, methanol, and water ~10 computed poses/solute with many more solvated geometries available 462K poses have computed PBE0/pcSseg-1 shieldings
sigma-pepper (part 1)	GDB molecules with ≤ 10 heavy atoms AIMNet2 stationary structures PBE0/pcSseg-1 shieldings
sigma-pepper (part 2)	GDB molecules with ≤ 9 heavy atoms AIMNet2 stationary structures WP04/pcSseg-2/gas (1H shieldings) ωB97X-D/pcSseg-2/gas (13C shieldings)
sigma-concentrate	50K random structures from sigma-shake PCM(chloroform) corrections at B3LYP/pcSseg-2

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 ¹H and ¹³C 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 ¹H nor ¹³C). 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:

site	MagNET	experiment (CDCl3)
carbonyl 13C	207.6	207.07
methyl 13C	30.7	30.92
methyl 1H	2.20	2.17

Notes:

• Passes and symmetry. n_passes (default 10) averages out equivariance error over the specified number of forward passes. If symmetrize (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).
• 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.

dataset	size	contents
sigma-shake	2.2 GB	4.4M GDB solutes, perturbed geometries, PBE0/pcSseg-1 shieldings
sigma-fresh	30 GB	10K solutes in explicit benzene, chloroform, methanol, and water
sigma-pepper	501 MB	AIMNet2 GDB structures used to train MagNET-Zero and MagNET-PCM
sigma-concentrate	17 MB	PCM corrections on 50K sigma-shake structures
delta22	2.5 GB	experimental 1H/13C shifts with matched DFT and MagNET predictions
dft8k	12 MB	external benchmark of 7111 organics
gdb_qcd	291 MB	quasiclassical dynamics of 2461 GDB molecules
magnet_test_predictions	822 MB	MagNET predictions and DFT targets
applications	5.1 GB	natural product MD geometries and experimental shifts
supertestset_magnet_x	1.2 MB	explicit solvent test set for MagNET-x

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.).