feat(neighbors): support global mode-2 sparse batches

CHANGELOG.md

@@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 ### Changed
 
+- Added canonical global packed mode 2 for batched sparse neighbor matrices. Callers with 3D local-index matrices must convert them with `aimnet.nbops.convert_mode2_local_to_global`; every system must reserve a final dummy atom, and full-3D periodic inputs must provide aligned lattice shifts.
 - Migrated Ewald Coulomb to the nvalchemiops (>=0.4) energy+autograd API: energies are bit-identical and forces/stress unchanged to kernel rounding (<=2e-06 eV/A), the upstream `DeprecationWarning` for Ewald is gone, and Ewald dense Hessians / HVPs are now true relaxed-charge autograd derivatives (comparable with simple/DSF) instead of fixed-charge finite differences. The `eps` argument of `hessian_vector_product` now applies to PME only. PME intentionally remains on the legacy explicit-terms path: its charge-gradient backward misbehaves inside the full calculator graph under `create_graph=True` (train mode), tracked in a dedicated issue; PME behavior is bit-identical to the previous release.
 - Hardened `make test`: the parallel run now hides CUDA and caps per-worker threads (`CUDA_VISIBLE_DEVICES="" OMP_NUM_THREADS=1`); previously xdist workers either all initialized the first GPU (OOM on CUDA boxes) or oversubscribed the CPU with per-worker torch thread pools. A new `make test-gpu` target runs the GPU-marked tests serially on CUDA.
 - Bumped `nvalchemi-toolkit-ops` to `>=0.4.0` and `warp-lang` to `>=1.13,<2` (installs 1.15). Energies, charges, and Hessians are bit-identical to 0.3.1; explicit force/virial outputs shift within float32 accumulation noise (max 6.7e-05 eV/A on a periodic system, 40x inside the project's cross-version acceptance of 1e-4 Hartree/A). The 0.4.0 direct-output flags used by the Ewald/PME path (`compute_forces`/`compute_virial`) are deprecated upstream and now emit `DeprecationWarning`; migrating to the autograd-based API is tracked as follow-up work.

aimnet/calculators/calculator.py

@@ -9,6 +9,7 @@
 import torch
 from torch import Tensor, nn
 
+from aimnet import nbops
 from aimnet.modules import DFTD3, LRCoulomb
 from aimnet.modules.lr import ExternalDerivativeTerms
 
@@ -112,12 +113,18 @@ class AIMNet2Calculator:
         "mol_idx": torch.int,
         "nbmat": torch.int,
         "nbmat_lr": torch.int,
+        "nbmat_coulomb": torch.int,
+        "nbmat_dftd3": torch.int,
         "nb_pad_mask": torch.bool,
         "nb_pad_mask_lr": torch.bool,
         "shifts": torch.float,
         "shifts_lr": torch.float,
+        "shifts_coulomb": torch.float,
+        "shifts_dftd3": torch.float,
         "cell": torch.float,
         "pbc": torch.bool,
+        "cutoff_coulomb": torch.float,
+        "cutoff_dftd3": torch.float,
     }
     keys_out: ClassVar[list[str]] = ["energy", "charges", "spin_charges", "forces", "hessian", "stress"]
     atom_feature_keys: ClassVar[list[str]] = ["coord", "numbers", "charges", "spin_charges", "forces"]
@@ -846,6 +853,14 @@ def eval(
             )
 
         if hessian:
+            probe = self.to_input_tensors(data)
+            primary_nbmat = probe.get("nbmat")
+            nbops.validate_neighbor_suffix_layout(probe)
+            if primary_nbmat is not None and primary_nbmat.ndim == 3:
+                nbops.normalize_mode2_periodic_geometry(probe, B=primary_nbmat.shape[0])
+                for suffix in nbops.NBMAT_SUFFIXES:
+                    if f"nbmat{suffix}" in probe or f"shifts{suffix}" in probe:
+                        nbops.validate_mode2_nbmat_raw(probe, suffix=suffix)
             subsystems = self._split_hessian_batch(data)
             if subsystems is not None:
                 stack = torch.as_tensor(data["coord"]).ndim == 3
