jaxdem.utils.neighbor_list_sizing#

Measure neighbor candidates and estimate capacity from relaxed bounding spheres.

Functions

estimate_neighbor_capacity(state, system, ...)

Estimate neighbor capacity using compressed, relaxed bounding spheres.

measure_neighbor_candidates(state, system, *)

Count candidates in a snapshot without allocating a neighbor-pair buffer.

Classes

CandidateStatistics(degree, ...)

Candidate counts for one configuration, before applying a capacity limit.

NeighborCapacityEstimate(max_neighbors, ...)

Capacity recommendations from a relaxed bounding-sphere configuration.
class jaxdem.utils.neighbor_list_sizing.CandidateStatistics(degree: NDArray[int64], directed_pair_count: int, max_degree: int, mean_degree: float, p95_degree: float, p99_degree: float, list_cutoff: float)#

Bases: object

Candidate counts for one configuration, before applying a capacity limit.

Variables:

  • degree (numpy.ndarray) – Directed candidate count for each constituent sphere, shape (N,).

  • directed_pair_count (int) – Sum of degree; each directed entry occupies one pooled cache slot.

  • p99_degree (max_degree, mean_degree, p95_degree,) – Maximum, mean, and degree percentiles across constituent spheres.

  • list_cutoff (float) – Effective search distance, including the absolute skin.

degree: NDArray[int64]#
directed_pair_count: int#
max_degree: int#
mean_degree: float#
p95_degree: float#
p99_degree: float#
list_cutoff: float#
class jaxdem.utils.neighbor_list_sizing.NeighborCapacityEstimate(max_neighbors: int, query_max_neighbors: int, samples: tuple[CandidateStatistics, ...], bounding_radii: NDArray[np.float64], box_size: NDArray[np.float64], packing_fraction: float, cutoff: float, skin: float, relaxation_steps: int, relaxation_info: MinimizeInfo)#

Bases: object

Capacity recommendations from a relaxed bounding-sphere configuration.

Variables:

  • max_neighbors (int) – Average per-sphere budget for a NeighborList pool of size N * max_neighbors, including the requested margin and rounding.

  • query_max_neighbors (int) – Per-row capacity for dense neighbor queries at the sampled list cutoff.

  • samples (tuple[CandidateStatistics, ...]) – Counts for the restored input orientations followed by random rotations.

  • bounding_radii (numpy.ndarray) – Bounding radius of each clump, ordered by increasing clump ID.

  • box_size (numpy.ndarray) – Probe box lengths in input length units, shape (dim,).

  • packing_fraction (float) – Sum of bounding-sphere volumes divided by the probe box volume.

  • skin (cutoff,) – Configured cutoff and absolute skin suitable for NeighborList.Create. The force model can enlarge the effective cutoff reported in samples.

  • relaxation_steps (int) – Number of FIRE updates performed on the bounding spheres.

  • relaxation_info (MinimizeInfo) – Force and torque convergence diagnostics in analogue units. The largest bounding radius, sphere mass, and contact stiffness are each one.

max_neighbors: int#
query_max_neighbors: int#
samples: tuple[CandidateStatistics, ...]#
bounding_radii: NDArray[np.float64]#
box_size: NDArray[np.float64]#
packing_fraction: float#
cutoff: float#
skin: float#
relaxation_steps: int#
relaxation_info: MinimizeInfo#
jaxdem.utils.neighbor_list_sizing.measure_neighbor_candidates(state: State, system: System, *, cutoff: float | None = None, skin: float | None = None, skin_fraction: float | None = None) → CandidateStatistics[source]#

Count candidates in a snapshot without allocating a neighbor-pair buffer.

