Add files using upload-large-folder tool

.gitattributes

@@ -1,35 +1,4 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
+*.hdf5 filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.mnova filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,37 @@
+# Large files live on the Hugging Face copy of this repo (ekwan16/MagNET), never committed here.
+*.hdf5
+*.h5
+*.npy
+*.npz
+*.pkl
+*.pickle
+*.ckpt
+*.mnova
+model_checkpoints/
+# Exception: the tiny (54 KB) natural-products reference used by magnet/test_magnet.py is a committed
+# test fixture, so those tests are self-contained (they need only the released checkpoints, not a
+# separate reference download).
+!magnet/test_data/magnet_zero_and_pcm_shieldings.h5
+# Exception: delta50 is tiny (33 KB), so it ships with the repo instead of a Hugging Face download.
+!data/delta50/delta50.hdf5
+
+# Regenerated figure outputs (the notebook redraws them and embeds them in its own cells).
+analysis/**/figures/
+analysis/**/documents/
+
+# macOS
+.DS_Store
+
+# Python
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.ipynb_checkpoints/
+*.egg-info/
+.venv/
+venv/
+dist/
+build/
+
+# Rendered API docs (users regenerate with `pdoc magnet -t pdoc_templates -o api_docs`).
+api_docs/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 The MagNET authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,215 @@
 ---
 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 | - for near-equilibrium geometries <br /> - trained on PBE0/pcSseg-1/gas shieldings |
+| **MagNET-Zero** | high-quality gas-phase solute shieldings | - use AIMNet2-optimized geometries <br /> - WP04/pcSseg-2 (<sup>1</sup>H shieldings) <br /> - ωB97X-D/pcSseg-2 (<sup>13</sup>C shieldings) |
+| **MagNET-PCM** | implicit solvent corrections | - shielding(PCM) - shielding(gas) <br />- use AIMNet2-optimized geometries <br /> - computed at B3LYP/pcSseg-2/chloroform |
+| **MagNET-x** | explicit solvent corrections | - shielding(solute+solvent) - shielding(solute) <br /> - supports chloroform, benzene, methanol, and water <br /> - use classical MD geometries <br /> - trained at PBE0/pcSseg-1 |
+
+All models use Equiformer-V2 and have approximately 10M weights. Separate weights are given for <sup>1</sup>H and <sup>13</sup>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 <br /> - stationary and quasiclassically perturbed structures at B3LYP-D3(BJ)/6-31G\* <br /> - PBE0/pcSseg-1/gas shieldings |
+| **_sigma_-fresh** | - 10K representative natural product, drug-like, sugar, and peptide solutes dissolved in benzene, chloroform, methanol, and water <br /> - ~10 computed poses/solute with many more solvated geometries available <br /> - 462K poses have computed PBE0/pcSseg-1 shieldings |
+| **_sigma_-pepper** (part 1) | - GDB molecules with ≤ 10 heavy atoms <br /> - AIMNet2 stationary structures <br /> - PBE0/pcSseg-1 shieldings |
+| **_sigma_-pepper** (part 2) | - GDB molecules with ≤ 9 heavy atoms <br /> - AIMNet2 stationary structures <br /> - WP04/pcSseg-2/gas (<sup>1</sup>H shieldings) <br /> - ωB97X-D/pcSseg-2/gas (<sup>13</sup>C shieldings) |
+| **_sigma_-concentrate** | - 50K random structures from *sigma*-shake <br /> - 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)**
+
+   ```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
+   ```
+   
+   **<i>During the review process</i>**, the GitHub repository is private, so use a read-only token for cloning:
+   
+   `git clone https://<token>@github.com/ekwan/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.
+   
+   **<i>During the review process</i>**, the HuggingFace repository is private, so use a read-only token for downloading:
+   
+   `hf download ekwan16/MagNET --local-dir . --include "model_checkpoints/*" --token hf_xxx`
+ 
+3. **Check the Installation**
+ 
+   `python -c "import magnet"`
+   
+   This should work with no errors.
+  
+### Your First Prediction
+
+Let's predict the <sup>1</sup>H and <sup>13</sup>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 <sup>1</sup>H nor <sup>13</sup>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 land on the measured shifts (experiment is acetone in chloroform, from the delta22
+dataset):
+
+| site | MagNET | experiment (CDCl<sub>3</sub>) |
+|---|---|---|
+| carbonyl <sup>13</sup>C | 207.6 | 207.07 |
+| methyl <sup>13</sup>C | 30.7 | 30.92 |
+| methyl <sup>1</sup>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](#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/<name>/*"` with one of the names below.
+
+| dataset | size | contents |
+|---|---|---|
+| `sigma-shake` | 2.2 GB | main training set: 4.4M GDB solutes, perturbed geometries, PBE0/pcSseg-1 shieldings |
+| `sigma-fresh` | 30 GB | explicit-solvent poses: 10K solutes in 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 <sup>1</sup>H/<sup>13</sup>C shifts with matched DFT and MagNET predictions |
+| `dft8k` | 12 MB | external benchmark of 7111 organics (Guan and Paton), two levels of theory |
+| `gdb_qcd` | 291 MB | quasiclassical-dynamics benchmark of 2461 GDB molecules, for rovibrational corrections |
+| `magnet_test_predictions` | 822 MB | MagNET's raw per-atom predictions and DFT targets on the test sets |
+| `applications` | 5.1 GB | natural-product evaluation set, with MD geometries and experimental shifts |
+| `supertestset_magnet_x` | 1.2 MB | explicit-solvent test set for MagNET-x: isolated vs solvated DFT shieldings |
+
+### 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 and do not need the downloaded files. Any test that does require the large model weights or datasets is skipped automatically when they are absent, and runs once they have been downloaded.
+
+### API Documentation
+
+You can render the API docs with [pdoc](https://pdoc.dev):
+
+    pip install pdoc
+    python build_api_docs.py
+
+Open `api_docs/index.html`. This documents the `magnet` package (`magnet.html` is the public
+API: `predict_shifts`, `predict_shieldings`, `implicit_solvent_correction`,
+`explicit_solvent_correction`) and every analysis module under `analysis/code`.
+
+### How to Cite
+
+"Chemical Shift Prediction Beyond the Electronic Structure Limit."
+
+Adams, K.; Wagen, C.C.W.; 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: the Goodman CP3 spreadsheet (`data/cp3/`), the Schattenberg and Kaupp NS372
+spreadsheet (`data/ns372/`), and the DFT8K benchmark (`data/dft8k/`, from Guan and Paton).

build_api_docs.py

@@ -0,0 +1,50 @@
+"""Render the MagNET API docs into api_docs/ with pdoc.
+
+    python build_api_docs.py
+
+Documents the magnet package's public modules and every analysis module under analysis/code,
+skipping test files. The analysis modules import their siblings and the dataset readers by bare
+name, so their directories are added to sys.path first.
+"""
+import glob
+import os
+import sys
+
+import pdoc
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+
+
+def _import_dirs():
+    """Directories the analysis modules import from: their own folder, shared/, and each
+    data/<dataset>/ folder that holds a reader module."""
+    dirs = [os.path.join(HERE, "analysis", "code"),
+            os.path.join(HERE, "analysis", "code", "shared")]
+    dirs += sorted(os.path.dirname(p) for p in glob.glob(os.path.join(HERE, "data", "*", "*_reader.py")))
+    return dirs
+
+
+def _modules():
+    """The .py files to document: the six public magnet modules, then the analysis modules."""
+    mods = [os.path.join(HERE, "magnet", name + ".py")
+            for name in ("__init__", "api", "inference", "model", "run_magnet", "scaling")]
+    analysis = glob.glob(os.path.join(HERE, "analysis", "code", "*.py"))
+    analysis += glob.glob(os.path.join(HERE, "analysis", "code", "shared", "*.py"))
+    for path in sorted(analysis):
+        name = os.path.basename(path)
+        if name.startswith("test_"):
+            continue
+        mods.append(path)
+    return mods
+
+
+def main():
+    for d in _import_dirs():
+        sys.path.insert(0, d)
+    pdoc.render.configure(template_directory=os.path.join(HERE, "pdoc_templates"))
+    pdoc.pdoc(*_modules(), output_directory=os.path.join(HERE, "api_docs"))
+    print("api docs written to", os.path.join(HERE, "api_docs", "index.html"))
+
+
+if __name__ == "__main__":
+    main()

conftest.py

@@ -0,0 +1,48 @@
+# The magnet/ model package needs the heavy PyTorch / PyTorch Geometric stack
+# (see magnet/requirements.txt), which the default dataset-reader test environment does not
+# install. When torch is absent, skip collecting the magnet tests entirely so the lightweight
+# CI stays green. They run locally once the torch stack from magnet/requirements.txt is present.
+#
+# This is needed (not just the importorskip guards inside test_magnet.py) because that test file
+# lives inside the magnet package, so pytest imports magnet/__init__.py during collection, which
+# imports torch before any in-test skip can run.
+import os
+import sys
+
+# The analysis code lives in analysis/code/ (importable modules) and analysis/code/shared/ (the
+# shared utilities). Put both on sys.path so tests anywhere resolve bare imports like `import
+# delta22` and `from stats import rmse`, no matter which folder the test file itself sits in.
+_HERE = os.path.dirname(os.path.abspath(__file__))
+for _p in (os.path.join(_HERE, "analysis", "code"),
+           os.path.join(_HERE, "analysis", "code", "shared")):
+    if _p not in sys.path:
+        sys.path.insert(0, _p)
+
+collect_ignore = []
+collect_ignore_glob = []
+try:
+    import torch  # noqa: F401
+except ImportError:
+    collect_ignore = ["magnet"]
+    collect_ignore_glob = ["magnet/*"]
+
+
+# The reproduce tests (test_reproduce.py) execute every notebook and take minutes; they are opt-in so
+# a plain `pytest` stays fast. Run them with `pytest --run-notebooks` (or `pytest -m reproduce`).
+def pytest_addoption(parser):
+    parser.addoption("--run-notebooks", action="store_true", default=False,
+                     help="run the reproduce tests that execute every analysis notebook")
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers", "reproduce: executes a full analysis notebook (slow, opt-in)")
+
+
+def pytest_collection_modifyitems(config, items):
+    import pytest
+    if config.getoption("--run-notebooks"):
+        return
+    skip = pytest.mark.skip(reason="notebook reproduce test; pass --run-notebooks to run")
+    for item in items:
+        if "reproduce" in item.keywords:
+            item.add_marker(skip)

data/applications/applications_reader.py

@@ -0,0 +1,260 @@
+"""Loader for the natural-products ("applications") release.
+
+applications.hdf5 stores shieldings and coordinates as whole numbers equal to the real value
+times 10,000 (scale 1e-4), with -2147483648 marking values that were never computed. Energies
+are stored as ordinary floating-point numbers. This module turns the stored numbers back into
+real values (leaving blanks where data is missing) and exposes one accessor per data layer, so
+notebooks never touch the raw file format. It is the single place where decoding happens.
+
+Solute names: vomicine, prednisone, peptide, flavone, dihydrotanshinone_I, and
+isomer_1E/1Z/2E/2Z/3E/3Z/4N/4O.
+Solvents: TIP4P (water), benzene, chloroform, methanol.
+"""
+import h5py
+import io
+import os
+
+import numpy as np
+import pandas as pd
+
+_SCALE = 1e4
+_MISSING_MARKER = -2147483648
+
+SOLVENTS = ["TIP4P", "benzene", "chloroform", "methanol"]
+# shell label = approximate number of heavy (non-hydrogen) atoms in the solute-plus-solvent cluster
+SHELL_SIZES = [50, 150, 250, 350, 450, 550, 650]
+
+
+def _decode(values):
+    """Turn stored whole numbers back into real values; blanks where data is missing."""
+    values = np.asarray(values)
+    if values.dtype == np.int32:
+        out = values.astype(np.float64) / _SCALE
+        out[values == _MISSING_MARKER] = np.nan
+        return out
+    return np.asarray(values, dtype=np.float64)
+
+
+class Applications:
+    """Read-only accessor for applications.hdf5 (+ the experimental spreadsheet)."""
+
+    def __init__(self, hdf5_path, experimental_xlsx=None, md_geometries_path=None):
+        self.hdf5_path = str(hdf5_path)
+        self.experimental_xlsx = str(experimental_xlsx) if experimental_xlsx else None
+        # the large MD-geometry companion is optional; default to the sibling file if present
+        if md_geometries_path is None:
+            sibling = self.hdf5_path.replace("applications.hdf5", "applications_md_geometries.hdf5")
+            md_geometries_path = sibling if (sibling != self.hdf5_path and os.path.exists(sibling)) else None
+        self.md_geometries_path = str(md_geometries_path) if md_geometries_path else None
+
+    def _open(self):
+        return h5py.File(self.hdf5_path, "r")
+
+    def _open_md(self):
+        if not self.md_geometries_path:
+            raise ValueError(
+                "MD-geometry companion not available. Download applications_md_geometries.hdf5 "
+                "next to applications.hdf5, or pass md_geometries_path=...")
+        return h5py.File(self.md_geometries_path, "r")
+
+    # ---- structure ----
+    _NON_SOLUTE_GROUPS = ("solvents", "composite_model")
+
+    def solutes(self):
+        """List the solute names stored in the release."""
+        with self._open() as f:
+            return [k for k in f.keys() if k not in self._NON_SOLUTE_GROUPS]
+
+    def atomic_numbers(self, solute):
+        """Return a solute's per-atom atomic numbers."""
+        with self._open() as f:
+            return np.asarray(f[solute]["atomic_numbers"][()])
+
+    # ---- MagNET-Zero (gas) + implicit-solvent (PCM) ----
+    def magnet_zero(self, solute):
+        """Boltzmann-weighted MagNET-Zero stationary shieldings and the PCM correction."""
+        with self._open() as f:
+            g = f[solute]["magnet_zero"]
+            return {
+                "stationary": _decode(g["stationary"][()]),
+                "pcm_correction": _decode(g["pcm_correction"][()]),
+            }
+
+    def magnet_zero_conformers(self, solute):
+        """Per-conformer dict: {conformer_name: {energy, geometry, magnet_zero, pcm}}."""
+        out = {}
+        with self._open() as f:
+            cg = f[solute]["magnet_zero"]["conformers"]
+            for ck in cg.keys():
+                c = cg[ck]
+                out[ck] = {
+                    "energy": float(c["energy"][()]),
+                    "geometry": _decode(c["geometry"][()]),
+                    "atomic_numbers": np.asarray(c["atomic_numbers"][()]),
+                    "magnet_zero": _decode(c["magnet_zero"][()]),
+                    "pcm": _decode(c["pcm"][()]),
+                }
+        return out
+
+    # ---- QCD (rovibrational) ----
+    def qcd(self, solute, shielding_only=False):
+        """qcd stationary (n_atoms,4)=[x,y,z,shielding] and trajectories (replica,frame,n_atoms,4).
+        With shielding_only=True returns just the shielding channel (column 3)."""
+        with self._open() as f:
+            g = f[solute]["qcd"]
+            stat = _decode(g["stationary"][()])
+            traj = _decode(g["trajectories"][()])
+        if shielding_only:
+            return {"stationary": stat[..., 3], "trajectories": traj[..., 3]}
+        return {"stationary": stat, "trajectories": traj}
+
+    # ---- MagNET-X (explicit solvent) ----
+    def magnet_x(self, solute, solvent, shell):
+        """Explicit-solvent shieldings (frames, n_atoms, 2)=[isolated, solvated] for one shell."""
+        with self._open() as f:
+            return _decode(f[solute]["magnet_x"][solvent][f"shell_{shell}"][()])
+
+    def magnet_x_stationary(self, solute, solvent):
+        """Explicit-solvent stationary-geometry shieldings for one solute/solvent pair."""
+        with self._open() as f:
+            return _decode(f[solute]["magnet_x"][solvent]["stationary"][()])
+
+    def available_shells(self, solute, solvent):
+        """List the shell sizes stored for one solute/solvent pair."""
+        with self._open() as f:
+            return sorted(
+                int(k.split("_")[1])
+                for k in f[solute]["magnet_x"][solvent].keys()
+                if k.startswith("shell_")
+            )
+
+    # ---- DFT reference (provided-only; PBE0(PBE1PBE)/pcSseg-3 GIAO, PCM per solvent) ----
+    def dft_reference(self, solute, solvent):
+        """Reference DFT shieldings for one solvent. Solvent keys here are 'gas', 'water',
+        'benzene', 'chloroform', 'methanol' -- note 'water' is the CPCM continuum-water model
+        (not the explicit 'TIP4P' key used by magnet_x), and 'gas' has no MagNET-x counterpart."""
+        with self._open() as f:
+            return _decode(f[solute]["dft_reference"][solvent]["shieldings"][()])
+
+    def dft_reference_geometries(self, solute):
+        """The DFT conformer geometries the reference shieldings were computed on, as a dict
+        {conformer_name: {coordinates (n_atoms, 3), atomic_numbers (n_atoms,), energy}}."""
+        out = {}
+        with self._open() as f:
+            g = f[solute]["dft_reference"].get("geometries")
+            if g is None:
+                return out
+            for ck in g.keys():
+                c = g[ck]
+                out[ck] = {
+                    "coordinates": _decode(c["coordinates"][()]),
+                    "atomic_numbers": np.asarray(c["atomic_numbers"][()]),
+                    "energy": float(c["energy"][()]),
+                }
+        return out
+
+    # ---- composite-model coefficients (fit on delta-22, applied to the natural products) ----
+    def _composite_csv(self, *path):
+        key = "composite_model/" + "/".join(path)
+        with self._open() as f:
+            if key not in f:
+                raise KeyError(f"{key} not in {self.hdf5_path}")
+            txt = f[key][()]
+        if isinstance(txt, bytes):
+            txt = txt.decode("utf-8")
+        return pd.read_csv(io.StringIO(txt))
+
+    def ols_coefficients(self, formula, nucleus):
+        """OLS composite-model coefficients (parameter rows x solvent columns) for one formula/nucleus."""
+        return self._composite_csv("ols_coefficients", formula, nucleus)
+
+    def bootstrap_coefficients(self, formula, nucleus):
+        """Bootstrap-resampled composite-model coefficients (one row per solvent x seed)."""
+        return self._composite_csv("bootstrap_coefficients", formula, nucleus)
+
+    def rmse_distribution(self, nucleus):
+        """delta-22 bootstrap RMSE distributions (solvent, nucleus, formula, seed, Bootstrap_RMSE)."""
+        return self._composite_csv("rmse_distributions", nucleus)
+
+    def pcm_conversion_factors(self, nucleus):
+        """Per-solvent scale factors applied to the chloroform-only PCM correction."""
+        return self._composite_csv("pcm_conversion_factors", nucleus)
+
+    def composite_formulas(self, kind="ols_coefficients"):
+        """List the formula keys available under a composite_model subgroup."""
+        with self._open() as f:
+            g = f.get(f"composite_model/{kind}")
+            return sorted(g.keys()) if g is not None else []
+
+    # ---- canonical solvent atom ordering ----
+    def solvent_atomic_numbers(self, solvent):
+        """Return one solvent molecule's atomic numbers, in the canonical atom order."""
+        with self._open() as f:
+            return np.asarray(f["solvents"][solvent]["atomic_numbers"][()])
+
+    # ---- MD solvent-shell geometries (the large companion file) ----
+    def available_md_shells(self, solute, solvent):
+        """List the shell sizes with MD geometries for one solute/solvent pair."""
+        with self._open_md() as f:
+            return sorted(int(k.split("_")[1]) for k in f[solute][solvent].keys()
+                          if k.startswith("geometries_"))
+
+    def md_geometry(self, solute, solvent, shell):
+        """Solvent-shell snapshot coordinates (frames, n_shell_atoms, 3) for one shell. The frame
+        axis aligns 1:1 with magnet_x(solute, solvent, shell)."""
+        with self._open_md() as f:
+            return _decode(f[solute][solvent][f"geometries_{shell}"][()])
+
+    def md_geometry_atomic_numbers(self, solute, solvent, shell):
+        """Atomic numbers for one MD solvent-shell snapshot's atoms."""
+        with self._open_md() as f:
+            return np.asarray(f[solute][solvent][f"atomic_numbers_{shell}"][()])
+
+    def md_stationary_geometry(self, solute):
+        """The solute's stationary-geometry coordinates from the OpenMM setup."""
+        with self._open_md() as f:
+            return _decode(f[solute]["openMM_stationary_geometry"][()])
+
+    # ---- experiment ----
+    @staticmethod
+    def _prefix_h_sites(df):
+        """Give each proton site a unique label per solute. The spreadsheet labels every H
+        site simply "H"; this prepends an incrementing two-digit number in row order per
+        solute ("01 H", "02 H", ...), matching the figure notebooks so per-site H rows are
+        distinguishable. Carbon sites already carry distinct labels and are left untouched."""
+        df = df.copy()
+        for solute in df["solute"].unique():
+            n = 1
+            mask = (df["solute"] == solute) & (df["nucleus"] == "H")
+            for idx in df[mask].index:
+                df.at[idx, "site"] = f"{n:02d} {df.at[idx, 'site']}"
+                n += 1
+        return df
+
+    def experiment(self, rename_water_to_tip4p=True, uuid_prefix_h_sites=False):
+        """Flat long table from the per-solute-sheet xlsx: one row per (solute, site).
+        With uuid_prefix_h_sites=True, proton sites are made unique per solute (see
+        _prefix_h_sites) as the figure notebooks require."""
+        if not self.experimental_xlsx:
+            raise ValueError("experimental_xlsx path not provided")
+        sheets = pd.read_excel(self.experimental_xlsx, sheet_name=None)
+        frames = []
+        for solute, df in sheets.items():
+            df = df.copy()
+            df.insert(0, "solute", solute)
+            frames.append(df)
+        out = pd.concat(frames, ignore_index=True)
+        if rename_water_to_tip4p and "water" in out.columns:
+            out = out.rename(columns={"water": "TIP4P"})
+        if uuid_prefix_h_sites:
+            out = self._prefix_h_sites(out)
+        return out
+
+    def site_atom_table(self, uuid_prefix_h_sites=False):
+        """Mapping of each (solute, nucleus, site) to its atom_numbers, used to aggregate
+        per-atom shieldings up to NMR sites. `atom_numbers` is returned as the raw string from
+        the spreadsheet (e.g. "17,19" or "21"); callers split/offset it themselves (atoms are
+        1-based in the spreadsheet, so subtract 1 to index the shielding arrays)."""
+        exp = self.experiment(rename_water_to_tip4p=False, uuid_prefix_h_sites=uuid_prefix_h_sites)
+        cols = ["solute", "site", "nucleus", "atom_numbers"]
+        return exp[cols].drop_duplicates().reset_index(drop=True)

data/applications/test_applications.py

@@ -0,0 +1,211 @@
+"""Tests for the applications (natural-products) loader.
+
+These build a tiny synthetic applications.hdf5 + experimental spreadsheet in a temp dir that
+mimic the real schema, so the loader and its decoding can be checked in CI without the multi-
+hundred-megabyte release files.
+"""
+import numpy as np
+import pandas as pd
+import h5py
+import pytest
+
+from applications_reader import Applications, _decode, SHELL_SIZES, SOLVENTS
+
+SCALE = 1e4
+MISS = np.int32(-2147483648)
+
+
+def _enc(a):
+    a = np.asarray(a, dtype=np.float64)
+    nan = np.isnan(a)
+    out = np.where(nan, 0.0, a)
+    out = np.round(out * SCALE).astype(np.int32)
+    out[nan] = MISS
+    return out
+
+
+@pytest.fixture
+def fake_release(tmp_path):
+    """A 2-solute synthetic release: vomicine (with a missing shell) + isomer_1E."""
+    h5 = tmp_path / "applications.hdf5"
+    xlsx = tmp_path / "applications_experimental.xlsx"
+    rng = np.random.default_rng(0)
+
+    def mk_solute(g, n_atoms, atomic_numbers, shells):
+        g.create_dataset("atomic_numbers", data=np.asarray(atomic_numbers, dtype=np.int32))
+        mz = g.create_group("magnet_zero")
+        # MagNET predicts H/C only -> heteroatoms read 0 (mirror the real data)
+        hetero = ~np.isin(np.asarray(atomic_numbers), [1, 6])
+        stat = rng.normal(100, 30, n_atoms); stat[hetero] = 0.0
+        pcm = rng.normal(0, 2, n_atoms); pcm[hetero] = 0.0
+        mz.create_dataset("stationary", data=_enc(stat))
+        mz.create_dataset("pcm_correction", data=_enc(pcm))
+        cf = mz.create_group("conformers").create_group("conformer_0")
+        cf.create_dataset("energy", data=np.float64(-1234.5))   # float, not encoded
+        cf.create_dataset("geometry", data=_enc(rng.normal(0, 5, (n_atoms, 3))))
+        cf.create_dataset("atomic_numbers", data=np.asarray(atomic_numbers, dtype=np.int32))
+        cf.create_dataset("magnet_zero", data=_enc(rng.normal(100, 30, n_atoms)))
+        cf.create_dataset("pcm", data=_enc(rng.normal(0, 2, n_atoms)))
+        q = g.create_group("qcd")
+        q.create_dataset("stationary", data=_enc(rng.normal(0, 50, (n_atoms, 4))))
+        q.create_dataset("trajectories", data=_enc(rng.normal(0, 50, (3, 4, n_atoms, 4))))
+        mx = g.create_group("magnet_x")
+        for sv in SOLVENTS:
+            svg = mx.create_group(sv)
+            for sh in shells:
+                svg.create_dataset(f"shell_{sh}", data=_enc(rng.normal(100, 30, (5, n_atoms, 2))))
+            svg.create_dataset("stationary", data=_enc(rng.normal(100, 30, n_atoms)))
+        dr = g.create_group("dft_reference")
+        # dft_reference uses gas + CPCM-continuum 'water' (not the explicit 'TIP4P'), like the real data
+        for sv in ["gas", "water", "benzene", "chloroform", "methanol"]:
+            dr.create_group(sv).create_dataset("shieldings", data=_enc(rng.normal(100, 40, n_atoms)))
+        gg = dr.create_group("geometries").create_group("conformer_001")
+        gg.create_dataset("coordinates", data=_enc(rng.normal(0, 5, (n_atoms, 3))))
+        gg.create_dataset("atomic_numbers", data=np.asarray(atomic_numbers, dtype=np.int32))
+        gg.create_dataset("energy", data=np.float64(-100.0))
+
+    with h5py.File(h5, "w") as f:
+        sg = f.create_group("solvents")
+        for sv in SOLVENTS:
+            sg.create_group(sv).create_dataset("atomic_numbers", data=np.array([6, 1, 1], dtype=np.int32))
+        # vomicine: drop one shell (mimic the isomer_1Z/chloroform/250 corruption)
+        mk_solute(f.create_group("vomicine"), 4, [8, 7, 6, 1], [s for s in SHELL_SIZES if s != 250])
+        mk_solute(f.create_group("isomer_1E"), 3, [6, 1, 1], SHELL_SIZES)
+        # composite_model group: fit tables stored verbatim as CSV text
+        cm = f.create_group("composite_model")
+        cm.create_group("pcm_conversion_factors").create_dataset(
+            "H", data="solvent,pcm_conversion_factor\nchloroform,1.0\nbenzene,0.6\n")
+        cm.create_group("ols_coefficients").create_group("stationary_plus_pcm").create_dataset(
+            "H", data="parameter,chloroform,benzene\nintercept,31.2,31.3\nstationary,-0.98,-0.99\n")
+        cm.create_group("bootstrap_coefficients").create_group("stationary_plus_pcm").create_dataset(
+            "H", data="solvent,seed,Intercept,stationary,pcm\nchloroform,0,31.2,-0.98,-1.0\nbenzene,0,31.3,-0.99,-0.6\n")
+        cm.create_group("rmse_distributions").create_dataset(
+            "H", data="solvent,nucleus,formula,seed,Bootstrap_RMSE\nchloroform,H,pcm2,0,0.07\nbenzene,H,pcm2,0,0.21\n")
+
+    # experiment xlsx: one sheet per solute
+    with pd.ExcelWriter(xlsx) as w:
+        for sol in ["vomicine", "isomer_1E"]:
+            pd.DataFrame({
+                "site": ["a_C", "b_H"], "nucleus": ["C", "H"], "atom_numbers": ["3", "4"],
+                "chloroform": [1.0, 2.0], "benzene": [1.1, 2.1],
+                "methanol": [1.2, 2.2], "water": [1.3, 2.3],
+            }).to_excel(w, sheet_name=sol, index=False)
+    return Applications(h5, xlsx)
+
+
+def test_decode_roundtrip_and_missing():
+    vals = np.array([1.2345, -50.0, 0.0], dtype=np.float64)
+    enc = _enc(vals)
+    out = _decode(enc)
+    assert np.allclose(out, vals, atol=5e-5)
+    enc_with_miss = np.append(enc, MISS).astype(np.int32)
+    assert np.isnan(_decode(enc_with_miss)[-1])
+
+
+def test_solutes_and_solvents_group_excluded(fake_release):
+    assert set(fake_release.solutes()) == {"vomicine", "isomer_1E"}
+
+
+def test_magnet_zero_heteroatoms_read_zero(fake_release):
+    # vomicine atomic_numbers [8, 7, 6, 1]: the O and N positions must decode to 0 (MagNET does H/C only)
+    mz = fake_release.magnet_zero("vomicine")
+    hetero = ~np.isin(fake_release.atomic_numbers("vomicine"), [1, 6])
+    assert np.allclose(mz["stationary"][hetero], 0.0)
+    assert np.allclose(mz["pcm_correction"][hetero], 0.0)
+    assert np.any(mz["stationary"][~hetero] != 0.0)  # H/C carry real values
+
+
+def test_qcd_shielding_only_drops_xyz(fake_release):
+    full = fake_release.qcd("isomer_1E")
+    just = fake_release.qcd("isomer_1E", shielding_only=True)
+    assert full["stationary"].shape == (3, 4)
+    assert just["stationary"].shape == (3,)
+    assert just["trajectories"].shape == (3, 4, 3)
+    assert np.allclose(just["stationary"], full["stationary"][..., 3], atol=5e-5)
+
+
+def test_magnet_x_and_shell_guard(fake_release):
+    mx = fake_release.magnet_x("isomer_1E", "benzene", 650)
+    assert mx.shape == (5, 3, 2)
+    # vomicine is missing shell 250 -> available_shells must exclude it, not crash
+    assert 250 not in fake_release.available_shells("vomicine", "chloroform")
+    assert fake_release.available_shells("isomer_1E", "chloroform") == SHELL_SIZES
+
+
+def test_dft_reference_and_conformers(fake_release):
+    assert fake_release.dft_reference("vomicine", "gas").shape == (4,)
+    conf = fake_release.magnet_zero_conformers("vomicine")
+    assert conf["conformer_0"]["energy"] == pytest.approx(-1234.5)
+    assert conf["conformer_0"]["geometry"].shape == (4, 3)
+
+
+def test_experiment_water_rename(fake_release):
+    df = fake_release.experiment()
+    assert "TIP4P" in df.columns and "water" not in df.columns
+    assert set(df["solute"].unique()) == {"vomicine", "isomer_1E"}
+
+
+def test_site_atom_table(fake_release):
+    tbl = fake_release.site_atom_table()
+    assert list(tbl.columns) == ["solute", "site", "nucleus", "atom_numbers"]
+    assert set(tbl["solute"].unique()) == {"vomicine", "isomer_1E"}
+    assert set(str(v) for v in tbl["atom_numbers"]) == {"3", "4"}
+
+
+def test_composite_model_accessors(fake_release):
+    conv = fake_release.pcm_conversion_factors("H")
+    assert list(conv.columns) == ["solvent", "pcm_conversion_factor"]
+    assert conv.set_index("solvent").loc["chloroform", "pcm_conversion_factor"] == 1.0
+    ols = fake_release.ols_coefficients("stationary_plus_pcm", "H")
+    assert "parameter" in ols.columns and "chloroform" in ols.columns
+    assert fake_release.composite_formulas("ols_coefficients") == ["stationary_plus_pcm"]
+    boot = fake_release.bootstrap_coefficients("stationary_plus_pcm", "H")
+    assert {"solvent", "seed", "Intercept", "stationary", "pcm"}.issubset(boot.columns)
+    rd = fake_release.rmse_distribution("H")
+    assert {"solvent", "formula", "Bootstrap_RMSE"}.issubset(rd.columns)
+
+
+def test_dft_reference_uses_water_not_tip4p(fake_release):
+    # the DFT reference is CPCM-continuum 'water'; the explicit 'TIP4P' key does not exist there
+    assert fake_release.dft_reference("vomicine", "water").shape == (4,)
+    with pytest.raises(KeyError):
+        fake_release.dft_reference("vomicine", "TIP4P")
+
+
+def test_dft_reference_geometries(fake_release):
+    geo = fake_release.dft_reference_geometries("vomicine")
+    assert "conformer_001" in geo
+    c = geo["conformer_001"]
+    assert c["coordinates"].shape == (4, 3)
+    assert c["energy"] == pytest.approx(-100.0)
+    assert list(c["atomic_numbers"]) == [8, 7, 6, 1]
+
+
+def test_md_geometry_companion(tmp_path):
+    md = tmp_path / "applications_md_geometries.hdf5"
+    with h5py.File(md, "w") as f:
+        g = f.create_group("vomicine")
+        g.create_dataset("openMM_stationary_geometry", data=_enc(np.zeros((4, 3))))
+        sv = g.create_group("chloroform")
+        sv.create_dataset("geometries_50", data=_enc(np.ones((5, 7, 3))))
+        sv.create_dataset("atomic_numbers_50", data=np.arange(7, dtype=np.int32))
+    app = Applications(tmp_path / "applications.hdf5", md_geometries_path=md)
+    assert app.available_md_shells("vomicine", "chloroform") == [50]
+    assert app.md_geometry("vomicine", "chloroform", 50).shape == (5, 7, 3)
+    assert app.md_geometry_atomic_numbers("vomicine", "chloroform", 50).shape == (7,)
+    assert app.md_stationary_geometry("vomicine").shape == (4, 3)
+
+
+def test_md_geometry_absent_raises_clearly(fake_release):
+    # the big companion is optional; accessing it without it gives a clear, actionable error
+    with pytest.raises(ValueError, match="companion"):
+        fake_release.available_md_shells("vomicine", "chloroform")
+
+
+def test_uuid_prefix_h_sites(fake_release):
+    # each solute's single H site (fixture site "b_H") becomes "01 b_H"; C sites untouched
+    tbl = fake_release.site_atom_table(uuid_prefix_h_sites=True)
+    h_sites = tbl[tbl["nucleus"] == "H"]["site"].tolist()
+    assert all(s.startswith("01 ") for s in h_sites)
+    c_sites = tbl[tbl["nucleus"] == "C"]["site"].tolist()
+    assert all(not s[:2].isdigit() for s in c_sites)

data/carbon_monoxide/co_bond_lengths_histogram.csv

@@ -0,0 +1,101 @@
+bin_left_angstrom,count
+0.6975547075,1
+0.7077355623,0
+0.7179164171,0
+0.7280972719,0
+0.7382781267,0
+0.7484589815,0
+0.7586398363,0
+0.7688206911,0
+0.7790015459,0
+0.7891824007,0
+0.7993632555,0
+0.8095441103,0
+0.8197249651,0
+0.8299058199,0
+0.8400866747,0
+0.8502675295,0
+0.8604483843,0
+0.8706292391,0
+0.8808100939,0
+0.8909909487,0
+0.9011718035,0
+0.9113526583,0
+0.9215335131,0
+0.9317143679,0
+0.9418952227,1
+0.9520760775,0
+0.9622569323,0
+0.9724377871,0
+0.9826186419,0
+0.9927994967,0
+1.002980351,0
+1.013161206,1
+1.023342061,0
+1.033522916,1
+1.043703771,0
+1.053884625,1
+1.06406548,0
+1.074246335,0
+1.08442719,5
+1.094608045,8
+1.104788899,2
+1.114969754,10
+1.125150609,34
+1.135331464,151
+1.145512319,768
+1.155693173,2689
+1.165874028,5616
+1.176054883,8403
+1.186235738,29687
+1.196416593,164521
+1.206597447,512713
+1.216778302,313754
+1.226959157,59816
+1.237140012,13908
+1.247320867,8563
+1.257501721,5611
+1.267682576,3237
+1.277863431,1824
+1.288044286,1890
+1.298225141,2919
+1.308405995,5201
+1.31858685,12030
+1.328767705,37567
+1.33894856,101935
+1.349129415,186377
+1.359310269,257009
+1.369491124,222621
+1.379671979,170158
+1.389852834,196816
+1.400033689,323785
+1.410214543,689662
+1.420395398,1028853
+1.430576253,790563
+1.440757108,503514
+1.450937963,247803
+1.461118817,124073
+1.471299672,59557
+1.481480527,28576
+1.491661382,15571
+1.501842237,9137
+1.512023091,5441
+1.522203946,3237
+1.532384801,1882
+1.542565656,1043
+1.552746511,552
+1.562927365,323
+1.57310822,170
+1.583289075,105
+1.59346993,45
+1.603650784,27
+1.613831639,20
+1.624012494,16
+1.634193349,1
+1.644374204,6
+1.654555058,5
+1.664735913,3
+1.674916768,3
+1.685097623,1
+1.695278478,4
+1.705459332,2

data/carbon_monoxide/co_magnet_vs_gaussian.csv

@@ -0,0 +1,82 @@
+bond_length_angstrom,bond_length_delta,gaussian_shielding_c_ppm,magnet_shielding_c_ppm,delta_c_ppm,gaussian_shielding_o_ppm,magnet_shielding_o_ppm,delta_o_ppm
+0.7000000476837158,0.0,100.323,-98.84709167480469,-199.17009167480467,118.92,0.0,-118.92
+0.7100000381469727,0.0,100.6646,-101.11131286621094,-201.77591286621094,119.3498,0.0,-119.3498
+0.7200000286102295,0.0,100.7627,-106.29923248291016,-207.06193248291015,119.4879,0.0,-119.4879
+0.7300000190734863,0.0,100.6302,-110.37581634521484,-211.00601634521485,119.3424,0.0,-119.3424
+0.7400000095367432,0.0,100.2785,-105.22541809082031,-205.50391809082032,118.9203,0.0,-118.9203
+0.75,0.0,99.7179,-82.99818420410156,-182.71608420410155,118.2274,0.0,-118.2274
+0.7599999904632568,0.0,98.9574,-58.89897537231445,-157.85637537231446,117.2682,0.0,-117.2682
+0.7699999809265137,0.0,98.0051,-41.90826416015625,-139.91336416015625,116.0464,0.0,-116.0464
+0.7799999713897705,0.0,96.8686,-30.977083206176758,-127.84568320617676,114.5651,0.0,-114.5651
+0.7900000214576721,0.0,95.5546,-24.81533432006836,-120.36993432006835,112.8267,0.0,-112.8267
+0.800000011920929,0.0,94.0699,-20.98507308959961,-115.05497308959961,110.8334,0.0,-110.8334
+0.8100000023841858,0.0,92.4205,-17.495792388916016,-109.91629238891602,108.5869,0.0,-108.5869
+0.8199999928474426,0.0,90.6124,-13.717061042785645,-104.32946104278564,106.089,0.0,-106.089
+0.8299999833106995,0.0,88.6517,-10.725614547729492,-99.3773145477295,103.3413,0.0,-103.3413
+0.8400000333786011,0.0,86.5441,-9.816015243530273,-96.36011524353027,100.3454,0.0,-100.3454
+0.8500000238418579,0.0,84.2955,-10.453042030334473,-94.74854203033448,97.103,0.0,-97.103
+0.8600000143051147,0.0,81.9118,-11.056601524353027,-92.96840152435303,93.6158,0.0,-93.6158
+0.8700000047683716,0.0,79.3986,-10.437362670898438,-89.83596267089844,89.8856,0.0,-89.8856
+0.8799999952316284,0.0,76.7615,-8.59598159790039,-85.35748159790039,85.9141,0.0,-85.9141
+0.8899999856948853,0.0,74.0062,-6.2477288246154785,-80.25392882461549,81.7031,0.0,-81.7031
+0.8999999761581421,0.0,71.1378,-3.5964534282684326,-74.73425342826843,77.2543,0.0,-77.2543
+0.9099999666213989,0.0,68.1616,-0.3521101176738739,-68.51371011767388,72.5691,0.0,-72.5691
+0.9199999570846558,0.0,65.0822,3.5139496326446533,-61.56825036735535,67.649,0.0,-67.649
+0.9300000071525574,0.0,61.9043,7.585774898529053,-54.31852510147095,62.495,0.0,-62.495
+0.9399999976158142,0.0,58.632,11.563019752502441,-47.06898024749756,57.1079,0.0,-57.1079
+0.949999988079071,0.0,55.269,15.541687965393066,-39.72731203460693,51.4883,0.0,-51.4883
+0.9599999785423279,0.0,51.8188,19.621856689453125,-32.19694331054688,45.6364,0.0,-45.6364
+0.9700000286102295,0.0,48.2845,23.571359634399414,-24.713140365600587,39.552,0.0,-39.552
+0.9800000190734863,0.0,44.6688,26.864168167114258,-17.80463183288574,33.2348,0.0,-33.2348
+0.9900000095367432,0.0,40.9739,29.0251407623291,-11.948759237670899,26.6838,0.0,-26.6838
+1.0,0.0,37.2017,30.12110710144043,-7.080592898559573,19.8976,0.0,-19.8976
+1.0099999904632568,0.0,33.3541,30.643943786621094,-2.7101562133789088,12.8748,0.0,-12.8748
+1.0199999809265137,0.0,29.4321,30.884471893310547,1.4523718933105485,5.6134,0.0,-5.6134
+1.0299999713897705,0.0,25.437,30.741764068603516,5.3047640686035145,-1.8888,0.0,1.8888
+1.0399999618530273,0.0,21.3696,30.09064292907715,8.72104292907715,-9.6345,0.0,9.6345
+1.0499999523162842,0.0,17.2304,29.05397605895996,11.823576058959961,-17.6265,0.0,17.6265
+1.0600000619888306,0.0,13.0196,27.87270736694336,14.853107366943359,-25.8681,0.0,25.8681
+1.0699999332427979,0.0,8.7375,26.578834533691406,17.841334533691406,-34.3628,0.0,34.3628
+1.0799999237060547,0.0,4.384,24.647645950317383,20.263645950317382,-43.1141,0.0,43.1141
+1.0899999141693115,0.0,-0.0411,21.391674041748047,21.432774041748047,-52.1261,0.0,52.1261
+1.0999999046325684,0.0,-4.5382,16.754995346069336,21.293195346069336,-61.403,0.0,61.403
+1.1100000143051147,0.0,-9.1078,11.06360149383545,20.171401493835447,-70.9492,0.0,70.9492
+1.1200000047683716,0.0,-13.7508,4.679224491119385,18.430024491119383,-80.7696,0.0,80.7696
+1.1299999952316284,0.0,-18.4679,-2.118180990219116,16.349719009780884,-90.8691,0.0,90.8691
+1.1399999856948853,0.0,-23.2601,-9.029415130615234,14.230684869384767,-101.253,0.0,101.253
+1.149999976158142,0.0,-28.1288,-15.528891563415527,12.599908436584471,-111.9268,0.0,111.9268
+1.159999966621399,0.0,-33.0753,-21.46068000793457,11.614619992065428,-122.8962,0.0,122.8962
+1.1700000762939453,0.0,-38.101,-27.291627883911133,10.809372116088866,-134.1672,0.0,134.1672
+1.1800000667572021,0.0,-43.2076,-33.76377868652344,9.443821313476562,-145.746,0.0,145.746
+1.190000057220459,0.0,-48.3971,-41.44232940673828,6.9547705932617205,-157.6393,0.0,157.6393
+1.2000000476837158,0.0,-53.6713,-50.427207946777344,3.2440920532226585,-169.8537,0.0,169.8537
+1.2100000381469727,0.0,-59.0325,-60.3585205078125,-1.3260205078125011,-182.3964,0.0,182.3964
+1.2200000286102295,0.0,-64.4831,-71.33944702148438,-6.856347021484382,-195.2749,0.0,195.2749
+1.2300000190734863,0.0,-70.0258,-83.5546875,-13.528887499999996,-208.4969,0.0,208.4969
+1.2400000095367432,0.0,-75.6632,-96.37044525146484,-20.70724525146484,-222.0705,0.0,222.0705
+1.25,0.0,-81.3984,-109.45909118652344,-28.060691186523442,-236.0042,0.0,236.0042
+1.2599999904632568,0.0,-87.2346,-123.06512451171875,-35.83052451171875,-250.3065,0.0,250.3065
+1.2699999809265137,0.0,-93.1752,-137.39169311523438,-44.21649311523437,-264.9868,0.0,264.9868
+1.2799999713897705,0.0,-99.2239,-152.3052215576172,-53.08132155761719,-280.0545,0.0,280.0545
+1.2899999618530273,0.0,-105.3845,-168.0139617919922,-62.629461791992185,-295.5198,0.0,295.5198
+1.2999999523162842,0.0,-111.6611,-185.2198944091797,-73.55879440917968,-311.3932,0.0,311.3932
+1.309999942779541,0.0,-118.0581,-204.2111053466797,-86.15300534667969,-327.6854,0.0,327.6854
+1.3199999332427979,0.0,-124.5801,-224.82327270507812,-100.24317270507812,-344.4081,0.0,344.4081
+1.3299999237060547,0.0,-131.2321,-247.2126007080078,-115.98050070800781,-361.5731,0.0,361.5731
+1.3399999141693115,0.0,-138.019,-271.37847900390625,-133.35947900390624,-379.1926,0.0,379.1926
+1.350000023841858,0.0,-144.9464,-297.06243896484375,-152.11603896484374,-397.2797,0.0,397.2797
+1.3600000143051147,0.0,-152.0198,-324.2068176269531,-172.18701762695312,-415.8481,0.0,415.8481
+1.3700000047683716,0.0,-159.2452,-352.2054748535156,-192.9602748535156,-434.9119,0.0,434.9119
+1.3799999952316284,0.0,-166.6289,-380.4760437011719,-213.8471437011719,-454.4861,0.0,454.4861
+1.3899999856948853,0.0,-174.1774,-408.8212585449219,-234.64385854492187,-474.5864,0.0,474.5864
+1.4000000953674316,0.0,-181.8975,-436.86309814453125,-254.96559814453124,-495.2291,0.0,495.2291
+1.4100000858306885,0.0,-189.7964,-463.40863037109375,-273.61223037109374,-516.4312,0.0,516.4312
+1.4200000762939453,0.0,-197.8817,-486.49700927734375,-288.6153092773437,-538.2106,0.0,538.2106
+1.4300000667572021,0.0,-206.1613,-504.5982360839844,-298.4369360839844,-560.5865,0.0,560.5865
+1.440000057220459,0.0,-214.6433,-517.5965576171875,-302.9532576171875,-583.5779,0.0,583.5779
+1.4500000476837158,0.0,-223.3363,-526.4462280273438,-303.10992802734376,-607.2056,0.0,607.2056
+1.4600000381469727,0.0,-232.2493,-532.2877807617188,-300.03848076171874,-631.4914,0.0,631.4914
+1.4700000286102295,0.0,-241.3916,-535.899169921875,-294.50756992187496,-656.4581,0.0,656.4581
+1.4800000190734863,0.0,-250.7732,-537.3216552734375,-286.5484552734375,-682.1295,0.0,682.1295
+1.4900000095367432,0.0,-260.4042,-535.8610229492188,-275.45682294921875,-708.5308,0.0,708.5308
+1.5,0.0,-270.2954,-530.5790405273438,-260.2836405273438,-735.6883,0.0,735.6883

data/delta22/delta22_reader.py

@@ -0,0 +1,1047 @@
+"""
+Data loading functions for delta22 DFT and NN data.
+
+Loads and processes the delta22 HDF5 file (DFT and MagNET shieldings, timings, experimental
+shifts) into the flat tables the analysis consumes. Read through analysis/code/delta22.py and
+the delta-22 figure notebooks (fig3, fig4, si_figure_s08_panelsBCD, ...).
+"""
+
+import csv
+import os
+import pandas as pd
+import numpy as np
+import h5py
+from tqdm import tqdm   # plain tqdm, not tqdm.auto: no ipywidgets/IProgress warning headless or in a notebook
+
+
+GEOMETRY_TYPES = ["aimnet2", "pbe0_tz"]
+IMPLICIT_SOLVATION_MODELS = ["pcm", "smd"]
+EXPLICIT_SOLVATION_MODELS = ["desmond", "openMM"]
+
+DEFAULT_PCM_REFERENCE_METHOD = "b3lyp_d3bj"
+DEFAULT_PCM_REFERENCE_BASIS = "pcSseg2"
+
+INDEX_COLUMNS = [
+    "solute",
+    "sap_geometry_type",
+    "sap_nmr_method",
+    "sap_basis",
+    "nucleus",
+    "site",
+    "solvent"
+]
+
+# This dataset stores shieldings and coordinates as whole numbers equal to the real
+# value times 10,000, with a reserved marker for values that were never computed.
+# These helpers turn them back into real values, leaving blanks where data is missing.
+_FIXED_POINT_SCALE = 1e4
+_MISSING_MARKER = -2147483648
+
+def _decode_fixed_point(values):
+    values = np.asarray(values)
+    if values.dtype == np.int32:
+        out = values.astype(np.float64) / _FIXED_POINT_SCALE
+        out[values == _MISSING_MARKER] = np.nan
+        return out
+    return np.asarray(values, dtype=np.float64)
+
+
+
+def _parse_method_name(s: str) -> tuple[str, str, str, str]:
+    """
+    Helper for parsing tuples with method details into a method name, basis, 
+    solvent model, and solvent name.
+    """
+    method, basis, solvent_model, solvent_name = next(csv.reader([s]))
+    return method, basis, solvent_model, solvent_name
+
+
+def _sort_and_parse_raw_shieldings(shield_df):
+    """
+    Helper for sorting the raw shieldings DataFrame to put 'none' solvents first
+    and for parsing the shielding lists into true lists from strings.
+    """
+    def _custom_sort(value):
+        """Helper function for sorting solvents with 'none' first."""
+        if pd.isna(value) or str(value).lower() == 'none':
+            return ('', str(value))
+        return (str(value).lower(), str(value))
+
+
+    def _parse_to_list(x):
+        """Parse shielding arrays from string format to float lists."""
+        value_list = x.strip("[]").replace('\n','').split(" ")
+        value_list = [v for v in value_list if v]
+        try:
+            return [float(i) for i in value_list]
+        except:
+            return []
+
+    def _to_float_list(x):
+        """Return the shieldings as a plain list of floats. Values loaded from the HDF5 are already
+        numeric arrays and are converted directly; the string form is only parsed when a caller hands
+        one in (e.g. a live-inference override), since converting an array through str() truncates to
+        print precision (~8 digits) instead of keeping full float64 precision."""
+        if isinstance(x, str):
+            return _parse_to_list(x)
+        return np.asarray(x, dtype=float).tolist()
+
+    sort_key = lambda x: x.map(_custom_sort) if x.name == "solvent" else x
+    shield_df.sort_values(["solute","solvent"],
+                                  key=sort_key,
+                                  inplace=True)
+    shield_df["shieldings"] = shield_df["shieldings"].apply(_to_float_list)
+    return shield_df
+
+
+def _fix_atom_numbers(atom_numbers, prefix="atom_", coerce_to=str, offset=0):
+    """Utility function for converting strings and ints to a common list of strings format."""
+    if isinstance(atom_numbers, int):
+        atom_numbers = [f"{prefix}{int(atom_numbers)+offset}"]
+    else:
+        atom_numbers = [f"{prefix}{int(i)+offset}" for i in atom_numbers.split(",")]
+    atom_numbers = [coerce_to(i) for i in atom_numbers]
+    return atom_numbers
+
+
+def _aggregate_shieldings_by_site(shieldings_data_df:pd.DataFrame, 
+                                  site_atoms_df:pd.DataFrame, 
+                                  verbose:bool=True):
+    """Aggregate DFT/NN shieldings by site."""
+
+    site_solutes = site_atoms_df.reset_index().solute.unique()
+
+    if verbose:
+        print("\033[93m", end="")
+        solute_iter = tqdm(site_solutes, 
+                           desc="Aggregating solute shieldings by site")
+    else:
+        solute_iter = site_solutes
+
+    rows = []
+    for solute in solute_iter:
+        sites_df = site_atoms_df.query(f"solute == '{solute}'")
+        data_df = shieldings_data_df.query(f"solute == '{solute}'").copy()
+
+        # Building the frame from the stacked arrays (every row of this solute has the same number of
+        # atoms) is far faster than .apply(pd.Series), which allocates a Series per row.
+        shieldings_columns = pd.DataFrame(np.stack(data_df.shieldings.to_numpy()),
+                                          index=data_df.index)
+        shieldings_columns.columns = [f"atom_{i}" for i in range(1, len(shieldings_columns.columns)+1)]
+        data_df = data_df.iloc[:,:-1]
+        data_df = data_df.join(shieldings_columns)
+        data_df.reset_index(inplace=True, drop=True)  # otherwise the index over the for loop won't work
+
+        # Metadata columns (geometry type, NMR method, basis, solvation model, solvent) are the same
+        # for every site of this solute; pulled out once as a plain object array so the per-site loop
+        # does not rebuild a pandas Series per row.
+        meta = data_df.iloc[:, 1:6].to_numpy(dtype=object)
+        # the first row is always the stationary (unsolvated) one
+        assert data_df.iloc[0, 5] == "none", \
+            f"Invalid first row: {data_df.iloc[0, 2]} for solute {solute}"
+
+        for (_, nucl, site), site_row in sites_df.iterrows():
+            atom_numbers = _fix_atom_numbers(site_row.iloc[0])
+            mean = data_df[atom_numbers].mean(axis=1).to_numpy()
+            for i in range(len(meta)):
+                geo_type, nmr_method, basis, implicit_solvation_model, solvent = meta[i]
+                rows.append([solute, geo_type, nmr_method, basis, implicit_solvation_model,
+                             solvent, nucl, site, mean[i]])
+
+    result_df = pd.DataFrame(rows, columns=[
+        "solute", 
+        "sap_geometry_type", 
+        "sap_nmr_method", 
+        "sap_basis", 
+        "implicit_solvation_model", 
+        "solvent", 
+        "nucleus", 
+        "site", 
+        "shielding"
+    ])
+
+    if verbose:
+        print(f"\033[92mRows of site-aggregated shieldings data: {len(result_df)}\033[0m")
+
+    return result_df
+
+def _create_stationary_df(shieldings_df, solvents):
+    """Create a dataframe for the stationary shieldings, tiled across."""
+    stationary_df = shieldings_df.query("solvent=='none'")
+    rows = []
+    for _, row in stationary_df.iterrows():
+        solute, geo_type, nmr_method, basis, _, _, nucl, site, stationary = row
+        for solvent in solvents:
+            row = [solute, geo_type, nmr_method, basis, nucl, site, solvent, stationary]
+            rows.append(row)
+
+    stationary_df = pd.DataFrame(rows, columns=[*INDEX_COLUMNS, "stationary"])
+    stationary_df.set_index(INDEX_COLUMNS, inplace=True)
+    stationary_df = stationary_df.sort_index()
+    return stationary_df
+
+def _create_dft_pcm_correction_df(sap_shieldings_df:pd.DataFrame,
+                                  stationary_df:pd.DataFrame,
+                                  pcm_reference_method:str=DEFAULT_PCM_REFERENCE_METHOD,
+                                  pcm_reference_basis:str=DEFAULT_PCM_REFERENCE_BASIS,
+                                  verbose:bool=True):
+    """
+    Create a dataframe for the DFT PCM corrections.
+    """
+    pcm_dft_shieldings_df = sap_shieldings_df.query("solvent != 'none'")
+
+    pcm_dft_shieldings_df = pd.pivot_table(pcm_dft_shieldings_df,
+                                           index=INDEX_COLUMNS,
+                                           columns="implicit_solvation_model",
+                                           values="shielding")
+
+    pcm_dft_shieldings_df = pcm_dft_shieldings_df.merge(stationary_df,
+                                                        on=INDEX_COLUMNS,
+                                                        how='left')
+
+    pcm_dft_corrections_df = pcm_dft_shieldings_df.copy()
+    pcm_dft_corrections_df["pcm"] = pcm_dft_corrections_df["pcm"] - \
+                                    pcm_dft_shieldings_df["stationary"]
+    pcm_dft_corrections_df.drop(columns=["stationary"], inplace=True)
+    pcm_dft_corrections_df.reset_index(inplace=True)
+
+    # methods with no PCM shielding data of their own (e.g. double hybrids) borrow the reference
+    # method's PCM correction below
+    missing_methods = set(sap_shieldings_df['sap_nmr_method'].unique()) - \
+                      set(pcm_dft_corrections_df['sap_nmr_method'].unique())
+
+    if verbose:
+        print("Methods without pcm shielding data:")
+        for method in missing_methods:
+            print(f" - {method}")
+
+        print("\nThe following reference PCM correction data will be used:")
+        print(f"  Reference Method: {pcm_reference_method}")
+        print(f"  Reference Basis: {pcm_reference_basis}")
+
+    # Get the data for the reference pcm corrections
+    reference_corrections = pcm_dft_corrections_df.query(
+        f"sap_nmr_method == '{pcm_reference_method}' and "
+        f"sap_basis == '{pcm_reference_basis}'"
+    ).copy().drop(columns=["sap_nmr_method","sap_basis"])
+
+    for method in missing_methods:
+        method_df = sap_shieldings_df.query(f"sap_nmr_method == '{method}'").copy()
+        method_df.drop(
+            columns=[
+                "implicit_solvation_model",
+                "solvent",
+                "shielding"
+            ], 
+            inplace=True
+        )
+
+        method_pcm_df = method_df.merge(reference_corrections,
+                                        on=[
+                                            "solute",
+                                            "sap_geometry_type",
+                                            "nucleus",
+                                            "site"
+                                        ],
+                                        how='left')
+
+        pcm_dft_corrections_df = pd.concat([pcm_dft_corrections_df, 
+                                            method_pcm_df],
+                                            ignore_index=True)
+
+    pcm_dft_corrections_df.set_index(INDEX_COLUMNS, inplace=True)
+
+    if verbose:
+        print("\nTotal rows in implicit corrections dataframe:", 
+            pcm_dft_corrections_df.shape[0])
+    
+    return pcm_dft_corrections_df
+
+def _create_nn_pcm_correction_df(zap_shieldings_df: pd.DataFrame,
+                                 solvents:list,
+                                 verbose: bool = True):
+    """
+    Create a DataFrame for NN PCM corrections.
+    """
+    pcm_corrections_df = zap_shieldings_df.query("solvent != 'none'").copy()
+
+    # MagNET only computes a PCM correction for chloroform; copy it unchanged to every other
+    # solvent (no per-solvent correction is computed here).
+    chloroform_df = pcm_corrections_df[pcm_corrections_df["solvent"] == "chloroform"]
+    new_rows = []
+    for nucleus in chloroform_df["nucleus"].unique():
+        nucleus_base = chloroform_df[chloroform_df["nucleus"] == nucleus]
+        for solvent in solvents:
+            if solvent == "chloroform":
+                continue
+            new_row = nucleus_base.copy()
+            new_row["solvent"] = solvent
+            new_rows.append(new_row)
+
+    if new_rows:
+        pcm_corrections_df = pd.concat([pcm_corrections_df, *new_rows], ignore_index=True)
+
+    pcm_corrections_df = pd.pivot_table(
+        pcm_corrections_df,
+        index=INDEX_COLUMNS,
+        columns="implicit_solvation_model",
+        values="shielding"
+    )
+
+    # renamed to match the DFT implicit-corrections dataframe's column name
+    pcm_corrections_df.columns = ["pcm"]
+
+    return pcm_corrections_df
+
+def _aggregate_qcd_by_site(qcd_raw_df:pd.DataFrame, 
+                           site_atoms_df:pd.DataFrame, 
+                           verbose:bool=True):
+    """Aggregate QCD corrections by site."""
+
+    if verbose:
+        print("\033[93m", end="")
+        solute_iter = tqdm(site_atoms_df.iterrows(), 
+                           desc="Aggregating solute QCD corrections by site",
+                           total=len(site_atoms_df))
+    else:
+        solute_iter = site_atoms_df.iterrows()
+
+    rows = []
+    for (solute, nucl, site), row in solute_iter:
+        atom_numbers, *_ = row
+        atom_numbers = _fix_atom_numbers(atom_numbers, prefix="", coerce_to=int)
+        query = qcd_raw_df.query(f"solute == '{solute}' and atom_number.isin({atom_numbers})")
+        correction = query.qcd_correction.mean()
+        stderr = query.stderr.mean()
+        rows.append([solute, nucl, site, correction, stderr])
+
+    result_df = pd.DataFrame(rows, columns=["solute", "nucleus", "site", "qcd", "qcd_stderr"])
+    if verbose:
+        print(f"\033[92mRows of site-aggregated QCD corrections: {result_df.shape[0]}\033[0m")
+    return result_df
+
+
+def _tile_qcd_corrections(qcd_df, solvents):
+    """Tile QCD corrections across solvents."""
+    rows = []
+    for _, row in qcd_df.iterrows():
+        solute, nucl, site, qcd, _ = row
+        for solvent in solvents:
+            rows.append([solute, nucl, site, solvent, qcd])
+    df = pd.DataFrame(rows, columns=["solute", "nucleus", "site", "solvent", "qcd"])
+    df.set_index(["solute", "nucleus", "site", "solvent"], inplace=True)
+    return df
+
+
+def _collate_explicit_corrections(df):
+    """
+    Collate desmond and openMM corrections for each (solute, nucleus, site, solvent).
+    Returns a DataFrame with columns ['desmond', 'openMM'].
+    """
+    def extract_corrections(group):
+        desmond = group.loc[group['explicit_solvation_model'] == 'desmond', 'correction']
+        openmm = group.loc[group['explicit_solvation_model'] == 'openMM', 'correction']
+        return pd.Series({
+            'desmond': desmond.iloc[0] if not desmond.empty else np.nan,
+            'openMM': openmm.iloc[0] if not openmm.empty else np.nan
+        })
+
+    result_df = df.groupby(["solute", "nucleus", "site", "solvent"]).apply(extract_corrections, include_groups=False)
+    result_df = pd.DataFrame(result_df)
+    return result_df
+
+
+def _build_combined_df(
+    experimental_df,
+    qcd_corrections_df,
+    explicit_corrections_df,
+    classical_df,
+    stationary_df,
+    implicit_corrections_df,
+    index_columns,
+):
+    """Build the final combined dataframe with all corrections."""
+    combined_corrections_df = pd.concat(
+        [experimental_df, qcd_corrections_df, explicit_corrections_df, classical_df], axis=1
+    )
+
+    stationary_and_pcm_df = pd.concat([stationary_df, implicit_corrections_df], axis=1)
+
+    # merged (not concatenated) so the non-nmr-specific data gets tiled across every nmr method
+    combined_corrections_df.reset_index(inplace=True)
+    stationary_and_pcm_df.reset_index(inplace=True)
+    combined_df = combined_corrections_df.merge(
+        stationary_and_pcm_df,
+        on=["solute", "nucleus", "site", "solvent"],
+        how='outer'
+    )
+    combined_df.set_index(index_columns, inplace=True)
+    combined_df = combined_df[
+        [
+            'experimental', 
+            'stationary',
+            'qcd',
+            'pcm',
+            'desmond',
+            'openMM',
+            'desmond_vib',
+            'openMM_vib'
+        ]
+    ]
+    return combined_df
+
+
+# Importable data-loading functions
+
+def load_experiment_data(filepath: str, verbose: bool = True) -> pd.DataFrame:
+    """
+    Reads the experimental data from an excel file at the specified path.
+    """
+    if verbose:
+        print("\033[93mLoading experimental data...\033[0m", end="", flush=True)
+    experimental_df = pd.read_excel(filepath)
+    experimental_df["nucleus"] = experimental_df.site.str.split("_").str[-1]
+    experimental_df.set_index(["solute","nucleus","site"], inplace=True)
+
+    # stack into column vectors indexed by solute/site/solvent, for vector regressions like
+    # y = c0 + c1 * x1 + ...
+    experimental_df = experimental_df.drop("atom_numbers", axis=1)
+    experimental_df.reset_index(inplace=True)
+    experimental_df.set_index(["solute","nucleus","site"], inplace=True)
+    experimental_df = pd.DataFrame(experimental_df.stack(future_stack=True))
+    experimental_df.columns = ["experimental"]
+    experimental_df.index.names = ["solute", "nucleus", "site", "solvent"]
+    experimental_df = experimental_df.sort_index()
+
+    if verbose:
+        print("\r\033[92mLoaded experimental data.    \033[0m")
+
+    return experimental_df
+
+def load_site_atom_data(filepath: str, verbose: bool = True) -> pd.DataFrame:
+    """
+    Loads atom number information for each solute/nucleus/site combination from
+    the experimental Excel file. 
+    
+    Returns a DataFrame indexed by solute, nucleus, and site, with a column 
+    'atom_numbers'.
+    """
+    if verbose:
+        print("\033[93mLoading site atom data...\033[0m", end="", flush=True)
+    experimental_df = pd.read_excel(filepath)
+    experimental_df["nucleus"] = experimental_df.site.str.split("_").str[-1]
+    result_df = experimental_df[["solute", "nucleus", "site", "atom_numbers"]]
+    result_df = result_df.drop_duplicates()
+    result_df.set_index(["solute", "nucleus", "site"], inplace=True)
+    if verbose:
+        print("\r\033[92mLoaded site atom data.       \033[0m")
+    return result_df
+
+def load_nmr_method_names(delta22_file_path:str) -> list[str]:
+    """
+    Returns the names of the nmr methods stored in the specified hdf5 file.
+    """
+    with h5py.File(delta22_file_path, 'r') as hdf5_file:
+        names_data = np.array(hdf5_file["conventional_nmr_method_names"])
+
+    nmr_method_tuples = [name.decode("ascii") for name in names_data]
+    parsed_methods = [_parse_method_name(t) for t in nmr_method_tuples]
+    columns = ["sap_nmr_method", "sap_basis", "solvent_model", "solvent_name"]
+    return pd.DataFrame(parsed_methods, columns=columns)
+
+def load_solvents_from_methods(delta22_file_path:str) -> list[str]:
+    """
+    Returns the list of solvents used in the nmr methods stored in the specified hdf5 file.
+    """
+    methods_df = load_nmr_method_names(delta22_file_path)
+    solvents = methods_df['solvent_name'].unique().tolist()
+    return [s for s in solvents if s.lower() != 'none']
+
+
+# Geometry and topology accessors.
+#
+# These return the raw stored arrays (atom coordinates and solvent atom orderings) that the
+# shielding loaders above do not expose. Coordinates are stored in the same fixed-point form as
+# the shieldings (whole numbers equal to the real value in angstroms times 10,000) and are decoded
+# back to angstroms here; the solvent atom orderings are plain atomic numbers and need no decoding.
+
+def load_solutes(hdf5_filepath: str) -> list[str]:
+    """
+    Returns the names of the solutes stored in the file, in file order.
+    """
+    with h5py.File(hdf5_filepath, "r") as hdf5_file:
+        return list(hdf5_file["solutes"].keys())
+
+
+def load_solvent_atomic_numbers(hdf5_filepath: str,
+                                solvent: str,
+                                explicit_solvation_model: str = "openMM") -> np.ndarray:
+    """
+    Returns the atomic numbers of one solvent molecule, in the atom order used by the explicit
+    solvent geometries. explicit_solvation_model is "openMM" (default) or "desmond"; the two
+    orderings are stored separately. This is the reference topology for reading the solvent atoms
+    in the solvated geometries from load_explicit_ensemble_geometries.
+    """
+    key = f"{explicit_solvation_model}_atomic_numbers"
+    with h5py.File(hdf5_filepath, "r") as hdf5_file:
+        return np.array(hdf5_file["solvents"][solvent][key])
+
+
+def load_stationary_geometries(hdf5_filepath: str, solute: str) -> dict:
+    """
+    Returns the two resting geometries the conventional (stationary and PCM) shieldings were
+    computed on, as a dict {geometry_type: coordinates (n_atoms, 3) in angstroms}. The geometry
+    types are "aimnet2" and "pbe0_tz", matching GEOMETRY_TYPES.
+    """
+    with h5py.File(hdf5_filepath, "r") as hdf5_file:
+        coords = _decode_fixed_point(hdf5_file["solutes"][solute]["stationary_and_pcm"]["geometries"][()])
+    return {geometry_type: coords[i] for i, geometry_type in enumerate(GEOMETRY_TYPES)}
+
+
+def load_qcd_geometries(hdf5_filepath: str, solute: str) -> dict:
+    """
+    Returns the gas-phase quasiclassical-dynamics geometries the QCD corrections were computed on,
+    as a dict with "unperturbed" (n_atoms, 3) = the resting geometry and "perturbed"
+    (n_trajectories, n_frames, n_atoms, 3) = the jiggled snapshots. Coordinates are in angstroms.
+    The perturbed frame axes align with the perturbed shieldings used by load_average_qcd_corrections.
+    """
+    with h5py.File(hdf5_filepath, "r") as hdf5_file:
+        qcd = hdf5_file["solutes"][solute]["qcd"]["gas_phase"]
+        return {
+            "unperturbed": _decode_fixed_point(qcd["unperturbed_geometry"][()]),
+            "perturbed": _decode_fixed_point(qcd["perturbed_geometries"][()]),
+        }
+
+
+def load_explicit_ensemble_geometries(hdf5_filepath: str,
+                                      solute: str,
+                                      solvent: str,
+                                      explicit_solvation_model: str = "openMM") -> np.ndarray:
+    """
+    Returns the explicit-solvent snapshot coordinates (n_frames, n_cluster_atoms, 3) in angstroms
+    for one solute/solvent under one engine ("openMM" or "desmond"). The solute atoms come first
+    (load_solutes order), then whole solvent molecules in the load_solvent_atomic_numbers order.
+    The frame axis aligns with the perturbed shieldings used by the explicit-correction loaders.
+    """
+    with h5py.File(hdf5_filepath, "r") as hdf5_file:
+        node = hdf5_file["solutes"][solute][explicit_solvation_model][solvent]
+        return _decode_fixed_point(node["perturbed_ensemble_geometries"][()])
+
+
+def load_explicit_solute_geometry(hdf5_filepath: str,
+                                  solute: str,
+                                  explicit_solvation_model: str = "openMM") -> np.ndarray:
+    """
+    Returns the isolated resting solute geometry (n_atoms, 3) in angstroms used as the gas-phase
+    reference for the explicit-solvent corrections under one engine ("openMM" or "desmond").
+    """
+    with h5py.File(hdf5_filepath, "r") as hdf5_file:
+        node = hdf5_file["solutes"][solute][explicit_solvation_model]["gas_phase"]
+        return _decode_fixed_point(node["unperturbed_solute_geometry"][()])
+
+
+def load_perturbed_shieldings(hdf5_filepath, solute, solvent,
+                              explicit_solvation_model="openMM", shield_type="dft"):
+    """The raw per-frame explicit-solvent shieldings (n_frames, n_atoms, 2) for one solute/solvent
+    under one engine ("openMM" or "desmond"). shield_type is "dft" or "nn". The last axis is
+    [isolated, solvated], so the per-frame solvent correction is solvated minus isolated. This is the
+    frame-resolved data behind the convergence and autocorrelation panels of Figure S8; the
+    site-aggregated correction over all frames comes from load_site_aggregated_explicit_shieldings.
+    Frames that were not computed read as NaN."""
+    key = f"perturbed_{shield_type}_shieldings"
+    with h5py.File(hdf5_filepath, "r") as hdf5_file:
+        node = hdf5_file["solutes"][solute][explicit_solvation_model][solvent]
+        return _decode_fixed_point(node[key][()])
+
+
+def load_stationary_and_pcm_dft_shieldings(hdf5_filepath:str,
+                                           verbose:bool=True):
+    """Load raw stationary and PCM DFT shieldings, one row per (solute, geometry, nmr method)."""
+    if verbose:
+        print("\033[93mLoading stationary and PCM DFT shieldings...\033[0m", end="", flush=True)
+    with h5py.File(hdf5_filepath, 'r') as hdf5_file:
+        nmr_method_names = np.array(hdf5_file["conventional_nmr_method_names"])
+        nmr_method_names = [i.decode("utf-8") for i in nmr_method_names]
+
+        dft_rows = []
+        nan_dft_rows = 0
+        for solute in hdf5_file["solutes"]:
+            solute_group = hdf5_file["solutes"][solute]["stationary_and_pcm"]
+
+            geometry_optimization_timings = np.array(solute_group["geometry_optimization_timings"])
+            dft_shieldings = _decode_fixed_point(solute_group["conventional_shieldings"][()])
+            dft_timings = np.array(solute_group["conventional_nmr_timings"])
+
+            for nmr_method_index, nmr_method_name in enumerate(nmr_method_names):
+                nmr_method, basis, solvent_model, solvent_name = _parse_method_name(nmr_method_name)
+
+                for geometry_index, geometry_type in enumerate(GEOMETRY_TYPES):
+                    method_geometry_time = geometry_optimization_timings[geometry_index]
+                    method_nmr_time = dft_timings[nmr_method_index, geometry_index]
+                    method_shieldings = dft_shieldings[nmr_method_index, geometry_index]
+                    if np.isnan(method_shieldings).any():
+                        nan_dft_rows += 1
+                    dft_row = [
+                        solute,
+                        geometry_type,
+                        nmr_method,
+                        basis,
+                        solvent_model,
+                        solvent_name,
+                        method_shieldings,
+                        method_geometry_time,
+                        method_nmr_time
+                    ]
+                    dft_rows.append(dft_row)
+
+        if nan_dft_rows > 0 and verbose:
+            print(f"Found {nan_dft_rows} total DFT rows with NaNs!")
+
+    columns = [
+        "solute",
+        "sap_geometry_type",
+        "sap_nmr_method",
+        "sap_basis",
+        "solvent_model",
+        "solvent",
+        "shieldings",
+        "geometry_time",
+        "nmr_time"
+    ]
+    dft_shieldings_df = pd.DataFrame(dft_rows, columns=columns)
+    dft_shieldings_df = _sort_and_parse_raw_shieldings(dft_shieldings_df)
+    if verbose:
+        print("\r\033[92mLoaded stationary and PCM DFT shieldings.    \033[0m")
+    return dft_shieldings_df
+
+def load_stationary_and_pcm_nn_shieldings(hdf5_filepath:str,
+                                          verbose:bool=True):
+    """Load raw stationary and PCM NN (MagNET) shieldings, one row per (solute, geometry, nmr method)."""
+    with h5py.File(hdf5_filepath, 'r') as hdf5_file:
+        nn_rows = []
+        nan_nn_rows = 0
+        for solute in hdf5_file["solutes"]:
+            atomic_numbers = np.array(hdf5_file["solutes"][solute]["atomic_numbers"])
+            ch_mask = np.where((atomic_numbers == 1) | (atomic_numbers == 6))[0]
+
+            solute_group = hdf5_file["solutes"][solute]["stationary_and_pcm"]
+            geometry_optimization_timings = np.array(solute_group["geometry_optimization_timings"])
+            nn_gas_shieldings = _decode_fixed_point(solute_group["nn_gas_shieldings"][()])
+            nn_pcm_corrections = _decode_fixed_point(solute_group["nn_pcm_corrections"][()])
+            nn_nmr_timings = np.array(solute_group["nn_nmr_timings"])
+
+            if np.isnan(nn_gas_shieldings[ch_mask]).any():
+                nan_nn_rows += 1
+            if np.isnan(nn_pcm_corrections[ch_mask]).any():
+                nan_nn_rows += 1
+            
+            nn_row_gas = [
+                solute, 
+                "aimnet2", 
+                "MagNET",
+                "N/A", 
+                "gas", 
+                "none", 
+                nn_gas_shieldings, 
+                geometry_optimization_timings[0], # AimNet2 timing
+                nn_nmr_timings[0]                 # MagNET-Zero timing 
+            ]
+            nn_row_pcm = [
+                solute, 
+                "aimnet2", 
+                "MagNET", 
+                "N/A", 
+                "pcm_correction", 
+                "chloroform", 
+                nn_pcm_corrections, 
+                geometry_optimization_timings[0], # AimNet2 timing
+                nn_nmr_timings[1]                 # MagNET-PCM timing 
+            ]
+            nn_rows.append(nn_row_gas)
+            nn_rows.append(nn_row_pcm)
+
+        if nan_nn_rows > 0 and verbose:
+            print(f"Found {nan_nn_rows} total NN rows with NaNs!")
+
+    columns = [
+        "solute",
+        "sap_geometry_type",
+        "sap_nmr_method",
+        "sap_basis",
+        "solvent_model",
+        "solvent",
+        "shieldings",
+        "geometry_time",
+        "nmr_time"
+    ]
+    nn_shieldings_df = pd.DataFrame(nn_rows, columns=columns)
+    nn_shieldings_df = _sort_and_parse_raw_shieldings(nn_shieldings_df)
+    return nn_shieldings_df
+
+def load_average_qcd_corrections(hdf5_filepath, 
+                                 shield_type="dft", 
+                                 verbose:bool=True):
+    """Loads and computes the average QCD corrections."""
+    if verbose:
+        print("\033[93mLoading QCD data...\033[0m", end="", flush=True)
+    rows = []
+    with h5py.File(hdf5_filepath, "r") as hdf5_file:
+        solutes = hdf5_file["solutes"].keys()
+        for solute in solutes:
+            qcd_group = hdf5_file["solutes"][solute]["qcd"]["gas_phase"]
+            if shield_type == "dft":
+                unperturbed = _decode_fixed_point(qcd_group["unperturbed_dft_shieldings"][()])
+                perturbed = _decode_fixed_point(qcd_group["perturbed_dft_shieldings"][()])
+            elif shield_type == "nn":
+                unperturbed = _decode_fixed_point(qcd_group["unperturbed_nn_shieldings"][()])
+                perturbed = _decode_fixed_point(qcd_group["perturbed_nn_shieldings"][()])
+            else:
+                raise ValueError("shield_type must be 'dft' or 'nn'")
+            means = np.mean(perturbed, axis=1)
+            n = means.shape[0]
+            stdevs = np.std(means, axis=0)
+            means = np.mean(means, axis=0)
+            corrections = means - unperturbed
+            stderrs = stdevs / np.sqrt(n)
+            for i, (correction, error) in enumerate(zip(corrections, stderrs)):
+                row = [solute, i + 1, correction, error]
+                rows.append(row)
+    df = pd.DataFrame(rows, columns=["solute", "atom_number", "qcd_correction", "stderr"])
+    if verbose:
+        print("\r\033[92mLoaded QCD data.    \033[0m")
+    return df
+
+def load_site_aggregated_explicit_shieldings(hdf5_filepath: str,
+                                             site_atoms_df: pd.DataFrame,
+                                             shield_type: str = "dft",
+                                             verbose: bool = True):
+    """Compute, per site/solvent/explicit-solvation-engine, the mean explicit-solvent correction
+    (solvated minus isolated) over all frames, with its standard error."""
+    if verbose:
+        print("\033[93mLoading explicit shieldings data...\033[0m", end="", flush=True)
+
+    new_rows = []
+    with h5py.File(hdf5_filepath, "r") as hdf5_file:
+        for (solute, nucl, site), row in site_atoms_df.iterrows():
+            atom_numbers, *_ = row
+            atom_indices = _fix_atom_numbers(atom_numbers, prefix="", coerce_to=int, offset=-1)
+
+            # desmond and openMM may not cover the same solvents; take the union
+            solvents = set()
+            for explicit_solvation_model in EXPLICIT_SOLVATION_MODELS:
+                group = hdf5_file[f"solutes/{solute}/{explicit_solvation_model}"]
+                solvents.update([
+                    k for k in group.keys()
+                    if k != "gas_phase" and f"perturbed_{shield_type}_shieldings" in group[k]
+                ])
+            solvents = sorted(solvents)
+
+            for explicit_solvation_model in EXPLICIT_SOLVATION_MODELS:
+                group = hdf5_file[f"solutes/{solute}/{explicit_solvation_model}"]
+                for solvent in solvents:
+                    dataset = group.get(f"{solvent}/perturbed_{shield_type}_shieldings")
+                    if dataset is None:
+                        new_row = [solute, solvent, explicit_solvation_model, nucl, site, np.nan, 0, np.nan]
+                        new_rows.append(new_row)
+                        continue
+                    dataset = _decode_fixed_point(dataset[()])
+
+                    dataset = dataset[:,atom_indices,:]
+                    dataset = np.mean(dataset, axis=1)
+                    corrections = dataset[:,1] - dataset[:,0]
+                    corrections = corrections[~np.isnan(corrections)]
+                    correction = np.mean(corrections)
+                    n = len(corrections)
+                    stderr = np.std(corrections)/np.sqrt(n)
+                    
+                    new_row = [solute, solvent, explicit_solvation_model, nucl, site, correction, n, stderr]
+                    new_rows.append(new_row)
+
+    explicit_df = pd.DataFrame(new_rows)
+    explicit_df.columns = [
+        "solute", 
+        "solvent", 
+        "explicit_solvation_model", 
+        "nucleus", 
+        "site",
+        "correction",
+        "n",
+        "stderr"
+    ]
+
+    if verbose:
+        print("\r\033[92mLoaded explicit shieldings data.    \033[0m")
+
+    return explicit_df
+
+def load_classical_vibrational_corrections(hdf5_filepath: str, 
+                                           site_atoms_df: pd.DataFrame, 
+                                           shield_type: str = "dft",
+                                           verbose: bool = True):
+    """Compute classical vibrational corrections for both DFT and NN shielding values."""
+    rows = []
+    with h5py.File(hdf5_filepath, "r") as hdf5_file:
+        solutes = hdf5_file["solutes"].keys()
+        for solute in solutes:
+            sites_df = site_atoms_df.query(f"solute == '{solute}'")
+            for (solute, nucl, site), site_row in sites_df.iterrows():
+                atom_numbers, *_ = site_row
+                atom_numbers = _fix_atom_numbers(atom_numbers, prefix="", coerce_to=int, offset=-1)
+
+                # desmond and openMM may not cover the same solvents; take the union
+                solvents = set()
+                for explicit_solvation_model in EXPLICIT_SOLVATION_MODELS:
+                    group = hdf5_file[f"solutes/{solute}/{explicit_solvation_model}"]
+                    solvents.update([
+                        k for k in group.keys()
+                        if k != "gas_phase" and f"perturbed_{shield_type}_shieldings" in group[k]
+                    ])
+                solvents = sorted(solvents)
+
+                for explicit_solvation_model in EXPLICIT_SOLVATION_MODELS:
+                    stationary_path = f"solutes/{solute}/{explicit_solvation_model}/gas_phase/unperturbed_{shield_type}_shieldings"
+                    shieldings = hdf5_file[stationary_path]
+                    shieldings = _decode_fixed_point(shieldings[()])
+                    classical_stationary = np.mean(shieldings[atom_numbers])
+
+                    for solvent in solvents:
+                        path = f"solutes/{solute}/{explicit_solvation_model}/{solvent}/perturbed_{shield_type}_shieldings"
+                        added = False
+                        if path in hdf5_file:
+                            shieldings = hdf5_file[path]
+                            shieldings = _decode_fixed_point(shieldings[()])
+                            shieldings = shieldings[:, atom_numbers, 0]  # solute shieldings
+
+                            if np.all(np.isnan(shieldings)):
+                                print(f"Found all NaNs in {path} for solute {solute}, nucleus {nucl}, site {site}, solvent {solvent}")
+                            shieldings = shieldings[~np.isnan(shieldings)]
+
+                            if len(shieldings) > 0:
+                                explicit_mean = np.mean(shieldings)
+                                correction = explicit_mean - classical_stationary
+                                new_row = [solute, nucl, site, solvent, explicit_solvation_model, correction]
+                                added = True
+                        if not added:
+                            new_row = [solute, nucl, site, solvent, explicit_solvation_model, np.nan]
+                        rows.append(new_row)
+
+    columns = ["solute", "nucleus", "site", "solvent", "flavor", "correction"]
+    df = pd.DataFrame(rows, columns=columns)
+    df = pd.pivot_table(df, values="correction", columns="flavor", index=["solute", "nucleus", "site", "solvent"])
+    df.columns = [f"{col}_vib" for col in df.columns]
+    return df
+
+def load_delta22_dft_data(delta22_file_path: str, 
+                          experimental_data_path: str,
+                          pcm_reference_method: str = DEFAULT_PCM_REFERENCE_METHOD,
+                          pcm_reference_basis: str = DEFAULT_PCM_REFERENCE_BASIS,
+                          verbose: bool = True) -> pd.DataFrame:
+    """
+    Load and process DFT shielding data from the delta22 HDF5 file into one combined DataFrame.
+
+    Returns a DataFrame indexed by solute, sap_geometry_type, sap_nmr_method, sap_basis, nucleus,
+    site, solvent, with columns experimental, stationary, qcd, pcm, desmond, openMM, desmond_vib,
+    openMM_vib.
+    """
+    experimental_df = load_experiment_data(experimental_data_path, verbose)
+    site_atoms_df = load_site_atom_data(experimental_data_path, verbose)
+    dft_solvents = load_solvents_from_methods(delta22_file_path)
+    sap_shieldings_df = load_stationary_and_pcm_dft_shieldings(delta22_file_path, verbose)
+    sap_shieldings_df = _aggregate_shieldings_by_site(sap_shieldings_df,
+                                                      site_atoms_df)
+
+    stationary_df = _create_stationary_df(sap_shieldings_df, dft_solvents)
+    pcm_dft_corrections_df = _create_dft_pcm_correction_df(sap_shieldings_df,
+                                                           stationary_df,
+                                                           pcm_reference_method,
+                                                           pcm_reference_basis,
+                                                           verbose=False)
+
+    qcd_data_df = load_average_qcd_corrections(delta22_file_path,
+                                               "dft",
+                                               verbose)
+    qcd_data_df = _aggregate_qcd_by_site(qcd_data_df, site_atoms_df, verbose)
+    qcd_corrections_df = _tile_qcd_corrections(qcd_data_df, dft_solvents)
+
+    explicit_data_df = load_site_aggregated_explicit_shieldings(
+        delta22_file_path,
+        site_atoms_df,
+        shield_type="dft",
+        verbose=verbose
+    )
+    explicit_corrections_df = _collate_explicit_corrections(explicit_data_df)
+
+    classical_df = load_classical_vibrational_corrections(
+        delta22_file_path,
+        site_atoms_df,
+        shield_type="dft",
+        verbose=verbose
+    )
+
+    combined_df = _build_combined_df(
+        experimental_df,
+        qcd_corrections_df,
+        explicit_corrections_df,
+        classical_df,
+        stationary_df,
+        pcm_dft_corrections_df,
+        INDEX_COLUMNS
+    )
+
+    return combined_df
+
+
+def load_delta22_nn_data(delta22_file_path: str,
+                         experimental_data_path: str,
+                         exclude_solutes=None,
+                         verbose: bool = True,
+                         nn_shieldings_override_df: pd.DataFrame = None) -> pd.DataFrame:
+    """
+    Load and process NN (MagNET) shielding data from the delta22 HDF5 file into one combined
+    DataFrame, indexed by solute, sap_geometry_type, sap_nmr_method, sap_basis, nucleus, site,
+    solvent, with columns experimental, stationary, qcd, pcm, desmond, openMM, desmond_vib,
+    openMM_vib.
+
+    MagNET only computes a PCM correction for chloroform; that value is copied unchanged to every
+    other solvent.
+
+    exclude_solutes drops solute names from the result, e.g. ["nitromethane"]. The composite-model
+    fits exclude nitromethane because it is a stark outlier for MagNET (its rovibrational/QCD
+    correction is anomalous).
+
+    nn_shieldings_override_df, if given, is used in place of
+    load_stationary_and_pcm_nn_shieldings(delta22_file_path) and must have the same columns/shape
+    (one "gas" row and one "pcm_correction" row per solute). This lets a caller substitute
+    freshly-computed shieldings (e.g. from live, symmetrized inference) for the ones stored in the
+    HDF5, without touching the rest of the pipeline (QCD, explicit-solvent, and
+    classical-vibrational corrections still come from the file).
+    """
+    experimental_df = load_experiment_data(experimental_data_path, verbose)
+    site_atoms_df = load_site_atom_data(experimental_data_path, verbose)
+
+    if nn_shieldings_override_df is not None:
+        # same row-sorting / shieldings-parsing load_stationary_and_pcm_nn_shieldings applies, so
+        # downstream code (_aggregate_shieldings_by_site) sees an identical shape regardless of
+        # whether the shieldings came from the HDF5 or from a live override
+        nn_shieldings_df = _sort_and_parse_raw_shieldings(nn_shieldings_override_df.copy())
+    else:
+        nn_shieldings_df = load_stationary_and_pcm_nn_shieldings(delta22_file_path,
+                                                                 verbose)
+
+    nn_shieldings_df = _aggregate_shieldings_by_site(nn_shieldings_df,
+                                                     site_atoms_df,
+                                                     verbose)
+
+    nn_solvents = load_solvents_from_methods(delta22_file_path)
+    stationary_df = _create_stationary_df(nn_shieldings_df, nn_solvents)
+    pcm_nn_corrections_df = _create_nn_pcm_correction_df(nn_shieldings_df,
+                                                         nn_solvents,
+                                                         verbose=True)
+
+    qcd_data_df = load_average_qcd_corrections(delta22_file_path, "nn", verbose)
+    qcd_data_df = _aggregate_qcd_by_site(qcd_data_df, site_atoms_df, verbose)
+    qcd_corrections_df = _tile_qcd_corrections(qcd_data_df, nn_solvents)
+
+    explicit_data_df = load_site_aggregated_explicit_shieldings(
+        delta22_file_path,
+        site_atoms_df,
+        shield_type="nn",
+        verbose=verbose
+    )
+    explicit_corrections_df = _collate_explicit_corrections(explicit_data_df)
+
+    classical_df = load_classical_vibrational_corrections(
+        delta22_file_path,
+        site_atoms_df,
+        shield_type="nn",
+        verbose=verbose
+    )
+
+    combined_df = _build_combined_df(
+        experimental_df,
+        qcd_corrections_df,
+        explicit_corrections_df,
+        classical_df,
+        stationary_df,
+        pcm_nn_corrections_df,
+        INDEX_COLUMNS
+    )
+
+    if exclude_solutes:
+        solute_index = combined_df.index.get_level_values("solute")
+        combined_df = combined_df[~solute_index.isin(set(exclude_solutes))]
+
+    return combined_df
+
+def _compute_timings(df: pd.DataFrame) -> pd.DataFrame:
+    """Sum a grouped shielding dataframe's geometry and NMR timing columns and return
+    geometry_time, nmr_time, and their total as a Series."""
+    geometry_time = df["geometry_time"].sum()
+    nmr_time = df["nmr_time"].sum()
+    total_time = geometry_time + nmr_time
+    return pd.Series({"geometry_time":geometry_time, 
+                      "nmr_time":nmr_time, 
+                      "total_time":total_time})
+
+def load_delta22_dft_timings(delta22_file_path: str) -> pd.DataFrame:
+    """
+    Sums DFT geometry-optimization and NMR computation times for every combination of geometry
+    type, NMR method, basis, solvent model, and solvent name.
+    """
+    sap_dft_shieldings = load_stationary_and_pcm_dft_shieldings(delta22_file_path)
+
+    groupings = [
+        "sap_geometry_type",
+        "sap_nmr_method",
+        "sap_basis",
+        "solvent_model",
+        "solvent",
+    ]
+    return sap_dft_shieldings.groupby(groupings).apply(_compute_timings, 
+                                       include_groups=False)
+
+def load_delta22_nn_timings(delta22_file_path: str) -> pd.DataFrame:
+    """
+    Sums NN (MagNET) geometry-optimization and NMR computation times for every combination of
+    geometry type, NMR method, basis, solvent model, and solvent name.
+    """
+    sap_nn_shieldings = load_stationary_and_pcm_nn_shieldings(delta22_file_path)
+
+    groupings = [
+        "sap_geometry_type",
+        "sap_nmr_method",
+        "sap_basis",
+        "solvent_model",
+        "solvent"
+    ]
+    return sap_nn_shieldings.groupby(groupings).apply(_compute_timings,
+                                       include_groups=False)
+
+
+if __name__ == "__main__":
+    current_file_path = os.path.abspath(__file__)
+    here = os.path.dirname(current_file_path)
+    hdf5_filepath = os.path.join(here, "delta22.hdf5")
+    experiment_data_filepath = os.path.join(here, "delta22_experimental.xlsx")
+
+    print("LOADING DFT DATA:")
+    print("="*60)
+    dft_data = load_delta22_dft_data(hdf5_filepath, experiment_data_filepath)
+
+    dft_data_display = dft_data.query(
+        "sap_nmr_method == 'wp04' and "
+        "sap_basis == 'pcSseg2' and "
+        "sap_geometry_type == 'aimnet2'"
+    ).copy().reset_index()
+    dft_data_display.drop(columns=["sap_nmr_method", "sap_basis", "sap_geometry_type"], inplace=True)
+    print("DFT Data Display:")
+    print(dft_data_display.head())
+    print("Solvents:")
+    print(dft_data_display.reset_index().solvent.unique())
+
+    print("LOADING NN DATA:")
+    print("="*60)
+    nn_data = load_delta22_nn_data(hdf5_filepath, experiment_data_filepath)
+    print(nn_data.head())
+
+

data/delta22/test_delta22.py

@@ -0,0 +1,260 @@
+"""
+Tests for the delta22 reader (delta22_reader.py).
+
+These build a tiny synthetic dataset in the *exact* on-disk format (so the fixture also
+serves as an executable spec of the layout) plus a matching experimental spreadsheet, then
+run the public loaders against them. No dependency on the multi-GB real file. One opt-in
+test runs against the real `delta22.hdf5` if it is present next to this file.
+
+Run:  pytest test_delta22.py -q
+Requires: pytest, numpy, h5py, pandas, openpyxl, tqdm.
+"""
+import os
+import numpy as np
+import pandas as pd
+import h5py
+import pytest
+
+import delta22_reader as R
+from delta22_reader import (
+    _decode_fixed_point,
+    _MISSING_MARKER,
+    load_delta22_dft_data,
+    load_delta22_nn_data,
+    load_delta22_dft_timings,
+    load_delta22_nn_timings,
+)
+
+SCALE = 1e4
+
+
+def _enc(arr):
+    """Encode floats the way the real file does: int32 = round(value * 1e4), NaN -> marker."""
+    a = np.asarray(arr, dtype=np.float64)
+    mask = np.isnan(a)
+    out = np.zeros(a.shape, dtype=np.int32)
+    out[~mask] = np.rint(a[~mask] * SCALE).astype(np.int32)
+    out[mask] = _MISSING_MARKER
+    return out
+
+
+# ---- known synthetic content -------------------------------------------------
+# One solute "moltest", 4 atoms (H, C, H, C); one solvent (chloroform); one method
+# (b3lyp_d3bj/pcSseg2, which is also the default PCM reference) present as a gas entry
+# and a chloroform-PCM entry; two geometries (aimnet2, pbe0_tz).
+SOLUTE = "moltest"
+Z = np.array([1, 6, 1, 6], np.uint8)          # 1-based H at 1,3 ; C at 2,4
+METHOD_NAMES = [
+    'b3lyp_d3bj,pcSseg2,gas,"none"',
+    'b3lyp_d3bj,pcSseg2,pcm,"chloroform"',
+]
+# conventional_shieldings[method, geometry, atom]
+GAS = np.array([[30.0, 150.0, 31.0, 151.0],    # aimnet2
+                [30.2, 150.4, 31.2, 151.4]])   # pbe0_tz
+PCM = np.array([[30.5, 151.0, 31.5, 152.0],    # aimnet2  (+0.5 H, +1.0 C vs gas)
+                [30.7, 151.4, 31.7, 152.4]])   # pbe0_tz
+CONV = np.stack([GAS, PCM])                     # (2 methods, 2 geom, 4 atoms)
+
+NN_GAS = np.array([30.1, 150.1, 31.1, 151.1])
+NN_PCM = np.array([0.4, 0.9, 0.4, 0.9])         # chloroform - gas
+
+# Expected stationary shift for site Ha_H (atoms 1,3) on the aimnet2 geometry:
+EXP_STATIONARY_HA_AIMNET2 = np.mean([30.0, 31.0])   # 30.5
+# Expected PCM correction for Ha_H aimnet2 = mean(PCM 1,3) - mean(GAS 1,3):
+EXP_PCM_HA_AIMNET2 = np.mean([30.5, 31.5]) - np.mean([30.0, 31.0])  # 0.5
+
+
+def _write_solute(grp):
+    grp.create_dataset("atomic_numbers", data=Z)
+
+    sap = grp.create_group("stationary_and_pcm")
+    sap.create_dataset("conventional_shieldings", data=_enc(CONV))
+    sap.create_dataset("geometries", data=_enc(np.zeros((2, 4, 3))))
+    sap.create_dataset("conventional_nmr_timings", data=np.ones((2, 2)) * 1.5)
+    sap.create_dataset("geometry_optimization_timings", data=np.array([2.0, 9.0]))
+    sap.create_dataset("nn_gas_shieldings", data=_enc(NN_GAS))     # (4,) flat
+    sap.create_dataset("nn_pcm_corrections", data=_enc(NN_PCM))    # (4,) flat
+    sap.create_dataset("nn_nmr_timings", data=np.array([0.1, 0.2]))
+
+    qcd = grp.create_group("qcd").create_group("gas_phase")
+    qcd.create_dataset("unperturbed_dft_shieldings", data=_enc(GAS[0]))
+    qcd.create_dataset("unperturbed_nn_shieldings", data=_enc(NN_GAS))
+    # (n_trajectories=2, n_frames=2, n_atoms=4); means equal the unperturbed -> 0 correction
+    qcd.create_dataset("perturbed_dft_shieldings",
+                       data=_enc(np.broadcast_to(GAS[0], (2, 2, 4)).copy()))
+    qcd.create_dataset("perturbed_nn_shieldings",
+                       data=_enc(np.broadcast_to(NN_GAS, (2, 2, 4)).copy()))
+    # geometry companions (coordinates, never consumed by the shielding loaders)
+    qcd.create_dataset("unperturbed_geometry", data=_enc(np.full((4, 3), 0.5)))
+    qcd.create_dataset("perturbed_geometries", data=_enc(np.full((2, 2, 4, 3), 0.5)))
+
+    for engine in ("desmond", "openMM"):
+        eg = grp.create_group(engine)
+        gp = eg.create_group("gas_phase")
+        gp.create_dataset("unperturbed_solute_geometry", data=_enc(np.zeros((4, 3))))
+        gp.create_dataset("unperturbed_dft_shieldings", data=_enc(GAS[0]))
+        gp.create_dataset("unperturbed_nn_shieldings", data=_enc(NN_GAS))
+        sv = eg.create_group("chloroform")
+        # (n_frames=3, n_atoms=4, 2) -> [:, :, 0] isolated, [:, :, 1] solvated
+        iso = np.broadcast_to(GAS[0], (3, 4)).copy()
+        solv = iso + 0.2                                  # +0.2 ppm solvent shift
+        solv[2, 0] = np.nan                               # exercise the missing-value path
+        dft = np.stack([iso, solv], axis=-1)
+        nn = np.stack([np.broadcast_to(NN_GAS, (3, 4)).copy(),
+                       np.broadcast_to(NN_GAS, (3, 4)).copy() + 0.2], axis=-1)
+        sv.create_dataset("perturbed_dft_shieldings", data=_enc(dft))
+        sv.create_dataset("perturbed_nn_shieldings", data=_enc(nn))
+        # solvated-cluster coordinates: 4 solute atoms + one 3-atom solvent block = 7
+        sv.create_dataset("perturbed_ensemble_geometries", data=_enc(np.full((3, 7, 3), 1.25)))
+
+
+@pytest.fixture(scope="module")
+def fixture(tmp_path_factory):
+    d = tmp_path_factory.mktemp("d22")
+    hdf5 = str(d / "mini.hdf5")
+    xlsx = str(d / "mini_experimental.xlsx")
+    with h5py.File(hdf5, "w") as f:
+        f.create_dataset("conventional_nmr_method_names",
+                         data=np.array(METHOD_NAMES, dtype=h5py.string_dtype()))
+        _write_solute(f.create_group("solutes").create_group(SOLUTE))
+        # top-level solvent topology: one 3-atom solvent block per engine ordering
+        cf = f.create_group("solvents").create_group("chloroform")
+        cf.create_dataset("desmond_atomic_numbers", data=np.array([6, 1, 17], np.uint8))
+        cf.create_dataset("openMM_atomic_numbers", data=np.array([6, 1, 17], np.uint8))
+    pd.DataFrame({
+        "solute": [SOLUTE, SOLUTE],
+        "site": ["Ha_H", "Ca_C"],
+        "atom_numbers": ["1,3", "2,4"],
+        "chloroform": [2.10, 30.0],
+    }).to_excel(xlsx, index=False)
+    return hdf5, xlsx
+
+
+# ---- the encode/decode round-trip (the one new piece of logic) ----
+def test_decode_scaling_and_missing():
+    enc = np.array([2070700, _MISSING_MARKER, -150000], np.int32)
+    out = _decode_fixed_point(enc)
+    assert out[0] == pytest.approx(207.07)
+    assert np.isnan(out[1])
+    assert out[2] == pytest.approx(-15.0)
+
+
+def test_decode_passes_floats_through():
+    f = np.array([1.5, 2.5], np.float64)
+    np.testing.assert_array_equal(_decode_fixed_point(f), f)
+
+
+# ---- the full reader against the synthetic file ----
+EXPECTED_COLUMNS = ["experimental", "stationary", "qcd", "pcm",
+                    "desmond", "openMM", "desmond_vib", "openMM_vib"]
+EXPECTED_INDEX = ["solute", "sap_geometry_type", "sap_nmr_method",
+                  "sap_basis", "nucleus", "site", "solvent"]
+
+
+def test_dft_loader_shape_and_decode(fixture):
+    hdf5, xlsx = fixture
+    df = load_delta22_dft_data(hdf5, xlsx, verbose=False)
+    assert list(df.index.names) == EXPECTED_INDEX
+    assert list(df.columns) == EXPECTED_COLUMNS
+    assert len(df) > 0
+    row = df.reset_index().query(
+        "site == 'Ha_H' and solvent == 'chloroform' and sap_geometry_type == 'aimnet2'"
+    )
+    assert len(row) == 1
+    assert row["stationary"].iloc[0] == pytest.approx(EXP_STATIONARY_HA_AIMNET2, abs=1e-3)
+    assert row["pcm"].iloc[0] == pytest.approx(EXP_PCM_HA_AIMNET2, abs=1e-3)
+    # explicit correction is +0.2 ppm; the seeded missing frame must not break it
+    assert row["desmond"].iloc[0] == pytest.approx(0.2, abs=1e-3)
+    assert np.isfinite(row["desmond"].iloc[0])
+
+
+def test_nn_loader_shape(fixture):
+    hdf5, xlsx = fixture
+    df = load_delta22_nn_data(hdf5, xlsx, verbose=False)
+    assert list(df.columns) == EXPECTED_COLUMNS
+    assert len(df) > 0
+    assert df.index.get_level_values("sap_nmr_method").unique().tolist() == ["MagNET"]
+
+
+def test_load_perturbed_shieldings(fixture):
+    hdf5, _ = fixture
+    raw = R.load_perturbed_shieldings(hdf5, SOLUTE, "chloroform", "openMM", "dft")
+    assert raw.shape == (3, 4, 2)                 # (n_frames, n_atoms, [isolated, solvated])
+    # the fixture set solvated = isolated + 0.2 with one missing value seeded
+    corr = raw[:, :, 1] - raw[:, :, 0]
+    assert np.nanmax(np.abs(corr - 0.2)) < 1e-3
+    assert np.isnan(raw[2, 0, 1])                 # the seeded missing solvated value
+
+
+def test_nn_loader_exclude_solutes(fixture):
+    hdf5, xlsx = fixture
+    full = load_delta22_nn_data(hdf5, xlsx, verbose=False)
+    assert SOLUTE in full.index.get_level_values("solute")
+    excluded = load_delta22_nn_data(hdf5, xlsx, exclude_solutes=[SOLUTE], verbose=False)
+    assert SOLUTE not in excluded.index.get_level_values("solute")
+    assert len(excluded) == 0   # the fixture has only the one solute
+
+
+def test_timings_loaders(fixture):
+    hdf5, _ = fixture
+    dft = load_delta22_dft_timings(hdf5)
+    nn = load_delta22_nn_timings(hdf5)
+    for t in (dft, nn):
+        assert list(t.columns) == ["geometry_time", "nmr_time", "total_time"]
+        assert (t["total_time"] >= 0).all()
+
+
+# ---- geometry + topology accessors (the raw stored arrays the loaders don't expose) ----
+def test_solutes_listing(fixture):
+    hdf5, _ = fixture
+    assert R.load_solutes(hdf5) == [SOLUTE]
+
+
+def test_geometry_accessors_shape_and_decode(fixture):
+    hdf5, _ = fixture
+    sg = R.load_stationary_geometries(hdf5, SOLUTE)
+    assert set(sg) == {"aimnet2", "pbe0_tz"}
+    assert sg["aimnet2"].shape == (4, 3)
+    q = R.load_qcd_geometries(hdf5, SOLUTE)
+    assert q["unperturbed"].shape == (4, 3)
+    assert q["perturbed"].shape == (2, 2, 4, 3)
+    assert q["unperturbed"][0, 0] == pytest.approx(0.5)        # coordinates decode from fixed point
+    assert q["perturbed"][0, 0, 0, 0] == pytest.approx(0.5)
+    e = R.load_explicit_ensemble_geometries(hdf5, SOLUTE, "chloroform", "openMM")
+    assert e.shape == (3, 7, 3)                                # 4 solute + 3 solvent atoms
+    assert e[0, 0, 0] == pytest.approx(1.25)
+    gs = R.load_explicit_solute_geometry(hdf5, SOLUTE, "desmond")
+    assert gs.shape == (4, 3)
+
+
+def test_solvent_atomic_numbers(fixture):
+    hdf5, _ = fixture
+    for engine in ("desmond", "openMM"):
+        an = R.load_solvent_atomic_numbers(hdf5, "chloroform", engine)
+        assert list(an) == [6, 1, 17]
+
+
+# ---------- opt-in: real file, if present ----------
+REAL = os.path.join(os.path.dirname(__file__), "delta22.hdf5")
+REAL_XLSX = os.path.join(os.path.dirname(__file__), "delta22_experimental.xlsx")
+
+
+@pytest.mark.skipif(not (os.path.exists(REAL) and os.path.exists(REAL_XLSX)),
+                    reason="real delta22.hdf5 / experimental xlsx not present")
+def test_real_file_smoke():
+    with h5py.File(REAL, "r") as f:
+        assert len(f["solutes"].keys()) == 22
+        assert len(f["conventional_nmr_method_names"]) == 473
+        any_solute = next(iter(f["solutes"].keys()))
+        desmond = [k for k in f["solutes"][any_solute]["desmond"].keys() if k != "gas_phase"]
+        assert len(desmond) == 12
+    df = load_delta22_dft_data(REAL, REAL_XLSX, verbose=False)
+    assert len(df) > 0
+    assert list(df.columns) == EXPECTED_COLUMNS
+    # the geometry + topology accessors reach the real stored arrays and decode to angstroms
+    sol = R.load_solutes(REAL)[0]
+    sg = R.load_stationary_geometries(REAL, sol)
+    assert sg["aimnet2"].ndim == 2 and sg["aimnet2"].shape[1] == 3
+    assert np.nanmax(np.abs(sg["aimnet2"])) < 100.0
+    an = R.load_solvent_atomic_numbers(REAL, "chloroform", "openMM")
+    assert an.ndim == 1 and len(an) > 0

data/delta50/decode_delta50.py

@@ -0,0 +1,118 @@
+"""
+decode_delta50.py - reader/decoder for the delta50 HDF5 dataset.
+
+DELTA50 is a small benchmark of 50 common organic molecules (reference 47 in the paper). This
+file holds MagNET's predicted NMR shieldings for each molecule on its AIMNet2 geometry, at three
+model chemistries:
+
+  nn_magnet_zero  MagNET-Zero (WP04 for 1H, wB97X-D for 13C; pcSseg-2; 20-pass), the gas-phase
+                  triple-zeta-quality prediction
+  nn_b3lyp        the gas-phase component of the MagNET-PCM correction (B3LYP/pcSseg-2)
+  nn_b3lyp_pcm    the chloroform component of the MagNET-PCM correction (B3LYP/pcSseg-2 + PCM)
+
+The MagNET-PCM implicit-solvent correction for a molecule is nn_b3lyp_pcm minus nn_b3lyp. The
+paper uses this set to probe outliers (nitromethane, nitroethane, and t-butyl nitrate are the
+notable ones).
+
+Storage format:
+  - Per-atom arrays are concatenated across molecules. Molecule i has `n_atoms[i]` atoms and owns
+    the rows [start[i], start[i+1]) where start is the cumulative sum of `n_atoms`.
+  - coordinates and shieldings are int32 fixed point: physical = stored * 1e-4, so reconstruction
+    error is <= 5e-5 (Angstrom or ppm). A missing value would be the marker -2147483648 -> NaN.
+  - molecules are identified by name (`molecule_names`).
+
+Requires: Python >= 3.7, numpy, h5py. No other dependencies.
+"""
+from __future__ import annotations
+import numpy as np
+import h5py
+
+_SCALE = 1e-4
+_MARKER = -2147483648
+
+
+def _decode(values: np.ndarray) -> np.ndarray:
+    """int32 fixed-point -> float64 physical units; marker -> NaN. No-op on float input."""
+    values = np.asarray(values)
+    if not np.issubdtype(values.dtype, np.integer):
+        return np.asarray(values, dtype=np.float64)
+    out = values.astype(np.float64) * _SCALE
+    out[values == _MARKER] = np.nan
+    return out
+
+
+class Delta50:
+    """Reader for delta50.hdf5. Use as a context manager; index molecules by position or by name."""
+
+    def __init__(self, path: str):
+        self.f = h5py.File(path, "r")
+        self.molecule_names = [str(n) for n in self.f["molecule_names"].asstr()[:]]
+        self.n_molecules = int(self.f.attrs["n_molecules"])
+        self._n_atoms = self.f["n_atoms"][:].astype(np.int64)
+        self._start = np.empty(self.n_molecules + 1, np.int64)
+        self._start[0] = 0
+        np.cumsum(self._n_atoms, out=self._start[1:])
+        self._name_to_index = {n: i for i, n in enumerate(self.molecule_names)}
+
+    def _check(self, index: int):
+        if not 0 <= index < self.n_molecules:
+            raise IndexError(f"molecule index {index} out of range [0, {self.n_molecules})")
+
+    def index_of(self, name: str) -> int:
+        """Position of the molecule named `name`; raises KeyError if absent."""
+        if name not in self._name_to_index:
+            raise KeyError(f"molecule {name!r} not in delta50")
+        return self._name_to_index[name]
+
+    def molecule(self, index: int) -> dict:
+        """Return one molecule's data (physical units) by position.
+
+        Keys: name; atomic_numbers (n,); coordinates (n, 3 Angstrom); and the three predicted
+        shieldings (nn_magnet_zero, nn_b3lyp, nn_b3lyp_pcm), each (n,) ppm.
+        """
+        self._check(index)
+        sl = slice(int(self._start[index]), int(self._start[index + 1]))
+        return {
+            "name": self.molecule_names[index],
+            "atomic_numbers": self.f["atomic_numbers"][sl],
+            "coordinates": _decode(self.f["coordinates"][sl]),
+            "nn_magnet_zero": _decode(self.f["nn_magnet_zero"][sl]),
+            "nn_b3lyp": _decode(self.f["nn_b3lyp"][sl]),
+            "nn_b3lyp_pcm": _decode(self.f["nn_b3lyp_pcm"][sl]),
+        }
+
+    def molecule_by_name(self, name: str) -> dict:
+        """Return one molecule's data (see `molecule`) by name."""
+        return self.molecule(self.index_of(name))
+
+    def pcm_correction(self, index: int) -> np.ndarray:
+        """The MagNET-PCM implicit-solvent correction (nn_b3lyp_pcm minus nn_b3lyp), (n,) ppm."""
+        m = self.molecule(index)
+        return m["nn_b3lyp_pcm"] - m["nn_b3lyp"]
+
+    def close(self):
+        """Close the underlying HDF5 file."""
+        self.f.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        self.close()
+
+    def __len__(self):
+        return self.n_molecules
+
+    def __repr__(self):
+        return f"<Delta50: {self.n_molecules} molecules>"
+
+
+if __name__ == "__main__":
+    import sys
+    path = sys.argv[1] if len(sys.argv) > 1 else "delta50.hdf5"
+    with Delta50(path) as ds:
+        print(ds)
+        m = ds.molecule(ds.index_of("nitromethane") if "nitromethane" in ds.molecule_names else 0)
+        c = m["atomic_numbers"] == 6
+        print(f"  {m['name']}: {len(m['atomic_numbers'])} atoms")
+        print(f"    13C MagNET-Zero[:3] = {np.round(m['nn_magnet_zero'][c][:3], 3)} ppm")

data/delta50/test_delta50.py

@@ -0,0 +1,103 @@
+"""Tests for data/delta50/decode_delta50.py.
+
+A tiny synthetic file in the stored int32 format exercises the reader, the per-molecule slicing,
+the fixed-point decode, and the PCM-correction helper, without the real file. An opt-in smoke test
+runs against the real file when present.
+"""
+import os
+import sys
+
+import numpy as np
+import pytest
+import h5py
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, HERE)
+
+import decode_delta50 as D  # noqa: E402
+
+SCALE = 1e-4
+MARKER = -2147483648
+REAL = os.path.join(HERE, "delta50.hdf5")
+
+
+def _encode(values):
+    v = np.asarray(values, dtype=np.float64)
+    out = np.round(v / SCALE).astype(np.int64)
+    out[~np.isfinite(v)] = MARKER
+    return out.astype(np.int32)
+
+
+def _build_synthetic(path):
+    """Two molecules (3 and 5 atoms) in the stored int32 format, plus the decoded truth to check against."""
+    rng = np.random.default_rng(0)
+    names = ["alpha", "beta"]
+    specs = [3, 5]  # n_atoms
+    n_atoms = np.array(specs, dtype=np.int32)
+    an, coords, zero, b3, pcm, truth = [], [], [], [], [], []
+    for na in specs:
+        a = rng.integers(1, 9, size=na).astype(np.uint8)
+        c = rng.normal(0, 2, size=(na, 3))
+        z = rng.normal(100, 30, size=na)
+        g = rng.normal(100, 30, size=na)
+        p = g + rng.normal(0, 0.2, size=na)
+        an.append(a); coords.append(_encode(c)); zero.append(_encode(z))
+        b3.append(_encode(g)); pcm.append(_encode(p))
+        truth.append((a, c, z, g, p))
+    opts = dict(compression="gzip", shuffle=True)
+    with h5py.File(path, "w") as f:
+        f.attrs["n_molecules"] = 2
+        f.attrs["scale"] = SCALE
+        f.attrs["missing_value_marker"] = MARKER
+        f.create_dataset("molecule_names", data=np.array(names, dtype=object),
+                         dtype=h5py.string_dtype("utf-8"))
+        f.create_dataset("n_atoms", data=n_atoms, **opts)
+        f.create_dataset("atomic_numbers", data=np.concatenate(an), **opts)
+        f.create_dataset("coordinates", data=np.concatenate(coords), **opts)
+        f.create_dataset("nn_magnet_zero", data=np.concatenate(zero), **opts)
+        f.create_dataset("nn_b3lyp", data=np.concatenate(b3), **opts)
+        f.create_dataset("nn_b3lyp_pcm", data=np.concatenate(pcm), **opts)
+    return names, specs, truth
+
+
+def test_reader_and_decode(tmp_path):
+    path = str(tmp_path / "tiny.hdf5")
+    names, specs, truth = _build_synthetic(path)
+    with D.Delta50(path) as ds:
+        assert len(ds) == 2
+        assert ds.molecule_names == names
+        for i, na in enumerate(specs):
+            m = ds.molecule(i)
+            a, c, z, g, p = truth[i]
+            assert m["name"] == names[i]
+            assert np.array_equal(m["atomic_numbers"], a)
+            assert m["coordinates"].shape == (na, 3)
+            assert np.abs(m["coordinates"] - c).max() <= 5e-5
+            assert np.abs(m["nn_magnet_zero"] - z).max() <= 5e-5
+            assert np.allclose(ds.pcm_correction(i), p - g, atol=1e-4)
+
+
+def test_lookup_and_errors(tmp_path):
+    path = str(tmp_path / "tiny.hdf5")
+    names, _, _ = _build_synthetic(path)
+    with D.Delta50(path) as ds:
+        assert ds.molecule_by_name("beta")["name"] == "beta"
+        with pytest.raises(KeyError):
+            ds.index_of("missing")
+        with pytest.raises(IndexError):
+            ds.molecule(9)
+
+
+def test_decode_marker_to_nan():
+    out = D._decode(np.array([0, 10000, MARKER], dtype=np.int32))
+    assert out[1] == pytest.approx(1.0) and np.isnan(out[2])
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real delta50.hdf5 not present")
+def test_real_file_smoke():
+    with D.Delta50(REAL) as ds:
+        assert len(ds) == 50
+        assert "nitromethane" in ds.molecule_names
+        m = ds.molecule_by_name("nitromethane")
+        assert np.isfinite(m["nn_magnet_zero"]).all()
+        assert ds.pcm_correction(ds.index_of("nitromethane")).shape == m["atomic_numbers"].shape

data/dft8k/dft8k_reader.py

@@ -0,0 +1,236 @@
+"""
+dft8k_reader.py - reader/decoder for the dft8k HDF5 dataset.
+
+DFT8k is an external NMR shielding benchmark of 7111 organic molecules (Guan, Paton and
+co-workers; reference 14 in the paper), used to evaluate the MagNET models out of sample.
+The file has TWO groups, one per level of theory. Both groups cover the SAME 7111 molecules
+and use the SAME integer molecule ids, but in different row order, so match a molecule across
+the two groups by id (use `molecule_by_id`), not by row position.
+
+  /aimnet2_wp04_wb97xd_pcsseg2   AIMNet2 geometries; per atom, three DFT shieldings plus the
+                                 two MagNET predictions. This is the file behind the paper's
+                                 DFT8k panels.
+  /b3lyp_pbe0_pcsseg1            B3LYP/pcSseg-1 geometries; three geometries per molecule with
+                                 one PBE0/pcSseg-1 shielding each, plus symmetry-equivalence
+                                 groups and per-atom averaging weights. This is the format the
+                                 model's own training/eval loader consumes.
+
+Storage format:
+  - Per-atom arrays are concatenated across the molecules in a group. Each molecule's atom
+    count is in `n_atoms`; molecule i owns the rows [start[i], start[i+1]) where start is the
+    cumulative sum of `n_atoms` (reconstructed once, on open).
+  - coordinates and shieldings are int32 fixed-point: physical = stored * scale (scale = 1e-4),
+    so reconstruction error is <= 5e-5 (Angstrom or ppm). A missing value (an atom a model does
+    not predict, or an absent shielding) is the marker -2147483648 and decodes to NaN.
+  - The molecule's integer `molecule_ids` entry is its identifier.
+
+Requires: Python >= 3.7, numpy, h5py. No other dependencies.
+"""
+from __future__ import annotations
+import numpy as np
+import h5py
+
+_SCALE = 1e-4
+_MARKER = -2147483648
+
+
+def _decode(values: np.ndarray) -> np.ndarray:
+    """int32 fixed-point -> float64 physical units; marker -> NaN. No-op on float input."""
+    values = np.asarray(values)
+    if not np.issubdtype(values.dtype, np.integer):
+        return np.asarray(values, dtype=np.float64)
+    out = values.astype(np.float64) * _SCALE
+    out[values == _MARKER] = np.nan
+    return out
+
+
+class _Group:
+    """Common molecule-indexing logic shared by both groups."""
+
+    def __init__(self, grp: h5py.Group):
+        self.g = grp
+        self.name = grp.name.lstrip("/")
+        self.molecule_ids = grp["molecule_ids"][:]
+        self.n_molecules = int(grp.attrs["n_molecules"])
+        self.n_atoms = int(grp.attrs["n_atoms"])
+        self.level_of_theory = str(grp.attrs["level_of_theory"])
+        self.geometry = str(grp.attrs["geometry"])
+        na = grp["n_atoms"][:].astype(np.int64)
+        self._start = np.empty(self.n_molecules + 1, np.int64)
+        self._start[0] = 0
+        np.cumsum(na, out=self._start[1:])
+        self._id_to_index = {int(m): i for i, m in enumerate(self.molecule_ids)}
+
+    def _check(self, index: int):
+        if not 0 <= index < self.n_molecules:
+            raise IndexError(f"molecule index {index} out of range [0, {self.n_molecules})")
+
+    def _rows(self, index: int) -> slice:
+        return slice(int(self._start[index]), int(self._start[index + 1]))
+
+    def index_of(self, molecule_id: int) -> int:
+        """Return the row position of a molecule id within this group."""
+        if int(molecule_id) not in self._id_to_index:
+            raise KeyError(f"molecule id {molecule_id} not in group {self.name}")
+        return self._id_to_index[int(molecule_id)]
+
+    def molecule_by_id(self, molecule_id: int) -> dict:
+        """Look up a molecule by its id instead of its row position."""
+        return self.molecule(self.index_of(molecule_id))
+
+    def __len__(self):
+        return self.n_molecules
+
+
+class AIMNet2Group(_Group):
+    """/aimnet2_wp04_wb97xd_pcsseg2 - AIMNet2 geometries, DFT and MagNET shieldings per atom."""
+
+    def __init__(self, grp: h5py.Group):
+        super().__init__(grp)
+        # one SMILES per molecule. 6780 of 7111 are real; the rest are the literal "none".
+        # Older files may lack the field, so guard for it.
+        self._smiles = grp["smiles"].asstr()[:] if "smiles" in grp else None
+
+    def smiles(self, index: int) -> str:
+        """Return molecule i's SMILES string (the literal "none" if it was not provided)."""
+        self._check(index)
+        if self._smiles is None:
+            return "none"
+        return str(self._smiles[index])
+
+    def molecule(self, index: int) -> dict:
+        """Return one molecule's atom data (physical units) by position.
+
+        Keys: id, smiles; atomic_numbers, coordinates (n,3 Angstrom); the three DFT shieldings
+        (pbe0_pcSseg1, wp04_pcSseg2, wb97xd_pcSseg2; each (n,) ppm); and the two MagNET
+        predictions (nn_wp04_pcSseg2, nn_wb97xd_pcSseg2; (n,) ppm). A MagNET prediction is
+        NaN at atoms that model does not predict (WP04 is the 1H model, wB97X-D the 13C).
+        """
+        self._check(index)
+        sl = self._rows(index)
+        return {
+            "id": int(self.molecule_ids[index]),
+            "smiles": self.smiles(index),
+            "atomic_numbers": self.g["atomic_numbers"][sl],
+            "coordinates": _decode(self.g["coordinates"][sl]),
+            "pbe0_pcSseg1": _decode(self.g["shielding_pbe0_pcSseg1"][sl]),
+            "wp04_pcSseg2": _decode(self.g["shielding_wp04_pcSseg2"][sl]),
+            "wb97xd_pcSseg2": _decode(self.g["shielding_wb97xd_pcSseg2"][sl]),
+            "nn_wp04_pcSseg2": _decode(self.g["nn_shielding_wp04_pcSseg2"][sl]),
+            "nn_wb97xd_pcSseg2": _decode(self.g["nn_shielding_wb97xd_pcSseg2"][sl]),
+        }
+
+    def all_shieldings(self) -> dict:
+        """Every atom in the dataset as flat arrays (physical units), for whole-dataset analyses
+        like the MagNET-vs-DFT residual histograms.
+
+        Keys: atomic_numbers (n_atoms_total,); the three DFT shieldings (pbe0_pcSseg1,
+        wp04_pcSseg2, wb97xd_pcSseg2) and the two MagNET predictions (nn_wp04_pcSseg2,
+        nn_wb97xd_pcSseg2), each (n_atoms_total,) ppm. A MagNET prediction is NaN at atoms that
+        model does not predict (WP04 is the 1H model, wB97X-D the 13C). To get one nucleus' residual,
+        difference the DFT and MagNET column for that model at that element and drop the NaNs.
+        """
+        return {
+            "atomic_numbers": self.g["atomic_numbers"][:],
+            "pbe0_pcSseg1": _decode(self.g["shielding_pbe0_pcSseg1"][:]),
+            "wp04_pcSseg2": _decode(self.g["shielding_wp04_pcSseg2"][:]),
+            "wb97xd_pcSseg2": _decode(self.g["shielding_wb97xd_pcSseg2"][:]),
+            "nn_wp04_pcSseg2": _decode(self.g["nn_shielding_wp04_pcSseg2"][:]),
+            "nn_wb97xd_pcSseg2": _decode(self.g["nn_shielding_wb97xd_pcSseg2"][:]),
+        }
+
+    def __repr__(self):
+        return (f"<AIMNet2Group /{self.name}: {self.n_molecules:,} molecules, "
+                f"{self.n_atoms:,} atoms>")
+
+
+class B3LYPGroup(_Group):
+    """/b3lyp_pbe0_pcsseg1 - three geometries per molecule, symmetry groups and weights."""
+
+    def __init__(self, grp: h5py.Group):
+        super().__init__(grp)
+        self.n_geometries = int(grp.attrs["n_geometries"])
+        # symmetry groups are stored flat: n_groups per molecule, each group's length, then
+        # all member atom indices. Reconstruct the per-molecule offsets once, on open.
+        self._sym_n = grp["symmetric_atoms_n_groups"][:].astype(np.int64)
+        self._sym_len = grp["symmetric_atoms_group_lengths"][:].astype(np.int64)
+        self._sym_mem = grp["symmetric_atoms_members"][:]
+        self._group_start = np.empty(self.n_molecules + 1, np.int64)
+        self._group_start[0] = 0
+        np.cumsum(self._sym_n, out=self._group_start[1:])
+        self._member_start = np.empty(self._sym_len.shape[0] + 1, np.int64)
+        self._member_start[0] = 0
+        np.cumsum(self._sym_len, out=self._member_start[1:])
+
+    def symmetric_atoms(self, index: int) -> list:
+        """Return molecule i's symmetry-equivalence groups as a list of lists of atom indices.
+
+        The indices are 0-based positions into this molecule's atom arrays (atomic_numbers,
+        coordinates), so a group can be used to index those arrays directly.
+        """
+        self._check(index)
+        groups = []
+        for gi in range(int(self._group_start[index]), int(self._group_start[index + 1])):
+            a, b = int(self._member_start[gi]), int(self._member_start[gi + 1])
+            groups.append([int(x) for x in self._sym_mem[a:b]])
+        return groups
+
+    def molecule(self, index: int) -> dict:
+        """Return one molecule's data (physical units) by position.
+
+        Keys: id, atomic_numbers; coordinates (n_geometries, n, 3 Angstrom); shielding
+        (n_geometries, n, ppm), where geometry 0 is a reference and 1, 2 are perturbations;
+        weights (n,) per-atom symmetry-averaging weight; symmetric_atoms (list of lists of
+        equivalent atoms, each a 0-based index into this molecule's atom arrays).
+        """
+        self._check(index)
+        sl = self._rows(index)
+        return {
+            "id": int(self.molecule_ids[index]),
+            "atomic_numbers": self.g["atomic_numbers"][sl],
+            "coordinates": _decode(self.g["coordinates"][:, sl, :]),
+            "shielding": _decode(self.g["shielding"][:, sl]),
+            "weights": self.g["weights"][sl].astype(np.float64),
+            "symmetric_atoms": self.symmetric_atoms(index),
+        }
+
+    def __repr__(self):
+        return (f"<B3LYPGroup /{self.name}: {self.n_molecules:,} molecules, "
+                f"{self.n_atoms:,} atoms, {self.n_geometries} geometries>")
+
+
+class DFT8k:
+    """Reader for dft8k.hdf5. Use as a context manager; access the two groups by attribute."""
+
+    def __init__(self, path: str):
+        self.f = h5py.File(path, "r")
+        self.aimnet2 = AIMNet2Group(self.f["aimnet2_wp04_wb97xd_pcsseg2"])
+        self.b3lyp = B3LYPGroup(self.f["b3lyp_pbe0_pcsseg1"])
+
+    def close(self):
+        """Close the underlying HDF5 file."""
+        self.f.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        self.close()
+
+    def __repr__(self):
+        return f"<DFT8k: {self.aimnet2!r}, {self.b3lyp!r}>"
+
+
+if __name__ == "__main__":
+    import sys
+    path = sys.argv[1] if len(sys.argv) > 1 else "dft8k.hdf5"
+    with DFT8k(path) as ds:
+        print(ds)
+        m = ds.aimnet2.molecule(0)
+        h = m["atomic_numbers"] == 1
+        print(f"  aimnet2 mol id={m['id']}: {len(m['atomic_numbers'])} atoms  smiles={m['smiles']}")
+        print(f"    1H WP04 DFT[:3]={np.round(m['wp04_pcSseg2'][h][:3], 3)}  "
+              f"MagNET[:3]={np.round(m['nn_wp04_pcSseg2'][h][:3], 3)}")
+        b = ds.b3lyp.molecule(0)
+        print(f"  b3lyp   mol id={b['id']}: coordinates {b['coordinates'].shape}, "
+              f"{len(b['symmetric_atoms'])} symmetry groups")

data/dft8k/test_dft8k.py

@@ -0,0 +1,224 @@
+"""
+Tests for the dft8k reader (dft8k_reader.py).
+
+Most tests build a tiny synthetic dft8k.hdf5 in the exact on-disk format (so the fixture
+doubles as a spec of what the reader expects) and need no large data. One opt-in test runs
+against the real dft8k.hdf5 if it is present (skipped otherwise, e.g. in CI).
+
+Run:  pytest test_dft8k.py -q
+Requires: pytest, numpy, h5py.
+"""
+import os
+import numpy as np
+import h5py
+import pytest
+
+from dft8k_reader import DFT8k, _decode, _MARKER
+
+SCALE = 1e4
+
+
+def _enc(arr):
+    """Encode floats the way the build script does: int32 = round(value * 1e4), NaN -> marker."""
+    a = np.asarray(arr, dtype=np.float64)
+    nan = np.isnan(a)
+    out = np.rint(np.where(nan, 0.0, a) * SCALE).astype(np.int32)
+    out[nan] = _MARKER
+    return out
+
+
+# ---------------------------------------------------------------------------
+# fixed-point decoding
+# ---------------------------------------------------------------------------
+def test_decode_round_trip_and_marker():
+    values = np.array([0.0, 207.07, -355.5, np.nan, 875.0], dtype=np.float64)
+    decoded = _decode(_enc(values))
+    assert np.isnan(decoded[3])
+    np.testing.assert_allclose(decoded[~np.isnan(decoded)],
+                               values[~np.isnan(values)], atol=5e-5)
+
+
+def test_decode_passes_through_float():
+    # a plain-decimal array must read back unchanged (decode is a no-op on floats)
+    values = np.array([1.0, 2.5, -7.25], dtype=np.float64)
+    np.testing.assert_array_equal(_decode(values), values)
+
+
+def test_decode_handles_int64():
+    # any integer dtype decodes the same way as int32
+    vals = np.array([10000, -3550000, _MARKER], dtype=np.int64)
+    out = _decode(vals)
+    np.testing.assert_allclose(out[:2], [1.0, -355.0])
+    assert np.isnan(out[2])
+
+
+# ---------------------------------------------------------------------------
+# synthetic dft8k.hdf5 in the real two-group layout
+# ---------------------------------------------------------------------------
+# two molecules: ids 100 and 200, with 2 and 3 atoms (atoms are H, C, ...)
+A_IDS = [100, 200]
+A_COUNTS = [2, 3]
+A_Z = [1, 6, 6, 1, 7]          # mol100: H C ; mol200: C H N
+# per-atom DFT shieldings (concatenated): pbe0, wp04, wb97xd
+A_PBE0 = [30.0, 150.0, 151.0, 31.0, 200.0]
+A_WP04 = [30.5, 150.5, 151.5, 31.5, 200.5]   # the 1H reference
+A_WB97 = [30.7, 150.7, 151.7, 31.7, 200.7]   # the 13C reference
+# MagNET predicts H with the WP04 model and C with the wB97X-D model; other atoms blank
+A_NN_WP04 = [30.6, np.nan, np.nan, 31.6, np.nan]    # only at H atoms (index 0, 3)
+A_NN_WB97 = [np.nan, 150.6, 151.6, np.nan, np.nan]  # only at C atoms (index 1, 2)
+
+
+def _make_synthetic_dft8k(path):
+    with h5py.File(path, "w") as f:
+        g = f.create_group("aimnet2_wp04_wb97xd_pcsseg2")
+        g.attrs["geometry"] = "AIMNet2"
+        g.attrs["level_of_theory"] = "PBE0/pcSseg-1; WP04/pcSseg-2; wB97X-D/pcSseg-2"
+        g.attrs["n_molecules"] = len(A_IDS)
+        g.attrs["n_atoms"] = sum(A_COUNTS)
+        g.attrs["scale"] = 1e-4
+        g.create_dataset("molecule_ids", data=np.array(A_IDS, np.uint32))
+        g.create_dataset("n_atoms", data=np.array(A_COUNTS, np.uint8))
+        # one SMILES per molecule; mol100 has a real one, mol200 carries the literal "none"
+        g.create_dataset("smiles", data=np.array(["CC", "none"], dtype=object),
+                         dtype=h5py.string_dtype(encoding="utf-8"))
+        g.create_dataset("atomic_numbers", data=np.array(A_Z, np.int8))
+        coords = np.arange(sum(A_COUNTS) * 3, dtype=np.float64).reshape(-1, 3) * 0.1
+        g.create_dataset("coordinates", data=_enc(coords))
+        g.create_dataset("shielding_pbe0_pcSseg1", data=_enc(A_PBE0))
+        g.create_dataset("shielding_wp04_pcSseg2", data=_enc(A_WP04))
+        g.create_dataset("shielding_wb97xd_pcSseg2", data=_enc(A_WB97))
+        g.create_dataset("nn_shielding_wp04_pcSseg2", data=_enc(A_NN_WP04))
+        g.create_dataset("nn_shielding_wb97xd_pcSseg2", data=_enc(A_NN_WB97))
+
+        # second group, SAME ids but reversed row order, three geometries per molecule
+        b = f.create_group("b3lyp_pbe0_pcsseg1")
+        b.attrs["geometry"] = "B3LYP/pcSseg-1"
+        b.attrs["level_of_theory"] = "PBE0/pcSseg-1 on B3LYP/pcSseg-1 geometries"
+        b.attrs["n_molecules"] = 2
+        b.attrs["n_atoms"] = 5
+        b.attrs["n_geometries"] = 3
+        b.attrs["scale"] = 1e-4
+        b.create_dataset("molecule_ids", data=np.array([200, 100], np.uint32))  # reversed
+        b.create_dataset("n_atoms", data=np.array([3, 2], np.uint8))
+        b.create_dataset("atomic_numbers", data=np.array([6, 1, 7, 1, 6], np.int8))
+        bcoords = np.arange(3 * 5 * 3, dtype=np.float64).reshape(3, 5, 3) * 0.01
+        b.create_dataset("coordinates", data=_enc(bcoords))
+        bshield = np.arange(3 * 5, dtype=np.float64).reshape(3, 5) + 100.0
+        b.create_dataset("shielding", data=_enc(bshield))
+        b.create_dataset("weights", data=np.array([1, 1, 1, 0.5, 0.5], np.float32))
+        # symmetry groups are 0-based atom indices, local to each molecule.
+        # mol200 (first row, 3 atoms): two groups [[1, 2]] and [[0]];
+        # mol100 (second row, 2 atoms): one group [[0, 1]].
+        b.create_dataset("symmetric_atoms_n_groups", data=np.array([2, 1], np.int32))
+        b.create_dataset("symmetric_atoms_group_lengths", data=np.array([2, 1, 2], np.int32))
+        b.create_dataset("symmetric_atoms_members", data=np.array([1, 2, 0, 0, 1], np.int32))
+
+
+def test_aimnet2_group(tmp_path):
+    p = tmp_path / "dft8k.hdf5"
+    _make_synthetic_dft8k(str(p))
+    with DFT8k(str(p)) as ds:
+        assert len(ds.aimnet2) == 2
+        m = ds.aimnet2.molecule(0)
+        assert m["id"] == 100
+        assert m["smiles"] == "CC"
+        assert ds.aimnet2.smiles(1) == "none"   # mol200 has no SMILES
+        np.testing.assert_array_equal(m["atomic_numbers"], [1, 6])
+        np.testing.assert_allclose(m["wp04_pcSseg2"], [30.5, 150.5], atol=5e-5)
+        # H atom (index 0) has a WP04 prediction; C atom (index 1) does not
+        assert not np.isnan(m["nn_wp04_pcSseg2"][0])
+        assert np.isnan(m["nn_wp04_pcSseg2"][1])
+        # C atom has a wB97X-D prediction; H atom does not
+        assert np.isnan(m["nn_wb97xd_pcSseg2"][0])
+        assert not np.isnan(m["nn_wb97xd_pcSseg2"][1])
+        np.testing.assert_allclose(m["nn_wb97xd_pcSseg2"][1], 150.6, atol=5e-5)
+
+
+def test_all_shieldings_flat(tmp_path):
+    p = tmp_path / "dft8k.hdf5"
+    _make_synthetic_dft8k(str(p))
+    with DFT8k(str(p)) as ds:
+        flat = ds.aimnet2.all_shieldings()
+        # one entry per atom across both molecules, in stored order
+        np.testing.assert_array_equal(flat["atomic_numbers"], [1, 6, 6, 1, 7])
+        np.testing.assert_allclose(flat["wp04_pcSseg2"], A_WP04, atol=5e-5)
+        # MagNET prediction is present only at the atoms its model covers
+        assert np.isnan(flat["nn_wp04_pcSseg2"]).tolist() == [False, True, True, False, True]
+        # the 1H residual (DFT minus MagNET) at the two H atoms is -0.1
+        Z = flat["atomic_numbers"]
+        h = (Z == 1) & np.isfinite(flat["nn_wp04_pcSseg2"])
+        np.testing.assert_allclose(flat["wp04_pcSseg2"][h] - flat["nn_wp04_pcSseg2"][h],
+                                   [-0.1, -0.1], atol=5e-5)
+
+
+def test_b3lyp_group_and_symmetry(tmp_path):
+    p = tmp_path / "dft8k.hdf5"
+    _make_synthetic_dft8k(str(p))
+    with DFT8k(str(p)) as ds:
+        b = ds.b3lyp
+        assert b.n_geometries == 3
+        m = b.molecule(0)               # first row is id 200
+        assert m["id"] == 200
+        assert m["coordinates"].shape == (3, 3, 3)   # (geometries, atoms, xyz)
+        assert m["shielding"].shape == (3, 3)
+        assert m["symmetric_atoms"] == [[1, 2], [0]]
+        m2 = b.molecule(1)              # id 100
+        assert m2["symmetric_atoms"] == [[0, 1]]
+        np.testing.assert_allclose(m2["weights"], [0.5, 0.5])
+        # the symmetry indices are valid 0-based positions into this molecule's atoms
+        for group in m["symmetric_atoms"] + m2["symmetric_atoms"]:
+            for atom_index in group:
+                assert 0 <= atom_index < len(m["atomic_numbers"]) or \
+                       0 <= atom_index < len(m2["atomic_numbers"])
+
+
+def test_match_across_groups_by_id(tmp_path):
+    p = tmp_path / "dft8k.hdf5"
+    _make_synthetic_dft8k(str(p))
+    with DFT8k(str(p)) as ds:
+        # rows are ordered differently in the two groups; matching by id must still work
+        a = ds.aimnet2.molecule_by_id(200)
+        b = ds.b3lyp.molecule_by_id(200)
+        assert a["id"] == b["id"] == 200
+        np.testing.assert_array_equal(a["atomic_numbers"], b["atomic_numbers"])
+
+
+# ---------------------------------------------------------------------------
+# opt-in test against the real file (skipped when absent, e.g. in CI)
+# ---------------------------------------------------------------------------
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_REAL = os.path.join(_HERE, "dft8k.hdf5")
+
+
+@pytest.mark.skipif(not os.path.exists(_REAL), reason="real dft8k.hdf5 not present")
+def test_real_dft8k():
+    with DFT8k(_REAL) as ds:
+        assert len(ds.aimnet2) == 7111
+        assert len(ds.b3lyp) == 7111
+        assert ds.b3lyp.n_geometries == 3
+        # the two groups carry the same id set
+        assert set(int(x) for x in ds.aimnet2.molecule_ids) == \
+               set(int(x) for x in ds.b3lyp.molecule_ids)
+        m = ds.aimnet2.molecule(0)
+        # SMILES: 6780 of the 7111 molecules have one; the rest carry the literal "none"
+        n_real_smiles = sum(1 for i in range(len(ds.aimnet2)) if ds.aimnet2.smiles(i) != "none")
+        assert n_real_smiles == 6780
+        assert isinstance(m["smiles"], str) and m["smiles"] != "none"
+        z = m["atomic_numbers"]
+        h = z == 1
+        c = z == 6
+        other = ~(h | c)
+        # the WP04 model predicts every hydrogen and nothing else; the wB97X-D model
+        # predicts every carbon and nothing else; non-target atoms read back blank
+        assert not np.isnan(m["nn_wp04_pcSseg2"][h]).any()
+        assert np.isnan(m["nn_wp04_pcSseg2"][~h]).all()
+        assert not np.isnan(m["nn_wb97xd_pcSseg2"][c]).any()
+        assert np.isnan(m["nn_wb97xd_pcSseg2"][~c]).all()
+        assert other.sum() == 0 or np.isnan(m["nn_wp04_pcSseg2"][other]).all()
+        # the same molecule resolves in the other group by id
+        b = ds.b3lyp.molecule_by_id(m["id"])
+        np.testing.assert_array_equal(m["atomic_numbers"], b["atomic_numbers"])
+        # every symmetry index is a valid 0-based position into this molecule's atoms
+        for group in b["symmetric_atoms"]:
+            for atom_index in group:
+                assert 0 <= atom_index < len(b["atomic_numbers"])

data/gdb_qcd/decode_gdb_qcd.py

@@ -0,0 +1,136 @@
+"""
+decode_gdb_qcd.py - reader/decoder for the gdb_qcd HDF5 dataset.
+
+gdb_qcd is a quasiclassical-dynamics benchmark: replicate molecular-dynamics trajectories for
+2461 small molecules drawn from GDB (the MagNET / sigma-shake test set). It is used to benchmark
+rovibrational corrections, which are computed by comparing NMR shieldings on the moving
+(displaced) geometries against the resting (stationary) geometry.
+
+For each molecule there are several trajectories; each trajectory has 33 frames; each frame gives
+every atom's position and its PBE0/pcSseg-1 NMR shielding. Trajectory counts vary by molecule
+(about 23 on average); the frame count is always 33. The resting (stationary) reference shieldings
+are not included here: they are the test-set entries of the MagNET training data (sigma-shake),
+matched by the same molecule id.
+
+Storage format:
+  - Per-atom arrays are concatenated across molecules. Molecule i has `n_atoms[i]` atoms and
+    `n_trajectories[i]` trajectories. Its atomic numbers occupy `n_atoms[i]` rows; its coordinates
+    and shieldings occupy `n_trajectories[i] * 33 * n_atoms[i]` rows, which reshape to
+    (n_trajectories, 33, n_atoms, ...). Offsets are reconstructed once, on open.
+  - coordinates and shieldings are int32 fixed point: physical = stored * 1e-4, so reconstruction
+    error is <= 5e-5 (Angstrom or ppm). A missing value would be the marker -2147483648 -> NaN
+    (none occur in this dataset).
+  - `molecule_ids` are integers and match the sigma-shake test-set ids.
+
+Requires: Python >= 3.7, numpy, h5py. No other dependencies.
+"""
+from __future__ import annotations
+import numpy as np
+import h5py
+
+_SCALE = 1e-4
+_MARKER = -2147483648
+
+
+def _decode(values: np.ndarray) -> np.ndarray:
+    """int32 fixed-point -> float64 physical units; marker -> NaN. No-op on float input."""
+    values = np.asarray(values)
+    if not np.issubdtype(values.dtype, np.integer):
+        return np.asarray(values, dtype=np.float64)
+    out = values.astype(np.float64) * _SCALE
+    out[values == _MARKER] = np.nan
+    return out
+
+
+class GDBQCD:
+    """Reader for gdb_qcd.hdf5. Use as a context manager; index molecules by position or by id."""
+
+    def __init__(self, path: str):
+        self.f = h5py.File(path, "r")
+        self.molecule_ids = self.f["molecule_ids"][:]
+        self.n_molecules = int(self.f.attrs["n_molecules"])
+        self.n_frames = int(self.f.attrs["n_frames"])
+        self.nmr_level_of_theory = str(self.f.attrs["nmr_level_of_theory"])
+        self._n_atoms = self.f["n_atoms"][:].astype(np.int64)
+        self._n_traj = self.f["n_trajectories"][:].astype(np.int64)
+        # atom-row offsets (atomic_numbers) and coordinate/shielding-row offsets
+        self._astart = np.empty(self.n_molecules + 1, np.int64)
+        self._astart[0] = 0
+        np.cumsum(self._n_atoms, out=self._astart[1:])
+        block = self._n_traj * self.n_frames * self._n_atoms
+        self._cstart = np.empty(self.n_molecules + 1, np.int64)
+        self._cstart[0] = 0
+        np.cumsum(block, out=self._cstart[1:])
+        self._id_to_index = {int(m): i for i, m in enumerate(self.molecule_ids)}
+
+    def _check(self, index: int):
+        if not 0 <= index < self.n_molecules:
+            raise IndexError(f"molecule index {index} out of range [0, {self.n_molecules})")
+
+    def index_of(self, molecule_id: int) -> int:
+        """Position of the molecule with this id; raises KeyError if absent."""
+        if int(molecule_id) not in self._id_to_index:
+            raise KeyError(f"molecule id {molecule_id} not in gdb_qcd")
+        return self._id_to_index[int(molecule_id)]
+
+    def n_trajectories(self, index: int) -> int:
+        """How many trajectories molecule i has."""
+        self._check(index)
+        return int(self._n_traj[index])
+
+    def molecule(self, index: int) -> dict:
+        """Return one molecule's trajectory data (physical units) by position.
+
+        Keys:
+          id               - integer molecule id (matches the sigma-shake test set)
+          atomic_numbers   - (n_atoms,) uint8
+          coordinates      - (n_trajectories, 33, n_atoms, 3) Angstrom
+          shieldings       - (n_trajectories, 33, n_atoms) ppm, PBE0/pcSseg-1
+        """
+        self._check(index)
+        na = int(self._n_atoms[index])
+        nt = int(self._n_traj[index])
+        asl = slice(int(self._astart[index]), int(self._astart[index + 1]))
+        csl = slice(int(self._cstart[index]), int(self._cstart[index + 1]))
+        coords = _decode(self.f["coordinates"][csl]).reshape(nt, self.n_frames, na, 3)
+        shields = _decode(self.f["shieldings"][csl]).reshape(nt, self.n_frames, na)
+        return {
+            "id": int(self.molecule_ids[index]),
+            "atomic_numbers": self.f["atomic_numbers"][asl],
+            "coordinates": coords,
+            "shieldings": shields,
+        }
+
+    def molecule_by_id(self, molecule_id: int) -> dict:
+        """Return one molecule's trajectory data (see `molecule`) by id."""
+        return self.molecule(self.index_of(molecule_id))
+
+    def close(self):
+        """Close the underlying HDF5 file."""
+        self.f.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        self.close()
+
+    def __len__(self):
+        return self.n_molecules
+
+    def __repr__(self):
+        return (f"<GDBQCD: {self.n_molecules:,} molecules, "
+                f"{int(self.f.attrs['n_trajectories']):,} trajectories, "
+                f"{self.n_frames} frames each, {self.nmr_level_of_theory}>")
+
+
+if __name__ == "__main__":
+    import sys
+    path = sys.argv[1] if len(sys.argv) > 1 else "gdb_qcd.hdf5"
+    with GDBQCD(path) as ds:
+        print(ds)
+        m = ds.molecule(0)
+        print(f"  molecule id={m['id']}: {len(m['atomic_numbers'])} atoms, "
+              f"coordinates {m['coordinates'].shape}, shieldings {m['shieldings'].shape}")
+        print(f"    trajectory 0, frame 0, atom 0: xyz={np.round(m['coordinates'][0,0,0], 4)}  "
+              f"shielding={m['shieldings'][0,0,0]:.4f} ppm")

data/gdb_qcd/test_gdb_qcd.py

@@ -0,0 +1,112 @@
+"""Tests for data/gdb_qcd/decode_gdb_qcd.py.
+
+A tiny synthetic file in the stored int32 format exercises the reader, the per-molecule block
+reconstruction (variable atoms and trajectories, fixed 33 frames), and the fixed-point decode,
+all without the multi-hundred-MB real file. An opt-in smoke test runs against the real file when
+it is present.
+"""
+import os
+import sys
+
+import numpy as np
+import pytest
+import h5py
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, HERE)
+
+import decode_gdb_qcd as D  # noqa: E402
+
+SCALE = 1e-4
+MARKER = -2147483648
+N_FRAMES = 33
+REAL = os.path.join(HERE, "gdb_qcd.hdf5")
+
+
+def _encode(values):
+    out = np.round(np.asarray(values, dtype=np.float64) / SCALE).astype(np.int64)
+    out[~np.isfinite(values)] = MARKER
+    return out.astype(np.int32)
+
+
+def _build_synthetic(path):
+    """Two molecules: 3 atoms / 2 trajectories, and 4 atoms / 1 trajectory."""
+    rng = np.random.default_rng(0)
+    specs = [(3, 2), (4, 1)]  # (n_atoms, n_trajectories)
+    mol_ids = np.array([11, 22], dtype=np.int64)
+    n_atoms = np.array([s[0] for s in specs], dtype=np.int32)
+    n_traj = np.array([s[1] for s in specs], dtype=np.int32)
+    atomic, coord_blocks, shield_blocks, truth = [], [], [], []
+    for na, nt in specs:
+        an = rng.integers(1, 9, size=na).astype(np.uint8)
+        coords = rng.normal(0, 2, size=(nt, N_FRAMES, na, 3))
+        shields = rng.normal(100, 30, size=(nt, N_FRAMES, na))
+        atomic.append(an)
+        coord_blocks.append(_encode(coords.reshape(-1, 3)))
+        shield_blocks.append(_encode(shields.reshape(-1)))
+        truth.append((an, coords, shields))
+    opts = dict(compression="gzip", shuffle=True)
+    with h5py.File(path, "w") as f:
+        f.attrs["n_molecules"] = 2
+        f.attrs["n_trajectories"] = int(n_traj.sum())
+        f.attrs["n_frames"] = N_FRAMES
+        f.attrs["nmr_level_of_theory"] = "PBE0/pcSseg-1"
+        f.attrs["scale"] = SCALE
+        f.attrs["missing_value_marker"] = MARKER
+        f.create_dataset("molecule_ids", data=mol_ids, **opts)
+        f.create_dataset("n_atoms", data=n_atoms, **opts)
+        f.create_dataset("n_trajectories", data=n_traj, **opts)
+        f.create_dataset("atomic_numbers", data=np.concatenate(atomic), **opts)
+        f.create_dataset("coordinates", data=np.concatenate(coord_blocks), **opts)
+        f.create_dataset("shieldings", data=np.concatenate(shield_blocks), **opts)
+    return mol_ids, specs, truth
+
+
+def test_reader_reconstructs_blocks(tmp_path):
+    path = str(tmp_path / "tiny.hdf5")
+    mol_ids, specs, truth = _build_synthetic(path)
+    with D.GDBQCD(path) as ds:
+        assert len(ds) == 2
+        assert ds.n_frames == 33
+        assert list(ds.molecule_ids) == list(mol_ids)
+        for i, (na, nt) in enumerate(specs):
+            m = ds.molecule(i)
+            an, coords, shields = truth[i]
+            assert m["id"] == int(mol_ids[i])
+            assert ds.n_trajectories(i) == nt
+            assert np.array_equal(m["atomic_numbers"], an)
+            assert m["coordinates"].shape == (nt, 33, na, 3)
+            assert m["shieldings"].shape == (nt, 33, na)
+            # fixed-point round-trip within the 5e-5 floor
+            assert np.abs(m["coordinates"] - coords).max() <= 5e-5
+            assert np.abs(m["shieldings"] - shields).max() <= 5e-5
+
+
+def test_lookup_by_id_and_errors(tmp_path):
+    path = str(tmp_path / "tiny.hdf5")
+    mol_ids, _, _ = _build_synthetic(path)
+    with D.GDBQCD(path) as ds:
+        assert ds.molecule_by_id(22)["id"] == 22
+        assert ds.index_of(11) == 0
+        with pytest.raises(KeyError):
+            ds.index_of(999)
+        with pytest.raises(IndexError):
+            ds.molecule(5)
+
+
+def test_decode_marker_to_nan():
+    vals = np.array([0, 10000, MARKER, -25000], dtype=np.int32)
+    out = D._decode(vals)
+    assert out[0] == 0.0 and out[1] == pytest.approx(1.0) and out[3] == pytest.approx(-2.5)
+    assert np.isnan(out[2])
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real gdb_qcd.hdf5 not present")
+def test_real_file_smoke():
+    with D.GDBQCD(REAL) as ds:
+        assert len(ds) == 2461
+        assert int(ds.f.attrs["n_trajectories"]) == 56702
+        m = ds.molecule(0)
+        assert m["coordinates"].shape[1] == 33
+        assert m["coordinates"].shape[0] == ds.n_trajectories(0)
+        assert np.isfinite(m["shieldings"]).all()

data/magnet_benchmark/performance_results.csv

@@ -0,0 +1,17 @@
+model,test_set,nucleus,median_ae,mae,rmse,n
+MagNET,stationary_internal,13C,0.23667734375001714,0.37584171551096685,0.7389307200428491,5259
+MagNET,stationary_internal,1H,0.020697677612304233,0.0311367929561522,0.054721767847735706,7655
+MagNET,vibrated_internal,13C,0.3584277770996209,0.5341938983186548,0.896683642076138,4264
+MagNET,vibrated_internal,1H,0.03247062377929666,0.047286319514361755,0.10046215637564167,6780
+MagNET,vibrated_external,13C,0.350236883544925,0.5368293357258397,1.088311121228106,12854
+MagNET,vibrated_external,1H,0.029520549011230912,0.04531780213565488,0.08093539176210796,15506
+MagNET,isolated_chloroform,13C,0.3621,0.596,1.2658,88550
+MagNET,isolated_chloroform,1H,0.0322,0.0745,0.3076,112460
+MagNET-x,stationary_internal,13C,0.2431782714843962,0.3895146201411134,0.7565858139938165,5259
+MagNET-x,stationary_internal,1H,0.020831005859374585,0.03165713338242734,0.05425704945428728,7655
+MagNET-x,vibrated_internal,13C,0.3772723861694356,0.5585515375031819,0.900760237056921,4264
+MagNET-x,vibrated_internal,1H,0.03227110290527335,0.046615542751976255,0.09691115895578069,6780
+MagNET-x,vibrated_external,13C,0.2872504028320293,0.43605742698970223,0.8808960500245364,12854
+MagNET-x,vibrated_external,1H,0.020489848327637006,0.03226036506553448,0.07018272591671143,15506
+MagNET-x,isolated_chloroform,13C,0.3062,0.4375,0.7232,88550
+MagNET-x,isolated_chloroform,1H,0.0224,0.0311,0.0462,112460

data/magnet_test_predictions/build_magnet_test_predictions.py

@@ -0,0 +1,132 @@
+"""Build magnet_test_predictions.hdf5 from the raw MagNET inference pickles.
+
+This regenerates the released archive so the packaging is reproducible end to end. It reads Keir
+Adams' per-test-set result pickles (each a {molecule_key: tuple} dict) and writes one canonical HDF5
+using the release's int32 fixed-point scheme (round(value * 1e4), NaN sentinel) with content-addressed
+dedup: any stored block that is byte-identical to one already written (the DFT targets are identical
+across a test set's model variants) is stored once and referenced.
+
+Usage:
+    python build_magnet_test_predictions.py <source_dir_of_pickles> <output.hdf5>
+
+The reader `magnet_test_predictions_reader.py` inverts this exactly (round-trip error <= 5e-5 ppm).
+"""
+import os
+import glob
+import pickle
+import hashlib
+import numpy as np
+import h5py
+
+SCALE = 1e4
+NAN_SENTINEL = np.iinfo(np.int32).min
+
+
+def _enc_float(a):
+    a = np.asarray(a, dtype=np.float64)
+    return np.where(np.isfinite(a), np.round(a * SCALE), NAN_SENTINEL).astype(np.int32)
+
+
+def _kind(elem):
+    if isinstance(elem, (str, np.str_)):
+        return "str"
+    a = np.asarray(elem)
+    if a.dtype == bool:
+        return "bool"
+    if np.issubdtype(a.dtype, np.floating):
+        return "float"
+    if np.issubdtype(a.dtype, np.integer):
+        return "int"
+    if a.dtype.kind in ("U", "S"):
+        return "str"
+    raise ValueError(f"unhandled element dtype {a.dtype}")
+
+
+class _Writer:
+    """Writes groups into an open h5py.File, sharing byte-identical blocks."""
+
+    def __init__(self, h):
+        self.h = h
+        self.store = {}
+        self.n_dedup = 0
+
+    def _put(self, grp, name, arr, **kw):
+        b = np.ascontiguousarray(arr)
+        key = hashlib.sha1(b.tobytes()).hexdigest() + f":{b.dtype.str}:{b.shape}"
+        if key in self.store:
+            grp.attrs[name + "_ref"] = self.store[key]
+            self.n_dedup += 1
+            return
+        ds = grp.create_dataset(name, data=b, **kw)
+        self.store[key] = ds.name
+
+    def encode_pickle(self, grp, data):
+        """Write one result pickle's {molecule_key: tuple} dict into `grp`, column by column."""
+        keys = list(data.keys())
+        n = len(keys)
+        tuples = [data[k] for k in keys]
+        tlen = len(tuples[0])
+        grp.attrs["n"] = n
+        grp.attrs["tuple_len"] = tlen
+        k0 = keys[0]
+        if isinstance(k0, tuple):
+            grp.attrs["key_kind"] = "tuple"
+            self._put(grp, "keys", np.asarray(keys, dtype=np.int64), compression="gzip", compression_opts=4)
+        elif isinstance(k0, (str, np.str_)):
+            grp.attrs["key_kind"] = "str"
+            grp.create_dataset("keys", data=np.asarray(keys, dtype=object),
+                               dtype=h5py.string_dtype(), compression="gzip", compression_opts=4)
+        else:
+            grp.attrs["key_kind"] = "int"
+            self._put(grp, "keys", np.asarray(keys, dtype=np.int64), compression="gzip", compression_opts=4)
+        kinds = [_kind(tuples[0][p]) for p in range(tlen)]
+        grp.attrs["pos_kinds"] = np.asarray(kinds, dtype=h5py.string_dtype())
+        for p in range(tlen):
+            kind = kinds[p]
+            grp.attrs[f"p{p}_kind"] = kind
+            if kind == "str":
+                vals = np.asarray([str(t[p]) for t in tuples], dtype=object)
+                grp.create_dataset(f"p{p}_str", data=vals, dtype=h5py.string_dtype(),
+                                   compression="gzip", compression_opts=4)
+                continue
+            arrs = [np.asarray(t[p]) for t in tuples]
+            shapes = np.asarray([a.shape for a in arrs], dtype=np.int32)
+            sizes = np.asarray([a.size for a in arrs], dtype=np.int64)
+            offsets = np.zeros(n + 1, dtype=np.int64)
+            offsets[1:] = np.cumsum(sizes)
+            flat = np.concatenate([a.ravel() for a in arrs]) if n else np.array([])
+            self._put(grp, f"p{p}_shapes", shapes, compression="gzip", compression_opts=4)
+            self._put(grp, f"p{p}_offsets", offsets, compression="gzip", compression_opts=4)
+            if kind == "float":
+                self._put(grp, f"p{p}_data", _enc_float(flat), compression="gzip", compression_opts=4, shuffle=True)
+            elif kind == "int":
+                self._put(grp, f"p{p}_data", flat.astype(np.int64), compression="gzip", compression_opts=4, shuffle=True)
+            elif kind == "bool":
+                self._put(grp, f"p{p}_data", flat.astype(np.int8), compression="gzip", compression_opts=4)
+
+
+def build_file(source_dir, outfile):
+    """Encode every *.pickle result file in `source_dir` into `outfile` as one HDF5 archive.
+
+    Returns (n_files_read, n_blocks_deduplicated).
+    """
+    files = sorted(glob.glob(os.path.join(source_dir, "*.pickle")))
+    with h5py.File(outfile, "w") as h:
+        h.attrs["format"] = "magnet_test_predictions v2 (dedup)"
+        h.attrs["scale"] = SCALE
+        h.attrs["nan_sentinel"] = NAN_SENTINEL
+        w = _Writer(h)
+        for f in files:
+            name = os.path.basename(f)[:-7]
+            if name.startswith("results_"):
+                name = name[len("results_"):]
+            w.encode_pickle(h.create_group(name), pickle.load(open(f, "rb")))
+    return len(files), w.n_dedup
+
+
+if __name__ == "__main__":
+    import sys
+    src, out = sys.argv[1], sys.argv[2]
+    nfiles, ndedup = build_file(src, out)
+    print(f"wrote {out} from {nfiles} pickles ({ndedup} shared blocks), "
+          f"{os.path.getsize(out) / 1e9:.3f} GB")

data/magnet_test_predictions/magnet_test_predictions_reader.py

@@ -0,0 +1,229 @@
+"""Reader for magnet_test_predictions.hdf5 - MagNET's raw per-atom predictions and DFT targets
+on every benchmark test set.
+
+This is the inference behind SI Tables S3 and S4 (proton and carbon performance of MagNET) and
+SI Table S5 (quasi-classical-dynamics corrections). Each test set was run through several models and
+geometries; this file stores, for every molecule, the predicted shieldings and the DFT reference so
+the accuracy statistics can be recomputed exactly.
+
+The file is a set of groups, one per original result file. List them with `available()`. Group names
+carry a solvent field only for the solvated MD sets:
+    <testset>_<nucleus>_<model>_<extra>                   (gas-phase and QCD sets)
+    <testset>_<solvent>_<nucleus>_<model>_<extra>         (solutesmd500 / solutesmd500isolated)
+for example  gasphaseinternal_C_pretrained_C_vib1  or  solutesmd500isolated_chloroform_H_chloroform_H.
+
+Test sets:
+    gasphaseinternal      internal GDB test set, gas phase
+    gasphasedft8k         external DFT8K benchmark (Guan/Paton), gas phase
+    qcdtraj2500           ~2500 molecules with quasi-classical dynamics trajectories
+    solutesmd500          500 solutes, solvated poses from classical MD
+    solutesmd500isolated  the same 500 solutes, isolated (single-molecule) poses
+    solutesmd50super      50 "super" test solutes (solvent-correction outputs)
+    solutesmd8superopenmm 8 "super" solutes (OpenMM)
+
+Models:
+    pretrained            MagNET foundation model (trained on vibrated geometries)
+    pretrained_stationary MagNET trained only on stationary geometries (the Figure S10 ablation)
+    chloroform/benzene/methanol/tip4p   MagNET-x solvent-specific models
+
+Geometry suffix on the gas-phase sets:
+    vib0                  stationary (equilibrium) test geometry
+    vib1                  vibrated test geometry
+
+Each group loads back with `load()` to the original {molecule_key: tuple} object. The tuple layout
+differs by family (this matters for computing errors correctly):
+    gasphaseinternal, gasphasedft8k : (prediction, target, atomic_numbers, mask)   mask selects the
+                                      scored nucleus; the error is over prediction[mask]-target[mask].
+    solutesmd500, solutesmd500isolated : (target, prediction)   already one nucleus, no mask.
+    solutesmd50super : (molecule_name, smiles, prediction, target)   these are solvent CORRECTIONS.
+    solutesmd8superopenmm : four float arrays of shape (n_frames, n_atoms); pairing is not documented,
+                                      so use load() and interpret them yourself.
+    qcdtraj2500 : (stationary_prediction, stationary_target, trajectory_prediction, trajectory_target)
+                                      where the trajectory arrays are (n_replicates, n_frames, n_scored).
+                                      The QCD *correction* (SI Table S5) is the trajectory mean minus the
+                                      stationary value; use qcd_correction_stats(), NOT stats().
+
+`stats(group)` gives the (median, MAE, RMSE, n) of the per-atom errors for the gas-phase and MD-solute
+families (SI Tables S3/S4). For the QCD family use `qcd_correction_stats(group)` (SI Table S5). Both
+raise on a family they do not handle rather than return a wrong number.
+"""
+import numpy as np
+import h5py
+
+SCALE = 1e4
+NAN_SENTINEL = np.iinfo(np.int32).min
+
+
+def _dec_float(a):
+    a = np.asarray(a)
+    return np.where(a == NAN_SENTINEL, np.nan, a / SCALE).astype(np.float32)
+
+
+def available(path):
+    """Return the sorted list of test-set group names in the file."""
+    with h5py.File(path, "r") as h:
+        return sorted(h.keys())
+
+
+def _fetch(h, g, name):
+    """Read dataset g/name, following a content-addressed dedup reference if present."""
+    if name in g:
+        return g[name][()]
+    ref = g.attrs.get(name + "_ref")
+    if ref is not None:
+        return h[ref][()]
+    raise KeyError(f"{name} missing in {g.name}")
+
+
+def _decode_group(h, g):
+    n = int(g.attrs["n"])
+    tlen = int(g.attrs["tuple_len"])
+    kk = g.attrs["key_kind"]
+    if kk == "str":
+        raw = g["keys"][()]
+        keys = [k.decode() if isinstance(k, bytes) else str(k) for k in raw]
+    else:
+        raw = _fetch(h, g, "keys")
+        if kk == "tuple":
+            keys = [tuple(int(x) for x in row) for row in raw]
+        else:
+            keys = [int(x) for x in raw]
+
+    cols = []
+    for p in range(tlen):
+        if f"p{p}_str" in g:
+            s = g[f"p{p}_str"][()]
+            cols.append(("str", [x.decode() if isinstance(x, bytes) else str(x) for x in s]))
+            continue
+        kind = g.attrs[f"p{p}_kind"]
+        data = _fetch(h, g, f"p{p}_data")
+        shapes = _fetch(h, g, f"p{p}_shapes")
+        offs = _fetch(h, g, f"p{p}_offsets")
+        if kind == "float":
+            data = _dec_float(data)
+        elif kind == "bool":
+            data = data.astype(bool)
+        else:
+            data = data.astype(np.int64)
+        vals = [data[offs[i]:offs[i + 1]].reshape(tuple(int(x) for x in shapes[i])) for i in range(n)]
+        cols.append((kind, vals))
+    return {keys[i]: tuple(cols[p][1][i] for p in range(tlen)) for i in range(n)}
+
+
+def load(path, group):
+    """Load one test-set group back to its original {molecule_key: tuple} dict."""
+    with h5py.File(path, "r") as h:
+        if group not in h:
+            raise KeyError(f"{group!r} not in file; see available()")
+        return _decode_group(h, h[group])
+
+
+def _require_aligned(h, g, group, *positions):
+    """The flat computation below pairs elements across columns by position, which is only equivalent
+    to the per-molecule form when the columns share the same per-molecule boundaries. Every column
+    stores those boundaries as `p{p}_offsets`; require them to match so malformed input (a molecule
+    whose prediction/target/mask arrays are not the same length) is refused loudly instead of pairing
+    the wrong atoms silently."""
+    offs = [_fetch(h, g, f"p{p}_offsets") for p in positions]
+    if not all(np.array_equal(offs[0], o) for o in offs[1:]):
+        raise ValueError(f"{group} has misaligned per-molecule lengths across columns "
+                         f"{positions}; the stored arrays are inconsistent. Use load() and inspect.")
+
+
+def _flat_error_array(h, g, group):
+    """Per-atom |prediction - target| for every scored atom in a group, computed straight from the
+    flat column arrays without rebuilding the per-molecule dict. Each column is stored concatenated in
+    molecule order with matching per-molecule boundaries (checked by _require_aligned), so a
+    whole-column subtraction (plus the stored mask, for the gas-phase family) gives the per-atom errors
+    for every recognized tuple layout. Returns None for an unrecognized layout, so the caller can
+    refuse rather than guess."""
+    tlen = int(g.attrs["tuple_len"])
+    if tlen == 2:
+        # (target, prediction), already one nucleus, no mask
+        _require_aligned(h, g, group, 0, 1)
+        target = _dec_float(_fetch(h, g, "p0_data"))
+        pred = _dec_float(_fetch(h, g, "p1_data"))
+        return np.abs(pred.ravel() - target.ravel())
+    if tlen == 4 and "p3_str" not in g and g.attrs.get("p3_kind") == "bool":
+        # (prediction, target, atomic_numbers, mask); the mask selects the scored nucleus
+        _require_aligned(h, g, group, 0, 1, 3)
+        pred = _dec_float(_fetch(h, g, "p0_data"))
+        target = _dec_float(_fetch(h, g, "p1_data"))
+        mask = _fetch(h, g, "p3_data").astype(bool)
+        return np.abs(pred[mask] - target[mask])
+    if tlen == 4 and "p0_str" in g:
+        # (name, smiles, prediction, target) - solvent-correction predictions
+        _require_aligned(h, g, group, 2, 3)
+        pred = _dec_float(_fetch(h, g, "p2_data"))
+        target = _dec_float(_fetch(h, g, "p3_data"))
+        return np.abs(pred.ravel() - target.ravel())
+    return None
+
+
+def errors(path, group):
+    """Flat array of per-atom absolute errors |prediction - target| for a per-atom group
+    (gas-phase and MD-solute families - SI Tables S3/S4).
+
+    Refuses the QCD family (raise), whose accuracy is a trajectory *correction*, not a per-atom
+    shielding error - use qcd_correction_errors() for that. Also refuses any tuple layout it does not
+    recognize (for example solutesmd8superopenmm, whose four float arrays have no documented pairing);
+    load() that group and interpret it yourself.
+    """
+    if group.startswith("qcdtraj"):
+        raise ValueError(f"{group} is a QCD trajectory set; its accuracy is a correction "
+                         f"(SI Table S5). Use qcd_correction_errors()/qcd_correction_stats().")
+    with h5py.File(path, "r") as h:
+        if group not in h:
+            raise KeyError(f"{group!r} not in file; see available()")
+        e = _flat_error_array(h, h[group], group)
+    if e is None:
+        raise ValueError(f"{group} has an unrecognized tuple layout for per-atom errors; "
+                         f"use load() and compute directly.")
+    return e
+
+
+def stats(path, group):
+    """(median_ae, mae, rmse, n) over all scored atoms in a per-atom group - the SI Table S3/S4
+    quantities. Raises for the QCD family (use qcd_correction_stats)."""
+    d = errors(path, group)
+    return float(np.median(d)), float(d.mean()), float((d ** 2).mean() ** 0.5), int(d.size)
+
+
+def qcd_correction_errors(path, group):
+    """Flat per-atom absolute errors in the QCD *correction* for a qcdtraj2500 group - the quantity
+    behind SI Table S5. The correction is the trajectory-averaged shielding minus the stationary
+    shielding; the error is |predicted_correction - DFT_correction|, per scored atom.
+
+    Each molecule stores (stationary_prediction, stationary_target, trajectory_prediction,
+    trajectory_target); the trajectory arrays are (n_replicates, n_frames, n_scored).
+    """
+    if not group.startswith("qcdtraj"):
+        raise ValueError(f"{group} is not a QCD trajectory set; use errors()/stats().")
+    d = load(path, group)
+    parts = []
+    for stat_pred, stat_target, traj_pred, traj_target in d.values():
+        pred_corr = np.asarray(traj_pred).mean(axis=(0, 1)) - np.asarray(stat_pred)
+        dft_corr = np.asarray(traj_target).mean(axis=(0, 1)) - np.asarray(stat_target)
+        parts.append(np.abs(pred_corr - dft_corr))
+    return np.concatenate(parts)
+
+
+def qcd_correction_stats(path, group):
+    """(median_ae, mae, rmse, n) of the QCD correction errors for a qcdtraj2500 group - SI Table S5."""
+    e = qcd_correction_errors(path, group)
+    return float(np.median(e)), float(e.mean()), float((e ** 2).mean() ** 0.5), int(e.size)
+
+
+if __name__ == "__main__":
+    import sys
+    p = sys.argv[1] if len(sys.argv) > 1 else "magnet_test_predictions.hdf5"
+    names = available(p)
+    print(f"{len(names)} groups.")
+    print("Tables S3/S4 (per-atom error):")
+    for g in ("gasphaseinternal_C_pretrained_C_vib1", "solutesmd500isolated_chloroform_H_chloroform_H"):
+        if g in names:
+            print(f"  {g}: median/MAE/RMSE/n =", stats(p, g))
+    print("Table S5 (QCD correction):")
+    for g in ("qcdtraj2500_H_pretrained_H", "qcdtraj2500_C_pretrained_C"):
+        if g in names:
+            print(f"  {g}: median/MAE/RMSE/n =", qcd_correction_stats(p, g))

data/magnet_test_predictions/test_magnet_test_predictions.py

@@ -0,0 +1,175 @@
+"""Tests for the magnet_test_predictions reader.
+
+The fixture builds a tiny archive with the real encoder (`build_magnet_test_predictions.build_file`)
+covering every tuple family, both key kinds, string columns, multi-dimensional trajectory arrays, and
+the content-addressed dedup path - so it is a faithful spec of the on-disk format, and the reader is
+tested end to end against it. Optional tests run against the real file when it is present.
+"""
+import os
+import pickle
+import numpy as np
+import pytest
+
+import magnet_test_predictions_reader as R
+import build_magnet_test_predictions as B
+
+
+def _make_source(dirpath):
+    """Write tiny pickles, one per family, mimicking the real result files (names drive the reader)."""
+    files = {}
+
+    # gas-phase 4-tuple (prediction, target, atomic_numbers, mask); NaN in a prediction
+    files["results_gasphaseinternal_H_pretrained_H_vib1.pickle"] = {
+        10: (np.array([1.0, 2.0, 3.0], np.float32), np.array([1.05, 2.10, 2.90], np.float32),
+             np.array([1, 1, 6]), np.array([True, True, False])),
+        11: (np.array([4.0, np.nan], np.float32), np.array([4.20, 5.00], np.float32),
+             np.array([1, 8]), np.array([True, False])),
+    }
+    # a SECOND gas-phase model on the SAME test set: identical target/atoms/mask -> must dedup
+    files["results_gasphaseinternal_H_chloroform_H_vib1.pickle"] = {
+        10: (np.array([1.1, 2.1, 3.1], np.float32), np.array([1.05, 2.10, 2.90], np.float32),
+             np.array([1, 1, 6]), np.array([True, True, False])),
+        11: (np.array([4.1, np.nan], np.float32), np.array([4.20, 5.00], np.float32),
+             np.array([1, 8]), np.array([True, False])),
+    }
+    # MD-solute 2-tuple (target, prediction), tuple keys
+    files["results_solutesmd500isolated_chloroform_C_chloroform_C.pickle"] = {
+        (1, 1): (np.array([10.0, 20.0], np.float32), np.array([10.1, 19.8], np.float32)),
+        (1, 2): (np.array([30.0], np.float32), np.array([30.5], np.float32)),
+    }
+    # super-corrections 4-tuple (name, smiles, prediction, target), 2-D arrays, int keys
+    files["results_solutesmd50super_corrections_TIP4P_C_tip4p_C.pickle"] = {
+        7: ("MOL1", "CCO", np.array([[-1.0, -2.0]], np.float32), np.array([[-1.2, -1.9]], np.float32)),
+    }
+    # QCD 4-tuple (stat_pred, stat_target, traj_pred(reps,frames,scored), traj_target)
+    tp = np.zeros((2, 3, 2), np.float32); tt = np.zeros((2, 3, 2), np.float32)
+    tp[:] = np.array([100.0, 50.0]); tt[:] = np.array([100.0, 50.0])
+    tp += 0.4  # predicted trajectory shifted +0.4 from stationary
+    tt += 0.3  # DFT trajectory shifted +0.3 -> correction error 0.1 per atom
+    files["results_qcdtraj2500_C_pretrained_C.pickle"] = {
+        100: (np.array([100.0, 50.0], np.float32), np.array([100.0, 50.0], np.float32), tp, tt),
+    }
+    # str-key family
+    files["results_solutesmd8superopenmm_TIP4P_C_tip4p_C.pickle"] = {
+        "MOLX": tuple(np.full((2, 2), v, np.float32) for v in (1.0, 1.1, 2.0, 2.1)),
+    }
+    for fn, d in files.items():
+        pickle.dump(d, open(os.path.join(dirpath, fn), "wb"))
+
+
+@pytest.fixture
+def archive(tmp_path):
+    src = tmp_path / "src"; src.mkdir()
+    _make_source(str(src))
+    out = str(tmp_path / "magnet_test_predictions.hdf5")
+    nfiles, ndedup = B.build_file(str(src), out)
+    assert nfiles == 6
+    assert ndedup > 0, "the two gas-phase models share a target/atoms/mask; dedup must trigger"
+    return out
+
+
+def test_available(archive):
+    got = R.available(archive)
+    assert "gasphaseinternal_H_pretrained_H_vib1" in got
+    assert "qcdtraj2500_C_pretrained_C" in got
+    assert len(got) == 6
+
+
+def test_roundtrip_all_families(archive):
+    """Every family decodes back to its source pickle exactly (floats within a quantum)."""
+    src = os.path.join(os.path.dirname(archive), "src")
+    for fn in os.listdir(src):
+        group = fn[len("results_"):-len(".pickle")]
+        orig = pickle.load(open(os.path.join(src, fn), "rb"))
+        dec = R.load(archive, group)
+        assert set(dec) == set(orig)
+        for k in orig:
+            for a, b in zip(orig[k], dec[k]):
+                if isinstance(a, str):
+                    assert a == b
+                else:
+                    a = np.asarray(a); b = np.asarray(b)
+                    assert a.shape == b.shape
+                    if a.dtype == bool:
+                        assert np.array_equal(a, b.astype(bool))
+                    elif np.issubdtype(a.dtype, np.integer):
+                        assert np.array_equal(a, b)
+                    else:
+                        m = np.isfinite(a)
+                        assert np.array_equal(np.isnan(a), np.isnan(np.asarray(b, float)))
+                        if m.any():
+                            assert np.abs(a[m] - np.asarray(b)[m]).max() < 1e-4
+
+
+def test_dedup_reference_resolves_to_identical_bytes(archive):
+    """The deduped target of the second model must equal the first model's target exactly."""
+    import h5py
+    with h5py.File(archive, "r") as h:
+        # exactly one of the two same-target groups stores p1_data; the other references it
+        g1 = h["gasphaseinternal_H_pretrained_H_vib1"]
+        g2 = h["gasphaseinternal_H_chloroform_H_vib1"]
+        has = [("p1_data" in g) for g in (g1, g2)]
+        ref = [("p1_data_ref" in g.attrs) for g in (g1, g2)]
+        assert sum(has) == 1 and sum(ref) == 1 and has != ref
+    a = R.load(archive, "gasphaseinternal_H_pretrained_H_vib1")
+    b = R.load(archive, "gasphaseinternal_H_chloroform_H_vib1")
+    for k in a:
+        assert np.array_equal(np.asarray(a[k][1]), np.asarray(b[k][1]))  # identical DFT target
+        assert not np.array_equal(np.asarray(a[k][0]), np.asarray(b[k][0]))  # different predictions
+
+
+def test_errors_apply_mask_and_stats(archive):
+    e = R.errors(archive, "gasphaseinternal_H_pretrained_H_vib1")
+    assert np.allclose(np.sort(e), [0.05, 0.10, 0.20], atol=1e-4)  # mol10 two atoms + mol11 one
+    med, mae, rmse, n = R.stats(archive, "gasphaseinternal_H_pretrained_H_vib1")
+    assert n == 3 and abs(mae - 0.35 / 3) < 1e-4
+
+
+def test_tuple_keys_roundtrip(archive):
+    d = R.load(archive, "solutesmd500isolated_chloroform_C_chloroform_C")
+    assert set(d) == {(1, 1), (1, 2)}
+
+
+def test_qcd_correction(archive):
+    # constructed so every atom's correction error is 0.1
+    med, mae, rmse, n = R.qcd_correction_stats(archive, "qcdtraj2500_C_pretrained_C")
+    assert abs(mae - 0.1) < 1e-3 and abs(rmse - 0.1) < 1e-3
+
+
+def test_stats_refuses_qcd(archive):
+    with pytest.raises(ValueError):
+        R.stats(archive, "qcdtraj2500_C_pretrained_C")
+
+
+def test_qcd_stats_refuses_nonqcd(archive):
+    with pytest.raises(ValueError):
+        R.qcd_correction_stats(archive, "gasphaseinternal_H_pretrained_H_vib1")
+
+
+def test_errors_refuses_unpaired_family(archive):
+    # solutesmd8superopenmm is four float arrays with no documented pairing -> must refuse, not guess
+    with pytest.raises(ValueError):
+        R.errors(archive, "solutesmd8superopenmm_TIP4P_C_tip4p_C")
+
+
+# ----- optional real-file checks (run locally when the archive is downloaded) -----
+REAL = os.path.join(os.path.dirname(__file__), "magnet_test_predictions.hdf5")
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real hdf5 not present (runs locally)")
+def test_real_tables_s3_s4():
+    # foundation 13C vibrated internal (SI: 0.31917763 / 0.47851510 / 7.13765463)
+    med, mae, rmse, n = R.stats(REAL, "gasphaseinternal_C_pretrained_C_vib1")
+    assert abs(med - 0.3192) < 1e-3 and abs(mae - 0.4785) < 1e-3 and abs(rmse - 7.1377) < 1e-2
+    # foundation 1H isolated chloroform (SI: 0.03216 / 0.07175 / 0.29437)
+    med, mae, rmse, n = R.stats(REAL, "solutesmd500isolated_chloroform_H_pretrained_H")
+    assert abs(med - 0.0322) < 1e-3 and abs(mae - 0.0718) < 1e-3 and abs(rmse - 0.2944) < 1e-2
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real hdf5 not present (runs locally)")
+def test_real_table_s5_qcd():
+    # SI Table S5: MagNET 1H 0.01526165/0.023047674/0.04020983 ; 13C 0.17198181/0.26223338/0.50198762
+    med, mae, rmse, n = R.qcd_correction_stats(REAL, "qcdtraj2500_H_pretrained_H")
+    assert abs(med - 0.0152617) < 1e-3 and abs(mae - 0.0230477) < 1e-3 and abs(rmse - 0.0402098) < 1e-3
+    med, mae, rmse, n = R.qcd_correction_stats(REAL, "qcdtraj2500_C_pretrained_C")
+    assert abs(med - 0.1719818) < 2e-3 and abs(mae - 0.2622334) < 1e-3 and abs(rmse - 0.5019876) < 1e-3

data/scaling_factors/scaling_factors_reader.py

@@ -0,0 +1,36 @@
+"""Reader for the symmetrized MagNET-Zero/MagNET-PCM linear-scaling tables in this folder.
+
+These two CSV files (one per nucleus) are a deployment-quality export: the per-solvent coefficients
+that turn a MagNET-Zero gas-phase shielding plus a MagNET-PCM chloroform correction into a predicted
+chemical shift, fit with reflection-symmetrized inference (n_passes=10, symmetrize=True). They are
+the "serving on new molecules" variant of Supporting Information Tables S10/S11; the paper's published
+tables use the unsymmetrized shieldings and live in analysis/code/scaling_factors.py.
+
+Each row is one solvent; columns are intercept, stationary (the MagNET-Zero slope), and pcm (the
+MagNET-PCM-correction slope). Predicted shift = intercept + stationary * shielding + pcm * correction.
+"""
+import os
+
+import pandas as pd
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+NUCLEI = ("H", "C")
+
+
+def _csv_path(nucleus):
+    if nucleus not in NUCLEI:
+        raise ValueError(f"nucleus must be one of {NUCLEI}, got {nucleus!r}")
+    return os.path.join(HERE, f"scaling_factors_symmetrized_{nucleus}.csv")
+
+
+def load_symmetrized_tables():
+    """Returns {"H": DataFrame, "C": DataFrame}, each indexed by solvent with columns
+    intercept / stationary / pcm."""
+    return {n: pd.read_csv(_csv_path(n)).set_index("solvent") for n in NUCLEI}
+
+
+def predict_shift(table, solvent, shielding, correction):
+    """Chemical shift for one site: intercept + stationary * shielding + pcm * correction, using
+    `table` (one nucleus' DataFrame from load_symmetrized_tables) and its `solvent` row."""
+    row = table.loc[solvent]
+    return float(row["intercept"] + row["stationary"] * shielding + row["pcm"] * correction)

data/scaling_factors/scaling_factors_symmetrized_C.csv

@@ -0,0 +1,13 @@
+solvent,intercept,stationary,pcm
+chloroform,171.726919,-0.924228,-0.936801
+tetrahydrofuran,171.052274,-0.918997,-1.069506
+dichloromethane,171.507154,-0.921165,-1.115013
+acetone,171.305850,-0.919608,-1.239500
+acetonitrile,171.877492,-0.922626,-1.287661
+dimethylsulfoxide,170.596065,-0.918672,-1.296497
+trifluoroethanol,174.235231,-0.939011,-1.290071
+methanol,172.423799,-0.927563,-1.288843
+TIP4P,173.692750,-0.937848,-1.342659
+benzene,171.965715,-0.927029,-0.593714
+toluene,171.689276,-0.924983,-0.618407
+chlorobenzene,171.073820,-0.921235,-0.997638

data/sigma-concentrate/decode_sigma_concentrate.py

@@ -0,0 +1,113 @@
+"""
+decode_sigma_concentrate.py - reader/decoder for the sigma-concentrate HDF5 dataset.
+
+sigma-concentrate holds B3LYP/pcSseg-2 NMR shieldings computed both in the gas phase and with
+an implicit chloroform solvent (PCM), on AIMNet2 stationary geometries for 50,000 molecules
+sampled from sigma-shake. It is the data behind the MagNET-PCM implicit-solvent model, which
+learns the solvent correction shielding_b3lyp_pcm - shielding_b3lyp_gas.
+
+The file is flat (a single molecule set, one geometry per molecule):
+
+  shielding_b3lyp_gas    isotropic sigma, B3LYP/pcSseg-2, gas phase
+  shielding_b3lyp_pcm    isotropic sigma, B3LYP/pcSseg-2, PCM, solvent = chloroform
+
+Storage format:
+  - Per-atom arrays are concatenated across all molecules. Each molecule's atom count is in
+    `n_atoms`; the row range for molecule i is [start[i], start[i+1]) where start = cumulative
+    sum of n_atoms (reconstructed once, on open).
+  - coordinates and shieldings are int32 fixed-point: physical = stored * scale (scale = 1e-4),
+    so reconstruction error is <= 5e-5 (Angstrom or ppm). Coordinates are mean-centered per
+    molecule (centroid at the origin), a translation only; interatomic distances are unchanged.
+  - The molecule's array index IS its identifier (row order is fixed to the source). `sigma_shake_id`
+    additionally gives each molecule's id in the sigma-shake dataset (same molecule there, at a
+    different level of theory), so the two datasets can be joined.
+  - `failed_indices` lists molecules whose shielding calculation failed (none in this release):
+    their shieldings would be sentinel zeros and must be ignored. There are no SMILES.
+
+Requires: Python >= 3.7, numpy, h5py. No other dependencies.
+"""
+from __future__ import annotations
+import numpy as np
+import h5py
+
+
+class SigmaConcentrate:
+    """Reader for the sigma-concentrate HDF5 file: gas/PCM shieldings on 50,000 sigma-shake molecules."""
+
+    def __init__(self, path: str):
+        self.f = h5py.File(path, "r")
+        self.n_molecules = int(self.f.attrs["n_molecules"])
+        self.n_atoms_total = int(self.f.attrs["n_atoms"])
+        self.level_of_theory = str(self.f.attrs["level_of_theory"])
+        self.geometry = str(self.f.attrs["geometry"])
+        # atom row offsets: start[i]..start[i+1] for molecule i (one cumsum, cached)
+        na = self.f["n_atoms"][:].astype(np.int64)
+        self._start = np.empty(self.n_molecules + 1, np.int64)
+        self._start[0] = 0
+        np.cumsum(na, out=self._start[1:])
+        self._sid = self.f["sigma_shake_id"][:]
+        self._failed = set(int(x) for x in self.f["failed_indices"][:])
+        self.shieldings = sorted(k for k in self.f.keys() if k.startswith("shielding_"))
+
+    def _check(self, index: int):
+        if not 0 <= index < self.n_molecules:
+            raise IndexError(f"molecule index {index} out of range [0, {self.n_molecules})")
+
+    def _scaled(self, name: str, sl: slice) -> np.ndarray:
+        d = self.f[name]
+        return d[sl].astype(np.float64) * float(d.attrs["scale"])
+
+    def sigma_shake_id(self, index: int) -> int:
+        """Return molecule `index`'s id in the sigma-shake dataset."""
+        self._check(index)
+        return int(self._sid[index])
+
+    def molecule(self, index: int) -> dict:
+        """Return one molecule's decoded data (physical units) by array index.
+
+        Keys: index, sigma_shake_id, atomic_numbers, coordinates (n,3 Angstrom),
+        shielding_b3lyp_gas (n, ppm), shielding_b3lyp_pcm (n, ppm), and `failed`. If `failed`
+        is True the shieldings are sentinel zeros and should be ignored; geometry is still valid.
+        """
+        self._check(index)
+        a, b = int(self._start[index]), int(self._start[index + 1])
+        sl = slice(a, b)
+        out = {
+            "index": index,
+            "sigma_shake_id": int(self._sid[index]),
+            "atomic_numbers": self.f["atomic_numbers"][sl],
+            "coordinates": self._scaled("coordinates", sl),   # (n,3) Angstrom
+            "failed": index in self._failed,
+        }
+        for s in self.shieldings:
+            out[s] = self._scaled(s, sl)                       # (n,) ppm
+        return out
+
+    def __len__(self):
+        return self.n_molecules
+
+    def close(self):
+        """Close the underlying HDF5 file."""
+        self.f.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        self.close()
+
+    def __repr__(self):
+        return (f"<SigmaConcentrate: {self.n_molecules:,} molecules, "
+                f"{self.n_atoms_total:,} atoms, {self.level_of_theory}>")
+
+
+if __name__ == "__main__":
+    import sys
+    path = sys.argv[1] if len(sys.argv) > 1 else "sigma-concentrate.hdf5"
+    with SigmaConcentrate(path) as ds:
+        print(ds)
+        m = ds.molecule(0)
+        print(f"  mol 0: sigma_shake_id={m['sigma_shake_id']}  Z={list(m['atomic_numbers'][:6])}  "
+              f"failed={m['failed']}")
+        print(f"    gas={np.round(m['shielding_b3lyp_gas'][:3], 4)}  "
+              f"pcm={np.round(m['shielding_b3lyp_pcm'][:3], 4)}")

data/sigma-concentrate/test_sigma_concentrate.py

@@ -0,0 +1,162 @@
+"""
+Tests for the sigma-concentrate reader (decode_sigma_concentrate.SigmaConcentrate).
+
+These build a tiny synthetic dataset in the *exact* on-disk format (so the fixture also serves as
+an executable spec of the layout) and exercise the decoder against it. No dependency on the real
+file. One opt-in test runs against the real `sigma-concentrate.hdf5` if it is present next to this
+file.
+
+Run:  pytest test_sigma_concentrate.py -q
+Requires: pytest, numpy, h5py.
+"""
+import os
+import numpy as np
+import h5py
+import pytest
+
+from decode_sigma_concentrate import SigmaConcentrate
+
+SCALE = 1e-4
+
+
+def _i32(x):
+    return np.round(np.asarray(x, float) / SCALE).astype(np.int32)
+
+
+def _write(f, *, n_atoms, znums, sids, coords, gas, pcm, failed):
+    assert len(znums) == int(np.sum(n_atoms))
+    f.attrs["dataset"] = "sigma-concentrate"
+    f.attrs["level_of_theory"] = "B3LYP/pcSseg-2"
+    f.attrs["geometry"] = "AIMNet2 stationary"
+    f.attrs["solvent"] = "gas + PCM chloroform"
+    f.attrs["n_molecules"] = len(n_atoms)
+    f.attrs["n_atoms"] = int(np.sum(n_atoms))
+    f.create_dataset("atomic_numbers", data=np.asarray(znums, np.int8))
+    f.create_dataset("n_atoms", data=np.asarray(n_atoms, np.uint8))
+    f.create_dataset("sigma_shake_id", data=np.asarray(sids, np.int32))
+    d = f.create_dataset("coordinates", data=_i32(coords)); d.attrs["scale"] = SCALE; d.attrs["units"] = "angstrom"
+    d = f.create_dataset("shielding_b3lyp_gas", data=_i32(gas)); d.attrs["scale"] = SCALE; d.attrs["units"] = "ppm"
+    d = f.create_dataset("shielding_b3lyp_pcm", data=_i32(pcm)); d.attrs["scale"] = SCALE; d.attrs["units"] = "ppm"
+    f.create_dataset("failed_indices", data=np.asarray(failed, np.int32))
+
+
+# 3 molecules (2,1,3 atoms = 6); heteroatoms exercise the passthrough; molecule 1 failed
+NATOMS = [2, 1, 3]
+Z = [6, 1,   8,   6, 7, 1]
+SIDS = [1000265, 42, 999808]
+COORDS = np.array([[0, 0, 0], [1.2345, 0, 0],
+                   [5, 5, 5],
+                   [3, 0, 0], [4.1, 0, 0], [3, 1.5, 0]], float)
+GAS = np.array([180.1234, 31.0,   0.0,   150.0, 151.5, 29.9])   # mol 1 (atom 2) = sentinel 0
+PCM = np.array([180.4567, 30.8,   0.0,   149.6, 151.2, 29.7])
+
+
+@pytest.fixture(scope="module")
+def mini_path(tmp_path_factory):
+    p = str(tmp_path_factory.mktemp("sc") / "mini.hdf5")
+    with h5py.File(p, "w") as f:
+        _write(f, n_atoms=NATOMS, znums=Z, sids=SIDS, coords=COORDS, gas=GAS, pcm=PCM, failed=[1])
+    return p
+
+
+@pytest.fixture
+def ds(mini_path):
+    with SigmaConcentrate(mini_path) as d:
+        yield d
+
+
+def test_dimensions(ds):
+    assert ds.n_molecules == 3 and ds.n_atoms_total == 6
+    assert ds.shieldings == ["shielding_b3lyp_gas", "shielding_b3lyp_pcm"]
+    assert len(ds) == 3
+
+
+def test_offsets_slice_each_molecule(ds):
+    assert len(ds.molecule(0)["atomic_numbers"]) == 2
+    assert len(ds.molecule(1)["atomic_numbers"]) == 1
+    assert len(ds.molecule(2)["atomic_numbers"]) == 3
+
+
+def test_values_roundtrip_and_id(ds):
+    m = ds.molecule(2)
+    assert m["sigma_shake_id"] == 999808
+    np.testing.assert_array_equal(m["atomic_numbers"], [6, 7, 1])   # heteroatom passthrough
+    np.testing.assert_allclose(m["coordinates"], COORDS[3:6], atol=5e-5)
+    np.testing.assert_allclose(m["shielding_b3lyp_gas"], GAS[3:6], atol=5e-5)
+    np.testing.assert_allclose(m["shielding_b3lyp_pcm"], PCM[3:6], atol=5e-5)
+    assert not m["failed"]
+
+
+def test_both_shieldings_distinct(ds):
+    m = ds.molecule(0)
+    # gas and pcm differ atom-by-atom (the solvent correction is the signal)
+    assert not np.allclose(m["shielding_b3lyp_gas"], m["shielding_b3lyp_pcm"])
+
+
+def test_sigma_shake_id_accessor(ds):
+    assert ds.sigma_shake_id(0) == 1000265
+    assert ds.sigma_shake_id(1) == 42
+
+
+def test_failed_flag(ds):
+    assert ds.molecule(1)["failed"] is True
+    assert ds.molecule(0)["failed"] is False
+    f = ds.molecule(1)
+    # failed molecule's shieldings are sentinel zeros; geometry/id still valid
+    assert np.all(f["shielding_b3lyp_gas"] == 0) and np.all(f["shielding_b3lyp_pcm"] == 0)
+    assert f["sigma_shake_id"] == 42
+    np.testing.assert_allclose(f["coordinates"], COORDS[2:3], atol=5e-5)
+
+
+def test_empty_failed_indices(tmp_path):
+    p = str(tmp_path / "nofail.hdf5")
+    with h5py.File(p, "w") as f:
+        _write(f, n_atoms=[2], znums=[6, 1], sids=[7],
+               coords=np.array([[0, 0, 0], [1.0, 0, 0]], float),
+               gas=np.array([150.0, 30.0]), pcm=np.array([149.5, 29.8]), failed=[])
+    with SigmaConcentrate(p) as d:
+        assert d.molecule(0)["failed"] is False
+
+
+def test_index_out_of_range_raises(ds):
+    with pytest.raises(IndexError):
+        ds.molecule(-1)
+    with pytest.raises(IndexError):
+        ds.molecule(3)
+
+
+def test_context_manager_closes(mini_path):
+    d = SigmaConcentrate(mini_path)
+    d.close()
+    assert not d.f.id.valid
+
+
+# ---------- opt-in: real file, if present ----------
+REAL = os.path.join(os.path.dirname(__file__), "sigma-concentrate.hdf5")
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real sigma-concentrate.hdf5 not present")
+def test_real_file():
+    with SigmaConcentrate(REAL) as ds:
+        assert ds.n_molecules == 50_000 and ds.n_atoms_total == 1_413_134
+        assert len(ds.f["failed_indices"]) == 0
+        # gas and pcm are distinct columns; PCM differs from gas on carbon
+        Z = ds.f["atomic_numbers"][:500_000]
+        gas = ds.f["shielding_b3lyp_gas"][:500_000].astype(float) * SCALE
+        pcm = ds.f["shielding_b3lyp_pcm"][:500_000].astype(float) * SCALE
+        C = Z == 6
+        assert gas[C].mean() != pcm[C].mean()
+        # sigma_shake_id is populated and plausibly a sigma-shake id (positive)
+        assert ds.sigma_shake_id(0) > 0
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real sigma-concentrate.hdf5 not present")
+def test_real_file_reproduces_table_s8_site_counts():
+    """SI Table S8's 'sigma-concentrate' row is this file exactly, and (together with
+    sigma-pepper's wp04_wb97xd_pcSseg2 group -- see test_sigma_pepper.py's matching test) these
+    50,000 reused molecules make up the other half of the 'MagNET-Zero (fine-tuning)' row:
+    3,287,698 + 754,437 = 4,042,135 H sites; 2,028,642 + 528,569 = 2,557,211 C sites."""
+    with SigmaConcentrate(REAL) as ds:
+        z = ds.f["atomic_numbers"][:]
+        assert int((z == 1).sum()) == 754_437
+        assert int((z == 6).sum()) == 528_569

data/sigma-pepper/decode_sigma_pepper.py

@@ -0,0 +1,134 @@
+"""
+decode_sigma_pepper.py - reader/decoder for the sigma-pepper HDF5 dataset.
+
+sigma-pepper holds computed gas-phase NMR shieldings on AIMNet2 stationary geometries for
+GDB-derived molecules; it is the data behind the MagNET-Zero model. The file has TWO groups,
+one per level of theory:
+
+  /pbe0_pcSseg1          GDB 1-10 heavy atoms, PBE0/pcSseg-1, one shielding per atom
+  /wp04_wb97xd_pcSseg2   GDB 1-9  heavy atoms, two shieldings per atom; the `shieldings` list is
+                         sorted alphabetically, so the order is shielding_wb97xd (13C reference)
+                         then shielding_wp04 (1H reference)
+
+Storage format:
+  - Per-atom arrays are concatenated across all molecules in a group. Each molecule's atom
+    count is in `n_atoms`; the row range for molecule i is [start[i], start[i+1]) where
+    start = cumulative sum of n_atoms (reconstructed once, on open).
+  - coordinates and shieldings are int32 fixed-point: physical = stored * scale (scale = 1e-4),
+    so reconstruction error is <= 5e-5 (Angstrom or ppm). Coordinates are mean-centered per
+    molecule (centroid at the origin), a translation only; interatomic distances are unchanged.
+  - SMILES are stored as a single concatenated UTF-8 buffer (`smiles_buffer`) plus
+    `smiles_offsets`; molecule i's SMILES is buffer[offsets[i]:offsets[i+1]].
+  - The molecule's array index IS its identifier (row order is fixed to the source).
+  - `failed_indices` lists molecules whose shielding calculation failed: their shieldings are
+    sentinel zeros and must be ignored. Geometry and SMILES remain valid.
+
+Requires: Python >= 3.7, numpy, h5py. No other dependencies.
+"""
+from __future__ import annotations
+import numpy as np
+import h5py
+
+
+class Component:
+    """One group of sigma-pepper (a single level of theory)."""
+
+    def __init__(self, grp: h5py.Group):
+        self.g = grp
+        self.name = grp.name.lstrip("/")
+        self.n_molecules = int(grp.attrs["n_molecules"])
+        self.n_atoms = int(grp.attrs["n_atoms"])
+        self.level_of_theory = str(grp.attrs["level_of_theory"])
+        self.geometry = str(grp.attrs["geometry"])
+        self.gdb_range = str(grp.attrs["gdb_range"])
+        # atom row offsets: start[i]..start[i+1] for molecule i (one cumsum, cached)
+        na = grp["n_atoms"][:].astype(np.int64)
+        self._start = np.empty(self.n_molecules + 1, np.int64)
+        self._start[0] = 0
+        np.cumsum(na, out=self._start[1:])
+        # smiles buffer (small) + offsets, read once
+        self._sbuf = grp["smiles_buffer"][:].tobytes()
+        self._soff = grp["smiles_offsets"][:]
+        self._failed = set(int(x) for x in grp["failed_indices"][:])
+        self.shieldings = sorted(k for k in grp.keys() if k.startswith("shielding_"))
+
+    def _check(self, index: int):
+        if not 0 <= index < self.n_molecules:
+            raise IndexError(f"molecule index {index} out of range [0, {self.n_molecules})")
+
+    def smiles(self, index: int) -> str:
+        """Return molecule `index`'s SMILES string."""
+        self._check(index)
+        return self._sbuf[self._soff[index]:self._soff[index + 1]].decode("utf-8")
+
+    def _scaled(self, name: str, sl: slice) -> np.ndarray:
+        d = self.g[name]
+        return d[sl].astype(np.float64) * float(d.attrs["scale"])
+
+    def molecule(self, index: int) -> dict:
+        """Return one molecule's decoded data (physical units) by array index.
+
+        Keys: index, smiles, atomic_numbers, coordinates (n,3 Angstrom), one entry per
+        shielding dataset (n, ppm), and `failed`. If `failed` is True the shieldings are
+        sentinel zeros and should be ignored; geometry and SMILES are still valid.
+        """
+        self._check(index)
+        a, b = int(self._start[index]), int(self._start[index + 1])
+        sl = slice(a, b)
+        out = {
+            "index": index,
+            "smiles": self.smiles(index),
+            "atomic_numbers": self.g["atomic_numbers"][sl],
+            "coordinates": self._scaled("coordinates", sl),   # (n,3) Angstrom
+            "failed": index in self._failed,
+        }
+        for s in self.shieldings:
+            out[s] = self._scaled(s, sl)                       # (n,) ppm
+        return out
+
+    def __len__(self):
+        return self.n_molecules
+
+    def __repr__(self):
+        return (f"<Component /{self.name}: {self.n_molecules:,} molecules, "
+                f"{self.n_atoms:,} atoms, {self.level_of_theory}>")
+
+
+class SigmaPepper:
+    """Reader for a sigma-pepper HDF5 file; exposes each level-of-theory group as a Component."""
+
+    def __init__(self, path: str):
+        self.f = h5py.File(path, "r")
+        self.components = {name: Component(self.f[name]) for name in self.f.keys()
+                          if isinstance(self.f[name], h5py.Group)}
+
+    @property
+    def group_names(self):
+        """List of level-of-theory group names in the file."""
+        return list(self.components)
+
+    def __getitem__(self, name: str) -> Component:
+        return self.components[name]
+
+    def close(self):
+        """Close the underlying HDF5 file."""
+        self.f.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        self.close()
+
+
+if __name__ == "__main__":
+    import sys
+    path = sys.argv[1] if len(sys.argv) > 1 else "sigma-pepper.hdf5"
+    with SigmaPepper(path) as ds:
+        for name in ds.group_names:
+            c = ds[name]
+            print(c)
+            m = c.molecule(0)
+            extra = "  ".join(f"{s}={np.round(m[s][:3], 4)}" for s in c.shieldings)
+            print(f"  mol 0: {m['smiles']}  Z={list(m['atomic_numbers'][:6])}  failed={m['failed']}")
+            print(f"         {extra}")

data/sigma-pepper/test_sigma_pepper.py

@@ -0,0 +1,197 @@
+"""
+Tests for the sigma-pepper reader (decode_sigma_pepper.SigmaPepper).
+
+These build a tiny synthetic two-group dataset in the *exact* on-disk format (so the fixture
+also serves as an executable spec of the layout) and exercise the decoder against it. No
+dependency on the multi-hundred-MB real file. One opt-in test runs against the real
+`sigma-pepper.hdf5` if it is present next to this file.
+
+Run:  pytest test_sigma_pepper.py -q
+Requires: pytest, numpy, h5py.
+"""
+import os
+import numpy as np
+import h5py
+import pytest
+
+from decode_sigma_pepper import SigmaPepper
+
+SCALE = 1e-4
+
+
+def _i32(x):
+    return np.round(np.asarray(x, float) / SCALE).astype(np.int32)
+
+
+def _write_group(f, name, *, n_atoms, znums, smiles, coords, shieldings, failed, attrs):
+    """shieldings: dict name->per-atom array; failed: list of molecule indices."""
+    g = f.create_group(name)
+    g.attrs.update(attrs)
+    g.attrs["n_molecules"] = len(n_atoms)
+    g.attrs["n_atoms"] = int(np.sum(n_atoms))
+    assert len(znums) == int(np.sum(n_atoms))
+    g.create_dataset("atomic_numbers", data=np.asarray(znums, np.int8))
+    g.create_dataset("n_atoms", data=np.asarray(n_atoms, np.uint8))
+    d = g.create_dataset("coordinates", data=_i32(coords)); d.attrs["scale"] = SCALE; d.attrs["units"] = "angstrom"
+    for sname, arr in shieldings.items():
+        d = g.create_dataset(sname, data=_i32(arr)); d.attrs["scale"] = SCALE; d.attrs["units"] = "ppm"
+    g.create_dataset("failed_indices", data=np.asarray(failed, np.int32))
+    enc = [s.encode("utf-8") for s in smiles]
+    off = np.concatenate([[0], np.cumsum([len(e) for e in enc])]).astype(np.int64)
+    g.create_dataset("smiles_buffer", data=np.frombuffer(b"".join(enc), np.uint8))
+    g.create_dataset("smiles_offsets", data=off)
+
+
+# group A: 3 molecules (2,1,3 atoms = 6), one shielding, molecule 1 failed
+A_NATOMS = [2, 1, 3]
+A_Z = [6, 1,   8,   6, 7, 1]            # heteroatoms exercise the passthrough
+A_SMILES = ["CC", "O", "CCC"]
+A_COORDS = np.array([[0, 0, 0], [1.2345, 0, 0],
+                     [5, 5, 5],
+                     [3, 0, 0], [4.1, 0, 0], [3, 1.5, 0]], float)
+A_PBE0 = np.array([180.1234, 31.0, 0.0, 150.0, 151.5, 29.9])   # mol 1 (atom idx 2) = sentinel 0
+# group B: 2 molecules (3,2 atoms = 5), two shieldings, molecule 0 failed
+B_NATOMS = [3, 2]
+B_Z = [6, 8, 1,   6, 7]
+B_SMILES = ["CCO", "CN"]
+B_COORDS = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0],
+                     [9, 9, 9], [9.5, 9, 9]], float)
+B_WB97 = np.array([0.0, 0.0, 0.0, 110.0, 28.0])    # mol 0 failed (all zero)
+B_WP04 = np.array([0.0, 0.0, 0.0, 98.0, 28.3])
+
+
+@pytest.fixture(scope="module")
+def mini_path(tmp_path_factory):
+    p = str(tmp_path_factory.mktemp("sp") / "mini.hdf5")
+    with h5py.File(p, "w") as f:
+        _write_group(f, "pbe0_pcSseg1", n_atoms=A_NATOMS, znums=A_Z, smiles=A_SMILES, coords=A_COORDS,
+                     shieldings={"shielding_pbe0": A_PBE0}, failed=[1],
+                     attrs=dict(level_of_theory="PBE0/pcSseg-1, gas",
+                                geometry="AIMNet2 stationary (fmax=0.01)", gdb_range="1-10 heavy"))
+        _write_group(f, "wp04_wb97xd_pcSseg2", n_atoms=B_NATOMS, znums=B_Z, smiles=B_SMILES, coords=B_COORDS,
+                     shieldings={"shielding_wb97xd": B_WB97, "shielding_wp04": B_WP04}, failed=[0],
+                     attrs=dict(level_of_theory="WP04 + wB97X-D /pcSseg-2, gas",
+                                geometry="AIMNet2 stationary (fmax=1e-4)", gdb_range="1-9 heavy"))
+    return p
+
+
+@pytest.fixture
+def ds(mini_path):
+    with SigmaPepper(mini_path) as d:
+        yield d
+
+
+def test_group_names(ds):
+    assert set(ds.group_names) == {"pbe0_pcSseg1", "wp04_wb97xd_pcSseg2"}
+
+
+def test_component_dimensions(ds):
+    a = ds["pbe0_pcSseg1"]
+    assert a.n_molecules == 3 and a.n_atoms == 6
+    b = ds["wp04_wb97xd_pcSseg2"]
+    assert b.n_molecules == 2 and b.n_atoms == 5
+    assert b.shieldings == ["shielding_wb97xd", "shielding_wp04"]
+
+
+def test_offsets_slice_each_molecule(ds):
+    a = ds["pbe0_pcSseg1"]
+    assert len(a.molecule(0)["atomic_numbers"]) == 2
+    assert len(a.molecule(1)["atomic_numbers"]) == 1
+    assert len(a.molecule(2)["atomic_numbers"]) == 3
+
+
+def test_values_roundtrip_and_smiles(ds):
+    m = ds["pbe0_pcSseg1"].molecule(2)
+    assert m["smiles"] == "CCC"
+    np.testing.assert_array_equal(m["atomic_numbers"], [6, 7, 1])   # heteroatom passthrough
+    np.testing.assert_allclose(m["coordinates"], A_COORDS[3:6], atol=5e-5)
+    np.testing.assert_allclose(m["shielding_pbe0"], A_PBE0[3:6], atol=5e-5)
+    assert not m["failed"]
+
+
+def test_two_shieldings_present(ds):
+    m = ds["wp04_wb97xd_pcSseg2"].molecule(1)
+    assert m["smiles"] == "CN"
+    np.testing.assert_allclose(m["shielding_wp04"], B_WP04[3:5], atol=5e-5)
+    np.testing.assert_allclose(m["shielding_wb97xd"], B_WB97[3:5], atol=5e-5)
+
+
+def test_failed_flag(ds):
+    assert ds["pbe0_pcSseg1"].molecule(1)["failed"] is True
+    assert ds["pbe0_pcSseg1"].molecule(0)["failed"] is False
+    b0 = ds["wp04_wb97xd_pcSseg2"].molecule(0)
+    assert b0["failed"] is True
+    # failed molecule's shieldings are sentinel zeros (the whole point of the flag)
+    assert np.all(b0["shielding_wp04"] == 0) and np.all(b0["shielding_wb97xd"] == 0)
+    # geometry/smiles still valid for a failed molecule
+    assert b0["smiles"] == "CCO"
+    np.testing.assert_allclose(b0["coordinates"], B_COORDS[0:3], atol=5e-5)
+
+
+def test_empty_failed_indices(tmp_path):
+    p = str(tmp_path / "nofail.hdf5")
+    with h5py.File(p, "w") as f:
+        _write_group(f, "pbe0_pcSseg1", n_atoms=[2], znums=[6, 1], smiles=["CO"],
+                     coords=np.array([[0, 0, 0], [1.0, 0, 0]], float),
+                     shieldings={"shielding_pbe0": np.array([150.0, 30.0])}, failed=[],
+                     attrs=dict(level_of_theory="x", geometry="x", gdb_range="x"))
+    with SigmaPepper(p) as d:
+        assert d["pbe0_pcSseg1"].molecule(0)["failed"] is False
+
+
+def test_smiles_accessor(ds):
+    assert ds["pbe0_pcSseg1"].smiles(0) == "CC"
+    assert ds["wp04_wb97xd_pcSseg2"].smiles(1) == "CN"
+
+
+def test_index_out_of_range_raises(ds):
+    with pytest.raises(IndexError):
+        ds["pbe0_pcSseg1"].molecule(-1)
+    with pytest.raises(IndexError):
+        ds["pbe0_pcSseg1"].molecule(3)
+
+
+def test_context_manager_closes(mini_path):
+    d = SigmaPepper(mini_path)
+    d.close()
+    assert not d.f.id.valid
+
+
+# ---------- opt-in: real file, if present ----------
+REAL = os.path.join(os.path.dirname(__file__), "sigma-pepper.hdf5")
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real sigma-pepper.hdf5 not present")
+def test_real_file():
+    with SigmaPepper(REAL) as ds:
+        a = ds["pbe0_pcSseg1"]
+        b = ds["wp04_wb97xd_pcSseg2"]
+        assert a.n_molecules == 2_167_210 and a.n_atoms == 45_418_890
+        assert b.n_molecules == 320_105 and b.n_atoms == 6_113_101
+        assert len(a.g["failed_indices"]) == 1
+        assert len(b.g["failed_indices"]) == 472
+        # column identity: WP04 must be lower than wB97XD on carbon
+        Z = b.g["atomic_numbers"][:2_000_000]
+        wp = b.g["shielding_wp04"][:2_000_000].astype(float) * SCALE
+        wb = b.g["shielding_wb97xd"][:2_000_000].astype(float) * SCALE
+        C = (Z == 6) & (wp != 0) & (wb != 0)
+        assert wp[C].mean() < wb[C].mean()
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real sigma-pepper.hdf5 not present")
+def test_real_file_reproduces_table_s8_site_counts():
+    """SI Table S8's 'MagNET-Zero (initial round)' row is this file's pbe0_pcSseg1 group exactly.
+    The 'MagNET-Zero (fine-tuning)' row is this group's wp04_wb97xd_pcSseg2 PLUS sigma-concentrate
+    (the 50,000 reused molecules) -- see test_sigma_concentrate.py's matching test for that half."""
+    with SigmaPepper(REAL) as ds:
+        a = ds["pbe0_pcSseg1"]
+        z_a = a.g["atomic_numbers"][:]
+        assert int((z_a == 1).sum()) == 24_047_999
+        assert int((z_a == 6).sum()) == 15_330_361
+
+        b = ds["wp04_wb97xd_pcSseg2"]
+        z_b = b.g["atomic_numbers"][:]
+        # this group alone is not a Table S8 row (fine-tuning adds sigma-concentrate's atoms too),
+        # but its counts are the fixed half of that sum, so lock them down here too.
+        assert int((z_b == 1).sum()) == 3_287_698
+        assert int((z_b == 6).sum()) == 2_028_642

data/sigma-shake/decode_sigma_shake.py

@@ -0,0 +1,123 @@
+"""
+decode_sigma_shake.py - reader/decoder for the sigma-shake HDF5 dataset.
+
+sigma-shake stores PBE0/pcSseg-1 isotropic NMR shieldings for the stationary and
+vibrationally-perturbed geometries of ~4.79M GDB-derived molecules. It is the training
+data for the MagNET foundation model.
+
+Storage format:
+  - Per-atom arrays are concatenated across all molecules; molecule i occupies
+    rows [atom_start[i], atom_end[i]).
+  - Coordinates, perturbation, and shieldings are stored as int32 fixed-point.
+    Physical value = stored_integer * scale, where `scale` (= 1e-4) is an attribute
+    on each dataset. Reconstruction error is <= 5e-5 (Angstrom or ppm).
+  - The perturbed geometry is NOT stored directly; it is
+        perturbed_coords = stationary_coords + perturbation
+  - Where `incomplete_mask` is True, the molecule has no perturbed calculation:
+    `perturbation` and `shielding_perturbed` are sentinel zeros and must be ignored.
+
+Requires: Python >= 3.7, numpy, h5py.  No other dependencies.
+"""
+from __future__ import annotations
+import numpy as np
+import h5py
+
+
+class SigmaShake:
+    """Reader for the sigma-shake HDF5 file: stationary and vibrationally-perturbed shieldings."""
+
+    def __init__(self, path: str):
+        self.f = h5py.File(path, "r")
+        self._id = self.f["molecule_id"][:]
+        self._start = self.f["atom_start"][:]
+        self._end = self.f["atom_end"][:]
+        self._id_to_row = None  # built lazily on first by_id()
+
+    # ---- scalars / metadata ----
+    @property
+    def n_molecules(self) -> int:
+        """Total number of molecules in the file."""
+        return self._id.shape[0]
+
+    @property
+    def n_atoms(self) -> int:
+        """Total number of atoms across all molecules."""
+        return int(self.f.attrs["n_atoms"])
+
+    def _scaled(self, name: str, sl: slice) -> np.ndarray:
+        d = self.f[name]
+        return d[sl].astype(np.float64) * float(d.attrs["scale"])
+
+    # ---- molecule access ----
+    def molecule(self, index: int) -> dict:
+        """Return one molecule's decoded data by ROW index (0..n_molecules-1).
+
+        For an incomplete molecule (no perturbed calculation) the returned dict OMITS the
+        'perturbation', 'perturbed_coords', and 'shielding_perturbed' keys.
+        """
+        if not 0 <= index < self.n_molecules:
+            raise IndexError(f"molecule index {index} out of range [0, {self.n_molecules})")
+        a, b = int(self._start[index]), int(self._end[index])
+        sl = slice(a, b)
+        Z = self.f["atomic_numbers"][sl]
+        stat = self._scaled("stationary_coords", sl)
+        incomplete = bool(self.f["incomplete_mask"][a])  # constant within a molecule
+        out = {
+            "molecule_id": int(self._id[index]),
+            "atomic_numbers": Z,
+            "stationary_coords": stat,                # (n,3) Angstrom
+            "shielding_stationary": self._scaled("shielding_stationary", sl),  # (n,) ppm
+            "incomplete": incomplete,
+        }
+        if not incomplete:
+            pert = self._scaled("perturbation", sl)
+            out["perturbed_coords"] = stat + pert     # reconstruct absolute geometry
+            out["perturbation"] = pert
+            out["shielding_perturbed"] = self._scaled("shielding_perturbed", sl)
+        return out
+
+    def by_id(self, molecule_id: int) -> dict:
+        """Return one molecule's decoded data (see `molecule`) looked up by molecule_id."""
+        if self._id_to_row is None:
+            self._id_to_row = {int(v): i for i, v in enumerate(self._id)}
+        try:
+            row = self._id_to_row[int(molecule_id)]
+        except KeyError:
+            raise KeyError(f"unknown molecule_id {molecule_id}") from None
+        return self.molecule(row)
+
+    # ---- splits ----
+    def split(self, which: str, effective: bool = True) -> np.ndarray:
+        """Return molecule_id values for 'train'/'val'/'test'.
+        effective=True drops remove_molecule_ids (the load-time filter), reproducing the
+        sizes reported in the SI (train 3,730,572 / val 329,154 / test 329,163)."""
+        ids = self.f[f"split_{which}_ids"][:]
+        if effective:
+            rm = self.f["remove_molecule_ids"][:]
+            ids = ids[~np.isin(ids, rm)]
+        return ids
+
+    def close(self):
+        """Close the underlying HDF5 file."""
+        self.f.close()
+
+    def __enter__(self): return self
+    def __exit__(self, *a): self.close()
+
+
+if __name__ == "__main__":
+    import sys
+    path = sys.argv[1] if len(sys.argv) > 1 else "sigma-shake.hdf5"
+    with SigmaShake(path) as ds:
+        print(f"molecules: {ds.n_molecules:,}   atoms: {ds.n_atoms:,}")
+        for w in ("train", "val", "test"):
+            print(f"  effective {w:5s}: {len(ds.split(w)):,}")
+        m = ds.molecule(0)
+        print(f"\nmolecule 0 (id={m['molecule_id']}): {len(m['atomic_numbers'])} atoms, "
+              f"incomplete={m['incomplete']}")
+        print("  Z         :", m["atomic_numbers"][:8])
+        print("  stat xyz  :", np.round(m["stationary_coords"][0], 4), "...")
+        print("  sigma_stat:", np.round(m["shielding_stationary"][:4], 4), "...")
+        if not m["incomplete"]:
+            print("  pert xyz  :", np.round(m["perturbed_coords"][0], 4), "...")
+            print("  sigma_pert:", np.round(m["shielding_perturbed"][:4], 4), "...")

data/sigma-shake/test_sigma_shake.py

@@ -0,0 +1,161 @@
+"""
+Tests for the sigma-shake reader (decode_sigma_shake.SigmaShake).
+
+These build a tiny synthetic dataset in the *exact* on-disk format (so the fixture also
+serves as an executable spec of the layout) and exercise the decoder against it. No
+dependency on the multi-GB real file. One opt-in test runs against the real
+`sigma-shake.hdf5` if it is present next to this file.
+
+Run:  pytest test_sigma_shake.py -q
+Requires: pytest, numpy, h5py.
+"""
+import os
+import numpy as np
+import h5py
+import pytest
+
+from decode_sigma_shake import SigmaShake
+
+SCALE = 1e-4
+
+# ---- known synthetic content: 3 molecules, 9 atoms; molecule 1 is incomplete ----
+MOL_ID = np.array([100, 200, 300], np.int32)
+ATOM_START = np.array([0, 3, 5], np.int32)
+ATOM_END = np.array([3, 5, 9], np.int32)
+Z = np.array([6, 1, 1,   8, 1,   6, 6, 1, 1], np.int8)
+STAT_XYZ = np.array([
+    [0.0000, 0.0000, 0.0000], [1.0000, 0.0000, 0.0000], [0.0000, 1.0000, 0.0000],   # mol 100
+    [2.0000, 2.0000, 2.0000], [2.5000, 2.0000, 2.0000],                             # mol 200 (incomplete)
+    [3.0000, 0.0000, 0.0000], [4.0000, 0.0000, 0.0000], [3.0000, 1.0000, 0.0000], [4.0000, 1.0000, 0.0000],
+], np.float64)
+STAT_SIG = np.array([150.1234, 30.5678, 31.0000,   200.0000, 25.0000,   140.0000, 141.0000, 29.0000, 29.5000])
+PERT_XYZ = np.array([
+    [0.0100, 0.0000, 0.0000], [-0.0050, 0.0000, 0.0000], [0.0000, 0.0200, 0.0000],  # mol 100
+    [0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000],                             # mol 200 (zeroed)
+    [0.0030, 0.0000, 0.0000], [-0.0030, 0.0000, 0.0000], [0.0000, 0.0050, 0.0000], [0.0000, -0.0050, 0.0000],
+], np.float64)
+PERT_SIG = np.array([151.0000, 30.9999, 31.5000,   0.0, 0.0,   140.5000, 141.5000, 29.2000, 29.7000])
+INCOMPLETE_MASK = np.array([0, 0, 0, 1, 1, 0, 0, 0, 0], bool)
+SPLIT = {"train": np.array([100, 200], np.int32), "val": np.array([300], np.int32),
+         "test": np.array([], np.int32)}
+REMOVE = np.array([200], np.int32)
+INCOMPLETE_IDS = np.array([200], np.int32)
+
+
+def _i32(x):
+    return np.round(np.asarray(x) / SCALE).astype(np.int32)
+
+
+@pytest.fixture(scope="module")
+def mini_path(tmp_path_factory):
+    p = str(tmp_path_factory.mktemp("ss") / "mini.hdf5")
+    with h5py.File(p, "w") as f:
+        f.attrs["n_atoms"] = len(Z)
+        f.attrs["n_molecules"] = len(MOL_ID)
+        for name, arr in [("stationary_coords", STAT_XYZ), ("perturbation", PERT_XYZ),
+                          ("shielding_stationary", STAT_SIG), ("shielding_perturbed", PERT_SIG)]:
+            d = f.create_dataset(name, data=_i32(arr))
+            d.attrs["scale"] = SCALE
+        f.create_dataset("atomic_numbers", data=Z)
+        f.create_dataset("incomplete_mask", data=INCOMPLETE_MASK)
+        f.create_dataset("molecule_id", data=MOL_ID)
+        f.create_dataset("atom_start", data=ATOM_START)
+        f.create_dataset("atom_end", data=ATOM_END)
+        f.create_dataset("remove_molecule_ids", data=REMOVE)
+        f.create_dataset("incomplete_molecule_ids", data=INCOMPLETE_IDS)
+        for k, v in SPLIT.items():
+            f.create_dataset(f"split_{k}_ids", data=v)
+    return p
+
+
+@pytest.fixture
+def ds(mini_path):
+    with SigmaShake(mini_path) as d:
+        yield d
+
+
+def test_dimensions(ds):
+    assert ds.n_molecules == 3
+    assert ds.n_atoms == 9
+
+
+def test_complete_molecule_values(ds):
+    m = ds.molecule(0)
+    assert m["molecule_id"] == 100
+    assert not m["incomplete"]
+    np.testing.assert_array_equal(m["atomic_numbers"], [6, 1, 1])
+    np.testing.assert_allclose(m["stationary_coords"], STAT_XYZ[0:3], atol=5e-5)
+    np.testing.assert_allclose(m["shielding_stationary"], STAT_SIG[0:3], atol=5e-5)
+    np.testing.assert_allclose(m["shielding_perturbed"], PERT_SIG[0:3], atol=5e-5)
+
+
+def test_perturbed_geometry_reconstruction(ds):
+    m = ds.molecule(0)
+    np.testing.assert_allclose(m["perturbed_coords"], STAT_XYZ[0:3] + PERT_XYZ[0:3], atol=1e-4)
+
+
+def test_incomplete_molecule_omits_perturbed_keys(ds):
+    m = ds.molecule(1)
+    assert m["molecule_id"] == 200
+    assert m["incomplete"] is True
+    for k in ("perturbation", "perturbed_coords", "shielding_perturbed"):
+        assert k not in m
+    # stationary data still valid
+    np.testing.assert_allclose(m["stationary_coords"], STAT_XYZ[3:5], atol=5e-5)
+    np.testing.assert_allclose(m["shielding_stationary"], STAT_SIG[3:5], atol=5e-5)
+
+
+def test_scale_roundtrip_is_within_tolerance(ds):
+    # every stored value should reconstruct to <= half the LSB (5e-5)
+    m = ds.molecule(2)
+    assert np.abs(m["stationary_coords"] - STAT_XYZ[5:9]).max() <= 5e-5
+    assert np.abs(m["shielding_stationary"] - STAT_SIG[5:9]).max() <= 5e-5
+
+
+def test_by_id(ds):
+    assert ds.by_id(100)["molecule_id"] == 100
+    assert ds.by_id(300)["molecule_id"] == 300
+
+
+def test_by_id_unknown_raises(ds):
+    with pytest.raises(KeyError):
+        ds.by_id(999999)
+
+
+def test_index_out_of_range_raises(ds):
+    with pytest.raises(IndexError):
+        ds.molecule(-1)
+    with pytest.raises(IndexError):
+        ds.molecule(3)
+
+
+def test_split_effective_filters_remove_ids(ds):
+    # raw train has 100 and 200; effective drops 200 (in remove_molecule_ids)
+    np.testing.assert_array_equal(np.sort(ds.split("train", effective=False)), [100, 200])
+    np.testing.assert_array_equal(ds.split("train"), [100])
+    np.testing.assert_array_equal(ds.split("val"), [300])
+    assert len(ds.split("test")) == 0
+
+
+def test_context_manager_closes(mini_path):
+    d = SigmaShake(mini_path)
+    d.close()
+    assert not d.f.id.valid
+
+
+# ---------- opt-in: real file, if present ----------
+REAL = os.path.join(os.path.dirname(__file__), "sigma-shake.hdf5")
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real sigma-shake.hdf5 not present")
+def test_real_file_matches_si():
+    with SigmaShake(REAL) as ds:
+        assert ds.n_molecules == 4_787_816
+        assert ds.n_atoms == 135_219_803
+        assert len(ds.split("train")) == 3_730_572
+        assert len(ds.split("val")) == 329_154
+        assert len(ds.split("test")) == 329_163
+        m = ds.molecule(0)
+        assert m["molecule_id"] == 24553
+        # spot round-trip: perturbed = stationary + perturbation, bond-length-scale accuracy
+        assert m["perturbed_coords"].shape == m["stationary_coords"].shape

data/supertestset_magnet_x/decode_supertest.py

@@ -0,0 +1,148 @@
+"""
+decode_supertest.py - reader/decoder for the supertestset_magnet_x HDF5 dataset.
+
+This is the explicit-solvent test set for MagNET-x. It holds DFT (Gaussian) NMR shieldings for
+natural-product solutes computed two ways: with the solute isolated, and with the solute sitting
+in an explicit box of solvent molecules sampled from an OpenMM molecular-dynamics simulation. The
+explicit-solvent correction at each atom is the solvated shielding minus the isolated shielding,
+which is what MagNET-x learns to predict.
+
+Each record is one computed structure for a (molecule, solvent, conformer) combination, in either
+the isolated or the solvated state:
+
+  isolated   the full system is just the solute
+  solvated   the full system is the solute (always the first atoms) followed by solvent molecules
+
+The shieldings are stored only for the solute atoms (the solvent carries no NMR signal of interest).
+There are 318 records covering 8 molecules across 4 solvents and 5 conformers. A separate catalog
+lists all 50 candidate molecules of the test set with their SMILES.
+
+Storage format:
+  - Per-atom arrays are concatenated across records. Record i has `n_full_atoms[i]` atoms in its
+    full system and `n_solute_atoms[i]` solute atoms (the leading atoms of the full system). Its
+    full atomic numbers and coordinates occupy `n_full_atoms[i]` rows; its solute shieldings occupy
+    `n_solute_atoms[i]` rows. Offsets are reconstructed once, on open.
+  - coordinates and shieldings are int32 fixed point: physical = stored * 1e-4, reconstruction
+    error <= 5e-5. A missing value would be the marker -2147483648 -> NaN.
+
+Requires: Python >= 3.7, numpy, h5py. No other dependencies.
+"""
+from __future__ import annotations
+import numpy as np
+import h5py
+
+_SCALE = 1e-4
+_MARKER = -2147483648
+
+
+def _decode(values: np.ndarray) -> np.ndarray:
+    """int32 fixed-point -> float64 physical units; marker -> NaN. No-op on float input."""
+    values = np.asarray(values)
+    if not np.issubdtype(values.dtype, np.integer):
+        return np.asarray(values, dtype=np.float64)
+    out = values.astype(np.float64) * _SCALE
+    out[values == _MARKER] = np.nan
+    return out
+
+
+class SuperTestMagNETX:
+    """Reader for supertestset_magnet_x.hdf5. Use as a context manager."""
+
+    def __init__(self, path: str):
+        self.f = h5py.File(path, "r")
+        self.n_records = int(self.f.attrs["n_records"])
+        self.record_names = [str(s) for s in self.f["record_names"].asstr()[:]]
+        self.record_solvents = [str(s) for s in self.f["record_solvents"].asstr()[:]]
+        self._conf = self.f["conformer"][:]
+        self._iso = self.f["is_isolated"][:].astype(bool)
+        self._n_sol = self.f["n_solute_atoms"][:].astype(np.int64)
+        self._n_full = self.f["n_full_atoms"][:].astype(np.int64)
+        self._fstart = np.empty(self.n_records + 1, np.int64)
+        self._fstart[0] = 0
+        np.cumsum(self._n_full, out=self._fstart[1:])
+        self._sstart = np.empty(self.n_records + 1, np.int64)
+        self._sstart[0] = 0
+        np.cumsum(self._n_sol, out=self._sstart[1:])
+
+    # ------------------------------------------------------------------ catalog
+    def catalog(self) -> dict:
+        """The 50-molecule candidate catalog as a {name: SMILES} dict."""
+        names = self.f["catalog_names"].asstr()[:]
+        smiles = self.f["catalog_smiles"].asstr()[:]
+        return {str(n): str(s) for n, s in zip(names, smiles)}
+
+    # ------------------------------------------------------------------ records
+    def _check(self, index: int):
+        if not 0 <= index < self.n_records:
+            raise IndexError(f"record index {index} out of range [0, {self.n_records})")
+
+    def record(self, index: int) -> dict:
+        """Return one record (physical units) by position.
+
+        Keys: name, solvent, conformer (int), is_isolated (bool); full_atomic_numbers (n_full,),
+        full_coordinates (n_full, 3 Angstrom); solute_atomic_numbers (n_solute,), solute_shieldings
+        (n_solute,) ppm. The solute atoms are the leading n_solute atoms of the full system.
+        """
+        self._check(index)
+        fsl = slice(int(self._fstart[index]), int(self._fstart[index + 1]))
+        ssl = slice(int(self._sstart[index]), int(self._sstart[index + 1]))
+        ns = int(self._n_sol[index])
+        full_an = self.f["full_atomic_numbers"][fsl]
+        return {
+            "name": self.record_names[index],
+            "solvent": self.record_solvents[index],
+            "conformer": int(self._conf[index]),
+            "is_isolated": bool(self._iso[index]),
+            "full_atomic_numbers": full_an,
+            "full_coordinates": _decode(self.f["full_coordinates"][fsl]),
+            "solute_atomic_numbers": full_an[:ns],
+            "solute_shieldings": _decode(self.f["solute_shieldings"][ssl]),
+        }
+
+    def find(self, name: str, solvent: str, conformer: int, is_isolated: bool) -> int:
+        """Index of the record matching (name, solvent, conformer, is_isolated), or raise KeyError."""
+        for i in range(self.n_records):
+            if (self.record_names[i] == name and self.record_solvents[i] == solvent
+                    and int(self._conf[i]) == conformer and bool(self._iso[i]) == is_isolated):
+                return i
+        raise KeyError(f"no record for {name!r}, {solvent!r}, conformer {conformer}, "
+                       f"isolated={is_isolated}")
+
+    def explicit_correction(self, name: str, solvent: str, conformer: int) -> np.ndarray:
+        """The explicit-solvent correction (solvated minus isolated solute shieldings), (n_solute,)
+        ppm, for one (molecule, solvent, conformer)."""
+        solvated = self.record(self.find(name, solvent, conformer, is_isolated=False))
+        isolated = self.record(self.find(name, solvent, conformer, is_isolated=True))
+        return solvated["solute_shieldings"] - isolated["solute_shieldings"]
+
+    def close(self):
+        """Close the underlying HDF5 file."""
+        self.f.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        self.close()
+
+    def __len__(self):
+        return self.n_records
+
+    def __repr__(self):
+        n_iso = int(self._iso.sum())
+        return (f"<SuperTestMagNETX: {self.n_records} records "
+                f"({n_iso} isolated, {self.n_records - n_iso} solvated), "
+                f"solvents {self.f.attrs['solvents']}>")
+
+
+if __name__ == "__main__":
+    import sys
+    path = sys.argv[1] if len(sys.argv) > 1 else "supertestset_magnet_x.hdf5"
+    with SuperTestMagNETX(path) as ds:
+        print(ds)
+        print(f"  catalog: {len(ds.catalog())} molecules")
+        r = ds.record(ds.find(ds.record_names[0], ds.record_solvents[0], int(ds._conf[0]), bool(ds._iso[0])))
+        corr = ds.explicit_correction(r["name"], r["solvent"], r["conformer"])
+        print(f"  {r['name']} in {r['solvent']} (conformer {r['conformer']}): "
+              f"{len(r['solute_atomic_numbers'])} solute atoms")
+        print(f"    explicit-solvent correction[:3] = {np.round(corr[:3], 4)} ppm")

data/supertestset_magnet_x/test_supertest.py

@@ -0,0 +1,110 @@
+"""Tests for data/supertestset_magnet_x/decode_supertest.py.
+
+A tiny synthetic file in the stored int32 format exercises the reader, the two-level block
+reconstruction (full system vs solute prefix), the isolated/solvated pairing, and the
+explicit-solvent correction, without the real file. An opt-in smoke test runs against the real
+file when present.
+"""
+import os
+import sys
+
+import numpy as np
+import pytest
+import h5py
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, HERE)
+
+import decode_supertest as D  # noqa: E402
+
+SCALE = 1e-4
+MARKER = -2147483648
+REAL = os.path.join(HERE, "supertestset_magnet_x.hdf5")
+
+
+def _encode(values):
+    v = np.asarray(values, dtype=np.float64)
+    out = np.round(v / SCALE).astype(np.int64)
+    out[~np.isfinite(v)] = MARKER
+    return out.astype(np.int32)
+
+
+def _build_synthetic(path):
+    """One molecule, one solvent, one conformer: an isolated record (3 solute atoms) and a
+    solvated record (3 solute + 4 solvent = 7 full atoms)."""
+    rng = np.random.default_rng(0)
+    solute_an = np.array([6, 1, 8], dtype=np.uint8)
+    iso_sh = np.array([100.0, 30.0, 250.0])
+    solv_sh = iso_sh + np.array([-0.3, 0.1, -0.5])     # known correction
+    full_an_solv = np.concatenate([solute_an, [1, 1, 1, 1]]).astype(np.uint8)
+    records = [
+        # name, solvent, conf, is_isolated, full_an, full_xyz, n_solute, solute_sh
+        ("MOLX", "benzene", 0, 1, solute_an, rng.normal(0, 2, (3, 3)), 3, iso_sh),
+        ("MOLX", "benzene", 0, 0, full_an_solv, rng.normal(0, 2, (7, 3)), 3, solv_sh),
+    ]
+    str_dt = h5py.string_dtype("utf-8")
+    opts = dict(compression="gzip", shuffle=True)
+    names, solvents, confs, iso, n_sol, n_full = [], [], [], [], [], []
+    an_b, co_b, sh_b = [], [], []
+    for name, solv, conf, is_iso, fan, fxyz, ns, ssh in records:
+        names.append(name); solvents.append(solv); confs.append(conf); iso.append(is_iso)
+        n_sol.append(ns); n_full.append(len(fan))
+        an_b.append(fan); co_b.append(_encode(fxyz)); sh_b.append(_encode(ssh))
+    with h5py.File(path, "w") as f:
+        f.attrs["n_records"] = len(records)
+        f.attrs["n_catalog"] = 1
+        f.attrs["solvents"] = "benzene"
+        f.attrs["scale"] = SCALE
+        f.attrs["missing_value_marker"] = MARKER
+        f.create_dataset("catalog_names", data=np.array(["MOLX"], dtype=object), dtype=str_dt)
+        f.create_dataset("catalog_smiles", data=np.array(["CCO"], dtype=object), dtype=str_dt)
+        f.create_dataset("record_names", data=np.array(names, dtype=object), dtype=str_dt)
+        f.create_dataset("record_solvents", data=np.array(solvents, dtype=object), dtype=str_dt)
+        f.create_dataset("conformer", data=np.asarray(confs, np.int32), **opts)
+        f.create_dataset("is_isolated", data=np.asarray(iso, np.uint8), **opts)
+        f.create_dataset("n_solute_atoms", data=np.asarray(n_sol, np.int32), **opts)
+        f.create_dataset("n_full_atoms", data=np.asarray(n_full, np.int32), **opts)
+        f.create_dataset("full_atomic_numbers", data=np.concatenate(an_b), **opts)
+        f.create_dataset("full_coordinates", data=np.concatenate(co_b), **opts)
+        f.create_dataset("solute_shieldings", data=np.concatenate(sh_b), **opts)
+
+
+def test_record_reconstruction_and_solute_prefix(tmp_path):
+    path = str(tmp_path / "tiny.hdf5")
+    _build_synthetic(path)
+    with D.SuperTestMagNETX(path) as ds:
+        assert len(ds) == 2
+        iso = ds.record(0)
+        solv = ds.record(1)
+        assert iso["is_isolated"] and not solv["is_isolated"]
+        # isolated full system is just the solute
+        assert len(iso["full_atomic_numbers"]) == 3
+        # solvated full system is solute (first 3) + solvent (4)
+        assert len(solv["full_atomic_numbers"]) == 7
+        assert np.array_equal(solv["solute_atomic_numbers"], solv["full_atomic_numbers"][:3])
+        assert solv["solute_shieldings"].shape == (3,)
+
+
+def test_explicit_correction(tmp_path):
+    path = str(tmp_path / "tiny.hdf5")
+    _build_synthetic(path)
+    with D.SuperTestMagNETX(path) as ds:
+        corr = ds.explicit_correction("MOLX", "benzene", 0)
+        assert np.allclose(corr, [-0.3, 0.1, -0.5], atol=1e-4)
+        assert ds.catalog() == {"MOLX": "CCO"}
+        with pytest.raises(KeyError):
+            ds.find("MOLX", "water", 0, is_isolated=True)
+
+
+def test_decode_marker_to_nan():
+    out = D._decode(np.array([0, 10000, MARKER], dtype=np.int32))
+    assert out[1] == pytest.approx(1.0) and np.isnan(out[2])
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real supertestset_magnet_x.hdf5 not present")
+def test_real_file_smoke():
+    with D.SuperTestMagNETX(REAL) as ds:
+        assert len(ds) == 318
+        assert len(ds.catalog()) == 50
+        corr = ds.explicit_correction("MOL176459", "benzene", 0)
+        assert corr.shape[0] == 39 and np.isfinite(corr).all()

magnet/__init__.py

@@ -0,0 +1,52 @@
+# ruff: noqa
+"""MagNET predicts NMR shieldings using equivariant neural networks.
+
+| function | what you get |
+|---|---|
+| `predict_shifts` | <sup>1</sup>H and <sup>13</sup>C chemical shifts (ppm) in a specific solvent using MagNET-Zero/MagNET-PCM |
+| `predict_shieldings` | <sup>1</sup>H and <sup>13</sup>C shieldings (ppm) in the gas phase |
+| `implicit_solvent_correction` | shieldings(PCM=chloroform) - shieldings(gas phase) |
+| `explicit_solvent_correction` | shieldings(solute+solvents) - shieldings(solute) |
+
+**Inputs and outputs**
+
+- One molecule: array of `atomic_numbers` (e.g. 6 for carbon) and
+  `coordinates` (an N-by-3 array of xyz positions in Angstrom, a numpy array or a nested list). You
+  get one numpy array back, one value per atom.
+- Many molecules: pass a list of each, and you get a list of arrays back.
+
+**Geometries** MagNET-Zero and MagNET-PCM expect AIMNet2-optimized geometries.
+
+**Supported Solvents**
+
+| solvent | `predict_shifts` | `implicit_solvent_correction` | `explicit_solvent_correction` |
+|---|:---:|:---:|:---:|
+| tetrahydrofuran | ✓ | | |
+| dichloromethane | ✓ | | |
+| chloroform | ✓ | ✓ | ✓ |
+| toluene | ✓ | | |
+| benzene | ✓ | | ✓ |
+| chlorobenzene | ✓ | | |
+| acetone | ✓ | | |
+| dimethylsulfoxide | ✓ | | |
+| acetonitrile | ✓ | | |
+| trifluoroethanol | ✓ | | |
+| methanol | ✓ | | ✓ |
+| water | ✓ | | ✓ |
+
+`implicit_solvent_correction` predicts only the chloroform correction; `predict_shifts` linearly scales it to the other solvents.
+
+**Shared options** (all four functions):
+
+- `n_passes` (default `10`): average out equivariance error over `n_passes` forward passes
+- `symmetrize` (default `True`): if True, average over `n_passes` on the input geometry and `n_passes` on the mirror image of the input geometry
+- `device` (default `None`): where to run, a torch device or a string like `"cpu"` or `"cuda"`; uses GPU if available
+
+"""
+__docformat__ = "google"
+
+from .api import (predict_shifts, predict_shieldings, implicit_solvent_correction,
+                  explicit_solvent_correction)
+
+__all__ = ["predict_shifts", "predict_shieldings", "implicit_solvent_correction",
+           "explicit_solvent_correction"]

magnet/api.py

@@ -0,0 +1,309 @@
+"""The MagNET user-facing API.
+
+Four entry points, each taking one molecule (`atomic_numbers` shaped `(N,)` and `coordinates` shaped
+`(N, 3)`) or a list of molecules, and returning the matching shape:
+
+| function | what you get |
+|---|---|
+| `predict_shifts` | <sup>1</sup>H and <sup>13</sup>C chemical shifts (ppm) in a specific solvent using MagNET-Zero/MagNET-PCM |
+| `predict_shieldings` | <sup>1</sup>H and <sup>13</sup>C shieldings (ppm) in the gas phase |
+| `implicit_solvent_correction` | shieldings(PCM=chloroform) - shieldings(gas phase) |
+| `explicit_solvent_correction` | shieldings(solute+solvents) - shieldings(solute) |
+
+**Shared options** (all four functions):
+
+- `n_passes` (default `10`): average out equivariance error over `n_passes` forward passes
+- `symmetrize` (default `True`): if True, average over `n_passes` on the input geometry and `n_passes` on the mirror image of the input geometry
+- `device` (default `None`): where to run, a torch device or a string like `"cpu"` or `"cuda"`; uses GPU if available
+- `checkpoints_dir` (default `None`): directory holding the released weights, if they are not in the current directory or the repo; may be the `model_checkpoints/` folder or its parent, and takes precedence over both
+"""
+__docformat__ = "google"
+
+from collections import Counter
+
+import numpy as np
+
+from . import run_magnet
+from . import scaling
+
+_MODELS = {
+    "MagNET": run_magnet.compute_MagNET_foundation_shieldings,
+    "MagNET-Zero": run_magnet.compute_MagNET_Zero_shieldings,
+}
+
+# Atomic-number counts in one molecule of each MagNET-x solvent, used to validate that the
+# non-solute atoms really are a whole number of molecules of the declared solvent.
+_SOLVENT_COMPOSITION = {
+    "chloroform": Counter({6: 1, 1: 1, 17: 3}),        # CHCl3
+    "benzene": Counter({6: 6, 1: 6}),                  # C6H6
+    "methanol": Counter({6: 1, 8: 1, 1: 4}),           # CH3OH
+    "water": Counter({8: 1, 1: 2}),                    # H2O (run as TIP4P)
+}
+
+
+def _check_passes(n_passes):
+    if not (isinstance(n_passes, (int, np.integer)) and n_passes >= 1):
+        raise ValueError(f"n_passes must be a positive integer; got {n_passes!r}")
+
+
+def _as_batch(atomic_numbers, coordinates):
+    """Normalize a single molecule or a list of molecules to parallel lists. Returns
+    (atomic_numbers_list, coordinates_list, was_single). Validates matching atom counts."""
+    if len(coordinates) == 0:
+        raise ValueError("empty input; pass one molecule ((N,) and (N, 3)) or a non-empty list")
+    single = np.ndim(np.asarray(coordinates[0])) == 1     # (N, 3): first row is length-3 -> single
+    if single:
+        an_list = [np.asarray(atomic_numbers)]
+        xyz_list = [np.asarray(coordinates, dtype=float)]
+    else:
+        an_list = [np.asarray(a) for a in atomic_numbers]
+        xyz_list = [np.asarray(c, dtype=float) for c in coordinates]
+    if len(an_list) != len(xyz_list):
+        raise ValueError(f"got {len(an_list)} atomic_numbers array(s) but {len(xyz_list)} geometry(ies)")
+    for a, c in zip(an_list, xyz_list):
+        if a.ndim != 1 or c.ndim != 2 or c.shape[1] != 3 or a.shape[0] != c.shape[0]:
+            raise ValueError("each molecule needs atomic_numbers shaped (N,) and coordinates shaped "
+                             f"(N, 3) with the same N; got {a.shape} and {c.shape}")
+    return an_list, xyz_list, single
+
+
+def predict_shifts(atomic_numbers, coordinates, solvent="chloroform", n_passes=10, symmetrize=True,
+                   device=None, return_components=False, checkpoints_dir=None):
+    """Predict the <sup>1</sup>H and <sup>13</sup>C chemical shifts of a molecule in a solvent.
+
+    This is the main entry point. It runs the MagNET-Zero and MagNET-PCM models and applies the paper's
+    per-solvent calibration, so the output is directly comparable to an experimental spectrum. Needs an
+    AIMNet2-optimized geometry.
+
+    Args:
+        atomic_numbers: element numbers, shape `(N,)`, for one molecule; or a list of such arrays.
+        coordinates: xyz positions in Angstrom, shape `(N, 3)`; or a list of them.
+        solvent: one of the 12 calibrated solvents: `"tetrahydrofuran"`, `"dichloromethane"`,
+            `"chloroform"`, `"toluene"`, `"benzene"`, `"chlorobenzene"`, `"acetone"`,
+            `"dimethylsulfoxide"`, `"acetonitrile"`, `"trifluoroethanol"`, `"methanol"`, `"water"`.
+        n_passes: average out equivariance error over `n_passes` forward passes (default `10`).
+        symmetrize: if True, also average over the mirror image of the input geometry (default `True`),
+            doubling the passes.
+        device: where to run, a torch device or a string like `"cpu"` or `"cuda"` (default `None`);
+            uses GPU if available.
+        return_components: return the numbers behind each shift (see Returns) instead of just the shifts.
+        checkpoints_dir: directory holding the released weights, if they are not in the current
+            directory or the repo (default `None`); may be the `model_checkpoints/` folder or its
+            parent, and takes precedence over both defaults.
+
+    Returns:
+        One chemical shift in ppm per atom, shape `(N,)`, with NaN at any atom that is not
+        <sup>1</sup>H or <sup>13</sup>C. Pass a list of molecules and you get a list of arrays back.
+
+        With `return_components=True` you get a dict instead (or a list of dicts):
+
+        - `shifts`: the shifts, as above.
+        - `zero_shielding`: the MagNET-Zero gas-phase shielding at each atom.
+        - `pcm_correction`: the MagNET-PCM solvent correction at each atom.
+        - `coefficients`: the calibration used, as
+          `{"H": {"intercept": ..., "stationary": ..., "pcm": ...}, "C": {...}}`, so that
+          `shift = intercept + stationary * zero_shielding + pcm * pcm_correction`.
+    """
+    _check_passes(n_passes)
+    tables = scaling.published_scaling_tables()
+    key = "TIP4P" if solvent == "water" else solvent
+    if key not in tables["C"] or key not in tables["H"]:
+        options = ["water" if s == "TIP4P" else s for s in tables["C"]]
+        raise ValueError(f"unknown solvent {solvent!r}; choose one of {options}")
+
+    an_list, xyz_list, single = _as_batch(atomic_numbers, coordinates)
+    zero = run_magnet.compute_MagNET_Zero_shieldings(an_list, xyz_list, n_passes=n_passes,
+                                                     symmetrize=symmetrize, device=device,
+                                                     checkpoints_dir=checkpoints_dir)
+    pcm = run_magnet.compute_MagNET_PCM_corrections(an_list, xyz_list, n_passes=n_passes,
+                                                    symmetrize=symmetrize, device=device,
+                                                    checkpoints_dir=checkpoints_dir)
+    results = []
+    for atoms, sigma, delta in zip(an_list, zero, pcm):
+        out = np.full(atoms.shape, np.nan, dtype=float)
+        for nucleus, atomic_number in (("H", 1), ("C", 6)):
+            mask = atoms == atomic_number
+            if mask.any():
+                out[mask] = scaling.predict_shift(tables[nucleus], key, sigma[mask], delta[mask])
+        if return_components:
+            results.append({"shifts": out, "zero_shielding": sigma, "pcm_correction": delta,
+                            "coefficients": {"H": dict(tables["H"][key]), "C": dict(tables["C"][key])}})
+        else:
+            results.append(out)
+    return results[0] if single else results
+
+
+def predict_shieldings(atomic_numbers, coordinates, model="MagNET", n_passes=10, symmetrize=True,
+                       device=None, checkpoints_dir=None):
+    """Predict gas-phase NMR shielding constants for a molecule.
+
+    A shielding constant is what the network outputs directly, before it is calibrated into a chemical
+    shift; if you want shifts, use `predict_shifts`. For solvent effects use
+    `implicit_solvent_correction` or `explicit_solvent_correction`.
+
+    Args:
+        atomic_numbers: element numbers, shape `(N,)`; or a list of such arrays.
+        coordinates: xyz positions in Angstrom, shape `(N, 3)`; or a list of them.
+        model: `"MagNET"`, the general foundation model, or `"MagNET-Zero"`, which is more accurate but
+            expects an AIMNet2-optimized geometry.
+        n_passes: average out equivariance error over `n_passes` forward passes (default `10`).
+        symmetrize: if True, also average over the mirror image of the input geometry (default `True`),
+            doubling the passes.
+        device: where to run, a torch device or a string like `"cpu"` or `"cuda"` (default `None`);
+            uses GPU if available.
+        checkpoints_dir: directory holding the released weights, if they are not in the current
+            directory or the repo (default `None`); may be the `model_checkpoints/` folder or its
+            parent, and takes precedence over both defaults.
+
+    Returns:
+        One shielding constant in ppm per atom, shape `(N,)`. The models are trained on <sup>1</sup>H
+        and <sup>13</sup>C, so the values are meaningful only at hydrogen and carbon atoms; atoms of
+        other elements come back as `0.0`, not a prediction. (`predict_shifts` instead returns `NaN`
+        there.) Pass a list of molecules and you get a list of arrays back.
+    """
+    redirect = {"MagNET-PCM": "implicit_solvent_correction", "MagNET-x": "explicit_solvent_correction"}
+    if model in redirect:
+        raise ValueError(f"{model!r} is a correction model; use magnet.{redirect[model]}(...), "
+                         f"not predict_shieldings().")
+    if model not in _MODELS:
+        raise ValueError(f"model must be one of {list(_MODELS)}; got {model!r}")
+    _check_passes(n_passes)
+    an_list, xyz_list, single = _as_batch(atomic_numbers, coordinates)
+    out = _MODELS[model](an_list, xyz_list, n_passes=n_passes, symmetrize=symmetrize, device=device,
+                         checkpoints_dir=checkpoints_dir)
+    return out[0] if single else out
+
+
+def implicit_solvent_correction(atomic_numbers, coordinates, n_passes=10, symmetrize=True, device=None,
+                                checkpoints_dir=None):
+    """Predict how a solvent changes a molecule's shieldings, using a fast continuum-solvent model.
+
+    The result is a per-atom shielding change: add it to a `predict_shieldings(..., model="MagNET-Zero")`
+    value to get the solvated shielding, or just use `predict_shifts`, which does this for you. Uses
+    MagNET-PCM, which only predicts the chloroform correction, so there is no solvent argument here;
+    `predict_shifts` is what reuses this one correction for other solvents, through its per-solvent
+    calibration. Needs an AIMNet2-optimized geometry.
+
+    Args:
+        atomic_numbers: element numbers, shape `(N,)`; or a list of such arrays.
+        coordinates: xyz positions in Angstrom, shape `(N, 3)`; or a list of them.
+        n_passes: average out equivariance error over `n_passes` forward passes (default `10`).
+        symmetrize: if True, also average over the mirror image of the input geometry (default `True`),
+            doubling the passes.
+        device: where to run, a torch device or a string like `"cpu"` or `"cuda"` (default `None`);
+            uses GPU if available.
+        checkpoints_dir: directory holding the released weights, if they are not in the current
+            directory or the repo (default `None`); may be the `model_checkpoints/` folder or its
+            parent, and takes precedence over both defaults.
+
+    Returns:
+        One shielding change in ppm per atom, shape `(N,)`. Meaningful only at hydrogen and carbon
+        atoms; atoms of other elements come back as `0.0`. Pass a list of molecules and you get a list
+        of arrays back.
+    """
+    _check_passes(n_passes)
+    an_list, xyz_list, single = _as_batch(atomic_numbers, coordinates)
+    out = run_magnet.compute_MagNET_PCM_corrections(an_list, xyz_list, n_passes=n_passes,
+                                                    symmetrize=symmetrize, device=device,
+                                                    checkpoints_dir=checkpoints_dir)
+    return out[0] if single else out
+
+
+def _validate_solvent(solvent, solvent_atomic_numbers):
+    """Confirm the non-solute atoms really are whole molecules of `solvent`, one after another.
+
+    The solvent atoms must be listed one molecule at a time: all atoms of the first molecule, then
+    all atoms of the next, and so on. This splits them into consecutive groups of the solvent's
+    molecule size and checks each group has the right number of each element (the order of atoms
+    within a group does not matter). Listing them this way is required by the model, which finds the
+    solvent molecules by splitting the atom list into equal-size groups the same way, and uses that
+    to drop solvent molecules that sit far from the solute."""
+    if solvent not in _SOLVENT_COMPOSITION:
+        raise ValueError(f"unknown solvent {solvent!r}; expected one of {sorted(_SOLVENT_COMPOSITION)}")
+    atoms = np.asarray(solvent_atomic_numbers).tolist()
+    if not atoms:
+        raise ValueError("no solvent atoms; solute_atoms cannot name every atom")
+    per_molecule = _SOLVENT_COMPOSITION[solvent]
+    size = run_magnet.N_ATOMS_PER_SOLVENT[solvent]
+    if len(atoms) % size != 0:
+        raise ValueError(f"got {len(atoms)} non-solute atoms, not a whole number of {solvent} "
+                         f"molecules ({size} atoms each); check solute_atoms and solvent")
+    for start in range(0, len(atoms), size):
+        group = Counter(atoms[start:start + size])
+        if group != per_molecule:
+            raise ValueError(f"the non-solute atoms are not whole {solvent} molecules listed one at "
+                             f"a time: the group of {size} atoms starting at position {start} has "
+                             f"element counts {dict(group)}, expected {dict(per_molecule)}. List each "
+                             f"solvent molecule's atoms together.")
+
+
+def _explicit_one(atoms, xyz, solute_atoms, solvent):
+    """Validate solute_atoms and the solvent composition, then reorder one snapshot to
+    solute-first-then-solvent. Returns (solute_atomic_numbers, reordered_atomic_numbers,
+    reordered_coordinates)."""
+    n = len(atoms)
+    solute = np.asarray(solute_atoms, dtype=int)
+    solute = np.where(solute < 0, solute + n, solute)          # normalize negative indices
+    if solute.size == 0:
+        raise ValueError("solute_atoms is empty; name at least one solute atom")
+    if np.any(solute < 0) or np.any(solute >= n):
+        raise ValueError(f"solute_atoms has an index out of range for a {n}-atom system")
+    if len(set(solute.tolist())) != solute.size:
+        raise ValueError("solute_atoms has duplicate indices")
+    solute_set = set(solute.tolist())
+    solvent_idx = np.array([i for i in range(n) if i not in solute_set], dtype=int)
+    _validate_solvent(solvent, atoms[solvent_idx])
+    order = np.concatenate([solute, solvent_idx])              # solute first, then solvent blocks
+    return atoms[solute], atoms[order], xyz[order]
+
+
+def explicit_solvent_correction(atomic_numbers, coordinates, solute_atoms, solvent="chloroform",
+                                n_passes=10, symmetrize=True, device=None,
+                                solvent_distance_threshold=12.0, checkpoints_dir=None):
+    """Predict a solvent's effect on shieldings from an MD snapshot with explicit solvent molecules.
+
+    Uses MagNET-x. Corrections are returned for solute atoms only. Corrections should be averaged
+    over multiple frames.
+
+    The expected atom ordering is [solute, n x solvents].
+
+    Args:
+        atomic_numbers: element numbers for the whole solute + solvent system, shape `(N,)`; or a list
+            of snapshots.
+        coordinates: xyz positions in Angstrom, shape `(N, 3)`; or a list of them.
+        solute_atoms: which atoms are the solute, as 0-based indices (a list or array of ints;
+            negative indices count from the end, like Python lists); the rest are treated as solvent.
+        solvent: `"chloroform"`, `"benzene"`, `"methanol"`, or `"water"`, the four MagNET-x supports
+            (default `"chloroform"`).
+        n_passes: average out equivariance error over `n_passes` forward passes (default `10`).
+        symmetrize: if True, also average over the mirror image of the input geometry (default `True`),
+            doubling the passes.
+        device: where to run, a torch device or a string like `"cpu"` or `"cuda"` (default `None`);
+            uses GPU if available.
+        solvent_distance_threshold: solvent molecules whose nearest atom is farther than this many
+            Angstrom from the solute are dropped automatically (default 12.0), so extra solvent in the
+            snapshot is harmless.
+        checkpoints_dir: directory holding the released weights, if they are not in the current
+            directory or the repo (default `None`); may be the `model_checkpoints/` folder or its
+            parent, and takes precedence over both defaults.
+
+    Returns:
+        One correction in ppm per solute atom, in the order you listed them in `solute_atoms`. Pass a
+        list of snapshots and you get a list of arrays back.
+    """
+    if solvent not in run_magnet.N_ATOMS_PER_SOLVENT:
+        raise ValueError(f"solvent must be one of {sorted(run_magnet.N_ATOMS_PER_SOLVENT)}; "
+                         f"got {solvent!r}")
+    _check_passes(n_passes)
+    an_list, xyz_list, single = _as_batch(atomic_numbers, coordinates)
+    solute_an_list, full_an_list, full_xyz_list = [], [], []
+    for atoms, xyz in zip(an_list, xyz_list):
+        solute_an, full_an, full_xyz = _explicit_one(atoms, xyz, solute_atoms, solvent)
+        solute_an_list.append(solute_an)
+        full_an_list.append(full_an)
+        full_xyz_list.append(full_xyz)
+    out = run_magnet.compute_MagNET_x_corrections(
+        solvent, solute_an_list, full_an_list, full_xyz_list,
+        n_passes=n_passes, symmetrize=symmetrize, device=device,
+        solvent_distance_threshold=solvent_distance_threshold, checkpoints_dir=checkpoints_dir)
+    return out[0] if single else out

magnet/inference.py

@@ -0,0 +1,154 @@
+import numpy as np
+import torch
+import torch_geometric
+import torch_scatter
+
+# The elements MagNET was trained on (the sigma datasets are GDB molecules plus F, S, and Cl
+# substituents): hydrogen, carbon, nitrogen, oxygen, fluorine, sulfur, chlorine. Predictions for any
+# other element are not meaningful, so predict_shieldings refuses them.
+SUPPORTED_ELEMENTS = frozenset({1, 6, 7, 8, 9, 16, 17})
+
+
+def yield_data(solute_atomic_numbers, geometries, atomic_numbers, shieldings = None, N_atoms_per_solvent = 3, solvent_distance_threshold = None, atom_type = 'H'):
+    data = torch_geometric.data.Data(
+        x = torch.as_tensor(atomic_numbers, dtype = torch.long), 
+    )
+    if shieldings is None:
+        data.shieldings = torch.zeros(data.x.shape[0], dtype = torch.float)
+    else:
+        data.shieldings = torch.as_tensor(shieldings, dtype = torch.float)
+
+    data.pos = torch.as_tensor(geometries, dtype = torch.float)
+    
+    solute = np.zeros(len(atomic_numbers), dtype = int)
+    solute[0:len(solute_atomic_numbers)] = 1
+    data.solute = torch.as_tensor(solute, dtype = torch.long)
+    
+    if atom_type == 'H':
+        atom_type_mask = (atomic_numbers == 1) & (solute == 1) # H shieldings
+    elif atom_type == 'C':
+        atom_type_mask = (atomic_numbers == 6) & (solute == 1) # C shieldings
+    data.atom_type_mask = torch.as_tensor(atom_type_mask, dtype = torch.bool)
+        
+    data.batch = torch.zeros(data.x.shape[0], dtype = torch.long)
+    data.natoms = torch.unique_consecutive(data.batch, return_counts = True)[1] 
+    
+    N_solute_atoms = sum(data.solute).item()
+    N_solvent_atoms = data.x.shape[0] - N_solute_atoms
+    N_solvent_molecules = N_solvent_atoms // N_atoms_per_solvent
+    unique_molecule_batch = np.zeros(N_solute_atoms, dtype = int)
+    if N_solvent_molecules > 0:
+        unique_molecule_batch = np.concatenate([
+            unique_molecule_batch,
+            np.concatenate([np.zeros(N_atoms_per_solvent, dtype = int) + 1 + i for i in range(N_solvent_molecules)], axis = 0), 
+        ], axis = 0)
+    data.unique_molecule_batch = torch.as_tensor(unique_molecule_batch, dtype = torch.long)
+    
+    if solvent_distance_threshold is not None:
+        min_distance_to_solute_peratom = torch.linalg.norm(data.pos - (data.pos[data.solute > 0])[:, None, ...], dim = -1).min(axis = 0).values
+        min_distance_to_solute_permolecule = torch_scatter.scatter_min(min_distance_to_solute_peratom, data.unique_molecule_batch)[0]
+        remove_atom = np.array([m in torch.where((min_distance_to_solute_permolecule > solvent_distance_threshold))[0] for m in data.unique_molecule_batch])
+        
+        data.x = data.x[~remove_atom]
+        if shieldings is None:
+            data.shieldings = data.shieldings[~remove_atom]
+        data.pos = data.pos[~remove_atom]
+        data.solute = data.solute[~remove_atom]
+        data.atom_type_mask = data.atom_type_mask[~remove_atom]
+        data.batch = torch.zeros(data.x.shape[0], dtype = torch.long)
+        data.natoms = torch.unique_consecutive(data.batch, return_counts = True)[1] 
+        N_solute_atoms = sum(data.solute).item()
+        N_solvent_atoms = data.x.shape[0] - N_solute_atoms
+        N_solvent_molecules = N_solvent_atoms // N_atoms_per_solvent
+        unique_molecule_batch = np.zeros(N_solute_atoms, dtype = int)
+        if N_solvent_molecules > 0:
+            unique_molecule_batch = np.concatenate([
+                unique_molecule_batch,
+                np.concatenate([np.zeros(N_atoms_per_solvent, dtype = int) + 1 + i for i in range(N_solvent_molecules)], axis = 0), 
+            ], axis = 0)
+        data.unique_molecule_batch = torch.as_tensor(unique_molecule_batch, dtype = torch.long)
+    
+    return data
+
+
+def _predict_once(model_H, model_C, solute_atomic_numbers, geometry, atomic_numbers, N_atoms_per_solvent, solvent_distance_threshold, device):
+    # one forward pass each for the 1H and 13C heads, combined into a single per-atom array
+    for atom_type in ['H', 'C']:
+        data = yield_data(
+            solute_atomic_numbers = solute_atomic_numbers,
+            geometries = geometry,
+            atomic_numbers = atomic_numbers,
+            shieldings = None,
+            N_atoms_per_solvent = N_atoms_per_solvent if N_atoms_per_solvent is not None else 3,
+            solvent_distance_threshold = solvent_distance_threshold,
+            atom_type = atom_type,
+        )
+        data = data.to(device)
+
+        with torch.no_grad():
+            if atom_type == 'H':
+                y_pred = model_H.forward(data).cpu()
+            elif atom_type == 'C':
+                y_pred = model_C.forward(data).cpu()
+            # the masks live on `device`; move them to CPU to match y_pred (a no-op on CPU, and
+            # required on CUDA, where indexing a CPU tensor with a device mask would raise)
+            atom_type_mask = data.atom_type_mask.cpu()
+            solute = data.solute.cpu()
+            y_pred[~atom_type_mask] = 0.
+            y_pred = y_pred[solute == 1].numpy()
+
+        # combine H and C shieldings into one array
+        if atom_type == 'H':
+            y_pred_combined = y_pred
+        elif atom_type == 'C':
+            y_pred_combined[y_pred != 0.0] = y_pred[y_pred != 0.0]
+
+    return y_pred_combined
+
+
+def predict_shieldings(model_H, model_C, solute_atomic_numbers, geometry, atomic_numbers = None, N_atoms_per_solvent = None, solvent_distance_threshold = None, device = 'cpu', n_passes = 1, symmetrize = False):
+    """Predict 1H/13C shieldings for the solute atoms.
+
+    A single forward pass is NOT deterministic: each edge picks a random local
+    reference frame (eqV2/edge_rot_mat.py), so with a finite spherical-harmonic grid
+    the output varies by ~0.01 ppm (13C) between passes. Set n_passes (e.g. 20, as in
+    the MagNET paper) to average that frame noise away.
+
+    symmetrize=True also averages the prediction over the molecule and its mirror image.
+    Isotropic shielding is parity-even (a molecule and its reflection have identical
+    shieldings), but the SO(3)-only model does not enforce this and can disagree by
+    ~0.3 ppm on 13C for large molecules; symmetrizing removes that spurious error.
+    """
+    solute_atomic_numbers = np.asarray(solute_atomic_numbers)
+    if atomic_numbers is None:
+        atomic_numbers = solute_atomic_numbers
+    else:
+        # coerce so the element comparisons in yield_data (atomic_numbers == 1) stay array-wise;
+        # a bare Python list would compare to a scalar False and silently mis-mask
+        atomic_numbers = np.asarray(atomic_numbers)
+
+    # MagNET was only ever trained on these elements. On anything else it would emit a confident but
+    # meaningless prediction, so refuse it rather than let bad numbers through (validate the full
+    # system, which for MagNET-x includes the solvent atoms).
+    unsupported = sorted(set(np.unique(atomic_numbers).tolist()) - SUPPORTED_ELEMENTS)
+    if unsupported:
+        raise ValueError(
+            f"MagNET supports only the elements {sorted(SUPPORTED_ELEMENTS)} "
+            f"(H, C, N, O, F, S, Cl); the input contains unsupported atomic numbers {unsupported}. "
+            f"Exclude molecules with these elements before predicting.")
+
+    geometry = np.asarray(geometry, dtype=float)
+    geometries = [geometry]
+    if symmetrize:
+        reflected = geometry.copy()
+        reflected[..., 0] = -reflected[..., 0]   # mirror across the yz-plane (an improper rotation)
+        geometries.append(reflected)
+
+    preds = []
+    for geom in geometries:
+        for _ in range(n_passes):
+            preds.append(np.squeeze(_predict_once(model_H, model_C, solute_atomic_numbers, geom,
+                                                  atomic_numbers, N_atoms_per_solvent,
+                                                  solvent_distance_threshold, device)))
+    # atleast_1d keeps a one-atom solute a (1,) array instead of a 0-d scalar after the per-pass squeeze
+    return np.atleast_1d(np.mean(preds, axis=0))

magnet/model.py

@@ -0,0 +1,236 @@
+import torch
+import pytorch_lightning as pl
+from .eqV2 import EquiformerV2_NMR
+
+class MagNET_Lightning(pl.LightningModule):
+    def __init__(self, params, precision = 32, filter_solvent_edges = None, solvent_edge_radius = None):
+        super().__init__()
+        
+        self.save_hyperparameters()
+        
+        self.precision = precision
+        
+        # if filter_solvent_edges and solvent_edge_radius are not provided as arguments, read from params file
+            # this is convenient, because we don't need to constantly provide filter_solvent_edges and solvent_edge_radius
+            # as arguments if they're already specified in params, while still allowing us to override their previous 
+            # values if we need to (e.g., when loading in a pretrained checkpoint)
+        if filter_solvent_edges is None:
+            assert solvent_edge_radius is None
+            filter_solvent_edges = False
+            solvent_edge_radius = params['max_radius']
+            if 'filter_solvent_edges' in params:
+                filter_solvent_edges = params['filter_solvent_edges']
+            if 'solvent_edge_radius' in params:
+                solvent_edge_radius = params['solvent_edge_radius']
+        self.filter_solvent_edges = filter_solvent_edges
+        self.solvent_edge_radius = solvent_edge_radius
+        
+        self.lr = params['lr']
+            
+        if 'min_lr' in params:
+            self.min_lr = params['min_lr']
+        else:
+            self.min_lr = self.lr
+        
+        # number of exponential decay steps (= # batches) from self.lr to self.min_lr
+        if 'lr_steps' in params:
+            self.lr_steps = params['lr_steps']
+        else:
+            self.lr_steps = 300000 
+        
+        self.model_type = params['model_type'] # in ['EquiformerV2']
+        
+        evidential_regression = False
+        if 'evidential_regression' in params:
+            evidential_regression = params['evidential_regression']
+        self.evidential_regression = evidential_regression
+        
+        if self.model_type == 'EquiformerV2':
+            assert self.precision == 32
+            
+            self.model = EquiformerV2_NMR(
+                None, None, None,
+                use_pbc=False, regress_forces=False,
+                
+                otf_graph = True,
+                max_neighbors = params['max_neighbors'],
+                max_radius = params['max_radius'],
+                max_num_elements = params['max_num_elements'],
+            
+                num_layers = params['num_layers'], 
+                sphere_channels = params['sphere_channels'],
+                attn_hidden_channels = params['attn_hidden_channels'],
+                num_heads = params['num_heads'], 
+                attn_alpha_channels = params['attn_alpha_channels'],
+                attn_value_channels = params['attn_value_channels'],
+                ffn_hidden_channels = params['ffn_hidden_channels'],
+                
+                norm_type='layer_norm_sh',
+                
+                lmax_list = params['lmax_list'],
+                mmax_list = params['mmax_list'],
+                grid_resolution = params['grid_resolution'], 
+            
+                num_sphere_samples=128,
+                
+                edge_channels=128,
+                use_atom_edge_embedding=True,
+                share_atom_edge_embedding=False,
+                use_m_share_rad=False,
+                distance_function="gaussian",
+                num_distance_basis=512, # not used; hard-coded by Equiformer-V2 to 600
+            
+                attn_activation='silu',
+                use_s2_act_attn=False, 
+                use_attn_renorm=True,
+                ffn_activation='silu',
+                use_gate_act=False,
+                use_grid_mlp=True, 
+                use_sep_s2_act=True,
+            
+                alpha_drop=0.0,
+                drop_path_rate=0.0, 
+                proj_drop=0.0, 
+            
+                weight_init = params['weight_init'],
+                
+                evidential_regression = self.evidential_regression,
+                filter_solvent_edges = self.filter_solvent_edges,
+                solvent_edge_radius = self.solvent_edge_radius,
+            )
+    
+    def forward(self, data):
+        if 'atomic_numbers' not in data.keys():
+            data.atomic_numbers = data.x 
+        if 'natoms' not in data.keys():
+            data.natoms = torch.unique_consecutive(data.batch, return_counts = True)[1]
+        
+        if self.model_type == 'EquiformerV2':
+            if self.evidential_regression:
+                mu,v,alpha,beta = self.model(data).squeeze(1).chunk(4, dim = 1)
+                v = torch.nn.functional.softplus(v) + 1e-8
+                alpha = torch.nn.functional.softplus(alpha) + 1. + 1e-8
+                beta = torch.nn.functional.softplus(beta) + 1e-8
+                return torch.cat([mu,v,alpha,beta], dim = 1)
+            else:
+                return self.model(data).squeeze(1)
+    
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr = self.lr)
+        
+        # exponential lr decay from self.lr to self.min_lr in self.lr_steps steps
+        gamma = (self.min_lr / self.lr) ** (1.0 / self.lr_steps)
+        func = lambda step: max(gamma**(step), self.min_lr / self.lr)
+        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda = func)
+        
+        lr_scheduler_config = {
+            "scheduler": scheduler,
+            "interval": "step",
+            "frequency": 1,
+            "strict": False,
+            "name": None,
+        }
+        
+        return {"optimizer": optimizer, "lr_scheduler": lr_scheduler_config}
+    
+    def training_step(self, train_batch, batch_idx):
+        data = train_batch
+        data.atomic_numbers = data.x 
+        data.natoms = torch.unique_consecutive(data.batch, return_counts = True)[1] 
+
+        out = self.forward(data)
+        
+        if self.precision == 64:
+            y = data.shieldings.double()
+        else:
+            y = data.shieldings
+            
+        if self.evidential_regression:
+            lam = 0.2
+            mu,v,alpha,beta = out.chunk(4, dim = 1)
+            mu = mu.squeeze(dim=1)[data.atom_type_mask]
+            v = v.squeeze(dim=1)[data.atom_type_mask]
+            alpha = alpha.squeeze(dim=1)[data.atom_type_mask]
+            beta = beta.squeeze(dim=1)[data.atom_type_mask]
+            
+            loss = 0.5 * (torch.log(torch.pi / v)) \
+                   - alpha * torch.log(2.*beta*(1. + v)) \
+                   + (alpha + 0.5) * torch.log(v*(y[data.atom_type_mask] - mu)**2. + 2.*beta*(1. + v)) \
+                   + torch.lgamma(alpha) - torch.lgamma(alpha + 0.5) \
+                   + lam * (torch.abs(y[data.atom_type_mask] - mu) * (2. * v + alpha) - 1e-4)
+            loss = loss.mean()
+        else:
+            loss = torch.mean(((out.squeeze(1) - y)[data.atom_type_mask])**2.0)
+        
+        self.log('train_loss', loss, batch_size = data.atom_type_mask.sum().item())
+        
+        # does not account for symmetric atoms 
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        data = val_batch
+        data.atomic_numbers = data.x 
+        data.natoms = torch.unique_consecutive(data.batch, return_counts = True)[1] 
+        
+        out = self.forward(data)
+        
+        if self.precision == 64:
+            y = data.shieldings.double()
+        else:
+            y = data.shieldings
+        
+        if self.evidential_regression:
+            mu,v,alpha,beta = out.chunk(4, dim = 1)
+            mu = mu.squeeze(dim=1)[data.atom_type_mask]
+            v = v.squeeze(dim=1)[data.atom_type_mask]
+            alpha = alpha.squeeze(dim=1)[data.atom_type_mask]
+            beta = beta.squeeze(dim=1)[data.atom_type_mask]
+            
+            lam = 0.2
+            loss = 0.5 * (torch.log(torch.pi / v)) \
+                   - alpha * torch.log(2.*beta*(1. + v)) \
+                   + (alpha + 0.5) * torch.log(v*(y[data.atom_type_mask] - mu)**2. + 2.*beta*(1. + v)) \
+                   + torch.lgamma(alpha) - torch.lgamma(alpha + 0.5) \
+                   + lam * (torch.abs(y[data.atom_type_mask] - mu) * (2. * v + alpha) - 1e-4)
+            loss = loss.mean()
+            mae = torch.mean(torch.abs(mu - y[data.atom_type_mask]))
+        else:
+            loss = torch.mean(((out.squeeze(1) - y)[data.atom_type_mask])**2.0)
+            mae = torch.mean(torch.abs(((out.squeeze(1) - y)[data.atom_type_mask])))
+            
+        self.log_dict({"val_loss": loss, "val_mae": mae}, batch_size = data.atom_type_mask.sum().item())
+        return loss
+    
+    def test_step(self, test_batch, batch_idx):
+        data = test_batch
+        data.atomic_numbers = data.x 
+        data.natoms = torch.unique_consecutive(data.batch, return_counts = True)[1] 
+        
+        out = self.forward(data)
+        
+        if self.precision == 64:
+            y = data.shieldings.double()
+        else:
+            y = data.shieldings
+        
+        if self.evidential_regression:
+            lam = 0.2
+            mu,v,alpha,beta = out.chunk(4, dim = 1)
+            mu = mu.squeeze(dim=1)[data.atom_type_mask]
+            v = v.squeeze(dim=1)[data.atom_type_mask]
+            alpha = alpha.squeeze(dim=1)[data.atom_type_mask]
+            beta = beta.squeeze(dim=1)[data.atom_type_mask]
+            
+            loss = 0.5 * (torch.log(torch.pi / v)) \
+                   - alpha * torch.log(2.*beta*(1. + v)) \
+                   + (alpha + 0.5) * torch.log(v*(y[data.atom_type_mask] - mu)**2. + 2.*beta*(1. + v)) \
+                   + torch.lgamma(alpha) - torch.lgamma(alpha + 0.5) \
+                   + lam * (torch.abs(y[data.atom_type_mask] - mu) * (2. * v + alpha) - 1e-4)
+            loss = loss.mean()
+            mae = torch.mean(torch.abs(mu - y[data.atom_type_mask]))
+        else:
+            loss = torch.mean(((out.squeeze(1) - y)[data.atom_type_mask])**2.0)
+            mae = torch.mean(torch.abs(((out.squeeze(1) - y)[data.atom_type_mask])))
+
+        self.log_dict({"test_loss": loss, "test_mae": mae}, batch_size = data.atom_type_mask.sum().item())
+        return loss

magnet/pyproject.toml

@@ -0,0 +1,49 @@
+[build-system]
+requires = ["setuptools>=61"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "magnet-nmr"
+version = "0.1.0"
+description = "MagNET: equivariant neural networks for NMR chemical-shift (shielding) prediction"
+readme = {text = "MagNET: equivariant neural networks for NMR chemical-shift (shielding) prediction. Models, datasets, install instructions, and usage examples are at https://github.com/ekwan/MagNET", content-type = "text/markdown"}
+requires-python = ">=3.9"
+license = "MIT"
+
+# torch_geometric is pinned to 2.6.x because 2.7+ requires pyg-lib for radius_graph.
+#
+# torch_scatter and torch_cluster are deliberately NOT listed here. They have no plain-PyPI
+# wheels for our platforms, so pip would try to build them from source and fail; they must come
+# from the PyG wheel index, keyed to the exact torch version:
+#     pip install torch_scatter torch_cluster --find-links https://data.pyg.org/whl/torch-2.5.0+cpu.html
+# The `pyg` extra records this so `pip install magnet-nmr[pyg]` documents the two, but the
+# --find-links index is still required for the install to resolve.
+dependencies = [
+    "torch==2.5.0",
+    "torch_geometric==2.6.1",
+    "pytorch_lightning>=2.0",
+    "e3nn>=0.5",
+    "numpy",
+]
+
+[project.optional-dependencies]
+examples = ["h5py"]
+pyg = ["torch_scatter", "torch_cluster"]
+
+# The importable package `magnet` lives in this directory (package-dir maps it to ".").
+[tool.setuptools]
+packages = [
+    "magnet",
+    "magnet.eqV2",
+    "magnet.eqV2.ocpmodels",
+    "magnet.eqV2.ocpmodels.common",
+    "magnet.eqV2.ocpmodels.models",
+    "magnet.eqV2.ocpmodels.models.scn",
+]
+
+[tool.setuptools.package-dir]
+"magnet" = "."
+
+# the Wigner-D rotation table that wigner.py loads at import time
+[tool.setuptools.package-data]
+"magnet.eqV2" = ["Jd.pt"]

magnet/requirements.txt

@@ -0,0 +1,32 @@
+# MagNET inference dependencies (CPU; tested on Apple Silicon arm64 and Linux x86-64).
+#
+# Two non-obvious points are baked in below:
+#   1. torch_scatter and torch_cluster are NOT on plain PyPI for macOS arm64. They come
+#      from the PyG wheel index, which is keyed to the EXACT torch version. The
+#      --find-links line points pip at the matching prebuilt wheels for torch 2.5.0 + CPU.
+#   2. torch_geometric is PINNED to 2.6.x. From 2.7 onward, radius_graph requires pyg-lib
+#      (which has no Apple-Silicon wheel); 2.6.x still dispatches radius_graph to
+#      torch_cluster, which is what the model needs.
+#
+# Install torch first, then this file, or just run this file twice if a fresh
+# environment resolves torch_scatter before torch:
+#   pip install torch==2.5.0
+#   pip install -r requirements.txt
+#
+# On a Linux machine with CUDA, `pip install torch==2.5.0` defaults to the multi-gigabyte GPU build.
+# This is a CPU stack, so install the CPU build explicitly first:
+#   pip install torch==2.5.0 --index-url https://download.pytorch.org/whl/cpu
+
+--find-links https://data.pyg.org/whl/torch-2.5.0+cpu.html
+
+torch==2.5.0
+torch_scatter
+torch_cluster
+torch_geometric==2.6.1
+pytorch_lightning>=2.0
+e3nn>=0.5
+numpy
+
+# Only needed to run the bundled examples/tests that read the .hdf5 datasets
+# (the model itself does not import h5py):
+h5py

magnet/run_magnet.py

@@ -0,0 +1,221 @@
+"""Worked examples for the MagNET model family.
+
+Each helper loads the released checkpoints for one model and runs `predict_shieldings`. All of them
+expose the two prediction options:
+
+  - `n_passes`: average over this many forward passes to remove the per-pass frame noise (the MagNET
+    paper uses 20).
+  - `symmetrize`: also average over the molecule and its mirror image. Isotropic shielding is
+    parity-even, but the SO(3)-only network does not enforce that, so symmetrizing removes a spurious
+    reflection error that grows to ~0.3 ppm on 13C for large molecules. It costs one extra evaluation.
+    The released reference shieldings are NOT symmetrized, so leave it False to reproduce them and turn
+    it on for best accuracy on large molecules.
+
+Geometries are the AIMNet2 (machine-learning-optimized) stationary geometries, in Angstrom.
+"""
+import os
+
+import numpy as np
+import torch
+
+from magnet.model import MagNET_Lightning
+from magnet.inference import predict_shieldings
+
+MODEL_CHECKPOINTS = {
+    # foundation model (predicts the gas-phase shielding the rovibrational/QCD analysis builds on)
+    "MagNET_H": "model_checkpoints/MagNET/MagNET_1H.ckpt",
+    "MagNET_C": "model_checkpoints/MagNET/MagNET_13C.ckpt",
+    # MagNET-Zero: gas-phase shielding (WP04 for 1H, wB97X-D for 13C, pcSseg-2)
+    "MagNET-Zero_H": "model_checkpoints/MagNET-Zero/MagNET-Zero_1H.ckpt",
+    "MagNET-Zero_C": "model_checkpoints/MagNET-Zero/MagNET-Zero_13C.ckpt",
+    # MagNET-PCM: B3LYP-D3(BJ)/pcSseg-2 gas-phase and chloroform-PCM shieldings; the correction is
+    # their difference
+    "MagNET-PCM-withPCM_H": "model_checkpoints/MagNET-PCM/MagNET-PCM-withPCM_1H.ckpt",
+    "MagNET-PCM-withPCM_C": "model_checkpoints/MagNET-PCM/MagNET-PCM-withPCM_13C.ckpt",
+    "MagNET-PCM-withoutPCM_H": "model_checkpoints/MagNET-PCM/MagNET-PCM-withoutPCM_1H.ckpt",
+    "MagNET-PCM-withoutPCM_C": "model_checkpoints/MagNET-PCM/MagNET-PCM-withoutPCM_13C.ckpt",
+    # MagNET-x: explicit-solvent shieldings, one model per solvent (see compute_MagNET_x_corrections).
+    # The released checkpoint folder is "MagNET-x" (lowercase x); paths are case-sensitive on Linux.
+    "MagNET-x-chloroform_H": "model_checkpoints/MagNET-x/MagNET-x-chloroform_1H.ckpt",
+    "MagNET-x-chloroform_C": "model_checkpoints/MagNET-x/MagNET-x-chloroform_13C.ckpt",
+    "MagNET-x-benzene_H": "model_checkpoints/MagNET-x/MagNET-x-benzene_1H.ckpt",
+    "MagNET-x-benzene_C": "model_checkpoints/MagNET-x/MagNET-x-benzene_13C.ckpt",
+    "MagNET-x-methanol_H": "model_checkpoints/MagNET-x/MagNET-x-methanol_1H.ckpt",
+    "MagNET-x-methanol_C": "model_checkpoints/MagNET-x/MagNET-x-methanol_13C.ckpt",
+    "MagNET-x-water_H": "model_checkpoints/MagNET-x/MagNET-x-water_1H.ckpt",
+    "MagNET-x-water_C": "model_checkpoints/MagNET-x/MagNET-x-water_13C.ckpt",
+}
+
+
+def _default_device():
+    return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+def _resolve_checkpoint(path, checkpoints_dir=None):
+    """Locate a checkpoint file given a MODEL_CHECKPOINTS-relative path like
+    "model_checkpoints/MagNET/MagNET_1H.ckpt". Resolution order for a relative path:
+      1. under `checkpoints_dir`, if given (it may be the model_checkpoints/ folder itself, or a
+         directory that contains one); this takes precedence over the defaults below,
+      2. as given (relative to the current directory),
+      3. relative to the release repo root (the parent of this magnet/ folder), where
+         model_checkpoints/ sits in the integrated Hugging Face checkout.
+    Absolute paths pass through unchanged. If nothing is found, return the original so the loader
+    raises a clear error on it."""
+    if checkpoints_dir:
+        # the stored path is "model_checkpoints/<rest>"; accept a checkpoints_dir that IS the
+        # model_checkpoints/ folder (join <rest>) or one that CONTAINS it (join the whole path)
+        rest = path[len("model_checkpoints/"):] if path.startswith("model_checkpoints/") else path
+        for candidate in (os.path.join(checkpoints_dir, rest), os.path.join(checkpoints_dir, path)):
+            if os.path.exists(candidate):
+                return candidate
+    if os.path.exists(path):
+        return path
+    repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))  # parent of magnet/
+    candidate = os.path.join(repo_root, path)
+    if os.path.exists(candidate):
+        return candidate
+    return path
+
+
+def load_model_to_device(path, device=None, checkpoints_dir=None):
+    """Load one checkpoint onto `device`. `checkpoints_dir`, if given, is where to find the released
+    weights: it may point at the model_checkpoints/ folder itself or at a directory that contains one,
+    and takes precedence over the current directory and the repo root (see _resolve_checkpoint). Pass
+    it when the weights live outside the repo; leave it None to use the in-repo/CWD defaults."""
+    if device is None:
+        device = _default_device()
+    path = _resolve_checkpoint(path, checkpoints_dir=checkpoints_dir)
+    return MagNET_Lightning.load_from_checkpoint(checkpoint_path=path, map_location=device).to(device).eval()
+
+
+def _predict_with(key_H, key_C, atomic_numbers_list, geometries_list,
+                  n_passes=20, symmetrize=False, device=None, checkpoints_dir=None, **predict_kwargs):
+    """Load one model pair and run predict_shieldings over a list of molecules. Returns a list of
+    per-atom shielding arrays. Extra keyword arguments are forwarded to predict_shieldings (used by
+    the explicit-solvent path for the solute/solvent split)."""
+    if device is None:
+        device = _default_device()
+    model_H = load_model_to_device(MODEL_CHECKPOINTS[key_H], device, checkpoints_dir=checkpoints_dir)
+    model_C = load_model_to_device(MODEL_CHECKPOINTS[key_C], device, checkpoints_dir=checkpoints_dir)
+    out = []
+    for entry in zip(atomic_numbers_list, geometries_list):
+        atomic_numbers, geometry = entry[0], entry[1]
+        out.append(np.atleast_1d(predict_shieldings(
+            model_H, model_C, solute_atomic_numbers=atomic_numbers, geometry=geometry,
+            device=device, n_passes=n_passes, symmetrize=symmetrize, **predict_kwargs).squeeze()))
+    return out
+
+
+def compute_MagNET_foundation_shieldings(atomic_numbers_list, geometries_list,
+                                         n_passes=20, symmetrize=False, device=None, checkpoints_dir=None):
+    """MagNET foundation-model gas-phase shieldings (1H and 13C) for AIMNet2 geometries."""
+    return _predict_with("MagNET_H", "MagNET_C", atomic_numbers_list, geometries_list,
+                         n_passes=n_passes, symmetrize=symmetrize, device=device,
+                         checkpoints_dir=checkpoints_dir)
+
+
+def compute_MagNET_Zero_shieldings(atomic_numbers_list, geometries_list,
+                                   n_passes=20, symmetrize=False, device=None, checkpoints_dir=None):
+    """MagNET-Zero gas-phase shieldings (1H and 13C) for a list of AIMNet2 geometries.
+
+    Args:
+        atomic_numbers_list: list of (N,) integer atomic-number arrays.
+        geometries_list: list of (N, 3) coordinate arrays, in Angstrom.
+        n_passes, symmetrize: prediction options (see module docstring).
+        checkpoints_dir: optional directory holding the released weights (see load_model_to_device).
+
+    Returns a list of (N,) shielding arrays (1H at H atoms, 13C at C atoms).
+    """
+    return _predict_with("MagNET-Zero_H", "MagNET-Zero_C", atomic_numbers_list, geometries_list,
+                         n_passes=n_passes, symmetrize=symmetrize, device=device,
+                         checkpoints_dir=checkpoints_dir)
+
+
+def compute_MagNET_PCM_corrections(atomic_numbers_list, geometries_list,
+                                   n_passes=20, symmetrize=False, device=None, checkpoints_dir=None):
+    """MagNET-PCM implicit-solvent (chloroform) corrections (1H and 13C) for AIMNet2 geometries.
+
+    The correction is the difference between the PCM-corrected shielding and the gas-phase shielding,
+    both from the B3LYP-D3(BJ)/pcSseg-2 MagNET-PCM checkpoints.
+
+    Returns a list of (N,) correction arrays.
+    """
+    gas = _predict_with("MagNET-PCM-withoutPCM_H", "MagNET-PCM-withoutPCM_C",
+                        atomic_numbers_list, geometries_list,
+                        n_passes=n_passes, symmetrize=symmetrize, device=device,
+                        checkpoints_dir=checkpoints_dir)
+    pcm = _predict_with("MagNET-PCM-withPCM_H", "MagNET-PCM-withPCM_C",
+                        atomic_numbers_list, geometries_list,
+                        n_passes=n_passes, symmetrize=symmetrize, device=device,
+                        checkpoints_dir=checkpoints_dir)
+    return [pcm[i] - gas[i] for i in range(len(pcm))]
+
+
+# Atoms per solvent molecule, for the explicit-solvent split (water is stored/run as TIP4P).
+N_ATOMS_PER_SOLVENT = {"chloroform": 5, "benzene": 12, "methanol": 6, "water": 3}
+
+
+def compute_MagNET_x_corrections(solvent, solute_atomic_numbers_list, atomic_numbers_list,
+                                 geometries_list, n_passes=20, symmetrize=False, device=None,
+                                 solvent_distance_threshold=12.0, checkpoints_dir=None):
+    """MagNET-x explicit-solvent corrections (solvated minus isolated) for one solvent.
+
+    Args:
+        solvent: 'chloroform' | 'benzene' | 'methanol' | 'water' (water uses the TIP4P model).
+        solute_atomic_numbers_list: list of (n_solute,) solute atomic-number arrays.
+        atomic_numbers_list: list of (n_total,) atomic numbers for the whole solute+solvent system.
+            The solute atoms must come first, followed by complete solvent molecules in contiguous
+            blocks of N_ATOMS_PER_SOLVENT[solvent] atoms.
+        geometries_list: list of (n_total, 3) coordinate arrays (same atom order), in Angstrom.
+        n_passes, symmetrize: prediction options (see module docstring).
+        solvent_distance_threshold: drop solvent molecules whose nearest atom is beyond this many
+            Angstrom of the solute (the released pipeline used 12.0). The isolated pass uses 0.0
+            internally, which strips all solvent.
+        checkpoints_dir: optional directory holding the released weights (see load_model_to_device).
+
+    Returns a list of (n_solute,) correction arrays. The published pipeline averages this over MD
+    frames; a single-frame value is only an estimate.
+    """
+    if solvent not in N_ATOMS_PER_SOLVENT:
+        raise ValueError(f"solvent must be one of {sorted(N_ATOMS_PER_SOLVENT)}; got {solvent!r}")
+    n_per = N_ATOMS_PER_SOLVENT[solvent]
+    if device is None:
+        device = _default_device()
+    model_H = load_model_to_device(MODEL_CHECKPOINTS[f"MagNET-x-{solvent}_H"], device, checkpoints_dir=checkpoints_dir)
+    model_C = load_model_to_device(MODEL_CHECKPOINTS[f"MagNET-x-{solvent}_C"], device, checkpoints_dir=checkpoints_dir)
+    out = []
+    for solute_Z, full_Z, geom in zip(solute_atomic_numbers_list, atomic_numbers_list, geometries_list):
+        kw = dict(model_H=model_H, model_C=model_C, solute_atomic_numbers=solute_Z, geometry=geom,
+                  atomic_numbers=full_Z, N_atoms_per_solvent=n_per, device=device,
+                  n_passes=n_passes, symmetrize=symmetrize)
+        solvated = np.atleast_1d(predict_shieldings(solvent_distance_threshold=solvent_distance_threshold, **kw).squeeze())
+        isolated = np.atleast_1d(predict_shieldings(solvent_distance_threshold=0.0, **kw).squeeze())
+        out.append(solvated - isolated)
+    return out
+
+
+if __name__ == "__main__":
+    atomic_numbers = np.array([6, 6, 6, 8, 8, 6, 1, 1, 1, 1, 1, 1])
+    coordinates = np.array(
+        [
+            [-2.07273936, 0.71782219, -0.27169511],
+            [-1.51154256, -0.47834334, -0.05051489],
+            [-0.06133485, -0.72471339, 0.10392854],
+            [0.38785097, -1.84332836, 0.30459866],
+            [0.65684897, 0.41875002, -0.00514369],
+            [2.0648613, 0.23182768, 0.13713977],
+            [-1.48494101, 1.6266042, -0.3569932],
+            [-3.14927745, 0.81001383, -0.37304503],
+            [-2.12480974, -1.36963022, 0.03013895],
+            [2.44585443, -0.42836204, -0.64841408],
+            [2.30088663, -0.16782625, 1.12843716],
+            [2.5483427, 1.20718563, 0.03357037],
+        ]
+    )
+
+    foundation = compute_MagNET_foundation_shieldings([atomic_numbers], [coordinates])
+    zero = compute_MagNET_Zero_shieldings([atomic_numbers], [coordinates])
+    pcm = compute_MagNET_PCM_corrections([atomic_numbers], [coordinates])
+    print("foundation gas-phase shieldings:", foundation[0])
+    print("MagNET-Zero gas-phase shieldings:", zero[0])
+    print("MagNET-PCM chloroform corrections:", pcm[0])

magnet/scaling.py

@@ -0,0 +1,78 @@
+"""Turn MagNET-Zero / MagNET-PCM shieldings into the chemical shifts a chemist measures, using the
+paper's published per-solvent linear scaling (SI Table S10 for proton, S11 for carbon).
+
+This is a tiny, self-contained copy (numpy only), so `from magnet import published_scaling_tables,
+predict_shift` works with no data download and no extra path setup. The full derivation of these
+tables from the delta-22 data lives in analysis/code/scaling_factors.py, and a test there checks the
+two copies agree so they cannot drift.
+
+The prediction equation, for one nucleus and one solvent, is
+
+    shift = intercept + stationary * sigma_zero + pcm * delta_pcm
+
+where sigma_zero is the MagNET-Zero gas-phase shielding (compute_MagNET_Zero_shieldings) and delta_pcm
+is the MagNET-PCM chloroform correction (compute_MagNET_PCM_corrections). You always pass the
+chloroform correction; the per-solvent scaling is already folded into the `pcm` coefficient.
+"""
+import csv
+import io
+
+import numpy as np
+
+# Published SI Tables S10 (proton) and S11 (carbon), verbatim. Columns: intercept, stationary, pcm.
+_PUBLISHED_TABLE_CSV = {
+    "H": """solvent,intercept,stationary,pcm
+tetrahydrofuran,31.321785,-0.980175,-0.786672
+dichloromethane,31.377310,-0.979870,-0.808573
+chloroform,31.294992,-0.975794,-0.852690
+toluene,31.716957,-0.996733,1.944856
+benzene,31.987569,-1.005289,2.236515
+chlorobenzene,31.681438,-0.993846,0.830236
+acetone,31.512157,-0.987261,-1.238414
+dimethylsulfoxide,31.598698,-0.991103,-1.355248
+acetonitrile,31.496110,-0.985716,-0.974465
+trifluoroethanol,30.876004,-0.959997,-0.958852
+methanol,31.256030,-0.976446,-1.250121
+TIP4P,31.359134,-0.978809,-1.478298
+""",
+    "C": """solvent,intercept,stationary,pcm
+tetrahydrofuran,171.054488,-0.919000,-1.069509
+dichloromethane,171.509389,-0.921167,-1.115015
+chloroform,171.728797,-0.924231,-0.936804
+toluene,171.690907,-0.924987,-0.618410
+benzene,171.967177,-0.927033,-0.593716
+chlorobenzene,171.075910,-0.921238,-0.997641
+acetone,171.308023,-0.919610,-1.239503
+dimethylsulfoxide,170.598425,-0.918674,-1.296501
+acetonitrile,171.879839,-0.922629,-1.287664
+trifluoroethanol,174.237671,-0.939012,-1.290073
+methanol,172.426161,-0.927566,-1.288847
+TIP4P,173.696379,-0.937854,-1.342668
+""",
+}
+
+
+def published_scaling_tables():
+    """The published scaling tables as {"H": table, "C": table}. Each table maps a solvent name to a
+    {"intercept", "stationary", "pcm"} dict of coefficients. Pass one table and a solvent name to
+    predict_shift. Water is named "TIP4P"."""
+    tables = {}
+    for nucleus, text in _PUBLISHED_TABLE_CSV.items():
+        table = {}
+        for row in csv.DictReader(io.StringIO(text.strip())):
+            solvent = row.pop("solvent")
+            table[solvent] = {name: float(value) for name, value in row.items()}
+        tables[nucleus] = table
+    return tables
+
+
+def predict_shift(table, solvent, magnet_zero_shielding, magnet_pcm_chloroform_correction):
+    """Chemical shift(s) from MagNET-Zero / MagNET-PCM outputs, for one nucleus and solvent.
+
+    table: one nucleus's table from published_scaling_tables(). solvent: a solvent name (water is
+    "TIP4P"). magnet_zero_shielding and magnet_pcm_chloroform_correction are scalars or arrays.
+    """
+    row = table[solvent]
+    return (row["intercept"]
+            + row["stationary"] * np.asarray(magnet_zero_shielding, dtype=float)
+            + row["pcm"] * np.asarray(magnet_pcm_chloroform_correction, dtype=float))

magnet/test_api.py

@@ -0,0 +1,463 @@
+"""Adversarial edge-case tests for the magnet public API (magnet/api.py).
+
+Pure-logic tests (single-vs-list detection, validation, solvent checking, solute selection) run
+whenever the torch stack is installed. Tests marked @needs_model load real checkpoints and are
+skipped unless the weights are present; they stay few and fast (n_passes=1, symmetrize=False, tiny
+molecules). Run the whole suite with `pytest`; to include the real-model tests, check out the
+Hugging Face copy so the weights land in model_checkpoints/.
+"""
+import os
+import warnings
+
+import numpy as np
+import pytest
+
+warnings.filterwarnings("ignore")
+
+import magnet.api as api
+from magnet import scaling
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_REPO = os.path.dirname(_HERE)                       # magnet/ -> repo root
+_CKPT_REL = os.path.join("model_checkpoints", "MagNET-Zero", "MagNET-Zero_1H.ckpt")
+_HAVE_MODELS = os.path.exists(os.path.join(_REPO, _CKPT_REL))
+needs_model = pytest.mark.skipif(not _HAVE_MODELS, reason="checkpoints not available")
+
+# ---- tiny geometries (validity irrelevant for API-contract tests) ----
+CH4_Z = np.array([6, 1, 1, 1, 1])
+CH4_XYZ = np.array([[0., 0., 0.],
+                    [0.629, 0.629, 0.629],
+                    [-0.629, -0.629, 0.629],
+                    [-0.629, 0.629, -0.629],
+                    [0.629, -0.629, -0.629]])
+# formaldehyde CH2O: exercises H, C, and an O (should be NaN in shifts)
+CH2O_Z = np.array([6, 8, 1, 1])
+CH2O_XYZ = np.array([[0., 0., 0.], [1.2, 0., 0.], [-0.5, 0.94, 0.], [-0.5, -0.94, 0.]])
+# one chloroform molecule, atoms in molecule order C H Cl Cl Cl
+CHCL3_Z = np.array([6, 1, 17, 17, 17])
+CHCL3_XYZ = np.array([[0., 0., 0.], [0., 0., 1.1], [1.7, 0., -0.4],
+                      [-0.85, 1.47, -0.4], [-0.85, -1.47, -0.4]])
+
+
+# =====================================================================
+# _as_batch: single-vs-list detection and validation
+# =====================================================================
+
+def test_single_ndarray_is_single():
+    an, xyz, single = api._as_batch(CH4_Z, CH4_XYZ)
+    assert single is True and len(an) == 1 and an[0].shape == (5,) and xyz[0].shape == (5, 3)
+
+def test_single_pythonlist_is_single():
+    an, xyz, single = api._as_batch([6, 1, 1, 1, 1], CH4_XYZ.tolist())
+    assert single is True and len(xyz) == 1
+
+def test_one_atom_molecule_is_single():
+    an, xyz, single = api._as_batch(np.array([1]), np.array([[0., 0., 0.]]))
+    assert single is True and xyz[0].shape == (1, 3)
+
+def test_list_of_molecules_is_batch():
+    an, xyz, single = api._as_batch([CH4_Z, CH2O_Z], [CH4_XYZ, CH2O_XYZ])
+    assert single is False and len(an) == 2
+
+def test_batch_size_one_list_is_batch():
+    an, xyz, single = api._as_batch([CH4_Z], [CH4_XYZ])
+    assert single is False and len(an) == 1
+
+def test_3d_ndarray_is_batch():
+    Z = np.stack([CH4_Z, CH4_Z]);  X = np.stack([CH4_XYZ, CH4_XYZ])  # (2,5) and (2,5,3)
+    an, xyz, single = api._as_batch(Z, X)
+    assert single is False and len(xyz) == 2 and xyz[0].shape == (5, 3)
+
+def test_ragged_batch_ok():
+    an, xyz, single = api._as_batch([CH4_Z, CH2O_Z], [CH4_XYZ, CH2O_XYZ])
+    assert xyz[0].shape == (5, 3) and xyz[1].shape == (4, 3)
+
+def test_empty_list_raises():
+    with pytest.raises(ValueError, match="empty input"):
+        api._as_batch([], [])
+
+def test_empty_ndarray_raises():
+    with pytest.raises(ValueError, match="empty input"):
+        api._as_batch(np.array([]), np.zeros((0, 3)))
+
+def test_mismatched_atom_counts_single_raises():
+    with pytest.raises(ValueError):
+        api._as_batch(np.array([6, 1, 1]), np.zeros((4, 3)))
+
+def test_wrong_last_dim_raises():
+    with pytest.raises(ValueError):
+        api._as_batch(np.array([6, 1]), np.zeros((2, 2)))
+
+def test_mismatched_list_lengths_raises():
+    with pytest.raises(ValueError):
+        api._as_batch([CH4_Z, CH4_Z], [CH4_XYZ, CH4_XYZ, CH4_XYZ])
+
+def test_coords_1d_single_atom_flat_is_misread():
+    # coordinates=(3,) flat for one atom: coordinates[0] is a scalar (ndim 0) -> not single ->
+    # treated as a 3-molecule batch of 0-d coords -> raises. Documents the (N,3) requirement.
+    with pytest.raises(ValueError):
+        api._as_batch(np.array([6]), np.array([0., 0., 0.]))
+
+
+# =====================================================================
+# _check_passes
+# =====================================================================
+
+@pytest.mark.parametrize("bad", [0, -1, 1.5, "3", None, 2.0])
+def test_check_passes_rejects(bad):
+    with pytest.raises(ValueError, match="n_passes"):
+        api._check_passes(bad)
+
+@pytest.mark.parametrize("good", [1, 10, np.int64(5)])
+def test_check_passes_accepts(good):
+    api._check_passes(good)  # no raise
+
+def test_check_passes_bool_quirk():
+    # bool is a subclass of int in Python; True passes as "1 pass", False is rejected.
+    api._check_passes(True)              # accepted (quirk, not a real bug)
+    with pytest.raises(ValueError):
+        api._check_passes(False)
+
+
+# =====================================================================
+# predict_shieldings argument validation (pure logic, before model load)
+# =====================================================================
+
+def test_predict_shieldings_bad_model():
+    with pytest.raises(ValueError, match="model must be one of"):
+        api.predict_shieldings(CH4_Z, CH4_XYZ, model="bogus")
+
+@pytest.mark.parametrize("m,fn", [("MagNET-PCM", "implicit_solvent_correction"),
+                                  ("MagNET-x", "explicit_solvent_correction")])
+def test_predict_shieldings_correction_redirect(m, fn):
+    with pytest.raises(ValueError, match=fn):
+        api.predict_shieldings(CH4_Z, CH4_XYZ, model=m)
+
+def test_predict_shieldings_bad_passes():
+    with pytest.raises(ValueError, match="n_passes"):
+        api.predict_shieldings(CH4_Z, CH4_XYZ, n_passes=0)
+
+
+# =====================================================================
+# predict_shifts solvent validation (pure logic, before model load)
+# =====================================================================
+
+def test_predict_shifts_unknown_solvent():
+    with pytest.raises(ValueError, match="unknown solvent"):
+        api.predict_shifts(CH4_Z, CH4_XYZ, solvent="ethanol")
+
+def test_predict_shifts_solvent_case_sensitive():
+    with pytest.raises(ValueError, match="unknown solvent"):
+        api.predict_shifts(CH4_Z, CH4_XYZ, solvent="Chloroform")
+
+def test_predict_shifts_water_not_rejected_by_validation():
+    # "water" must map to TIP4P and NOT trip the unknown-solvent guard.
+    tables = scaling.published_scaling_tables()
+    assert "TIP4P" in tables["C"] and "TIP4P" in tables["H"]
+    # the error message offers "water" (not "TIP4P") as an option:
+    options = ["water" if s == "TIP4P" else s for s in tables["C"]]
+    assert "water" in options and "TIP4P" not in options
+
+
+# =====================================================================
+# _validate_solvent: atom-type-count check against a named solvent
+# =====================================================================
+
+def test_validate_chloroform():
+    api._validate_solvent("chloroform", CHCL3_Z)  # no raise
+
+def test_validate_two_chloroforms():
+    api._validate_solvent("chloroform", np.concatenate([CHCL3_Z, CHCL3_Z]))
+
+def test_validate_water():
+    api._validate_solvent("water", np.array([8, 1, 1]))
+
+def test_validate_benzene():
+    api._validate_solvent("benzene", np.array([6] * 6 + [1] * 6))
+
+def test_validate_methanol():
+    api._validate_solvent("methanol", np.array([6, 8, 1, 1, 1, 1]))
+
+def test_validate_order_within_block_independent():
+    # order within one molecule's block does not matter
+    api._validate_solvent("water", np.array([1, 8, 1]))
+
+def test_validate_empty_raises():
+    with pytest.raises(ValueError, match="no solvent atoms"):
+        api._validate_solvent("chloroform", np.array([], dtype=int))
+
+def test_validate_wrong_count_raises():
+    # 7 atoms is not a whole number of chloroforms (5 atoms each)
+    with pytest.raises(ValueError, match="not a whole number"):
+        api._validate_solvent("chloroform", np.array([6, 6, 1, 1, 17, 17, 17]))
+
+def test_validate_wrong_composition_raises():
+    # 5 atoms (one chloroform-sized block) but the wrong elements
+    with pytest.raises(ValueError, match="not whole chloroform"):
+        api._validate_solvent("chloroform", np.array([6, 6, 6, 6, 6]))
+
+def test_validate_grouped_by_element_raises():
+    # two chloroforms' worth of atoms (correct TOTALS: 2 C, 2 H, 6 Cl) but grouped by element, so the
+    # contiguous 5-atom blocks are not whole molecules. Total-count check would miss this; block does not.
+    grouped = np.array([6, 6, 1, 1, 17, 17, 17, 17, 17, 17])
+    with pytest.raises(ValueError, match="not whole chloroform"):
+        api._validate_solvent("chloroform", grouped)
+
+def test_validate_two_chloroforms_contiguous_ok():
+    # same atoms as above but laid out as two whole contiguous molecules: fine
+    api._validate_solvent("chloroform", np.concatenate([CHCL3_Z, CHCL3_Z]))
+
+def test_validate_mismatched_solvent_raises():
+    # a benzene block passed as methanol (both divisible cases would still fail composition)
+    with pytest.raises(ValueError, match="not whole methanol"):
+        api._validate_solvent("methanol", np.array([6, 6, 6, 6, 6, 6]))
+
+
+# =====================================================================
+# supported-solvent tables stay in sync with the code (validation)
+# =====================================================================
+
+EXPLICIT_SOLVENTS = ["chloroform", "benzene", "methanol", "water"]
+PREDICT_SHIFTS_SOLVENTS = [
+    "tetrahydrofuran", "dichloromethane", "chloroform", "toluene", "benzene", "chlorobenzene",
+    "acetone", "dimethylsulfoxide", "acetonitrile", "trifluoroethanol", "methanol", "water",
+]
+
+def _one_molecule(solvent):
+    """Build one molecule of `solvent` as an atomic-number array from the composition table."""
+    return np.array([element for element, count in api._SOLVENT_COMPOSITION[solvent].items()
+                     for _ in range(count)])
+
+def test_explicit_solvent_set_matches_run_magnet():
+    # composition table, the model's block-size table, and the documented list all agree
+    from magnet import run_magnet
+    assert sorted(api._SOLVENT_COMPOSITION) == sorted(run_magnet.N_ATOMS_PER_SOLVENT)
+    assert sorted(api._SOLVENT_COMPOSITION) == sorted(EXPLICIT_SOLVENTS)
+
+def test_solvent_composition_sums_to_block_size():
+    from magnet import run_magnet
+    for solvent, comp in api._SOLVENT_COMPOSITION.items():
+        assert sum(comp.values()) == run_magnet.N_ATOMS_PER_SOLVENT[solvent]
+
+@pytest.mark.parametrize("solvent", EXPLICIT_SOLVENTS)
+def test_validate_accepts_one_molecule(solvent):
+    api._validate_solvent(solvent, _one_molecule(solvent))  # must not raise
+
+@pytest.mark.parametrize("solvent", EXPLICIT_SOLVENTS)
+def test_validate_accepts_two_contiguous_molecules(solvent):
+    one = _one_molecule(solvent)
+    api._validate_solvent(solvent, np.concatenate([one, one]))  # must not raise
+
+@pytest.mark.parametrize("solvent", EXPLICIT_SOLVENTS)
+def test_validate_rejects_dropping_one_atom(solvent):
+    # one atom short of a whole molecule -> wrong count
+    short = _one_molecule(solvent)[:-1]
+    with pytest.raises(ValueError, match="not a whole number"):
+        api._validate_solvent(solvent, short)
+
+def test_predict_shifts_supports_the_12_documented_solvents():
+    tables = scaling.published_scaling_tables()
+    for nucleus in ("H", "C"):
+        keys = {"water" if s == "TIP4P" else s for s in tables[nucleus]}
+        assert keys == set(PREDICT_SHIFTS_SOLVENTS)
+
+def test_explicit_rejects_predict_only_solvent():
+    # acetone is valid for predict_shifts but not for explicit corrections
+    atoms = np.concatenate([np.array([6, 1, 1, 1]), CHCL3_Z])
+    xyz = np.zeros((len(atoms), 3))
+    with pytest.raises(ValueError, match="solvent must be one of"):
+        api.explicit_solvent_correction(atoms, xyz, solute_atoms=[0, 1, 2, 3], solvent="acetone",
+                                        n_passes=1, symmetrize=False)
+
+
+# =====================================================================
+# _explicit_one: solute selection, reordering, negative/dup/range
+# =====================================================================
+
+def _system():  # methane solute (0..3) + one chloroform (4..8)
+    atoms = np.concatenate([np.array([6, 1, 1, 1]), CHCL3_Z])
+    xyz = np.concatenate([CH4_XYZ[:4] + 10.0, CHCL3_XYZ])  # solute far, solvent near origin
+    return atoms, xyz
+
+def test_explicit_one_basic():
+    atoms, xyz = _system()
+    solute_an, full_an, full_xyz = api._explicit_one(atoms, xyz, [0, 1, 2, 3], "chloroform")
+    assert np.array_equal(solute_an, np.array([6, 1, 1, 1]))
+    assert np.array_equal(full_an, atoms)  # already solute-first
+
+def test_explicit_one_reorders_solvent_first_input():
+    # input has chloroform first, methane second; solute indices point at the methane
+    atoms = np.concatenate([CHCL3_Z, np.array([6, 1, 1, 1])])
+    xyz = np.concatenate([CHCL3_XYZ, CH4_XYZ[:4] + 10.0])
+    solute_an, full_an, full_xyz = api._explicit_one(atoms, xyz, [5, 6, 7, 8], "chloroform")
+    assert np.array_equal(solute_an, np.array([6, 1, 1, 1]))
+    # reordered: solute first, then the solvent block
+    assert np.array_equal(full_an[:4], np.array([6, 1, 1, 1]))
+    assert np.array_equal(np.sort(full_an[4:]), np.sort(CHCL3_Z))
+
+def test_explicit_one_output_order_follows_solute_atoms():
+    atoms, xyz = _system()
+    solute_an, full_an, _ = api._explicit_one(atoms, xyz, [3, 2, 1, 0], "chloroform")  # reversed
+    assert np.array_equal(solute_an, np.array([1, 1, 1, 6]))  # matches requested order
+    assert np.array_equal(full_an[:4], np.array([1, 1, 1, 6]))
+
+def test_explicit_one_negative_indices():
+    atoms, xyz = _system()  # n=9
+    solute_an, _, _ = api._explicit_one(atoms, xyz, [-9, -8, -7, -6], "chloroform")  # -> 0,1,2,3
+    assert np.array_equal(solute_an, np.array([6, 1, 1, 1]))
+
+def test_explicit_one_wrong_solvent_raises():
+    atoms, xyz = _system()  # solvent is chloroform
+    with pytest.raises(ValueError, match="not .* benzene"):
+        api._explicit_one(atoms, xyz, [0, 1, 2, 3], "benzene")
+
+def test_explicit_one_duplicate_raises():
+    atoms, xyz = _system()
+    with pytest.raises(ValueError, match="duplicate"):
+        api._explicit_one(atoms, xyz, [0, 0, 1, 2], "chloroform")
+
+def test_explicit_one_negative_normalizes_to_duplicate_raises():
+    atoms, xyz = _system()  # n=9; index 2 and -7 both -> 2
+    with pytest.raises(ValueError, match="duplicate"):
+        api._explicit_one(atoms, xyz, [2, -7], "chloroform")
+
+def test_explicit_one_out_of_range_raises():
+    atoms, xyz = _system()
+    with pytest.raises(ValueError, match="out of range"):
+        api._explicit_one(atoms, xyz, [0, 1, 2, 9], "chloroform")
+
+def test_explicit_one_very_negative_raises():
+    atoms, xyz = _system()
+    with pytest.raises(ValueError, match="out of range"):
+        api._explicit_one(atoms, xyz, [-100], "chloroform")
+
+def test_explicit_one_all_atoms_solute_raises():
+    atoms, xyz = _system()
+    with pytest.raises(ValueError, match="no solvent atoms"):
+        api._explicit_one(atoms, xyz, list(range(len(atoms))), "chloroform")
+
+def test_explicit_one_empty_solute_raises():
+    atoms, xyz = _system()
+    with pytest.raises(ValueError, match="empty"):
+        api._explicit_one(atoms, xyz, [], "chloroform")
+
+
+# =====================================================================
+# explicit_solvent_correction: pure-logic validation before any model call
+# =====================================================================
+
+def test_explicit_unknown_solvent_raises():
+    atoms, xyz = _system()
+    with pytest.raises(ValueError, match="solvent must be one of"):
+        api.explicit_solvent_correction(atoms, xyz, solute_atoms=[0, 1, 2, 3], solvent="ethanol",
+                                        n_passes=1, symmetrize=False)
+
+def test_explicit_solvent_mismatch_raises():
+    # declare benzene but the snapshot has chloroform -> count check fails before any model load
+    atoms, xyz = _system()
+    with pytest.raises(ValueError, match="not .* benzene"):
+        api.explicit_solvent_correction(atoms, xyz, solute_atoms=[0, 1, 2, 3], solvent="benzene",
+                                        n_passes=1, symmetrize=False)
+
+
+# =====================================================================
+# REAL MODEL TESTS (few, tiny, n_passes=1, symmetrize=False)
+# =====================================================================
+
+@needs_model
+def test_real_predict_shieldings_single_returns_array():
+    out = api.predict_shieldings(CH4_Z, CH4_XYZ, n_passes=1, symmetrize=False)
+    assert isinstance(out, np.ndarray) and out.shape == (5,)
+    assert np.all(np.isfinite(out))  # H and C all predicted, none NaN
+
+@needs_model
+def test_real_predict_shieldings_list_returns_list():
+    out = api.predict_shieldings([CH4_Z], [CH4_XYZ], n_passes=1, symmetrize=False)
+    assert isinstance(out, list) and len(out) == 1 and out[0].shape == (5,)
+
+@needs_model
+def test_real_predict_shieldings_batch_shapes():
+    out = api.predict_shieldings([CH4_Z, CH2O_Z], [CH4_XYZ, CH2O_XYZ],
+                                 n_passes=1, symmetrize=False)
+    assert isinstance(out, list) and out[0].shape == (5,) and out[1].shape == (4,)
+
+@needs_model
+def test_real_unsupported_element_rejected():
+    # PH3-like: phosphorus (15) is out of vocab
+    Z = np.array([15, 1, 1, 1])
+    X = np.array([[0., 0., 0.], [1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])
+    with pytest.raises(ValueError, match="unsupported atomic numbers"):
+        api.predict_shieldings(Z, X, n_passes=1, symmetrize=False)
+
+@needs_model
+def test_real_predict_shifts_components_reproduce():
+    d = api.predict_shifts(CH2O_Z, CH2O_XYZ, solvent="chloroform", n_passes=1,
+                           symmetrize=False, return_components=True)
+    assert set(d) == {"shifts", "zero_shielding", "pcm_correction", "coefficients"}
+    shifts, sigma, delta = d["shifts"], d["zero_shielding"], d["pcm_correction"]
+    coef = d["coefficients"]
+    # coefficients must match chloroform rows
+    tbl = scaling.published_scaling_tables()
+    assert coef["H"] == tbl["H"]["chloroform"] and coef["C"] == tbl["C"]["chloroform"]
+    # reproduce shift = intercept + stationary*zero + pcm*pcm for H (idx 2,3) and C (idx 0)
+    for idx, nuc in [(0, "C"), (2, "H"), (3, "H")]:
+        expect = coef[nuc]["intercept"] + coef[nuc]["stationary"] * sigma[idx] \
+                 + coef[nuc]["pcm"] * delta[idx]
+        assert np.isclose(shifts[idx], expect, atol=1e-6), (idx, nuc, shifts[idx], expect)
+    # O atom (idx 1) must be NaN, H/C finite
+    assert np.isnan(shifts[1])
+    assert np.all(np.isfinite([shifts[0], shifts[2], shifts[3]]))
+
+@needs_model
+def test_real_predict_shifts_acetone_matches_readme():
+    # the README "Your First Prediction" worked example: acetone in chloroform on an AIMNet2
+    # geometry, default n_passes/symmetrize. Guards the documented numbers against model drift.
+    Z = np.array([6, 6, 8, 6, 1, 1, 1, 1, 1, 1])
+    X = np.array([
+        [1.2913, -0.5947, -0.0016], [0.0029, 0.1931, -0.0010], [-0.0174, 1.3994, -0.0003],
+        [-1.2743, -0.6189, 0.0004], [-1.0822, -1.6899, 0.0014], [-1.8597, -0.3513, 0.8791],
+        [-1.8605, -0.3530, -0.8783], [1.3415, -1.2208, 0.8916], [1.3170, -1.2669, -0.8614],
+        [2.1415, 0.0788, -0.0298],
+    ])
+    shifts = api.predict_shifts(Z, X, solvent="chloroform")
+    carbonyl = shifts[[1]].mean()
+    methyl_C = shifts[[0, 3]].mean()
+    methyl_H = shifts[[4, 5, 6, 7, 8, 9]].mean()
+    assert np.isclose(carbonyl, 207.6, atol=1.0), carbonyl
+    assert np.isclose(methyl_C, 30.7, atol=0.5), methyl_C
+    assert np.isclose(methyl_H, 2.20, atol=0.15), methyl_H
+
+@needs_model
+def test_real_predict_shifts_water_maps_to_tip4p():
+    d = api.predict_shifts(CH4_Z, CH4_XYZ, solvent="water", n_passes=1,
+                           symmetrize=False, return_components=True)
+    tbl = scaling.published_scaling_tables()
+    assert d["coefficients"]["H"] == tbl["H"]["TIP4P"]
+    assert d["coefficients"]["C"] == tbl["C"]["TIP4P"]
+
+@needs_model
+def test_real_single_vs_batch_same_shapes():
+    single = api.predict_shieldings(CH4_Z, CH4_XYZ, n_passes=1, symmetrize=False)
+    batch = api.predict_shieldings([CH4_Z, CH4_Z], [CH4_XYZ, CH4_XYZ],
+                                   n_passes=1, symmetrize=False)
+    assert single.shape == batch[0].shape == batch[1].shape
+    # values are close (stochastic frames -> not identical); documents non-determinism
+    assert np.allclose(single, batch[0], atol=0.5)
+
+@needs_model
+@pytest.mark.xfail(reason="known: zero-edge graph (lone atom) gives a cryptic torch error, not a "
+                          "clean ValueError; a lone atom is not a valid NMR input",
+                   raises=RuntimeError)
+def test_real_single_atom_solute():
+    # isolated single atom: no neighbors -> empty edge tensor -> deep torch.min() crash.
+    api.predict_shieldings(np.array([6]), np.array([[0., 0., 0.]]),
+                           n_passes=1, symmetrize=False)
+
+@needs_model
+@pytest.mark.xfail(reason="known: atoms all beyond the model cutoff -> zero edges -> torch "
+                          "RuntimeError, not a clean ValueError",
+                   raises=RuntimeError)
+def test_real_atoms_beyond_cutoff_crash():
+    # two atoms 50 A apart: no edges within cutoff -> same zero-edge crash as one atom.
+    api.predict_shieldings(np.array([6, 1]), np.array([[0., 0., 0.], [0., 0., 50.]]),
+                           n_passes=1, symmetrize=False)

magnet/test_magnet.py

@@ -0,0 +1,262 @@
+"""Tests for the magnet package.
+
+The synthetic tests build a tiny EquiformerV2 model (no checkpoint, no dataset) and run a
+forward pass, so they exercise the whole package wiring and run in CI. The tests that need the
+released checkpoints / reference data are gated with skipif and run only when those files are
+present in the repo (model_checkpoints/ and data/<name>/, carried by the Hugging Face checkout).
+"""
+import os
+import numpy as np
+import pytest
+
+# The model needs the heavy PyTorch / PyG stack (see requirements.txt). When it is not
+# installed (e.g. the dataset-reader CI, which installs only numpy/h5py/pandas), skip the whole
+# module cleanly instead of erroring at import.
+pytest.importorskip("torch")
+pytest.importorskip("torch_scatter")
+pytest.importorskip("torch_cluster")
+pytest.importorskip("torch_geometric")
+pytest.importorskip("magnet")
+
+from magnet.model import MagNET_Lightning
+from magnet.inference import predict_shieldings
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_REPO = os.path.abspath(os.path.join(_HERE, ".."))
+_CKPT_ROOT = os.path.join(_REPO, "model_checkpoints")  # absent on GitHub (git-ignored); every skipif below is then False
+CKPT = os.path.join(_CKPT_ROOT, "MagNET-Zero")
+# The natural-products reference shieldings are a tiny (54 KB) committed test fixture (test_data/), so
+# the reproduce/symmetrize tests need only the released checkpoints, not a separate reference download.
+NP_H5 = os.path.join(_HERE, "test_data", "magnet_zero_and_pcm_shieldings.h5")
+
+
+def _dataset_file(name):
+    """Resolve a released dataset hdf5 the way analysis/code/paths.py does, but without importing it
+    (this package is standalone): the in-repo data/<name>/<name>.hdf5, carried by the Hugging Face
+    checkout via Git LFS."""
+    return os.path.join(_REPO, "data", name, name + ".hdf5")
+
+TINY_PARAMS = dict(
+    model_type="EquiformerV2", max_neighbors=20, max_radius=5.0, max_num_elements=20,
+    num_layers=1, sphere_channels=8, attn_hidden_channels=8, num_heads=1,
+    attn_alpha_channels=8, attn_value_channels=8, ffn_hidden_channels=16,
+    lmax_list=[2], mmax_list=[2], grid_resolution=14, weight_init="uniform", lr=1e-4,
+)
+
+# a small tetrahedral-ish CH4 so the graph has edges
+METHANE_Z = np.array([6, 1, 1, 1, 1])
+METHANE_XYZ = np.array([[0, 0, 0], [0.63, 0.63, 0.63], [-0.63, -0.63, 0.63],
+                        [-0.63, 0.63, -0.63], [0.63, -0.63, -0.63]], dtype=float)
+
+
+def test_predict_rejects_unsupported_elements():
+    """MagNET only knows H, C, N, O, F, S, Cl; predicting on anything else must raise, not guess.
+    The check runs before the model, so dummy models are fine."""
+    # phosphorus (15) in the solute
+    with pytest.raises(ValueError, match="unsupported atomic numbers"):
+        predict_shieldings(None, None, np.array([6, 1, 1, 15]), np.zeros((4, 3)))
+    # an unsupported atom in the solvent block (the full atomic_numbers, MagNET-x path)
+    with pytest.raises(ValueError, match="unsupported"):
+        predict_shieldings(None, None, np.array([6, 1]), np.zeros((4, 3)),
+                           atomic_numbers=np.array([6, 1, 35, 1]))   # bromine
+    # a supported-only molecule passes the element check (it may fail later for lack of a real
+    # model, but not with an "unsupported" element error)
+    try:
+        predict_shieldings(None, None, METHANE_Z, METHANE_XYZ)
+    except ValueError as e:
+        assert "unsupported" not in str(e)
+    except Exception:
+        pass
+
+
+def test_resolve_checkpoint_honors_checkpoints_dir(tmp_path):
+    """checkpoints_dir points at the released weights when they live outside the repo. It may be the
+    model_checkpoints/ folder itself or a directory that contains one, and takes precedence."""
+    from magnet.run_magnet import _resolve_checkpoint
+    rel = "model_checkpoints/MagNET-Zero/MagNET-Zero_1H.ckpt"
+
+    # (a) checkpoints_dir IS the model_checkpoints/ folder -> the leading path segment is stripped
+    mc = tmp_path / "mc"
+    (mc / "MagNET-Zero").mkdir(parents=True)
+    inside = mc / "MagNET-Zero" / "MagNET-Zero_1H.ckpt"
+    inside.write_bytes(b"")
+    assert _resolve_checkpoint(rel, checkpoints_dir=str(mc)) == str(inside)
+
+    # (b) checkpoints_dir CONTAINS a model_checkpoints/ folder (an `hf download --local-dir` target)
+    parent = tmp_path / "download"
+    contained = parent / "model_checkpoints" / "MagNET-Zero" / "MagNET-Zero_1H.ckpt"
+    contained.parent.mkdir(parents=True)
+    contained.write_bytes(b"")
+    assert _resolve_checkpoint(rel, checkpoints_dir=str(parent)) == str(contained)
+
+    # (c) a checkpoints_dir without the file is not fabricated (falls through to the defaults)
+    empty = tmp_path / "empty"
+    empty.mkdir()
+    assert not _resolve_checkpoint(rel, checkpoints_dir=str(empty)).startswith(str(empty))
+
+    # (d) an absolute path resolves regardless of checkpoints_dir
+    assert _resolve_checkpoint(str(inside), checkpoints_dir=str(empty)) == str(inside)
+
+
+@pytest.fixture(scope="module")
+def tiny_model():
+    return MagNET_Lightning(TINY_PARAMS).eval()
+
+
+def test_predict_shape_and_finite(tiny_model):
+    out = predict_shieldings(tiny_model, tiny_model, solute_atomic_numbers=METHANE_Z,
+                             geometry=METHANE_XYZ, device="cpu")
+    out = np.asarray(out)
+    assert out.shape == (METHANE_Z.shape[0],)
+    assert np.isfinite(out).all()
+
+
+def test_n_passes_averages(tiny_model):
+    out = predict_shieldings(tiny_model, tiny_model, solute_atomic_numbers=METHANE_Z,
+                             geometry=METHANE_XYZ, device="cpu", n_passes=3)
+    assert np.asarray(out).shape == (METHANE_Z.shape[0],)
+    assert np.isfinite(out).all()
+
+
+def test_symmetrize_runs_and_is_reflection_consistent(tiny_model):
+    # the symmetrized prediction averages over G and its mirror image, so evaluating it on G
+    # and on reflect(G) must give the same per-atom values (up to per-pass frame noise).
+    refl = METHANE_XYZ.copy(); refl[:, 0] = -refl[:, 0]
+    a = predict_shieldings(tiny_model, tiny_model, solute_atomic_numbers=METHANE_Z,
+                           geometry=METHANE_XYZ, device="cpu", n_passes=8, symmetrize=True)
+    b = predict_shieldings(tiny_model, tiny_model, solute_atomic_numbers=METHANE_Z,
+                           geometry=refl, device="cpu", n_passes=8, symmetrize=True)
+    # both are reflection-symmetrized estimates of the same molecule
+    assert np.allclose(np.asarray(a), np.asarray(b), atol=0.5)
+
+
+def test_solvent_path_runs():
+    """Regression test for the explicit-solvent (MagNET-x) code path. That branch runs only when
+    filter_solvent_edges=True (the MagNET-x checkpoints set it), so the gas-phase tests never reach
+    it; a torch_geometric >= 2.6 change (data.keys became a method) once broke it. Build a tiny
+    solvent-aware model and feed a solute plus two solvent molecules."""
+    model = MagNET_Lightning(dict(TINY_PARAMS, filter_solvent_edges=True)).eval()
+    solute_Z = np.array([6, 1, 1, 1, 1])
+    solvent_mol = np.array([[0.0, 0.0, 0.0], [0.0, 0.96, 0.0], [0.93, -0.24, 0.0]])  # 3-atom solvent
+    full_Z = np.concatenate([solute_Z, [8, 1, 1], [8, 1, 1]])
+    geom = np.concatenate([METHANE_XYZ, solvent_mol + [3.5, 0, 0], solvent_mol + [0, 3.5, 0]])
+    out = np.asarray(predict_shieldings(model, model, solute_atomic_numbers=solute_Z, geometry=geom,
+                                        atomic_numbers=full_Z, N_atoms_per_solvent=3,
+                                        solvent_distance_threshold=12.0, device="cpu"))
+    assert out.shape == (solute_Z.shape[0],)
+    assert np.isfinite(out).all()
+
+
+@pytest.mark.skipif(not os.path.exists(CKPT), reason="released MagNET-Zero checkpoints not present")
+def test_real_magnet_zero_reproduces_reference():
+    import h5py
+    dev = "cpu"
+    mH = MagNET_Lightning.load_from_checkpoint(os.path.join(CKPT, "MagNET-Zero_1H.ckpt"), map_location=dev).eval()
+    mC = MagNET_Lightning.load_from_checkpoint(os.path.join(CKPT, "MagNET-Zero_13C.ckpt"), map_location=dev).eval()
+    with h5py.File(NP_H5, "r") as f:
+        name = list(f.keys())[0]
+        g = f[name]
+        z = g["atomic_numbers"][:]; xyz = g["geometry"][:].astype(np.float64)
+        # This reference file predates the MagNET casing standardization, so its dataset is stored
+        # under the old "MagNet_..." spelling. Read whichever spelling the file actually uses.
+        ref_key = "MagNET_Zero_shieldings" if "MagNET_Zero_shieldings" in g else "MagNet_Zero_shieldings"
+        ref = g[ref_key][:]
+    pred = predict_shieldings(mH, mC, solute_atomic_numbers=z, geometry=xyz, device=dev, n_passes=20)
+    tgt = (z == 1) | (z == 6)
+    # un-symmetrized 20-pass reproduces the (un-symmetrized) released numbers well under 0.1 ppm
+    assert np.abs(pred[tgt] - ref[tgt]).mean() < 0.05
+
+
+@pytest.mark.skipif(not os.path.exists(CKPT), reason="released MagNET-Zero checkpoints not present")
+def test_real_symmetrize_changes_13C_on_large_molecule():
+    import h5py
+    dev = "cpu"
+    mH = MagNET_Lightning.load_from_checkpoint(os.path.join(CKPT, "MagNET-Zero_1H.ckpt"), map_location=dev).eval()
+    mC = MagNET_Lightning.load_from_checkpoint(os.path.join(CKPT, "MagNET-Zero_13C.ckpt"), map_location=dev).eval()
+    with h5py.File(NP_H5, "r") as f:
+        # pick a large molecule (>=50 atoms) where the parity error is significant
+        big = [k for k in f.keys() if f[k]["atomic_numbers"].shape[0] >= 50][0]
+        g = f[big]
+        z = g["atomic_numbers"][:]; xyz = g["geometry"][:].astype(np.float64)
+    plain = predict_shieldings(mH, mC, solute_atomic_numbers=z, geometry=xyz, device=dev, n_passes=12)
+    sym = predict_shieldings(mH, mC, solute_atomic_numbers=z, geometry=xyz, device=dev, n_passes=12, symmetrize=True)
+    c = z == 6
+    # symmetrization removes a real 13C parity error: it should shift carbons by > 0.05 ppm RMS
+    assert np.sqrt(((plain[c] - sym[c]) ** 2).mean()) > 0.05
+
+
+CKPT_FOUNDATION = os.path.join(_CKPT_ROOT, "MagNET")
+CKPT_PCM = os.path.join(_CKPT_ROOT, "MagNET-PCM")
+
+# a small, valid molecule (methyl acetate) to smoke-run the foundation and PCM checkpoints
+_SMOKE_Z = np.array([6, 6, 6, 8, 8, 6, 1, 1, 1, 1, 1, 1])
+_SMOKE_XYZ = np.array([
+    [-2.07273936, 0.71782219, -0.27169511], [-1.51154256, -0.47834334, -0.05051489],
+    [-0.06133485, -0.72471339, 0.10392854], [0.38785097, -1.84332836, 0.30459866],
+    [0.65684897, 0.41875002, -0.00514369], [2.0648613, 0.23182768, 0.13713977],
+    [-1.48494101, 1.6266042, -0.3569932], [-3.14927745, 0.81001383, -0.37304503],
+    [-2.12480974, -1.36963022, 0.03013895], [2.44585443, -0.42836204, -0.64841408],
+    [2.30088663, -0.16782625, 1.12843716], [2.5483427, 1.20718563, 0.03357037]])
+
+
+@pytest.mark.skipif(not os.path.exists(CKPT_FOUNDATION),
+                    reason="released MagNET (foundation) checkpoints not present")
+def test_real_foundation_loads_and_runs():
+    """The foundation checkpoints load and produce finite per-atom shieldings of the right shape.
+    Regression for the checkpoint-family load/run path (the MagNET-X case bug was this class)."""
+    dev = "cpu"
+    mH = MagNET_Lightning.load_from_checkpoint(os.path.join(CKPT_FOUNDATION, "MagNET_1H.ckpt"), map_location=dev).eval()
+    mC = MagNET_Lightning.load_from_checkpoint(os.path.join(CKPT_FOUNDATION, "MagNET_13C.ckpt"), map_location=dev).eval()
+    pred = predict_shieldings(mH, mC, solute_atomic_numbers=_SMOKE_Z, geometry=_SMOKE_XYZ, device=dev, n_passes=2)
+    assert pred.shape == (_SMOKE_Z.shape[0],)
+    assert np.isfinite(pred[(_SMOKE_Z == 1) | (_SMOKE_Z == 6)]).all()
+
+
+@pytest.mark.skipif(not os.path.exists(CKPT_PCM),
+                    reason="released MagNET-PCM checkpoints not present")
+def test_real_pcm_correction_is_finite_and_nonzero():
+    """The MagNET-PCM gas and with-PCM checkpoints load, run, and genuinely differ: the chloroform
+    correction (with-PCM minus gas) is finite and not identically zero over the H/C atoms."""
+    dev = "cpu"
+    mH = MagNET_Lightning.load_from_checkpoint(os.path.join(CKPT_PCM, "MagNET-PCM-withoutPCM_1H.ckpt"), map_location=dev).eval()
+    mC = MagNET_Lightning.load_from_checkpoint(os.path.join(CKPT_PCM, "MagNET-PCM-withoutPCM_13C.ckpt"), map_location=dev).eval()
+    pH = MagNET_Lightning.load_from_checkpoint(os.path.join(CKPT_PCM, "MagNET-PCM-withPCM_1H.ckpt"), map_location=dev).eval()
+    pC = MagNET_Lightning.load_from_checkpoint(os.path.join(CKPT_PCM, "MagNET-PCM-withPCM_13C.ckpt"), map_location=dev).eval()
+    gas = predict_shieldings(mH, mC, solute_atomic_numbers=_SMOKE_Z, geometry=_SMOKE_XYZ, device=dev, n_passes=4)
+    pcm = predict_shieldings(pH, pC, solute_atomic_numbers=_SMOKE_Z, geometry=_SMOKE_XYZ, device=dev, n_passes=4)
+    correction = pcm - gas
+    hc = (_SMOKE_Z == 1) | (_SMOKE_Z == 6)
+    assert np.isfinite(correction[hc]).all()
+    assert np.abs(correction[hc]).max() > 1e-3   # the two checkpoints are not the same model
+
+
+CKPT_X = os.path.join(_CKPT_ROOT, "MagNET-x")
+SIGMA_FRESH = _dataset_file("sigma-fresh")  # in-repo data/sigma-fresh/, git-ignored, via the Hugging Face checkout
+
+
+@pytest.mark.skipif(not (os.path.exists(CKPT_X) and os.path.exists(SIGMA_FRESH)),
+                    reason="MagNET-x checkpoints or sigma-fresh data not present")
+def test_real_magnet_x_reproduces_dft_correction():
+    """The MagNET-x explicit-solvent correction (solvated minus isolated) should approximate the DFT
+    correction stored in sigma-fresh for a single solute+solvent pose."""
+    import sys
+    sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data", "sigma-fresh"))
+    from decode_sigma_fresh import SigmaFresh
+    dev = "cpu"
+    mH = MagNET_Lightning.load_from_checkpoint(os.path.join(CKPT_X, "MagNET-x-chloroform_1H.ckpt"), map_location=dev).eval()
+    mC = MagNET_Lightning.load_from_checkpoint(os.path.join(CKPT_X, "MagNET-x-chloroform_13C.ckpt"), map_location=dev).eval()
+    with SigmaFresh(SIGMA_FRESH) as ds:
+        p = ds.pose("chloroform", 3, 1)
+    ns = p["n_solute_atoms"]; nper = 5; cut = ns + p["n_solvents_partial"] * nper
+    Zf = p["atomic_numbers"][:cut]; geom = p["coordinates"][:cut]; Zs = p["atomic_numbers"][:ns]
+    ref = (p["shielding_solvated"] - p["shielding_isolated"])[:ns]
+
+    def run(threshold):
+        return np.asarray(predict_shieldings(
+            mH, mC, solute_atomic_numbers=Zs, geometry=geom, atomic_numbers=Zf,
+            N_atoms_per_solvent=nper, solvent_distance_threshold=threshold, device=dev, n_passes=8)).squeeze()
+
+    correction = run(12.0) - run(0.0)
+    keep = ((Zs == 1) | (Zs == 6)) & np.isfinite(ref)
+    rms = np.sqrt(np.mean((correction[keep] - ref[keep]) ** 2))
+    assert rms < 0.3   # measured ~0.085 ppm on this pose; model error, not exact

pdoc_templates/module.html.jinja2

@@ -0,0 +1,9 @@
+{% extends "default/module.html.jinja2" %}
+{% block style %}
+    {{ super() }}
+    <style>
+        /* Widen the content column so wide tables (e.g. the supported-solvent table) are not clipped. */
+        main, header { width: calc(78rem + var(--sidebar-width)); }
+        .pdoc table { width: 100%; }
+    </style>
+{% endblock %}

reproduce.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+"""Regenerate every figure and table in the paper from the committed data.
+
+Runs each analysis notebook headless (no display, no widgets). Every notebook reads the in-repo
+data/<name>/ files and writes its figures to <notebook>/figures/ and its tables to
+<notebook>/documents/. The notebooks are stored without outputs, so this script is how the figures
+and tables are produced. No environment variables and no network are needed.
+
+Usage:
+    python reproduce.py                 # run all notebooks
+    python reproduce.py fig3 s10        # run only notebooks whose path contains one of these strings
+"""
+import argparse
+import glob
+import os
+import sys
+import time
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__,
+                                 formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("filters", nargs="*", help="only run notebooks whose path contains one of these")
+    ap.add_argument("--timeout", type=int, default=900, help="per-cell timeout in seconds (default 900)")
+    args = ap.parse_args()
+
+    try:
+        import nbformat
+        from nbconvert.preprocessors import ExecutePreprocessor
+    except ImportError:
+        sys.exit("reproduce.py needs the analysis dependencies: pip install -r requirements.txt")
+
+    notebooks = sorted(glob.glob(os.path.join(HERE, "analysis", "*", "*.ipynb")))
+    if args.filters:
+        notebooks = [p for p in notebooks if any(f in p for f in args.filters)]
+    if not notebooks:
+        sys.exit("no matching notebooks")
+
+    print("running {} notebooks headless".format(len(notebooks)))
+    n_ok = n_fail = 0
+    for path in notebooks:
+        rel = os.path.relpath(path, HERE)
+        start = time.time()
+        notebook = nbformat.read(path, as_version=4)
+        # run each notebook from its own directory so its in-repo data and figures/ paths resolve
+        executor = ExecutePreprocessor(timeout=args.timeout)
+        try:
+            executor.preprocess(notebook, {"metadata": {"path": os.path.dirname(path)}})
+            print("  OK   ({:5.0f}s) {}".format(time.time() - start, rel))
+            n_ok += 1
+        except Exception as error:
+            print("  FAIL ({:5.0f}s) {}: {}: {}".format(
+                time.time() - start, rel, type(error).__name__, str(error)[:120]))
+            n_fail += 1
+
+    print("\n{} ok, {} failed".format(n_ok, n_fail))
+    sys.exit(1 if n_fail else 0)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,38 @@
+# Dependencies for the dataset readers, the analysis code, and reproducing the figures/tables
+# (reproduce.py). This is NOT the inference stack: `pip install magnet-nmr` and the Option B checkout
+# pull torch from magnet/pyproject.toml and magnet/requirements.txt instead, so nothing here reaches
+# an inference-only install.
+# Loose floors so a future release can't silently break CI; tests also need pytest.
+# numpy + h5py cover the sigma readers; the delta22 reader additionally uses
+# pandas (DataFrame assembly), openpyxl (reads the experimental .xlsx), and tqdm.
+numpy>=1.20
+h5py>=3.0
+pandas>=1.3
+openpyxl>=3.0
+tqdm>=4.0
+# the leveling-effect analysis (analysis/code/leveling.py) and its tests use scikit-learn
+# for the principal-components step. matplotlib + adjustText are needed only to redraw its
+# figures, not to run the tests, so they are not listed here.
+scikit-learn>=1.0
+# the applications analysis (analysis/code/applications.py) fits the composite linear model with
+# statsmodels. Its figure notebooks also draw boxplots with matplotlib + seaborn; those are needed
+# only to redraw the figures (not to run the readers or the tests), and seaborn pulls matplotlib in
+# as a dependency.
+statsmodels>=0.13
+seaborn>=0.12
+# the DFT8K functional-group analysis (analysis/code/dft8k_residuals.py's
+# functional_group_errors/find_extreme_residual, used by fig5b_dft8k.ipynb) parses each molecule's
+# SMILES and matches functional-group substructures with RDKit. Not needed for the rest of the
+# repo's readers or tests.
+rdkit>=2023.9
+
+# Reproducing the figures and tables (reproduce.py) executes the analysis notebooks, which draw with
+# matplotlib (adjustText places the leveling-figure labels) and use scipy directly. matplotlib and
+# scipy already arrive transitively via seaborn and scikit-learn; they are listed explicitly because
+# the notebooks import them. nbconvert + ipykernel are the headless notebook-execution stack
+# reproduce.py drives; they are the only reason a plain reader/test run does not need them.
+matplotlib>=3.5
+scipy>=1.7
+adjustText>=1.0
+nbconvert>=7.0
+ipykernel>=6.0

test_reproduce.py

@@ -0,0 +1,31 @@
+"""Reproduction check: every analysis notebook executes headless.
+
+This follows the repo's real-data test pattern: it SKIPS when the datasets are not present (a fresh
+CI checkout has none of the big `.hdf5` files) and RUNS locally once they are downloaded. So it never
+slows or breaks CI, but a `pytest` on a machine with the data verifies the whole figure/table pipeline.
+The human-facing entrypoint is reproduce.py; this is the pytest wrapper.
+"""
+import glob
+import os
+
+import pytest
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+NOTEBOOKS = sorted(glob.glob(os.path.join(HERE, "analysis", "*", "*.ipynb")))
+
+# Sentinel for "the datasets have been downloaded": delta22 is read by many notebooks. os.path.exists
+# follows symlinks, so a working in-repo symlink into the big-file source counts as present, while a
+# fresh CI checkout (no file, or a dangling link) does not.
+_HAVE_DATA = os.path.exists(os.path.join(HERE, "data", "delta22", "delta22.hdf5"))
+
+
+@pytest.mark.reproduce
+@pytest.mark.skipif(not _HAVE_DATA, reason="datasets not downloaded (hf download ekwan16/MagNET)")
+@pytest.mark.parametrize("notebook", NOTEBOOKS, ids=[os.path.relpath(p, HERE) for p in NOTEBOOKS])
+def test_notebook_reproduces(notebook):
+    nbformat = pytest.importorskip("nbformat")
+    nbconvert = pytest.importorskip("nbconvert.preprocessors")  # requires Jupyter
+    executor = nbconvert.ExecutePreprocessor(timeout=900)
+    nb = nbformat.read(notebook, as_version=4)
+    # run from the notebook's own directory so its in-repo data and figures/ paths resolve
+    executor.preprocess(nb, {"metadata": {"path": os.path.dirname(notebook)}})