@@ -984,6 +999,13 @@ def prepare_input(self, data: dict[str, Any], *, hessian: bool = False) -> dict[
         caller_had_mol_idx = raw_data.get("mol_idx") is not None
         caller_had_nbmat = raw_data.get("nbmat") is not None
         data = self.to_input_tensors(data)
+        primary_nbmat = data.get("nbmat")
+        nbops.validate_neighbor_suffix_layout(data)
+        if primary_nbmat is not None and primary_nbmat.ndim == 3:
+            nbops.normalize_mode2_periodic_geometry(data, B=primary_nbmat.shape[0])
+            for suffix in nbops.NBMAT_SUFFIXES:
+                if f"nbmat{suffix}" in data or f"shifts{suffix}" in data:
+                    nbops.validate_mode2_nbmat_raw(data, suffix=suffix)
         data = self.mol_flatten(data, hessian=hessian)
         if data.get("cell") is not None and self._coulomb_method == "simple":
             warnings.warn(
@@ -1219,22 +1241,36 @@ def _select_for_structure(value: Any, index: int, n_struct: int) -> Any:
     def _split_batch_dim(self, data: dict[str, Any], B: int) -> list[dict[str, Any]]:
         """Slice a (B, ...) batched input into B single-structure dicts."""
         subs: list[dict[str, Any]] = []
+        primary_nbmat = data.get("nbmat")
+        mode2 = primary_nbmat is not None and torch.as_tensor(primary_nbmat).ndim == 3
         for b in range(B):
             sub: dict[str, Any] = {}
             for k, v in data.items():
                 if v is None:
                     continue
                 if k in ("coord", "numbers"):
-                    sub[k] = torch.as_tensor(v)[b]
+                    value = torch.as_tensor(v)
+                    sub[k] = value[b : b + 1] if mode2 else value[b]
                 elif k in ("charge", "mult"):
                     sub[k] = self._select_for_structure(v, b, B)
                 elif k == "cell":
                     t = torch.as_tensor(v)
-                    sub[k] = t[b] if t.ndim == 3 else t
-                elif k.startswith(("nbmat", "shifts")) or k == "mol_idx":
-                    # Precomputed neighbor-list keys must not be shared unsliced
-                    # across subsystems (they'd be wrong per-structure); the
-                    # recursive eval rebuilds them.
+                    sub[k] = t[b : b + 1] if mode2 and t.ndim == 3 else (t[b] if t.ndim == 3 else t)
+                elif k == "pbc":
+                    t = torch.as_tensor(v)
+                    sub[k] = t[b : b + 1] if mode2 and t.ndim == 2 else (t[b] if t.ndim == 2 else t)
+                elif k.startswith("nbmat"):
+                    t = torch.as_tensor(v)[b : b + 1]
+                    if t.ndim == 3:
+                        sentinel = B * t.shape[1]
+                        singleton_sentinel = t.shape[1]
+                        usable = t != sentinel
+                        t = torch.where(usable, t - b * t.shape[1], torch.full_like(t, singleton_sentinel))
+                    sub[k] = t
+                elif k.startswith("shifts"):
+                    t = torch.as_tensor(v)
+                    sub[k] = t[b : b + 1] if t.ndim >= 3 else t
+                elif k == "mol_idx":
                     continue
                 else:
                     sub[k] = v
@@ -1404,11 +1440,22 @@ def to_input_tensors(self, data: dict[str, Any]) -> dict[str, Tensor]:
             if not (isinstance(data[k], Tensor) and data[k].requires_grad):
                 t = t.detach()
             ret[k] = t
+        neighbor_memo: list[tuple[Any, Tensor]] = []
+        neighbor_keys = {f"nbmat{suffix}" for suffix in nbops.NBMAT_SUFFIXES}
         for k in self.keys_in_optional:
             if k in data and data[k] is not None:
-                t = torch.as_tensor(data[k], device=self.device, dtype=self.keys_in_optional[k])
-                if not (isinstance(data[k], Tensor) and data[k].requires_grad):
-                    t = t.detach()
+                if k in neighbor_keys:
+                    source = data[k]
+                    t = next((value for previous, value in neighbor_memo if source is previous), None)
+                    if t is None:
+                        t = torch.as_tensor(source, device=self.device)
+                        if not (isinstance(source, Tensor) and source.requires_grad):
+                            t = t.detach()
+                        neighbor_memo.append((source, t))
+                else:
+                    t = torch.as_tensor(data[k], device=self.device, dtype=self.keys_in_optional[k])
+                    if not (isinstance(data[k], Tensor) and data[k].requires_grad):
+                        t = t.detach()
                 ret[k] = t
         # Ensure all tensors have at least 1D shape for consistent batch processing
         for k, v in ret.items():
@@ -1421,6 +1468,11 @@ def mol_flatten(self, data: dict[str, Tensor], *, hessian: bool = False) -> dict
         Will not flatten for batched input and molecule size below threshold.
         """
         ndim = data["coord"].ndim
+        explicit_mode2 = data.get("nbmat") is not None and data["nbmat"].ndim == 3
+        if explicit_mode2:
+            self._batch = None
+            self._max_mol_size = data["coord"].shape[1]
+            return data
         if ndim == 2:
             self._batch = None
             if "mol_idx" not in data:
@@ -1884,28 +1936,29 @@ def _hessian_vector_product_impl(
             prepared, forces=False, stress=False, hessian=external_hessian
         )
 
-        coord = self._saved_for_grad["coord"]  # (N+1, 3), requires_grad
+        coord = self._saved_for_grad["coord"]
+        coord_flat = coord.reshape(-1, 3)
+        real_indices = (~prepared["mask_i"].reshape(-1)).nonzero(as_tuple=False).flatten()
+        n_real = real_indices.numel()
         tot_energy = prepared["energy"].sum()
         # Differentiable part of the forces only. The detached ``coulomb_terms``
         # forces are a constant w.r.t. coord (zero second derivative), so they are
         # intentionally excluded from the vjp; the periodic curvature they would
         # have carried is supplied by the directional FD helper instead.
-        forces_diff = -torch.autograd.grad(tot_energy, coord, create_graph=True)[0]  # (N+1, 3)
-        N = coord.shape[0] - 1
+        forces_diff = -torch.autograd.grad(tot_energy, coord, create_graph=True)[0]
 
         device = coord.device
         vecs = torch.as_tensor(vectors, device=device)
         single = vecs.ndim == 2
         if single:
             vecs = vecs.unsqueeze(0)
-        if vecs.shape[-2:] != (N, 3):
-            raise ValueError(f"vectors must have trailing shape ({N}, 3); got {tuple(vecs.shape)}")
+        if vecs.shape[-2:] != (n_real, 3):
+            raise ValueError(f"vectors must have trailing shape ({n_real}, 3); got {tuple(vecs.shape)}")
 
         outs = []
         for k in range(vecs.shape[0]):
             v = vecs[k].to(forces_diff.dtype)
-            v_full = torch.zeros_like(coord)
-            v_full[:N] = v
+            v_full = torch.zeros_like(coord_flat).index_copy(0, real_indices, v).reshape_as(coord)
             # autograd Hv = -d(forces . v)/dcoord = d^2E/dr^2 . v
             # (NN + short-range + dsf/simple charge-response + dftd3)
             hv_full = -torch.autograd.grad(
@@ -1916,7 +1969,7 @@ def _hessian_vector_product_impl(
                 create_graph=create_graph,
                 allow_unused=True,
             )[0]
-            hv = hv_full[:N]
+            hv = hv_full.reshape(-1, 3).index_select(0, real_indices)
             if method == "pme" and self.external_coulomb is not None:
                 # Full-periodic fixed-position curvature (directional FD).
                 # Ewald needs no FD block: its energy is in the autograd graph

aimnet/calculators/derivatives.py

@@ -69,7 +69,13 @@ def set_grad_tensors(
         coord_unstrained = data["coord"]
         cell = data["cell"]
         cell_unstrained = cell
-        if cell.ndim == 2:
+        if coord_unstrained.ndim == 3:
+            B = coord_unstrained.shape[0]
+            scaling = torch.eye(3, dtype=cell.dtype, device=cell.device).unsqueeze(0).repeat(B, 1, 1)
+            scaling.requires_grad_(True)
+            data["coord"] = torch.einsum("bni,bij->bnj", coord_unstrained, scaling)
+            data["cell"] = cell @ scaling
+        elif cell.ndim == 2:
             # Single system: (3, 3) scaling
             scaling = torch.eye(3, requires_grad=True, dtype=cell.dtype, device=cell.device)
             data["coord"] = data["coord"] @ scaling
@@ -136,7 +142,10 @@ def get_derivatives(
                 volume = torch.linalg.det(cell).abs().unsqueeze(-1).unsqueeze(-1)  # (B, 1, 1)
             data["stress"] = dedc / volume
     if hessian:
-        H = calculate_hessian(data["forces"], saved_for_grad["coord"])
+        real_atom_mask = None
+        if saved_for_grad["coord"].ndim == 3 and "numbers" in data:
+            real_atom_mask = data["numbers"].ne(0)
+        H = calculate_hessian(data["forces"], saved_for_grad["coord"], real_atom_mask=real_atom_mask)
         if coulomb_terms is not None and getattr(coulomb_terms, "hessian", None) is not None:
             # The LR coulomb hessian is computed in float64 via finite
             # differences. Accumulate in that (higher) precision rather than
@@ -146,7 +155,7 @@ def get_derivatives(
     return data
 
 
-def calculate_hessian(forces: Tensor, coord: Tensor) -> Tensor:
+def calculate_hessian(forces: Tensor, coord: Tensor, real_atom_mask: Tensor | None = None) -> Tensor:
     """Dense ``(N, 3, N, 3)`` Hessian of the energy w.r.t. real-atom coordinates.
 
     Autograd contract (IMPORTANT):
@@ -169,6 +178,42 @@ def calculate_hessian(forces: Tensor, coord: Tensor) -> Tensor:
     periodic Ewald/PME long-range block remains a fixed-charge FD term in
     either case).
     """
+    if real_atom_mask is not None:
+        coord_for_grad = coord
+        if coord.ndim == 3:
+            if coord.shape[0] != 1 or forces.shape[0] != 1 or real_atom_mask.shape[0] != 1:
+                raise ValueError("real_atom_mask Hessian calculation requires a singleton mode-2 system.")
+            coord_real = coord[0][real_atom_mask[0]]
+            forces = forces[0]
+            real_atom_mask = real_atom_mask[0]
+        else:
+            coord_real = coord[real_atom_mask]
+        if real_atom_mask.shape != (coord_for_grad.shape[-2],):
+            raise ValueError("real_atom_mask must align with the coordinate atom dimension.")
+        if forces.shape[0] == coord_for_grad.shape[-2]:
+            forces_real = forces[real_atom_mask]
+        elif forces.shape[0] == coord_real.shape[0]:
+            forces_real = forces
+        else:
+            raise ValueError("forces must contain either all atoms or all real atoms.")
+        n = forces_real.numel()
+        eye = torch.eye(n, device=forces_real.device, dtype=forces_real.dtype)
+
+        def vjp_real(go: Tensor) -> Tensor:
+            grad = torch.autograd.grad(
+                forces_real.flatten(),
+                coord_for_grad,
+                grad_outputs=go,
+                retain_graph=True,
+                allow_unused=True,
+            )[0]
+            if grad is None:
+                return torch.zeros_like(coord_real)
+            grad_real = grad[0][real_atom_mask] if grad.ndim == 3 else grad[real_atom_mask]
+            return -grad_real
+
+        return torch.func.vmap(vjp_real, 0)(eye).view(-1, 3, coord_real.shape[0], 3)
+
     # Coord includes padding atom (shape N+1), forces only for real atoms (shape N).
     # Hessian computed only for actual atoms: (N, 3, N, 3).
     #