Parameters:

  • state – One configuration and its domain, force model, and interaction rules. Positions, forces, and contact history are unchanged.

  • system – One configuration and its domain, force model, and interaction rules. Positions, forces, and contact history are unchanged.

  • cutoff – Configured scalar broad-phase cutoff. Defaults to the collider’s cutoff if present, otherwise twice the largest force-model search radius. Pair counts are additionally filtered by the sum of their individual force-model search radii.

  • skin – Absolute buffer distance. Defaults to the collider’s stored skin if present, otherwise 0.05 * cutoff.

  • skin_fraction – Buffer as a fraction of the configured cutoff; mutually exclusive with skin. Overrides the collider’s stored skin.

Returns:

Directed per-sphere counts with the domain’s distance convention and the collider’s clump and bond exclusions. Empty states have zero counts.

Return type:

CandidateStatistics

Raises:

RuntimeError – The spatial hash overflows, so counts would be incomplete.

Notes

This is a host-side utility. A JIT-compiled search counts candidates without evaluating forces or energy. A pair is retained through the skin distance using distance <= radius_i + radius_j + skin. At most 64 particles use the all-pairs backend to avoid grid setup. Larger configurations use a cell search, except custom domains without a declared cell-search geometry, which use all pairs.

jaxdem.utils.neighbor_list_sizing.estimate_neighbor_capacity(state: State, system: System, packing_fraction: float, *, cutoff: float | None = None, skin: float | None = None, skin_fraction: float | None = None, max_steps: int = 1000000, force_tol: float = 1e-08, orientation_samples: int = 3, safety_factor: float = 1.2, extra_neighbors: int = 2, round_to: int = 8, seed: int = 0) → NeighborCapacityEstimate[source]#

Estimate neighbor capacity using compressed, relaxed bounding spheres.

Parameters:

  • state – One periodic 2D or 3D configuration of spheres or rigid clumps. Bodies are identified by clump_id. All analogue bodies can move, including those marked fixed in the input. Input arrays and history are unchanged.

  • system – One periodic 2D or 3D configuration of spheres or rigid clumps. Bodies are identified by clump_id. All analogue bodies can move, including those marked fixed in the input. Input arrays and history are unchanged.

  • packing_fraction – Positive target sum of bounding-sphere volumes divided by box volume. Values above one are permitted for overlapping soft bounding spheres.

  • cutoff – Search settings as in measure_neighbor_candidates().

  • skin – Search settings as in measure_neighbor_candidates().

  • skin_fraction – Search settings as in measure_neighbor_candidates().

  • max_steps – FIRE update budget for one analogue relaxation. Default: 1,000,000.

  • force_tol – Maximum force norm tolerance in analogue units: the largest bounding radius, sphere mass, and spring stiffness are each one. Default: 1e-8.

  • orientation_samples – Total orientation samples, including the input orientations first. Later samples independently rotate each clump. Default: three. Single-sphere states require only one sample.

  • safety_factor – Multiplier >= 1 applied to the largest observed mean and maximum degree.

  • extra_neighbors – Minimum additive margin on each of those statistics. Default: two. The larger of the additive and multiplicative margins is used.

  • round_to – Round each recommendation upward to this integer multiple. Default: eight.

  • seed – Seed for the initial small center perturbation and random orientations.

Returns:

Pooled and dense-query capacities, sample counts, probe geometry, and analogue convergence diagnostics. Capacities are Python integers.

Return type:

NeighborCapacityEstimate

Raises:

  • ValueError – Inputs are invalid, empty, or nonperiodic.

  • RuntimeError – Analogue relaxation fails to converge or candidate counting overflows.

Notes

Body centers and box lengths are scaled once to the requested packing fraction, preserving the box aspect ratio and periodic shear. A small seeded perturbation separates coincident centers before relaxation. Frictionless spheres are relaxed with FIRE and a CellList. Only force evaluations are needed during its iterations. Clump geometry is then restored at the relaxed centers; orientation samples reuse that relaxation.

Counts are exact for the sampled configurations. The recommendations are estimates for subsequent motion, so production overflow checks remain necessary. Dense queries with a different cutoff require separate sizing. This host-side preparation utility cannot be called inside jax.jit.