You can deploy this model and start earning money today!

ekwan16/
MagNET

chemistry

nmr

chemical-shift-prediction

molecular-property-prediction

equiformer

mit

Views

No views yet

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.

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 (¹H shieldings) ωB97X-D/pcSseg-2 (¹³C 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 (¹H shieldings) ωB97X-D/pcSseg-2/gas (¹³C shieldings)
sigma-concentrate	50K random structures from sigma-shake PCM(chloroform) corrections at B3LYP/pcSseg-2

Installing MagNET

Install MagNET and its dependencies.

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

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

Option B: Full Installation

bash
1git clone https://github.com/ekwan/MagNET 2cd MagNET 3pip install torch==2.5.0 --index-url https://download.pytorch.org/whl/cpu 4pip install -r magnet/requirements.txt 5pip install --no-deps ./magnet
Fetch Model Weights

bash
1pip install "huggingface_hub[cli]" 2hf 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.
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.

python

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:

python

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 (CDCl₃)
carbonyl ¹³C	207.6	207.07
methyl ¹³C	30.7	30.92
methyl ¹H	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 ¹H/¹³C 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/.

Use the Option B checkout.
Download the datasets.
Install the analysis dependencies: pip install -r requirements.txt
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

bash

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

ekwan16/MagNET