SFI.utils.neighbors module

Cell-list neighbor builder for truncated-range pair interactions.

Provides build_neighbor_csr() which constructs a sparse CSR neighbor list from particle positions and a cutoff radius, using scipy.spatial.cKDTree. The returned (indptr, indices) arrays plug directly into dispatch_pairs_from_extras.

All routines run on the host (pure NumPy) and are meant to be called between JIT-compiled simulation chunks, not inside jax.lax.scan.

SFI.utils.neighbors.build_neighbor_csr(positions, cutoff, box=None, *, exclude_self=True)

Build a CSR neighbor list using scipy.spatial.cKDTree.

Parameters:

• positions (ndarray, shape (N, d)) – Particle positions (spatial coordinates only).

• cutoff (float) – Cutoff radius. Pairs with r_ij > cutoff are excluded.

• box (ndarray, shape (d,), optional) – Periodic box lengths. If None, open (non-periodic) boundaries.

• exclude_self (bool) – If True (default), self-pairs (i, i) are never included.

Returns:

• indptr (ndarray, shape (N + 1,), dtype int32) – CSR row pointers.

• indices (ndarray, shape (nnz,), dtype int32) – CSR column indices (neighbour particle indices).

Return type:

Tuple[ndarray, ndarray]

SFI.utils.neighbors.make_neighbor_extras(positions, cutoff, box=None, *, indptr_key='indptr', indices_key='indices', exclude_self=True)

Build a CSR neighbor list and return it as an extras dict.

Convenience wrapper around build_neighbor_csr(). The returned dict is ready to be merged into extras_global for a process that uses dispatch_pairs_from_extras(indptr_key, indices_key).

Parameters:

• positions (ndarray) – Forwarded to build_neighbor_csr().

• cutoff (float) – Forwarded to build_neighbor_csr().

• box (ndarray | None) – Forwarded to build_neighbor_csr().

• exclude_self (bool) – Forwarded to build_neighbor_csr().

• indptr_key (str) – Keys under which CSR arrays are stored.

• indices_key (str) – Keys under which CSR arrays are stored.

Returns:

{indptr_key: indptr, indices_key: indices}

Return type:

dict

SFI.utils.neighbors.pad_neighbor_csr(indptr, indices, max_nnz, *, fill_index=0)

Pad a CSR neighbor list to a fixed max_nnz.

JAX JIT recompiles when array shapes change. Padding the indices array to a fixed length avoids recompilation across simulation chunks with fluctuating neighbor counts.

Excess entries are filled with fill_index (default 0). Because indptr keeps the true lengths, the dispatcher will only iterate over the real neighbours — the padded entries are never evaluated.

Note

This only pads indices. indptr is left unchanged (always N + 1 long). If the actual nnz exceeds max_nnz, a ValueError is raised.

Parameters:

• indptr (ndarray) – As returned by build_neighbor_csr().

• indices (ndarray) – As returned by build_neighbor_csr().

• max_nnz (int) – Target length for indices.

• fill_index (int) – Index used to fill padded entries.

Returns:

Same indptr, padded indices of length max_nnz.

Return type:

indptr, indices_padded