sync code from github@fbebb603b47d270f0d13b4146970b99ea374ddd7

README.md

@@ -130,7 +130,7 @@ The predictions are close to experiment:
 
 **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`).
+- **Passes and mirror averaging.** `n_passes` (default 10) averages out equivariance error over the specified number of forward passes. If `mirror_average` (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`). The passes are run as a single batch per head rather than one forward pass each, so raising `n_passes` costs far less than it used to; `max_batch_graphs` caps how many run at once and does not change the answer. `mirror_average` was called `symmetrize` before, and that name still works but warns: the option averages over the mirror image, not over symmetry-equivalent nuclei.
 - **Geometry.** MagNET-Zero and MagNET-PCM require [AIMNet2](https://github.com/isayevlab/AIMNet2)-optimized geometries. Do not use other geometries.
 - **Components.** Pass `return_components=True` to also get the MagNET-Zero shielding, the MagNET-PCM correction, and the scaling coefficients behind each shift, as a `dict`.
 - **Other models.** For raw shieldings and solvent corrections, see the [API documentation](#api-documentation).

magnet/api.py

@@ -13,7 +13,8 @@ Four entry points, each taking one molecule (`atomic_numbers` shaped `(N,)` and
 **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
+- `mirror_average` (default `True`): if True, average over `n_passes` on the input geometry and `n_passes` on the mirror image of the input geometry.
+  (Formerly `symmetrize`, which still works but warns: it averages over the mirror image, not over symmetry-equivalent nuclei.)
 - `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
 """
@@ -24,6 +25,7 @@ from collections import Counter
 import numpy as np
 
 from . import run_magnet
+from .inference import resolve_mirror_average
 from . import scaling
 
 _MODELS = {
@@ -67,7 +69,7 @@ def _as_batch(atomic_numbers, coordinates):
     return an_list, xyz_list, single
 
 
-def predict_shifts(atomic_numbers, coordinates, solvent="chloroform", n_passes=10, symmetrize=True,
+def predict_shifts(atomic_numbers, coordinates, solvent="chloroform", n_passes=10, mirror_average=True, symmetrize=None,
                    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.
 
@@ -82,7 +84,7 @@ def predict_shifts(atomic_numbers, coordinates, solvent="chloroform", n_passes=1
             `"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`),
+        mirror_average: 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.
@@ -104,6 +106,7 @@ def predict_shifts(atomic_numbers, coordinates, solvent="chloroform", n_passes=1
           `{"H": {"intercept": ..., "stationary": ..., "pcm": ...}, "C": {...}}`, so that
           `shift = intercept + stationary * zero_shielding + pcm * pcm_correction`.
     """
+    mirror_average = resolve_mirror_average(mirror_average, symmetrize, "predict_shifts")
     _check_passes(n_passes)
     tables = scaling.published_scaling_tables()
     key = "TIP4P" if solvent == "water" else solvent
@@ -113,10 +116,10 @@ def predict_shifts(atomic_numbers, coordinates, solvent="chloroform", n_passes=1
 
     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,
+                                                     mirror_average=mirror_average, 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,
+                                                    mirror_average=mirror_average, device=device,
                                                     checkpoints_dir=checkpoints_dir)
     results = []
     for atoms, sigma, delta in zip(an_list, zero, pcm):
@@ -133,7 +136,7 @@ def predict_shifts(atomic_numbers, coordinates, solvent="chloroform", n_passes=1
     return results[0] if single else results
 
 
-def predict_shieldings(atomic_numbers, coordinates, model="MagNET", n_passes=10, symmetrize=True,
+def predict_shieldings(atomic_numbers, coordinates, model="MagNET", n_passes=10, mirror_average=True, symmetrize=None,
                        device=None, checkpoints_dir=None):
     """Predict gas-phase NMR shielding constants for a molecule.
 
@@ -147,7 +150,7 @@ def predict_shieldings(atomic_numbers, coordinates, model="MagNET", n_passes=10,
         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`),
+        mirror_average: 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.
@@ -161,6 +164,7 @@ def predict_shieldings(atomic_numbers, coordinates, model="MagNET", n_passes=10,
         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.
     """
+    mirror_average = resolve_mirror_average(mirror_average, symmetrize, "predict_shieldings")
     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]}(...), "
@@ -169,12 +173,12 @@ def predict_shieldings(atomic_numbers, coordinates, model="MagNET", n_passes=10,
         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,
+    out = _MODELS[model](an_list, xyz_list, n_passes=n_passes, mirror_average=mirror_average, 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,
+def implicit_solvent_correction(atomic_numbers, coordinates, n_passes=10, mirror_average=True, symmetrize=None, device=None,
                                 checkpoints_dir=None):
     """Predict how a solvent changes a molecule's shieldings, using a fast continuum-solvent model.
 
@@ -188,7 +192,7 @@ def implicit_solvent_correction(atomic_numbers, coordinates, n_passes=10, symmet
         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`),
+        mirror_average: 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.
@@ -201,10 +205,11 @@ def implicit_solvent_correction(atomic_numbers, coordinates, n_passes=10, symmet
         atoms; atoms of other elements come back as `0.0`. Pass a list of molecules and you get a list
         of arrays back.
     """
+    mirror_average = resolve_mirror_average(mirror_average, symmetrize, "implicit_solvent_correction")
     _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,
+                                                    mirror_average=mirror_average, device=device,
                                                     checkpoints_dir=checkpoints_dir)
     return out[0] if single else out
 
@@ -258,7 +263,7 @@ def _explicit_one(atoms, xyz, solute_atoms, solvent):
 
 
 def explicit_solvent_correction(atomic_numbers, coordinates, solute_atoms, solvent="chloroform",
-                                n_passes=10, symmetrize=True, device=None,
+                                n_passes=10, mirror_average=True, symmetrize=None, 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.
 
@@ -276,7 +281,7 @@ def explicit_solvent_correction(atomic_numbers, coordinates, solute_atoms, solve
         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`),
+        mirror_average: 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.
@@ -291,6 +296,7 @@ def explicit_solvent_correction(atomic_numbers, coordinates, solute_atoms, solve
         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.
     """
+    mirror_average = resolve_mirror_average(mirror_average, symmetrize, "explicit_solvent_correction")
     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}")
@@ -304,6 +310,6 @@ def explicit_solvent_correction(atomic_numbers, coordinates, solute_atoms, solve
         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,
+        n_passes=n_passes, mirror_average=mirror_average, device=device,
         solvent_distance_threshold=solvent_distance_threshold, checkpoints_dir=checkpoints_dir)
     return out[0] if single else out

magnet/inference.py

@@ -1,3 +1,5 @@
+import warnings
+
 import numpy as np
 import torch
 import torch_geometric
@@ -8,6 +10,36 @@ import torch_scatter
 # other element are not meaningful, so predict_shieldings refuses them.
 SUPPORTED_ELEMENTS = frozenset({1, 6, 7, 8, 9, 16, 17})
 
+# How many graphs go through one forward pass at most. The passes of one prediction are batched
+# together, and a large solvated system at a high n_passes would otherwise have to fit in memory
+# all at once. Chunking changes nothing about the answer.
+MAX_BATCH_GRAPHS = 64
+
+
+def resolve_mirror_average(mirror_average, symmetrize, caller):
+    """Take whichever of the two names the caller used, and warn if it was the old one.
+
+    `symmetrize` was renamed to `mirror_average`: the option averages a prediction over the
+    molecule and its mirror image, where the old name read as though it averaged
+    symmetry-equivalent nuclei. The old name still works.
+
+    Args:
+        mirror_average: what the caller passed under the new name.
+        symmetrize: what the caller passed under the old name, or None if it did not.
+        caller: the function name to name in the warning.
+
+    Returns:
+        The setting to use.
+    """
+    if symmetrize is None:
+        return mirror_average
+    warnings.warn(
+        f"{caller}(symmetrize=...) is deprecated; use mirror_average=... instead. The option "
+        f"averages the prediction over the molecule and its mirror image, which 'symmetrize' "
+        f"read as though it meant averaging symmetry-equivalent nuclei.",
+        DeprecationWarning, stacklevel = 3)
+    return symmetrize
+
 
 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(
@@ -106,7 +138,60 @@ def _predict_once(model_H, model_C, solute_atomic_numbers, geometry, atomic_numb
     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):
+def _predict_batch(model_H, model_C, solute_atomic_numbers, geometries, atomic_numbers,
+                   N_atoms_per_solvent, solvent_distance_threshold, device):
+    """Run every geometry in `geometries` through both heads, one forward pass per head.
+
+    Each geometry becomes one graph, and the graphs are concatenated into a single
+    torch_geometric Batch, so a list of geometries costs two forward passes rather than two per
+    geometry. The model recomputes `natoms` from `batch.batch`, so the per-graph `batch` and
+    `natoms` that yield_data writes are dropped before concatenating and PyG assigns its own.
+
+    Returns one (n_solute,) array per geometry, in the order given.
+    """
+    n = len(geometries)
+    combined = [None] * n
+    for atom_type in ['H', 'C']:
+        data_list = []
+        for geom in geometries:
+            data = yield_data(
+                solute_atomic_numbers = solute_atomic_numbers,
+                geometries = geom,
+                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,
+            )
+            # PyG owns `batch` on a Batch, and the model rebuilds `natoms` from it; keeping the
+            # single-graph versions here would have them concatenated as ordinary attributes.
+            del data.batch, data.natoms
+            data_list.append(data)
+        batch = torch_geometric.data.Batch.from_data_list(data_list).to(device)
+
+        with torch.no_grad():
+            y_pred = (model_H if atom_type == 'H' else model_C).forward(batch).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 = batch.atom_type_mask.cpu()
+        solute = batch.solute.cpu()
+        graph_of_atom = batch.batch.cpu()
+        y_pred[~atom_type_mask] = 0.
+        keep = solute == 1
+        y_solute = y_pred[keep]
+        graph_of_solute = graph_of_atom[keep]
+
+        for i in range(n):
+            y = np.squeeze(y_solute[graph_of_solute == i].numpy())
+            # combine H and C shieldings into one array, exactly as the unbatched path does
+            if atom_type == 'H':
+                combined[i] = y
+            elif atom_type == 'C':
+                combined[i][y != 0.0] = y[y != 0.0]
+    return 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, mirror_average = False, symmetrize = None, max_batch_graphs = MAX_BATCH_GRAPHS):
     """Predict 1H/13C shieldings for the solute atoms.
 
     A single forward pass is NOT deterministic: each edge picks a random local
@@ -114,11 +199,22 @@ def predict_shieldings(model_H, model_C, solute_atomic_numbers, geometry, atomic
     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.
+    mirror_average=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.
+    ~0.3 ppm on 13C for large molecules; averaging over the mirror image removes that
+    spurious error.
+
+    The passes are independent of one another, so they are run as one batch per head rather than
+    one forward pass each: n_passes with mirror_average costs 2*n_passes graphs and two forward
+    passes, not 4*n_passes. `max_batch_graphs` caps how many graphs go through at once, so a large
+    solvated system at a high n_passes does not have to fit in memory all at once. The answer does
+    not depend on it.
+
+    `symmetrize` is the old name for `mirror_average` and still works, with a DeprecationWarning.
     """
+    mirror_average = resolve_mirror_average(mirror_average, symmetrize, "predict_shieldings")
+
     solute_atomic_numbers = np.asarray(solute_atomic_numbers)
     if atomic_numbers is None:
         atomic_numbers = solute_atomic_numbers
@@ -139,16 +235,18 @@ def predict_shieldings(model_H, model_C, solute_atomic_numbers, geometry, atomic
 
     geometry = np.asarray(geometry, dtype=float)
     geometries = [geometry]
-    if symmetrize:
+    if mirror_average:
         reflected = geometry.copy()
         reflected[..., 0] = -reflected[..., 0]   # mirror across the yz-plane (an improper rotation)
         geometries.append(reflected)
 
+    # every pass of every geometry, run together rather than one at a time
+    repeated = [geom for geom in geometries for _ in range(n_passes)]
+
     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)))
+    for start in range(0, len(repeated), max_batch_graphs):
+        preds.extend(_predict_batch(model_H, model_C, solute_atomic_numbers,
+                                    repeated[start:start + max_batch_graphs], 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/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "magnet-nmr"
-version = "0.1.0"
+version = "0.2.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"

magnet/run_magnet.py

@@ -5,10 +5,10 @@ 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
+  - `mirror_average`: 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
+    The released reference shieldings are NOT mirror-averaged, 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.
@@ -19,7 +19,7 @@ import numpy as np
 import torch
 
 from magnet.model import MagNET_Lightning
-from magnet.inference import predict_shieldings
+from magnet.inference import predict_shieldings, resolve_mirror_average
 
 MODEL_CHECKPOINTS = {
     # foundation model (predicts the gas-phase shielding the rovibrational/QCD analysis builds on)
@@ -89,10 +89,11 @@ def load_model_to_device(path, device=None, checkpoints_dir=None):
 
 
 def _predict_with(key_H, key_C, atomic_numbers_list, geometries_list,
-                  n_passes=20, symmetrize=False, device=None, checkpoints_dir=None, **predict_kwargs):
+                  n_passes=20, mirror_average=False, symmetrize=None, 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)."""
+    mirror_average = resolve_mirror_average(mirror_average, symmetrize, "_predict_with")
     if device is None:
         device = _default_device()
     model_H = load_model_to_device(MODEL_CHECKPOINTS[key_H], device, checkpoints_dir=checkpoints_dir)
@@ -102,37 +103,39 @@ def _predict_with(key_H, key_C, 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()))
+            device=device, n_passes=n_passes, mirror_average=mirror_average, **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):
+                                         n_passes=20, mirror_average=False, symmetrize=None, device=None, checkpoints_dir=None):
     """MagNET foundation-model gas-phase shieldings (1H and 13C) for AIMNet2 geometries."""
+    mirror_average = resolve_mirror_average(mirror_average, symmetrize, "compute_MagNET_foundation_shieldings")
     return _predict_with("MagNET_H", "MagNET_C", atomic_numbers_list, geometries_list,
-                         n_passes=n_passes, symmetrize=symmetrize, device=device,
+                         n_passes=n_passes, mirror_average=mirror_average, 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):
+                                   n_passes=20, mirror_average=False, symmetrize=None, 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).
+        n_passes, mirror_average: 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).
     """
+    mirror_average = resolve_mirror_average(mirror_average, symmetrize, "compute_MagNET_Zero_shieldings")
     return _predict_with("MagNET-Zero_H", "MagNET-Zero_C", atomic_numbers_list, geometries_list,
-                         n_passes=n_passes, symmetrize=symmetrize, device=device,
+                         n_passes=n_passes, mirror_average=mirror_average, 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):
+                                   n_passes=20, mirror_average=False, symmetrize=None, 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,
@@ -140,13 +143,14 @@ def compute_MagNET_PCM_corrections(atomic_numbers_list, geometries_list,
 
     Returns a list of (N,) correction arrays.
     """
+    mirror_average = resolve_mirror_average(mirror_average, symmetrize, "compute_MagNET_PCM_corrections")
     gas = _predict_with("MagNET-PCM-withoutPCM_H", "MagNET-PCM-withoutPCM_C",
                         atomic_numbers_list, geometries_list,
-                        n_passes=n_passes, symmetrize=symmetrize, device=device,
+                        n_passes=n_passes, mirror_average=mirror_average, 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,
+                        n_passes=n_passes, mirror_average=mirror_average, device=device,
                         checkpoints_dir=checkpoints_dir)
     return [pcm[i] - gas[i] for i in range(len(pcm))]
 
@@ -156,7 +160,7 @@ 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,
+                                 geometries_list, n_passes=20, mirror_average=False, symmetrize=None, device=None,
                                  solvent_distance_threshold=12.0, checkpoints_dir=None):
     """MagNET-x explicit-solvent corrections (solvated minus isolated) for one solvent.
 
@@ -167,7 +171,7 @@ def compute_MagNET_x_corrections(solvent, solute_atomic_numbers_list, atomic_num
             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).
+        n_passes, mirror_average: 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.
@@ -176,6 +180,7 @@ def compute_MagNET_x_corrections(solvent, solute_atomic_numbers_list, atomic_num
     Returns a list of (n_solute,) correction arrays. The published pipeline averages this over MD
     frames; a single-frame value is only an estimate.
     """
+    mirror_average = resolve_mirror_average(mirror_average, symmetrize, "compute_MagNET_x_corrections")
     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]
@@ -187,7 +192,7 @@ def compute_MagNET_x_corrections(solvent, solute_atomic_numbers_list, atomic_num
     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)
+                  n_passes=n_passes, mirror_average=mirror_average)
         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)

magnet/scaling.py

@@ -20,7 +20,7 @@ import io
 import numpy as np
 
 # Published SI Tables S10 (proton) and S11 (carbon), verbatim (reflection-symmetrized: n_passes=10,
-# symmetrize=True). Columns: intercept, stationary, pcm. Kept byte-for-byte in sync with
+# mirror_average=True, the option this code called symmetrize when the tables were made). Columns: intercept, stationary, pcm. Kept byte-for-byte in sync with
 # data/scaling_factors/scaling_factors_symmetrized_{H,C}.csv by test_scaling_factors.py.
 _PUBLISHED_TABLE_CSV = {
     "H": """solvent,intercept,stationary,pcm

tests/test_magnet.py

@@ -260,3 +260,88 @@ def test_real_magnet_x_reproduces_dft_correction():
     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
+
+
+# =====================================================================
+# mirror_average (formerly symmetrize), and the batching behind the passes
+# =====================================================================
+
+def test_symmetrize_is_deprecated_but_still_works():
+    """The old name keeps working and says so, since magnet-nmr is published under it."""
+    import warnings
+    from magnet.inference import resolve_mirror_average
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        assert resolve_mirror_average(False, True, "predict_shieldings") is True
+    assert len(caught) == 1
+    assert issubclass(caught[0].category, DeprecationWarning)
+    assert "mirror_average" in str(caught[0].message)
+
+
+def test_mirror_average_passes_without_warning():
+    """The new name is the quiet one."""
+    import warnings
+    from magnet.inference import resolve_mirror_average
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        assert resolve_mirror_average(True, None, "predict_shieldings") is True
+        assert resolve_mirror_average(False, None, "predict_shieldings") is False
+    assert caught == []
+
+
+def test_symmetrize_wins_when_both_are_given():
+    """A caller still passing the old name gets what it asked for, whatever the new one defaults to."""
+    import warnings
+    from magnet.inference import resolve_mirror_average
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        assert resolve_mirror_average(True, False, "predict_shieldings") is False
+
+
+@pytest.mark.skipif(not os.path.exists(CKPT), reason="released MagNET-Zero checkpoints not present")
+def test_mirror_average_matches_symmetrize():
+    """The rename is a rename: both names give the same numbers."""
+    import warnings
+    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()
+    z = np.array([6, 1, 1, 1, 1])
+    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]])
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", DeprecationWarning)
+        old = predict_shieldings(mH, mC, solute_atomic_numbers=z, geometry=xyz, device=dev,
+                                 n_passes=60, symmetrize=True)
+    new = predict_shieldings(mH, mC, solute_atomic_numbers=z, geometry=xyz, device=dev,
+                             n_passes=60, mirror_average=True)
+    # both average the same frame noise away, so they agree to well inside it
+    assert np.max(np.abs(old - new)) < 0.05
+
+
+@pytest.mark.skipif(not os.path.exists(CKPT), reason="released MagNET-Zero checkpoints not present")
+def test_chunking_does_not_change_the_answer():
+    """max_batch_graphs bounds memory and nothing else: the passes are averaged either way."""
+    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()
+    z = np.array([6, 1, 1, 1, 1])
+    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]])
+    kw = dict(solute_atomic_numbers=z, geometry=xyz, device=dev, n_passes=40, mirror_average=True)
+    whole = predict_shieldings(mH, mC, max_batch_graphs=1000, **kw)
+    chunked = predict_shieldings(mH, mC, max_batch_graphs=3, **kw)
+    # same passes, same average, only the number of forward calls differs
+    assert np.max(np.abs(whole - chunked)) < 0.05
+
+
+@pytest.mark.skipif(not os.path.exists(CKPT), reason="released MagNET-Zero checkpoints not present")
+def test_one_atom_solute_stays_one_dimensional():
+    """The per-graph squeeze must not turn a one-atom solute into a scalar."""
+    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()
+    z = np.array([6, 1, 1, 1, 1])
+    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]])
+    out = predict_shieldings(mH, mC, solute_atomic_numbers=z, geometry=xyz, device=dev, n_passes=2)
+    assert out.ndim == 1 and out.shape == (5,)