neighbors

Neighbor-list construction for periodic MD.

import mlx_atomistic.neighbors

Classes

NeighborBlocks

class NeighborBlocks
    def __init__(left: mx.array, right: mx.array, valid_mask: mx.array, block_size: int, candidate_count: int, compact_pair_count: int)

Fixed-shape candidate pair blocks for MLX-side cutoff filtering.

Parameters

Name	Type	Default	Description
left	mx.array
right	mx.array
valid_mask	mx.array
block_size	int
candidate_count	int
compact_pair_count	int

Properties

• block_count int — Number of fixed-size candidate blocks.
• candidate_waste_count int — Number of emitted block candidates outside the neighbor radius.
• estimated_bytes int — Estimated storage bytes for block indices and validity mask.
• padded_candidate_count int — Number of candidate slots including padding.

NeighborList

class NeighborList
    def __init__(pairs: mx.array | None, cutoff: float, skin: float = 0.0, stats: PairListStats | None = None, blocks: NeighborBlocks | None = None, tiles: NeighborTiles | None = None, interaction32: _DeviceFusedHalfSchedule32 | None = None, *, sort_diagnostic_pairs: bool = False, reference_positions: mx.array | None = None, cell: Cell | None = None)

Neighbor interactions for pairwise potentials.

Parameters

Name	Type	Default	Description
pairs	mx.array | None
cutoff	float
skin	float	0.0
stats	PairListStats | None	None
blocks	NeighborBlocks | None	None
tiles	NeighborTiles | None	None
interaction32	_DeviceFusedHalfSchedule32 | None	None
sort_diagnostic_pairs	bool	False
reference_positions	mx.array | None	None
cell	Cell | None	None

Properties

• adaptation_reason str | None — Reason the active backend selected a different internal schedule.
• backend str — Pair-construction backend name.
• candidate_count int | None — Number of candidate interactions tested before cutoff filtering.
• candidate_waste_count int | None — Number of candidate interactions rejected by compaction/filtering.
• candidate_waste_fraction float | None — Fraction of emitted candidates rejected by compaction/filtering.
• compact_pair_count int — Number of compact pairs accepted by the neighbor search radius.
• compaction_backend str | None — Backend used to compact candidates into explicit pairs, if any.
• diagnostic_pairs mx.array — Return the compact cutoff-plus-skin pairs for this generation.
• diagnostic_pairs_materialized bool — Whether the optional exact-pair diagnostic currently occupies memory.
• estimated_candidate_bytes int — Estimated bytes for backend candidate testing arrays.
• estimated_cell_list_bytes int — Estimated bytes for cell-list construction arrays.
• estimated_compact_pair_bytes int — Estimated bytes for compact int32 pairs accepted by the search radius.
• estimated_pair_bytes int — Estimated bytes for the compact int32 pair array.
• fallback_reason str | None — Reason an accelerated representation fell back or used a hybrid step.
• interactions mx.array | NeighborBlocks | NeighborTiles — Return the active force-evaluation representation.
• materialized_diagnostic_pairs mx.array | None — Return resident diagnostic pairs without triggering materialization.
• pair_count int — Number of unique pairs or candidate block entries.
• pairs mx.array — Return compact pairs, materializing deferred tile diagnostics once.
• representation_kind str — Neighbor interaction representation shape.
• supports_async_force_submission bool — Whether the active structured force route supports asynchronous submission.
• supports_tile_diagnostics bool — Whether exact diagnostic tiles are resident or can be built lazily.

Methods

diagnostic_force_candidates

def diagnostic_force_candidates(*, prefer_tiles: bool) -> mx.array | NeighborTiles

Return an exact diagnostic representation for this generation.

Parameters

Name	Type	Default	Description
prefer_tiles	bool

Returns

• mx.array | NeighborTiles

force_candidates

def force_candidates(*, prefer_tiles: bool) -> mx.array | NeighborTiles | _DeviceFusedHalfSchedule32

Select compact pairs or exact tiles for force binding.

Parameters

Name	Type	Default	Description
prefer_tiles	bool		Return exact tiles when they are present. Otherwise, return the compact pair representation.

