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

[zubatyuk]

AIMNetCentral Pull Request

Description

This change adds canonical global-indexed mode-2 sparse neighbor matrices for batches containing molecules of different sizes. It lets callers supply padded batched neighbor lists while preserving correct short-range, long-range, periodic, derivative, and Hessian-vector-product behavior.

Mode-2 inputs are now validated before computation, including padded-center sentinels, per-system ownership, shift alignment, and neighbor-matrix suffix layout. This prevents invalid global indices from reaching CUDA kernels or periodic backends, where failures otherwise can be difficult to diagnose.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that changes existing supported behavior)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Changes Made

• Added global packed mode-2 neighbor-list handling across calculator, model, AEV, DFT-D3, Coulomb, and periodic long-range paths.
• Validated raw mode-2 neighbor matrices and periodic geometry before flattening, including global index bounds, same-system neighbors, sentinel-only padding, and aligned lattice shifts.
• Added convert_mode2_local_to_global for callers migrating legacy local-index 3D matrices to the canonical global representation.
• Preserved shared neighbor-list tensor identity where suffixes alias the same input and avoided re-masking freshly produced mode-1 ConvSV output.
• Made Hessian-vector products select real atoms from the prepared padding mask, so vectors and results are correct for padded mode-2 batches.
• Added CPU, CUDA subprocess, periodic-backend, kernel, calculator, derivative, compile, and documentation coverage for the new layout.

Compatibility

• Existing flattened mode-1 inputs and canonical global mode-2 inputs remain supported.
• Legacy 3D matrices that use per-system local atom indices must be converted with aimnet.nbops.convert_mode2_local_to_global before use.
• Every mode-2 system must reserve its final padded dummy atom; full 3D periodic inputs must supply lattice shifts aligned with each neighbor matrix.

Testing

• Unit, focused, or full tests pass locally.
• Ruff and repository quality checks pass.
• New tests cover the changed behavior.
• Documentation builds without warnings.

Checklist

• I have performed a self-review of the code and documentation.
• I have checked that the code and documentation follow the project style.
• I have checked that the documentation is clear and readable.
• I have updated the changelog when the change is user-visible.

[isayev]

Reviewed the full diff with empirical verification on an L40S (torch 2.9.1+cu128). The design direction is right — validation-first, clean kernel generalization, good CUDA-subprocess assert testing, and the CUDA compiled mode-2 forward is clean (0 graph breaks, eager parity ~1e-5). Legacy B>1 local-index inputs are rejected loudly on every entry path; no silent-wrong path through the calculator. But one measured physics bug and several majors need fixing before merge.

Blocker — Ewald/PME energies depend on the padding width. _mode2_backend_inputs builds batch_idx over all B*N rows (dummies included), and nvalchemiops auto-estimates the Ewald splitting parameter from the per-system atom count when only accuracy is passed. Padding inflates the count, shifts alpha, and the caller-truncated real-space sum no longer cancels. Measured on a 4-atom cation (12 A cell, ewald_accuracy=1e-6, identical geometry and neighbor list): mode-1 flat 1.99822 eV; mode-2 with 1 pad 2.00033; 16 pads 2.00442 — 2.1e-3 eV mode disagreement and 4.1e-3 eV drift across padding widths (~0.1 kcal/mol), unchanged at accuracy=1e-4. The same molecule gets a different energy depending on the largest molecule in its batch. Fix: estimate alpha/k-cutoff from real per-system counts (mol_sizes) or pass explicit parameters. Note the mandatory trailing dummy makes this bias always nonzero in mode 2.

Major — shared (3,3) cell silently corrupts DSF/DFT-D3 for systems b>0. _mode2_backend_inputs replaced main's expand(B,-1,-1) with unsqueeze(0); the batched kernels index cell[system_id] unchecked. Measured with two translation-equivalent systems: DSF returns [-0.8738, -1.3215] instead of two identical values; DFT-D3 likewise; Ewald/PME at least fail loudly. The calculator path is protected by normalize_mode2_periodic_geometry, but LRCoulomb/DFTD3 are public modules driven directly by training code. Restore the expand or validate cell.shape[0] == B inside _mode2_backend_inputs.

Major — ConvSV einsum fallback can leak features across systems. The masked gather index is 0 — system 0's first real atom in the flattened layout — so the CPU/float64 branch is only correct because g_sv happens to be pre-masked by its producer. Verified: with an unmasked g_sv, a real center picks up system-0 features. The CUDA kernel branch is safe (sentinel break). Mask a_j/avf_sv with mask_ij in the fallback rather than relying on a second module's masking discipline; the new mask_i_ on the output only covers padded centers, not padded neighbor slots of real centers.

Major — non-finite dummy coordinates poison Ewald (NaN energy; PME/DSF unaffected). Mode 2 keeps every dummy row in the backend call, unlike mode 1 which strips them. validate_mode2_nbmat_raw already asserts shift finiteness — add coordinate finiteness and state the requirement in the mode-2 contract docs.

Major — _mode2_check reintroduces the CPU graph breaks #109 just removed. condition.item() runs inside models/base.py:prepare_input, which is inside the compiled forward: a CPU compile_model=True mode-2 eval accumulates 20 measured graph breaks (all Tensor.item() from nbops), and fullgraph=True hard-fails. CUDA is clean — _assert_async traces and even survives inductor fusion. Fix: torch._check(condition) under torch.compiler.is_compiling(), or route the CPU branch through _assert_async (it has an eager CPU kernel).

