sync code from github@2603923ed5412e943ea2128d491e9c54f907332f

magnet/inference.py

@@ -138,29 +138,35 @@ def _predict_once(model_H, model_C, solute_atomic_numbers, geometry, atomic_numb
     return y_pred_combined
 
 
-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.
+def _predict_batch(model_H, model_C, graphs, device):
+    """Run every graph in `graphs` through both heads, one forward pass per head.
+
+    Each entry of `graphs` is one geometry to predict for, as
+    `(solute_atomic_numbers, geometry, atomic_numbers, N_atoms_per_solvent,
+    solvent_distance_threshold)`. They need not be the same molecule: the graphs are concatenated
+    into a single torch_geometric Batch, and the answers are split apart again by which graph each
+    atom came from, so molecules of different sizes batch together as readily as repeated passes
+    of one.
+
+    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.
+    `unique_molecule_batch` rides along unread: yield_data consumes it for solvent filtering
+    before this point, and no model looks at it, so nothing has to be offset across graphs.
+
+    Returns one (n_solute,) array per entry, in the order given.
     """
-    n = len(geometries)
+    n = len(graphs)
     combined = [None] * n
     for atom_type in ['H', 'C']:
         data_list = []
-        for geom in geometries:
+        for solute_atomic_numbers, geometry, atomic_numbers, per_solvent, threshold in graphs:
             data = yield_data(
                 solute_atomic_numbers = solute_atomic_numbers,
-                geometries = geom,
+                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,
+                N_atoms_per_solvent = per_solvent if per_solvent is not None else 3,
+                solvent_distance_threshold = threshold,
                 atom_type = atom_type,
             )
             # PyG owns `batch` on a Batch, and the model rebuilds `natoms` from it; keeping the
@@ -191,8 +197,37 @@ def _predict_batch(model_H, model_C, solute_atomic_numbers, geometries, atomic_n
     return combined
 
 
+def _check_elements(atomic_numbers):
+    """Refuse a system holding an element MagNET was never trained on.
+
+    MagNET was only ever trained on SUPPORTED_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).
+
+    Raises:
+        ValueError: naming every atomic number that is not supported.
+    """
+    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.")
+
+
+def _geometries_of(geometry, mirror_average, n_passes):
+    """Say which geometries one prediction runs over, the mirror image and the passes included."""
+    geometry = np.asarray(geometry, dtype=float)
+    geometries = [geometry]
+    if mirror_average:
+        reflected = geometry.copy()
+        reflected[..., 0] = -reflected[..., 0]   # mirror across the yz-plane (an improper rotation)
+        geometries.append(reflected)
+    return [geom for geom in geometries for _ in range(n_passes)]
+
+
 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.
+    """Predict 1H/13C shieldings for the solute atoms of one molecule.
 
     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
@@ -206,47 +241,89 @@ def predict_shieldings(model_H, model_C, solute_atomic_numbers, geometry, atomic
     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.
+    one forward pass each. `predict_shieldings_batch` batches whole molecules together as well,
+    and is what to call for more than one.
 
     `symmetrize` is the old name for `mirror_average` and still works, with a DeprecationWarning.
     """
     mirror_average = resolve_mirror_average(mirror_average, symmetrize, "predict_shieldings")
+    return predict_shieldings_batch(
+        model_H, model_C, [solute_atomic_numbers], [geometry],
+        atomic_numbers_list = None if atomic_numbers is None else [atomic_numbers],
+        N_atoms_per_solvent = N_atoms_per_solvent,
+        solvent_distance_threshold = solvent_distance_threshold, device = device,
+        n_passes = n_passes, mirror_average = mirror_average,
+        max_batch_graphs = max_batch_graphs)[0]
 
-    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.")
+def predict_shieldings_batch(model_H, model_C, solute_atomic_numbers_list, geometries_list, atomic_numbers_list = 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 many molecules at once.
 
-    geometry = np.asarray(geometry, dtype=float)
-    geometries = [geometry]
-    if mirror_average:
-        reflected = geometry.copy()
-        reflected[..., 0] = -reflected[..., 0]   # mirror across the yz-plane (an improper rotation)
-        geometries.append(reflected)
+    Every molecule's passes, and its mirror image where `mirror_average` is set, go through as one
+    batch rather than one forward pass each, and molecules share those batches with one another.
+    A 31-atom graph occupies very little of a GPU, so what costs the time is the number of forward
+    passes and not the size of any one of them: batching across molecules is what fills the card.
+
+    The molecules need not be the same size or the same shape. Each answer is split out by which
+    graph its atoms came from, so a list of molecules comes back as a list of per-atom arrays in
+    the order given, exactly as calling `predict_shieldings` on each would have.
+
+    Args:
+        model_H, model_C: the 1H and 13C models to run.
+        solute_atomic_numbers_list: one array of solute atomic numbers per molecule.
+        geometries_list: one (n, 3) coordinate array per molecule, in the same order.
+        atomic_numbers_list: the whole system per molecule where it differs from the solute (the
+            explicit-solvent path), or None where every system is its own solute.
+        N_atoms_per_solvent, solvent_distance_threshold: the explicit-solvent options, applied to
+            every molecule.
+        device: where to run.
+        n_passes: how many forward passes to average per geometry.
+        mirror_average: whether to average over the mirror image as well.
+        symmetrize: the old name for `mirror_average`, which still works and warns.
+        max_batch_graphs: how many graphs go through one forward pass at most. It bounds memory
+            and changes no answer.
 
-    # every pass of every geometry, run together rather than one at a time
-    repeated = [geom for geom in geometries for _ in range(n_passes)]
+    Returns:
+        One (n_solute,) array per molecule, in the order given.
+
+    Raises:
+        ValueError: a molecule holds an element MagNET was not trained on, or the lists differ in
+            length.
+    """
+    mirror_average = resolve_mirror_average(mirror_average, symmetrize, "predict_shieldings_batch")
 
-    preds = []
-    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))
+    if len(solute_atomic_numbers_list) != len(geometries_list):
+        raise ValueError(
+            f"got {len(solute_atomic_numbers_list)} solutes and {len(geometries_list)} geometries; "
+            f"they name the same molecules and must be the same length")
+    if atomic_numbers_list is not None and len(atomic_numbers_list) != len(geometries_list):
+        raise ValueError(
+            f"got {len(atomic_numbers_list)} systems and {len(geometries_list)} geometries; "
+            f"they name the same molecules and must be the same length")
+
+    # every graph to run, and which molecule each one belongs to
+    graphs, molecule_of_graph = [], []
+    for i, (solute, geometry) in enumerate(zip(solute_atomic_numbers_list, geometries_list)):
+        solute = np.asarray(solute)
+        if atomic_numbers_list is None:
+            whole = solute
+        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
+            whole = np.asarray(atomic_numbers_list[i])
+        _check_elements(whole)
+        for geom in _geometries_of(geometry, mirror_average, n_passes):
+            graphs.append((solute, geom, whole, N_atoms_per_solvent, solvent_distance_threshold))
+            molecule_of_graph.append(i)
+
+    per_graph = []
+    for start in range(0, len(graphs), max_batch_graphs):
+        per_graph.extend(_predict_batch(model_H, model_C, graphs[start:start + max_batch_graphs],
+                                        device))
+
+    # average each molecule's own passes, and nothing else's
+    gathered = [[] for _ in geometries_list]
+    for prediction, molecule in zip(per_graph, molecule_of_graph):
+        gathered[molecule].append(prediction)
     # 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))
+    return [np.atleast_1d(np.mean(passes, axis=0)) for passes in gathered]

magnet/pyproject.toml

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

@@ -19,7 +19,8 @@ import numpy as np
 import torch
 
 from magnet.model import MagNET_Lightning
-from magnet.inference import predict_shieldings, resolve_mirror_average
+from magnet.inference import (predict_shieldings, predict_shieldings_batch,
+                              resolve_mirror_average)
 
 MODEL_CHECKPOINTS = {
     # foundation model (predicts the gas-phase shielding the rovibrational/QCD analysis builds on)
@@ -98,13 +99,10 @@ def _predict_with(key_H, key_C, atomic_numbers_list, geometries_list,
         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, mirror_average=mirror_average, **predict_kwargs).squeeze()))
-    return out
+    # one batch across the molecules as well as their passes, rather than a call per molecule
+    return [np.atleast_1d(one.squeeze()) for one in predict_shieldings_batch(
+        model_H, model_C, list(atomic_numbers_list), list(geometries_list),
+        device=device, n_passes=n_passes, mirror_average=mirror_average, **predict_kwargs)]
 
 
 def compute_MagNET_foundation_shieldings(atomic_numbers_list, geometries_list,

tests/test_magnet.py

@@ -345,3 +345,75 @@ def test_one_atom_solute_stays_one_dimensional():
                     [-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,)
+
+
+# =====================================================================
+# batching across molecules
+# =====================================================================
+
+def test_batch_refuses_mismatched_lists():
+    """Two lists naming the same molecules have to be the same length."""
+    from magnet.inference import predict_shieldings_batch
+    z = np.array([6, 1, 1, 1, 1])
+    xyz = np.zeros((5, 3))
+    with pytest.raises(ValueError, match="same length"):
+        predict_shieldings_batch(None, None, [z, z], [xyz])
+    with pytest.raises(ValueError, match="same length"):
+        predict_shieldings_batch(None, None, [z], [xyz], atomic_numbers_list=[z, z])
+
+
+def test_batch_refuses_unsupported_element():
+    """A molecule out of vocabulary is refused before any model runs, as the single path does."""
+    from magnet.inference import predict_shieldings_batch
+    good = np.array([6, 1, 1, 1, 1])
+    bad = np.array([15, 1, 1, 1])          # phosphorus
+    with pytest.raises(ValueError, match="unsupported atomic numbers"):
+        predict_shieldings_batch(None, None, [good, bad], [np.zeros((5, 3)), np.zeros((4, 3))])
+
+
+@pytest.mark.skipif(not os.path.exists(CKPT), reason="released MagNET-Zero checkpoints not present")
+def test_batch_matches_one_at_a_time():
+    """Molecules batched together answer as they do apart, whatever their sizes.
+
+    The molecules are deliberately of two different atom counts, since what splits the answers
+    apart is which graph each atom came from and not any assumption that they match.
+    """
+    from magnet.inference import predict_shieldings_batch
+    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()
+    methane_z = np.array([6, 1, 1, 1, 1])
+    methane_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]])
+    ethanol_z = np.array([6, 6, 8, 1, 1, 1, 1, 1, 1])
+    ethanol_xyz = np.array([[1.17, -0.24, 0.], [0., 0.55, 0.], [-1.16, -0.25, 0.],
+                            [2.08, 0.37, 0.], [1.17, -0.87, 0.89], [1.17, -0.87, -0.89],
+                            [0.01, 1.19, 0.88], [0.01, 1.19, -0.88], [-1.90, 0.35, 0.]])
+    zs, xyzs = [methane_z, ethanol_z], [methane_xyz, ethanol_xyz]
+    kw = dict(device=dev, n_passes=30, mirror_average=True)
+    apart = [predict_shieldings(mH, mC, solute_atomic_numbers=z, geometry=x, **kw)
+             for z, x in zip(zs, xyzs)]
+    together = predict_shieldings_batch(mH, mC, zs, xyzs, **kw)
+    assert [one.shape for one in together] == [one.shape for one in apart]
+    for a, b in zip(apart, together):
+        # both average the same frame noise away, so they agree to well inside it
+        assert np.max(np.abs(a - b)) < 0.05
+
+
+@pytest.mark.skipif(not os.path.exists(CKPT), reason="released MagNET-Zero checkpoints not present")
+def test_batch_chunking_does_not_change_the_answer():
+    """max_batch_graphs splits the work and nothing else, across molecules as within one."""
+    from magnet.inference import predict_shieldings_batch
+    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(device=dev, n_passes=20, mirror_average=True)
+    whole = predict_shieldings_batch(mH, mC, [z, z, z], [xyz, xyz + 0.01, xyz - 0.01],
+                                     max_batch_graphs=1000, **kw)
+    chunked = predict_shieldings_batch(mH, mC, [z, z, z], [xyz, xyz + 0.01, xyz - 0.01],
+                                       max_batch_graphs=3, **kw)
+    for a, b in zip(whole, chunked):
+        assert np.max(np.abs(a - b)) < 0.05