Preserve radius-graph connectivity under neighbour truncation

[teerthsharma]

Summary

get_max_neighbors_mask keeps the nearest max_neighbors edges per atom. That
is a purely local rule, so it carries no guarantee about the graph as a whole:
where two dense regions touch through a single bridging contact, both endpoints
can rank that contact outside their own budget and drop it. The bridge
disappears by local agreement and a physically connected structure reaches the
model as two disconnected components, which nothing downstream detects.

This adds preserve_connectivity, off by default, plumbed through
generate_graph and get_max_neighbors_mask. When set, it re-adds the shortest
dropped edges needed to restore the component count of the untruncated radius
graph. Component labelling (hooking plus pointer jumping) and the bridge
selection (Boruvka) are both vectorized with scatter_reduce_, so there is no
Python loop over edges and both run on GPU.

eSCNMDBackbone takes preserve_connectivity and forwards it into graph
generation, so a model can reach the flag rather than it stopping at
generate_graph. That constructor keyword is the only public-API addition; it
defaults to False and no shipped config, checkpoint, or YAML sets it. Default
behaviour is unchanged: with the flag off, the mask, the edge list and the
reported neighbour counts are the same objects as before, and there is no head,
loss, or training-unit change.

The correction is strictly an improvement rather than a different answer. Across
16 structures x 4 budgets the retained edge set is a superset of the flag-off set
in 64/64 configurations and the component count equals the untruncated count in
64/64, so an edge the budget kept is never removed or replaced, and the repair
never overshoots into gluing a real vacuum gap.

The failure mode

Two fcc Cu grains 3.2 A apart, cutoff 6 A, 64 atoms:

max_neighbors	components, untruncated	components, truncated
8	1	2
12	1	2
20	1	2
30 (shipped default)	1	2

Across 16 bridged and unbridged structures x 4 budgets, 55 of 64 configurations
lose connectivity with the flag off and 0 of 64 with it on. Uniformly random
dense periodic boxes lose it in 0 of 400 trials, so the defect requires bridge
topology — adsorbates, grain boundaries, cluster contacts — and random stress
testing does not surface it.

Edge growth is 0% on unbridged bulk and, on bridged fcc contacts, 2.86-4.76% at
max_neighbors: 30, rising to 13.33% at 8. Those are perfectly symmetric fcc
grains, so the degenerate shell that selection deliberately keeps whole is at its
widest; treat them as an upper bound rather than a typical case.

Each Boruvka round claims every edge that is shortest for either endpoint.
Requiring both would be a mutual-minimum rule, under which a component whose
shortest edge points at a component with an even shorter one elsewhere claims
nothing, no round is guaranteed to merge anything, and the round bound stops
being a bound. A chain of blocks with monotonically increasing gaps exhibits
exactly that and is covered by test_deep_chain_is_fully_reconnected.

All three radius-graph implementations (v1, v2, nvidia) call the shared mask, so
all three are covered by the one change.

Guarantees the selection keeps

Selection reads only distances and keeps the whole degenerate shell, the same
way the non-strict neighbour budget already does, so the reserved edges do not
depend on atom ordering: 0 of 54 permutation trials changed the reserved edge
set. This does not change the atom-order invariance of the existing selection —
enforce_max_strictly=False is already invariant, enforce_max_strictly=True
is not and stays that way.

Structures that are genuinely disconnected stay disconnected. The pass restores
the component count of the untruncated radius graph and never goes below it, so
a real vacuum gap is not bridged.

Cost

The repair is skipped when the retained graph already has one component per
system. No edge joins two systems and dropping edges can only add components, so
that is the floor, and reaching it proves equality with the untruncated graph.
That check is not free: it costs one component-labelling pass over the retained
graph, which is what the unbridged rows below measure.

RTX 4060 Laptop GPU, CUDA events, 200 timed iterations after 20 warmup, same
batch with the flag off as the control, at PR head:

batch	v1	v2
16 systems / 1728 atoms, nothing bridged	6.008 -> 6.485 ms, 1.079x	16.806 -> 17.652 ms, 1.050x
16 systems / 1552 atoms, 4 bridged	5.253 -> 8.044 ms, 1.532x	17.566 -> 24.129 ms, 1.374x

Contracting the retained components to the quotient graph bounds the repair by
the component count rather than the atom count, and bounds the merge rounds by
the deepest system rather than the batch. The 4-bridged batch runs 1 merge round;
a single 6-block chain, whose components all sit in one system, runs 3. Against
the atom-count formulation, dispatched aten operations fall from 111 to 80 and
device synchronizations from 10 to 5, and the no-fracture fast path is 25
operations and 3 synchronizations either way. Those operation counts were taken
before the either-endpoint correction above, which changes which edges are
claimed but not which operations are dispatched.

Rejected alternatives

