nvalchemi.distributed.SpatialPartitioner

class nvalchemi.distributed.SpatialPartitioner(config, cell_matrix, pbc)

Assigns atoms to spatial sub-domains on a Cartesian grid.

The partitioner divides the simulation cell into axis-aligned blocks and maps each atom to the rank that owns its block.

Parameters:

• config (DomainConfig) – Domain decomposition configuration.

• cell_matrix (torch.Tensor) – Cell / box matrix describing the simulation domain. Accepts either the Batch convention (1, 3, 3) or the raw (3, 3) shape; leading batch dimensions are squeezed internally.

• pbc (torch.Tensor) – Periodic boundary conditions per axis. Accepts either (1, 3) (Batch convention) or (3,); leading batch dimensions are squeezed internally.

assign_atoms_to_ranks(positions)

Assign each atom to a rank based on its position.

Parameters:

positions (torch.Tensor) – (N, 3) atom positions in Cartesian coordinates.

Returns:

(N,) integer tensor of rank assignments.

Return type:

torch.Tensor

static balance_cells_for_ranks(cells_per_dim, rank_grid)

Round each axis’s cell count to a multiple of its rank-grid factor.

The cell->rank block assignment (cx = ceil(Nx / Px)) gives the first ranks cx cells and the last rank the remainder, so when Nx is not a multiple of Px the domains are unequal — e.g. 3 cells across 2 ranks splits 2:1. Making Nx a multiple of Px gives every rank Nx / Px equal-width cells.

Rounds DOWN to the nearest multiple of Pi (floored at Pi so each rank keeps at least one cell). Rounding down — never up — keeps each cell at least as wide as the cutoff-derived size, so cells stay >= cutoff. Axes with a single rank (Pi == 1) are unchanged.

Parameters:

• cells_per_dim (tuple[int, int, int])

• rank_grid (tuple[int, int, int])

Return type:

tuple[int, int, int]

cell_to_rank(ix, iy, iz)

Map cell indices (ix, iy, iz) to the owning rank.

Works with both scalar ints and batched torch.Tensor inputs. Uses ceiling-division block assignment.

Parameters:

• ix (int | Tensor)

• iy (int | Tensor)

• iz (int | Tensor)

Return type:

int | Tensor

static compute_rank_grid(cells_per_dim, world_size)

Compute (Px, Py, Pz) rank grid minimizing surface area.

Enumerates all 3-factor factorizations of world_size and picks the one that minimizes the surface-area proxy 2 * (dx*dy + dy*dz + dx*dz) where dx = Nx/Px, etc.

Parameters:

• cells_per_dim (tuple[int, int, int])

• world_size (int)

Return type:

tuple[int, int, int]

domain_widths()

Physical width of one rank’s domain along each axis, in Angstroms.

Each rank owns ceil(N_d / P_d) grid cells of width face_distance_d / N_d.

Returns:

Per-axis domain width.

Return type:

tuple[float, float, float]

get_neighbor_ranks(rank)

Return precomputed neighbor ranks for rank.

Parameters:

rank (int)

Return type:

list[int]

keeps_owner(positions, owner_rank, hysteresis)

Hysteresis-aware ownership test.

An atom keeps its owner until it drifts more than hysteresis (Cartesian Angstrom) past the owner’s domain boundary — i.e. it stays while inside the owner’s domain expanded by hysteresis on every axis (PBC-wrapped on periodic axes). This stops the per-step migration thrashing of atoms that merely vibrate across a domain plane.

Parameters:

• positions (torch.Tensor) – [N, 3] atom positions in Cartesian coordinates.

• owner_rank (int) – The rank whose ownership is being tested.

• hysteresis (float) – Cartesian margin (Angstrom) an atom must exceed past the boundary before it loses its owner.

Returns:

[N] bool, True where an atom currently owned by owner_rank should keep that owner.

Return type:

torch.Tensor

Notes

Correctness relies on hysteresis <= skin / 2 (enforced by DomainConfig): a deferred atom plus inter-rebuild drift stays within the owner’s halo (ghost_width = cutoff + skin), so the owner still has all the atom’s neighbors and the neighbor still ghosts the atom. Uses the same fractional-bounds + reciprocal-norm geometry as the halo ghost region it must stay inside.

neighbor_span()

Rank offsets the ghost shell reaches along each axis.

ceil(ghost_width / domain_width) — 1 while the ghost fits inside one domain, more once it does not (many ranks, a small cell, or a barostat contracting the box). Communicating only +/-1 in that case silently drops every interaction that spans two or more domains.

Returns:

Per-axis offset range.

Return type:

tuple[int, int, int]

Raises:

ValueError – If the shell would have to wrap onto the rank’s own periodic image, which the halo exchange cannot express.

rank_to_cell_bounds(rank)

Return cell index bounds (lo, hi) owned by rank.

lo is inclusive, hi is exclusive.

Parameters:

rank (int)

Return type:

tuple[tuple[int, int, int], tuple[int, int, int]]

rank_to_grid_coords(rank)

Decompose a linear rank index into (rx, ry, rz) grid coords.

Parameters:

rank (int)

Return type:

tuple[int, int, int]

static refine_grid_for_ranks(cells_per_dim, world_size)

Subdivide cells until there are at least world_size cells.

Doubles the smallest dimension iteratively. Warns if total cells remain less than world_size after 64 iterations (safety cap).

Parameters:

• cells_per_dim (tuple[int, int, int])

• world_size (int)

Return type:

tuple[int, int, int]

topology_version: int = 0

Incremented whenever the neighbour-rank set changes, so anything derived from it can tell that its copy is stale. Class-level so a partitioner built without __init__ still carries one.

update_cell(cell_matrix)

Refresh the physical cell when a barostat (NPT/NPH) deforms the box.

Recomputes cell_matrix and its cached inverse; the fractional cell grid (cells_per_dim) and rank layout are intentionally kept fixed so rank assignment stays consistent as the box scales — only the physical size of each grid cell changes. Halo regions (rank_to_cell_bounds → cartesian via cell_matrix) and the fractional ghost width scale with the updated cell automatically. Without this, the partitioner keeps the partition-time box: as the cell grows, wrapped positions fall outside it (fractional coords >= 1) and assign_atoms_to_ranks misroutes atoms.

Contraction narrows every domain, so the ghost shell can come to span more than one of them; the neighbour-rank set is rebuilt whenever that happens, and neighbor_span() raises if the geometry has contracted past what the halo can express.

Parameters:

cell_matrix (torch.Tensor) – The barostat’s updated cell, (3, 3) or (1, 3, 3).

Return type:

None