Returns

• mx.array | NeighborTiles | _DeviceFusedHalfSchedule32 — Exact tiles only when requested and available, or compact pairs.

NeighborListManager

class NeighborListManager
    def __init__(cell: Cell, cutoff: float, skin: float = 0.3, check_interval: int = 1, sort_pairs: bool = False, max_workers: int | None = None, backend: NeighborBackend = 'auto', max_mlx_dense_atoms: int = DEFAULT_MLX_DENSE_PAIR_LIMIT, block_size: int = DEFAULT_MLX_CELL_BLOCK_SIZE, displacement_check_backend: NeighborCheckBackend = 'numpy', interaction32_exclusion_pairs: object = (), interaction32_one_four_pairs: object = (), interaction32_ordinary_tiles_per_group: int = 3, neighbor_list: NeighborList | None = None, reference_positions: mx.array | None = None, rebuild_count: int = 0, last_max_displacement: float = 0.0, updates_since_check: int = 0, rebuild_wall_seconds: float = 0.0, update_wall_seconds: float = 0.0)

Manage Verlet neighbor-list rebuilds during an MD trajectory.

Parameters

Name	Type	Default	Description
cell	Cell
cutoff	float
skin	float	0.3
check_interval	int	1
sort_pairs	bool	False
max_workers	int | None	None
backend	NeighborBackend	'auto'
max_mlx_dense_atoms	int	DEFAULT_MLX_DENSE_PAIR_LIMIT
block_size	int	DEFAULT_MLX_CELL_BLOCK_SIZE
displacement_check_backend	NeighborCheckBackend	'numpy'
interaction32_exclusion_pairs	object	()
interaction32_one_four_pairs	object	()
interaction32_ordinary_tiles_per_group	int	3
neighbor_list	NeighborList | None	None
reference_positions	mx.array | None	None
rebuild_count	int	0
last_max_displacement	float	0.0
updates_since_check	int	0
rebuild_wall_seconds	float	0.0
update_wall_seconds	float	0.0

Properties

• rebuild_threshold float — Maximum displacement before the Verlet list must be rebuilt.

Methods

build_cell_candidate

def build_cell_candidate(positions, cell: Cell) -> NeighborListManager

Build an isolated neighbor state for a proposed periodic cell.

Parameters

Name	Type	Default	Description
positions			Candidate particle positions.
cell	Cell		Candidate periodic cell.

Returns

• NeighborListManager — A distinct manager with one neighbor list built for the candidate state. This manager does not mutate the current manager.

commit_cell_candidate

def commit_cell_candidate(candidate: NeighborListManager) -> None

Replace the current cell-bound state with a compatible candidate.

Parameters

Name	Type	Default	Description
candidate	NeighborListManager		Isolated candidate returned by build_cell_candidate.

Returns

• None

Raises

• ValueError — If the candidate uses different neighbor-list policy.

needs_rebuild

def needs_rebuild(positions) -> bool

Return true when positions have moved too far from the reference frame.

Parameters

Name	Type	Default	Description
positions

Returns

• bool

rebuild

def rebuild(positions) -> NeighborList

Force a neighbor-list rebuild from current positions.

Parameters

Name	Type	Default	Description
positions

Returns

• NeighborList

update

def update(positions) -> NeighborList

Return a current neighbor list, rebuilding if needed.

Parameters

Name	Type	Default	Description
positions

Returns

• NeighborList

NeighborTiles

class NeighborTiles
    def __init__(atom_blocks: mx.array, tile_blocks: mx.array, member_mask: mx.array, exact_pair_count: int, raw_candidate_count: int, force_columns: mx.array | None = None, force_group_starts: mx.array | None = None, force_group_counts: mx.array | None = None, generation: int = 0, block_size: int = DEFAULT_MLX_CELL_TILE_BLOCK_SIZE)

Exact Verlet membership encoded over fixed four-atom block tiles.

