sync code from github@e853181b4c59e84169cc94394b0c9bbe8a58b077

conftest.py

@@ -1,11 +1,11 @@
-# The magnet/ model package needs the heavy PyTorch / PyTorch Geometric stack
+# The magnet/ package tests live in tests/ and need 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.
+# install. When torch is absent, skip collecting 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.
+# This is needed (not just the importorskip guards inside the test files) because tests/test_api.py
+# and tests/test_magnet.py import torch (and magnet, which imports torch) at module scope, before
+# any in-test skip can run.
 import os
 import sys
 
@@ -23,8 +23,8 @@ collect_ignore_glob = []
 try:
     import torch  # noqa: F401
 except ImportError:
-    collect_ignore = ["magnet"]
-    collect_ignore_glob = ["magnet/*"]
+    collect_ignore = ["magnet", "tests"]
+    collect_ignore_glob = ["magnet/*", "tests/*"]
 
 
 # The reproduce tests (test_reproduce.py) execute every notebook and take minutes; they are opt-in so

tests/test_api.py

@@ -0,0 +1,485 @@
+"""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 torch
+
+import magnet.api as api
+from magnet import scaling
+from magnet.eqV2.edge_rot_mat import init_edge_rot_mat
+
+_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)
+
+
+# ---- degenerate-geometry guard in the local-frame construction (no weights needed) ----
+def test_edge_rot_mat_rejects_coincident_atoms():
+    # an edge between two atoms at the same position has ~zero length; the old code printed a
+    # warning and then divided by ~zero, silently producing NaN frames that poison the whole
+    # forward pass. It must raise instead.
+    edge_vecs = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
+    with pytest.raises(ValueError, match="overlapping atoms"):
+        init_edge_rot_mat(edge_vecs)
+
+
+def test_edge_rot_mat_accepts_separated_atoms():
+    edge_vecs = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.5, 0.0], [0.0, 0.0, 2.0]])
+    rot = init_edge_rot_mat(edge_vecs)
+    assert rot.shape == (3, 3, 3)
+    # each frame is a proper rotation: R @ R^T == I
+    identity = torch.eye(3).expand(3, 3, 3)
+    assert torch.allclose(torch.bmm(rot, rot.transpose(1, 2)), identity, atol=1e-4)

tests/test_data/magnet_zero_and_pcm_shieldings.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32963207bef34c0ea4246042d9f8d8377d934654c20fdcdd10700643382cfa0c
+size 55384

tests/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