jaxdem.colliders.neighbor_list

Neighbor List Collider implementation.

Classes

NeighborList(cell_list, neighbor_list, ...)

Verlet Neighbor List collider.

class jaxdem.colliders.neighbor_list.NeighborList(cell_list: DynamicCellList, neighbor_list: Array, old_pos: Array, n_build_times: int, cutoff: Array, skin: Array, overflow: Array, max_neighbors: int)

Bases: Collider

Verlet Neighbor List collider.

This collider caches a list of neighbors for every particle. It only rebuilds the list when particles have moved more than half the ‘skin’ distance.

Performance Note: You must provide a non-zero skin (e.g., 0.1 * radius) for this collider to be efficient. If skin=0, it rebuilds every step.

cell_list

The underlying spatial partitioner used to build the list.

Type:

DynamicCellList

neighbor_list

Shape (N, max_neighbors). Contains the IDs of neighboring particles. padded with -1.

Type:

jax.Array

old_pos

Shape (N, dim). Positions of particles at the last build time.

Type:

jax.Array

n_build_times

Counter for how many times the list has been rebuilt.

Type:

int

cutoff

The interaction radius (force cutoff).

Type:

float

skin

Buffer distance. The list is built with radius = cutoff + skin and rebuilt when max_displacement > skin / 2.

Type:

float

overflow

Boolean flag indicating if the neighbor list overflowed during build.

Type:

jax.Array

max_neighbors

Static buffer size for the neighbor list.

Type:

int

cell_list: DynamicCellList

neighbor_list: Array

old_pos: Array

n_build_times: int

cutoff: Array

skin: Array

overflow: Array

max_neighbors: int

classmethod Create(state: State, cutoff: float, box_size: jax.Array | None = None, skin: float = 0.05, max_neighbors: int | None = None, number_density: float = 1.0, safety_factor: float = 1.2, cell_size: float | None = None) → Self

Creates a NeighborList collider.

Parameters:

• state (State) – Initial simulation state.

• cutoff (float) – The physical interaction cutoff radius.

• box_size (jax.Array, optional) – The size of the periodic box, if used.

• skin (float, default 0.05) – The buffer distance. Must be > 0.0 for performance.

• max_neighbors (int, optional) – Maximum neighbors to store per particle. If not provided, it is estimated from the number_density.

• number_density (float, default 1.0) – Number density for the state used to calculate max_neighbors, if not provided. Assumed to be 1.0.

• safety_factor (float, default 1.2) – Used to adjust the max_neighbors value calculated from number_density. Empirically obtained

• cell_size (float, optional) – Override for the underlying cell list size.

static create_neighbor_list(state: State, system: System, cutoff: float, max_neighbors: int) → Tuple[State, System, jax.Array, jax.Array]

Return the cached neighbor list from this collider.

Notes

• This method does not rebuild the neighbor list. It simply returns the last cached neighbor_list and overflow stored in system.collider.

• The returned neighbor indices refer to the collider’s internal particle ordering at the time the cache was last updated (i.e., after the most recent rebuild inside compute_force()).

• The cutoff and max_neighbors arguments are accepted for API compatibility but are currently ignored; the cache was built using this collider’s configured cutoff + skin and max_neighbors.

static compute_force(state: State, system: System) → Tuple[State, System]

static compute_potential_energy(state: State, system: System) → jax.Array