Major — validation is triplicated and measurably expensive. The identical 6-line gate is copy-pasted in calculator.prepare_input, the eval hessian probe, and models/base.py:prepare_input, so a plain eval validates twice (three times with hessian=True) with no dedup across aliased suffixes: measured 3.6-8.3 ms against a 20.7 ms total eager forward at B=32/N=30/M=60 — 17-40% overhead, launch-bound, with 8 identical validations of the same tensor object. Consolidate to one nbops.validate_mode2_input() chokepoint with an identity memo and a validated marker (keeping standalone-model users covered).

Major — the parity tests cannot see the blocker. Every periodic parity test compares mode-2 against mode-2 (B=2 vs B=1); the only genuine cross-mode check is non-periodic DSF, which is immune because its parameters are explicit. Please add: mode-1 vs mode-2 Ewald/PME energy/force parity (would have caught the blocker immediately), a padding-width-invariance test (same system, N_real+1 vs N_real+8 padding), a CPU mixed-size periodic case, and a compiled mode-2 forward test asserting zero graph breaks plus eager parity.

Major — breaking-change bookkeeping. The semver-breaking input contract gets a single CHANGELOG line under ### Changed with no BREAKING label. And the break is wider than the "local to global indices" story: every mode-2 input — including previously-unambiguous B=1 — must now reserve numbers[:, -1] == 0 as a dummy row, which convert_mode2_local_to_global explicitly cannot repair (it requires the dummy row to already exist). Label the entry BREAKING and document the dummy-row requirement as an explicit migration step.

Also: the branch now conflicts with main (CHANGELOG.md, aimnet/calculators/calculator.py, tests/test_calculator.py) after #108 merged — needs a rebase regardless.

Suggested order: fix the Ewald/PME parameter estimation first (with the parity + padding-invariance tests pinning it), then the cell expand and the compile-safe check; consolidate validation while touching those lines; rebase; relabel the changelog.

[isayev]

The global packed mode-2 design is right and the kernel work is high quality — sentinel contract preserved, double-backward capacity fix correct, the no-host-sync regression test is exactly what this path needs. No numerical defects found. Requesting changes on five points that need to land before merge:

1. Periodic parity tests never leave mode 2. _single_mode2_periodic_data in tests/test_mode2_periodic_backends.py builds a dict named mode1 but unsqueezes into a 3-D nbmat — a mode-2 singleton — and the same pattern repeats in the model/calculator/GPU parity tests, while the one pre-existing cross-mode calculator test was repurposed. A uniform mode-2 bug (shift sign, cell selection) passes the whole suite. Please add one flat mode-1 vs mode-2 anchor per backend (DSF/D3/Ewald/PME) with nonzero shifts — the flat periodic branch in _coul_nvalchemi is still there to serve as the reference.

2. cutoff_coulomb/cutoff_dftd3 are accepted but never consumed or checked. Mode 2 makes caller-supplied neighbor lists the only path, yet DSF still evaluates with self.dsf_rc, D3 with its smoothing radius, Ewald with its accuracy-derived cutoff. A list built with a smaller cutoff passes all structural validation and silently drops pairs — kcal/mol-scale. Consume the keys and check them against the module radii at mode-2 entry (and either validate or hard-document the s ⇒ −s edge-symmetry requirement).

3. The DSF Hessian claim is untested. The FD reference in test_global_mode2_periodic_hessian_diagonal_matches_independent differentiates the module energy directly, bypassing calculate_hessian, and _periodic_fd_setup rejects dsf — while docs/long_range.md now says DSF Hessians work. Since DSF's explicit term is autograd-detached, nothing proves its curvature isn't silently omitted. Add an absolute FD-vs-analytic test through the real calculator path, or change the table to No.

4. Rebase onto current main needs care beyond the 4 textual conflicts (CHANGELOG, aev.py, calculator.py, tests/test_calculator.py). Main's AEV dispatch now gates the Warp kernel on WARP_CUDA_AVAILABLE (CUDA torch + CPU-only warp-lang is a legal conda-forge solve); the new mode-2 branch checks only device+dtype, so a naive rebase crashes in wp.launch on those hosts. Fold the gate + one-time warning into the mode-2 d2features sub-branch, take main's kernels/__init__.py and _init_warp_quietly() wholesale, and add main's weights marker to tests/test_mode2_cuda_validation.py (its subprocess loads registry weights). Full suite + GPU run after rebase — models/base.py auto-merges cleanly over large unrelated changes, so a clean merge proves nothing.

5. Hot-path validation cost. Full mode-2 validation runs in AIMNet2Calculator.prepare_input, again in AIMNet2Base.prepare_input, and a third time in the Hessian probe — which also does a throwaway to_input_tensors device transfer for every Hessian request, mode-2 or not. With the usual all-suffixes-aliased input that's hundreds of extra kernel launches per periodic MD step. Validate once per call chain, memoize by tensor identity, gate the O(B·N·M) content checks behind a trust/debug flag. Fine as a fast-follow if you prefer, but flagging it now since large-system MD is the point of this feature.