--- license: mit tags: - chemistry - nmr - chemical-shift-prediction - molecular-property-prediction - equiformer --- # 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 | | | **MagNET-Zero** | high-quality gas-phase solute shieldings | | | **MagNET-PCM** | implicit solvent corrections | | | **MagNET-x** | explicit solvent corrections | | 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 | Dataset | Details | |---|---| | **_sigma_-shake** | | | **_sigma_-fresh** | | | **_sigma_-pepper** (part 1) | | | **_sigma_-pepper** (part 2) | | | **_sigma_-concentrate** | | ### Installing MagNET 1. Install MagNET and its dependencies. **Option A: Model Only via PyPI (won't work until publication)** ```bash 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** ```bash 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** ```bash 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. ```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 1H nor 13C). 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 (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](https://github.com/isayevlab/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](#api-documentation). ### Getting the datasets The supporting datasets for this paper are large (42 GB) and are [archived on Hugging Face](https://huggingface.co/ekwan16/MagNET). To fetch everything: hf download ekwan16/MagNET --local-dir . If you only want to download a single dataset, add `--include "data//*"` 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](#getting-the-datasets) 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 ```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](https://pdoc.dev): 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`](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.*).