approach	result
mutual-kNN (keep an edge only if each endpoint is in the other's top k)	56/64 fracture, against 55/64 for plain truncation on the same edge lists and budgets — strictly worse, since requiring both endpoints drops bridges the budget alone would have kept
union-kNN (keep an edge if either endpoint has the other in its top k)	55/64; identical to plain truncation, so relaxing the rule in that direction is not a fix either
torch.topk(k+1) instead of the full sort in the mask	1.33-1.45x on some shapes, 0.69-0.81x on others
building and sorting only rows over budget	bitwise identical, 0.69-0.78x; the sort is ~5% of the mask's runtime, so removing sort work cannot pay for the compaction
skipping the dense [num_atoms, global_max_degree] padding	occupancy is 61-100% even for heterogeneous batches, ceiling ~1.6x on one sub-step

Validation

tests/core/graph/test_beta0_neighbor_selection.py adds 52 tests, all passing:
the component labeller against an independent union-find, the fracture itself at
four budgets, restoration at four budgets, both v1 and v2, the no-op fast path on
dense bulk, disjointness of reserved and retained edges, bidirectionality of a
restored bridge, the merge-round bounds on chains of 2 to 6 bridged blocks under
two gap orderings, and the branches the two-grain case does not reach: strict
truncation, a mixed-PBC batch, a single-atom system in the batch, and the
backbone keyword actually arriving at graph generation.

Five properties of the reserved edge set are asserted directly, since selection
reads only distances: invariance under rigid motion (3 random rotations plus
translations), equivariance under uniform scaling of positions and cutoff
together (0.5x, 2x, 10x), invariance under atom relabelling, superset of the
flag-off edge set at four budgets, and bitwise reproducibility across repeated
calls on CPU and, under a gpu marker, on CUDA.

An eSCN-MD forward on an RTX 4060 Laptop GPU takes two fcc grains from 2 graph
components to 1 and a 4-block chain from 4 to 1, leaves a plain bulk control
edge-for-edge identical at 1120 edges, and keeps node embeddings and position
gradients finite in both settings. Weights are random, since facebook/UMA is a
gated repository, so this exercises graph generation and the equivariant blocks
rather than any trained prediction.

Correctness is measured by mutation rather than by coverage. Nine targeted
defects were injected into radius_graph_pbc.py one at a time and the suite rerun
against each:

injected defect	killed by
mutual-minimum instead of either-endpoint claim	test_deep_chain_is_fully_reconnected
absolute squared-distance threshold	10 tests
nondeterministic tie-break in the degenerate shell	6 tests
component-index tie-break instead of distance	4 tests
one Boruvka round regardless of depth	2 tests
fast path taken unconditionally	9 tests
reserved edges leaking into the retained set	2 tests
pointer jumping removed from the labeller	survives
one pointer jump per Boruvka round	survives

The two survivors reduce pointer-jump counts only; the labeller's convergence
loop and the Boruvka round count respectively absorb the shortfall on every
input tested, so neither is shown to change an output. They are not shown to be
safe in general either — see Limits.

tests/core/graph fails 31/190 on this branch and 31/138 on upstream main at
a73c95b, and tests/core/models/uma fails 3/239 on both, with identical
failure sets — no test regresses in either.

Limits

The 31 shared failures are a local Windows environment, not a clean run:
nvalchemiops is not installable there, which fails the radius_pbc_version=3
tests, and the checkpoint-downloading tests fail alongside them. They fail the
same way on upstream main, which is why the failure sets are diffed rather than
counted. The timing figures are from a single RTX 4060 laptop GPU and a single
batch shape; wall-clock ratios for this change moved between 1.29x and 3.35x on
that machine for identical code, because the mechanism is kernel-launch and
synchronization count rather than arithmetic, so the operation counts above are
the load-bearing measurement and the timings are indicative. Under exact
geometric degeneracy the minimum spanning forest is not unique (two symmetric
fcc grains gave 112 bitwise-equal candidate bridges); the whole degenerate shell
is kept rather than one representative, so this is not a minimal forest, and it
is what makes the 13.33% worst-case growth figure a property of symmetric test
structures rather than of the algorithm. The two surviving mutants show that the
pointer-jump counts are not tightly bounded by any current test: reducing them
changes no output on the structures tested, so the counts are sufficient but not
demonstrated to be necessary or minimal. The flag is not wired into any shipped
config — enabling it for UMA inference is a separate decision. The
radius_pbc_version=3 path is covered by reading rather than by execution:
nvalchemiops does not install on the machine used here, so the nvidia
wiring — where c_index and n_index share the global atom index space and the
neighbour counts are recomputed from the masked c_index — has not been run.
Graph parallelism with the flag on is likewise untested; v2 leaves
neighbor_index out under a node partition and v1 reaches the shared mask before
generate_graph filters by partition, but neither has been exercised. The
eSCN-MD forward itself is not bitwise reproducible run to run on CUDA with the
flag off or on, so that nondeterminism sits upstream of this change; the graph it
consumes is reproducible.