Fix non-periodic cell inflation causing PolarMACE k-space OOM

[samblau]

Summary

Replace the max(abs(positions)) * 5 * cutoff fictitious cell formula for non-periodic
directions with an extent-based cell (extent + 2*cutoff + 1). This fixes GPU OOM
errors when running PolarMACE on non-periodic molecules.

Problem

get_neighborhood() constructs a fictitious cell for non-periodic directions using
max(abs(positions)) * 5 * cutoff. This depends on the absolute coordinate origin
(not molecular extent) and creates unnecessarily large cells. The cell propagates into
PolarMACE's compute_k_vectors_flat(), where k-vector count scales with cell volume,
causing OOM on 80 GB GPUs for ordinary molecules (50-300 atoms) that happen to have
large absolute coordinates.

Example: a 130-atom molecule at coordinates ~(150, 50, 80) A gets a 4,779 A cubic cell
(volume 1.5x10^11 A^3) instead of the ~40 A cell it actually needs.

Fix

Use per-dimension extent + cutoff padding:

extent_x = positions[:, 0].max() - positions[:, 0].min()
cell[0, :] = (extent_x + 2 * cutoff + 1) * identity[0, :]

Verification

• Neighbor lists: Identical (matscipy gives same edge_index with either cell for non-periodic systems)
• PolarMACE energies: 0.000000 eV difference
• PolarMACE forces: <0.000004 eV/A max difference (float32 noise)
• Memory: 220,000x volume reduction on worst case (5,307 A -> 208 A cell), from >80 GB OOM to 0.05 GB

Tested on MACE-POLAR-1-S against the full OMol25 S2EF test set (2.8M non-periodic structures).
The fix resolves all 12,172 structures that previously OOMed on A100-80GB.

Note

This issue was identified by Axon, an AI assistant built atop Claude by Mirror Physics. Axon also determined the fix and drafted the PR. Thus, it should be carefully reviewed by the maintainers!