The representation is geometry-only. Each bit in member_mask records one pair that was inside cutoff + skin at rebuild time. Empty tiles are omitted. Each force-column descriptor packs its four membership bits above the tile and right-column index so the recurring Metal force kernel does not reload the tile membership word. Materializing explicit pairs remains an opt-in diagnostic action.

Parameters

Name	Type	Default	Description
atom_blocks	mx.array
tile_blocks	mx.array
member_mask	mx.array
exact_pair_count	int
raw_candidate_count	int
force_columns	mx.array | None	None
force_group_starts	mx.array | None	None
force_group_counts	mx.array | None	None
generation	int	0
block_size	int	DEFAULT_MLX_CELL_TILE_BLOCK_SIZE

Properties

• active_column_count int — Number of non-empty right-atom columns in the direct-force schedule.
• block_count int — Number of fixed-width atom blocks.
• estimated_bytes int — Estimated persistent bytes for block, tile, and membership arrays.
• force_group_count int — Number of same-left-block column groups in the direct-force schedule.
• lanes_per_tile int — Number of atom-pair lanes in one padded tile.
• mask_word_count int — Number of 32-bit membership words stored for each tile.
• padded_lane_count int — Number of scheduled tile lanes including inactive padding.
• padding_waste_count int — Number of scheduled lanes that are not exact Verlet members.
• padding_waste_fraction float — Fraction of scheduled tile lanes outside exact Verlet membership.
• scheduled_column_count int — Number of dispatched SIMD column lanes including group-tail padding.
• tile_count int — Number of retained non-empty block-pair tiles.

Methods

materialize_pairs

def materialize_pairs(*, sort: bool = True) -> mx.array

Decode exact Verlet members for diagnostics and tests.

Parameters

Name	Type	Default	Description
sort	bool	True	Whether to return canonical lexicographic pair order.

Returns

• mx.array — Exact unique atom pairs with shape (exact_pair_count, 2).

Functions

build_neighbor_list

def build_neighbor_list(positions, cell: Cell, *, cutoff: float, skin: float = 0.3, sort_pairs: bool = True, max_workers: int | None = None, backend: NeighborBackend = 'periodic_cell_list', max_mlx_dense_atoms: int = DEFAULT_MLX_DENSE_PAIR_LIMIT, block_size: int = DEFAULT_MLX_CELL_BLOCK_SIZE, generation: int = 0) -> NeighborList

Build a periodic cell-list neighbor list with unique i < j pairs.

Parameters

Name	Type	Default	Description
positions			Particle coordinates.
cell	Cell		Periodic cell owning the minimum-image convention.
cutoff	float		Interaction cutoff distance.
skin	float	0.3	Verlet-list skin added to the cutoff. Defaults to 0.3.
sort_pairs	bool	True	Order the compact pairs by left then right atom. Defaults to True.
max_workers	int | None	None	Worker count for the threaded CPU builder. Defaults to None.
backend	NeighborBackend	'periodic_cell_list'	Neighbor construction backend. Defaults to "periodic_cell_list".
max_mlx_dense_atoms	int	DEFAULT_MLX_DENSE_PAIR_LIMIT	Largest atom count the dense MLX pair backend accepts. Defaults to DEFAULT_MLX_DENSE_PAIR_LIMIT.
block_size	int	DEFAULT_MLX_CELL_BLOCK_SIZE	Atom-block width for block-structured backends. Defaults to DEFAULT_MLX_CELL_BLOCK_SIZE.
generation	int	0	Rebuild generation stamped onto tile geometry, so a caller never has to rebuild the tiles through dataclasses.replace and repeat their schedule validation. Defaults to 0.

Returns

• NeighborList — The constructed neighbor list.

validate_neighbor_backend

def validate_neighbor_backend(backend: str) -> NeighborBackend

Validate and normalize a neighbor-list construction backend.

Parameters

Name	Type	Default	Description
backend	str

Returns

• NeighborBackend

validate_neighbor_check_backend

def validate_neighbor_check_backend(backend: str) -> NeighborCheckBackend

Validate and normalize a neighbor-list displacement check backend.

Parameters

Name	Type	Default	Description
backend	str

Returns

• NeighborCheckBackend