nvalchemi.distributed.COOPadder

class nvalchemi.distributed.COOPadder

Built-in GraphPadder for COO edge_index graphs.

The inferred default: a model whose halo-padded graph is an ordinary Batch with per-atom fields and a per-edge neighbor_list (COO endpoints) declares nothing. It works on the abstract Batch storage groups, so it is model-agnostic.

Atom / edge counts are knowable up front, so pad() resolves both caps from cap_state before padding: "atoms" (1.15 initial headroom, strict_gt=False since the dead node sits at the last slot) and "edges" (1.35 initial headroom — edge count climbs ~25% through equilibration). Both regrow x1.30 on overflow, stride 16.

Layout:

• Per-atom fields -> n_cap: appended rows carry zeros and join the last graph (their node outputs are dropped by the owned-only consolidation).

• Per-edge fields -> e_cap: invalid edges (sentinel rows with endpoint >= n_real) and the fill are routed to an isolated dead node (the last row, n_cap - 1) as self-loops. neighbor_list_shifts for those rows uses the [1, 0, 0] image so the edge vector is non-degenerate (a zero vector NaNs through spherical harmonics). The dead node is referenced by no real edge and masked out of every owned output.

unpad() is a no-op: the owned-only output consolidation already drops the dead / ghost rows. The framework restores the transient padded storage separately (the padded Batch is reused in place across MD steps).

pad(data, cap_state, cap_atoms=True)

Resolve atom / edge caps from cap_state and pad the halo-padded Batch to them. Mutates data in place and returns it; None is a safe no-op.

Parameters:

• data (Batch)

• cap_state (dict[str, int])

• cap_atoms (bool)

Return type:

Batch

unpad(output, n_real=None)

No-op: the owned-only output consolidation drops the dead rows.

Parameters:

• output (Any)

• n_real (int | None)

Return type:

Any