Add optional NVIDIA electrostatics backend for PolarMACE

[ilyes319]

Summary

• add an optional nvalchemiops electrostatics backend for PolarMACE reciprocal-space multipole features and energy
• reuse one Toolkit-Ops SCF geometry cache across both spin channels and all fixed-point iterations
• expose backend selection through PolarMACE, MACECalculator, mace_polar, training CLI configuration, and model reconstruction
• preserve graph_longrange as the default and as the automatic fallback for molecular, unsupported partial-periodic, LAMMPS, and TorchScript paths
• support z-normal slabs (TTF) by combining NVIDIA multipole kernels with the existing analytic slab feature and energy corrections
• add the optional dependency, capability-gated CPU/GPU CI, and direct/batched/full-model parity coverage

Compatibility

Toolkit-Ops 0.4.0 requires Python 3.11+, PyTorch 2.8+, and float64 kernel inputs. The adapter promotes kernel inputs internally and casts outputs back so float32 and float64 PolarMACE models are both supported. The backend is currently eager-only. NVIDIA direct-k-space kernels handle fully periodic and z-normal slab calculations; exact molecular calculations continue through the graph_longrange fallback.

FFF PME benchmark

A float64, l=1 prototype was measured on the self-hosted NVIDIA A100 40 GB runner. It used 12 A vacuum, a 0.5 A mesh, and derived receiver features from the PME energy gradient. The PolarMACE estimate is 2R * feature + energy for R recursion steps.

Atoms	Exact FFF, R=1	Padded PME, R=1	Speedup	Feature memory, exact / PME
64	6.59 ms	28.93 ms	0.23x	0.014 / 0.030 GiB
256	6.73 ms	28.72 ms	0.23x	0.102 / 0.049 GiB
512	10.92 ms	29.14 ms	0.38x	0.383 / 0.049 GiB
1024	29.86 ms	30.00 ms	1.00x	1.508 / 0.078 GiB

At 256 atoms, increasing vacuum from 4 A to 24 A reduced the energy error from 7.63e-2 to 5.63e-4 eV/atom. The feature discrepancy plateaued near 1.19% and did not improve when refining the mesh from 1.0 A to 0.25 A. A direct diagnostic showed PME and NVIDIA direct-k features agree to 3.3e-5 relative error; both differ from the molecular graph_longrange path because that path uses displaced-charge finite differences rather than analytic multipoles.

The result supports PME for systems beyond roughly 1,000 atoms and for memory-constrained inference, but not as a transparent FFF replacement. Preserving existing model semantics needs a displaced-charge PME/receiver-feature implementation or a model trained with analytic multipoles, so exact molecular fallback remains in this PR.

Validation

• Toolkit-Ops 0.4.0 CPU suite: 7 passed (features, energy, batched systems, periodic and slab full-model parity, forces, virials/stress, first/second derivatives, float32/float64, molecular fallback, calculator guard)

• existing PolarMACE and argument-parser regressions: 36 passed; the subprocess training test was rerun outside the filesystem sandbox because OpenMP shared memory is unavailable inside it

• released MACEPOLAR1Mmodel checkpoint through MACECalculator: periodic float32 energy matched and maximum force difference was 3.19e-5

• Black, isort, pylint, bytecode compilation, workflow YAML/JSON parsing, and git diff --check

• final PR head: core, extension (including polar-nvalchemiops), integration, and GPU workflows all passed

• final NVIDIA A100 full suite: 28 passed, 24 skipped, 1 xfailed in 8m16s (periodic float32/float64 and slab full-model parity, compile, cuEquivariance, LES, schedule-free, and calculator GPU paths)

• derivative outputs use a focused 5e-4 full-model tolerance; direct feature/energy kernels and non-derivative outputs retain their stricter tolerances

• the single-GPU job runs pytest serially to prevent concurrent torch.compile workers from exhausting the A100

---

Note

Medium Risk
Changes core PolarMACE long-range energy/feature paths and introduces a heavy optional GPU dependency; fallbacks limit blast radius but incorrect routing or parity drift could affect forces and energies on periodic/slab systems.

Overview
Adds an optional nvalchemiops long-range electrostatics path for PolarMACE alongside the default graph_longrange implementation, wired through a new mace/modules/polar_backends.py adapter around NVIDIA ALCHEMI Toolkit-Ops multipole SCF APIs (shared geometry cache across spin channels and fixed-point steps).

PolarMACE gains set_electrostatics_backend and branches in forward to use NVIDIA kernels for fully periodic and z-normal slab (TTF) systems when not scripting/LAMMPS; molecules, other partial PBC, and unsupported cases still use graph_longrange, with slab analytic corrections layered on non-fully-periodic NVIDIA paths. Selection is exposed via MACECalculator, mace_polar, training CLI --electrostatics_backend, and model save/load config; compile_mode is rejected when this backend is active.

Packaging adds mace-torch[nvalchemiops], requirements, pytest nvalchemiops capability, a dedicated CI job, GPU fleet extras, README usage notes, and parity tests (direct, batched, full-model, slab, GPU).

^{Reviewed by Cursor Bugbot for commit 64185bb. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 4bad3d5. Configure here.}

[cursor]

Slab feature corrections unmasked

Medium Severity

The nval path allows mixed fully periodic and z-slab graphs in one batch, but feature corrections are gated only with a global torch.all(pbc) check while energy applies an explicit per-graph is_slab mask. If non_periodic_correction_terms does not zero fully periodic graphs itself, those systems get incorrect SCF features.

 

^{Reviewed by Cursor Bugbot for commit 4bad3d5. Configure here.}