kups.core.neighborlist ¶

@dataclass
class AdaptiveNeighborList(NeighborList[Literal[2]]):
    """Neighbor list that dispatches to the cheapest backing implementation.

    Holds ``(implementation, cost)`` pairs and, on each call, picks the
    implementation whose :class:`NeighborListCost` is smallest for that call's
    counts. Augment by appending pairs to :attr:`implementations`; the seeded
    implementations share the ``UniversalNeighborlistParameters`` capacities, so
    growing one grows the shared state.

    Attributes:
        implementations: Candidate :class:`NeighborListCandidate` pairs. Ties
            resolve to the earlier entry.

    Example:
        ```python
        nl = AdaptiveNeighborList.from_state(state, cutoffs)
        edges = nl(particles, systems)
        ```
    """

    implementations: tuple[NeighborListCandidate, ...]

    @classmethod
    def new[S](
        cls,
        state: S,
        lens: Lens[S, IsUniversalNeighborlistParams],
        cutoffs: Table[SystemId, Array],
    ) -> AdaptiveNeighborList:
        """Seed the library implementations with their default cost guesses.

        Args:
            state: Object exposing ``UniversalNeighborlistParameters`` via ``lens``.
            lens: Lens focusing the shared ``IsUniversalNeighborlistParams``.
            cutoffs: Per-system cutoffs bound onto each implementation.

        Returns:
            An ``AdaptiveNeighborList`` over dense, cell-list, and all-dense.
        """
        return cls(
            (
                NeighborListCandidate(
                    DenseNearestNeighborList.new(state, lens, cutoffs), dense_cost
                ),
                NeighborListCandidate(
                    CellListNeighborList.new(state, lens, cutoffs), cell_list_cost
                ),
                NeighborListCandidate(
                    AllDenseNearestNeighborList.new(state, lens, cutoffs),
                    all_dense_cost,
                ),
            )
        )

    @classmethod
    def from_state(
        cls,
        state: IsNeighborListState[IsUniversalNeighborlistParams],
        cutoffs: Table[SystemId, Array],
    ) -> AdaptiveNeighborList:
        """Seed from a state exposing ``neighborlist_params``."""
        return cls.new(state, lens(lambda s: s.neighborlist_params), cutoffs)

    def _choose(self, num_particles: int, num_systems: int) -> NeighborList[Literal[2]]:
        """Return the cheapest implementation for the given counts."""
        return min(
            self.implementations,
            key=lambda candidate: candidate.cost(num_particles, num_systems),
        ).neighborlist

    def pair_candidates(
        self,
        keys: Table[ParticleId, NeighborListPoints],
        systems: Table[SystemId, NeighborListSystems],
        *,
        inclusion: bool,
        exclusion: bool,
    ) -> CandidateBatch[Literal[2]]:
        """Full neighbors with shared group masks and retained image flags.

        Apply group masks shared by every pair term before compaction. Terms
        with differing policies apply their remaining masks during evaluation.
        Uses the same selectors and capacity repair as normal graph construction.
        The selected implementation must satisfy ``SelectableNeighborList``.
        """
        selected = self._choose(keys.size, systems.size)
        if not isinstance(selected, SelectableNeighborList):
            raise TypeError("Pair candidates require a SelectableNeighborList")
        masks: list[Mask[Literal[2]]] = [InBoundsMask()]
        if inclusion:
            masks.append(InclusionMatchMask())
        if exclusion:
            masks.append(ExclusionMask())
        else:
            # Always exclude a particle's zero-shift interaction with itself.
            masks.append(
                lambda batch, ctx: (
                    jnp.not_equal(*Index.match(batch.key_idx, batch.query_idx))
                    | ~batch.is_minimum_image
                )
            )
        masks.append(DistanceCutoffMask(selected.cutoffs))
        batch, keep, ctx = select_candidates(
            keys, systems, selected.selector(keys.size, systems), tuple(masks)
        )
        return ReduceCompactor(selected.avg_edges.multiply(keys.size)).compact_batch(
            keep, batch, ctx
        )

    @overload
    def __call__(
        self,
        keys: Table[ParticleId, NeighborListPoints],
        systems: Table[SystemId, NeighborListSystems],
        *,
        queries: Table[ParticleId, NeighborListPoints],
    ) -> Edges[Literal[2]]: ...
    @overload
    def __call__(
        self,
        keys: Table[ParticleId, NeighborListPoints],
        systems: Table[SystemId, NeighborListSystems],
        *,
        queried_keys: Index[ParticleId] | None = None,
    ) -> Edges[Literal[2]]: ...
    def __call__(
        self,
        keys: Table[ParticleId, NeighborListPoints],
        systems: Table[SystemId, NeighborListSystems],
        *,
        queries: Table[ParticleId, NeighborListPoints] | None = None,
        queried_keys: Index[ParticleId] | None = None,
    ) -> Edges[Literal[2]]:
        nl = self._choose(keys.size, systems.size)
        if queries is not None:
            return nl(keys, systems, queries=queries)
        return nl(keys, systems, queried_keys=queried_keys)

@dataclass
class AllDenseNearestNeighborList:
    """Dense O(N²) neighbor list considering all pairs across all systems.

    This implementation generates all possible particle pairs without spatial
    optimization. It is only suitable for very small systems or testing.

    **Warning**: This crosses system boundaries! Only use for single-system
    simulations. For multiple systems, use
    [DenseNearestNeighborList][kups.core.neighborlist.DenseNearestNeighborList]
    instead.

    Complexity: O(N²) where N is the total number of particles across all systems.

    Attributes:
        avg_edges: Capacity manager for edge array.
        avg_image_candidates: Capacity manager for image candidate pairs.
        cutoffs: Per-system cutoff distances used by this neighbor list.

    Example:
        ```python
        # Construct from state and a lens to the neighbor list parameters:
        nl = AllDenseNearestNeighborList.new(state, lens(lambda s: s.nl_params), cutoffs)

        # Or, if the state implements IsNeighborListState:
        nl = AllDenseNearestNeighborList.from_state(state, cutoffs)

        edges = nl(particles, systems)
        ```
    """

    avg_edges: Capacity[int]
    avg_image_candidates: Capacity[int]
    cutoffs: Table[SystemId, Array]

    @classmethod
    def new[S](
        cls,
        state: S,
        lens: Lens[S, IsAllDenseNeighborListParams],
        cutoffs: Table[SystemId, Array],
    ) -> AllDenseNearestNeighborList:
        params = lens.get(state)
        return AllDenseNearestNeighborList(
            avg_edges=LensCapacity(params.avg_edges, lens.focus(lambda x: x.avg_edges)),
            avg_image_candidates=LensCapacity(
                params.avg_image_candidates,
                lens.focus(lambda x: x.avg_image_candidates),
            ),
            cutoffs=cutoffs,
        )

    @classmethod
    def from_state(
        cls,
        state: IsNeighborListState[IsAllDenseNeighborListParams],
        cutoffs: Table[SystemId, Array],
    ) -> AllDenseNearestNeighborList:
        return cls.new(state, lens(lambda s: s.neighborlist_params), cutoffs)

    def selector(
        self, query_size: int, systems: Table[SystemId, NeighborListSystems]
    ) -> AllDenseSelector:
        """Build the candidate selector shared by graph and pair evaluation."""
        return AllDenseSelector(
            cutoffs=Table.broadcast_to(self.cutoffs, systems),
            max_image_candidates=self.avg_image_candidates.multiply(query_size),
        )

    @overload
    def __call__(
        self,
        keys: Table[ParticleId, NeighborListPoints],
        systems: Table[SystemId, NeighborListSystems],
        *,
        queries: Table[ParticleId, NeighborListPoints],
    ) -> Edges[Literal[2]]: ...
    @overload
    def __call__(
        self,
        keys: Table[ParticleId, NeighborListPoints],
        systems: Table[SystemId, NeighborListSystems],
        *,
        queried_keys: Index[ParticleId] | None = None,
    ) -> Edges[Literal[2]]: ...
    @jit
    def __call__(
        self,
        keys: Table[ParticleId, NeighborListPoints],
        systems: Table[SystemId, NeighborListSystems],
        *,
        queries: Table[ParticleId, NeighborListPoints] | None = None,
        queried_keys: Index[ParticleId] | None = None,
    ) -> Edges[Literal[2]]:
        if keys.data.inclusion.num_labels >= 2:
            logging.warning(
                "AllDenseNearestNeighborList is intended for single-system simulations. "
                "Performance may be degraded when using multiple systems. "
                "Consider using DenseNearestNeighborList or CellListNeighborList instead."
            )
        return build_radius_graph(
            self, keys, systems, queries=queries, queried_keys=queried_keys
        )

@dataclass
class AllDenseSelector:
    """Selector that emits every ``(i, j)`` pair across all systems."""

    cutoffs: Table[SystemId, Array]
    max_image_candidates: Capacity[int]

    def __call__(self, ctx: PipelineContext) -> CandidateBatch[Literal[2]]:
        query = ctx.query_table
        candidates = _all_subselect(ctx.keys, query, ctx.systems)
        candidates = lift_query_candidates(candidates, ctx)
        return replicate_for_images(
            candidates,
            ctx.keys,
            ctx.edge_query_table,
            ctx.systems,
            self.cutoffs,
            self.max_image_candidates,
        )

@dataclass
class CandidateBatch[D: int]:
    """Candidate set of degree ``D`` carried through the pipeline.

    Reuses [`Edges[D]`][kups.core.neighborlist.edges.Edges] for the
    `(indices, shifts)` layout (`indices` shape `(n, D)`,
    `shifts` shape `(n, D-1, 3)`); adds the
    ``is_minimum_image`` flag that
    [`ExclusionMask`][kups.core.neighborlist.masks.ExclusionMask] needs to
    keep non-minimum periodic copies of excluded pairs.

    Attributes:
        edges: Candidate edges (indices + fractional shifts), keyed by
            [`PipelineContext.edge_keys`][kups.core.neighborlist.types.PipelineContext.edge_keys].
        is_minimum_image: ``(n,)`` bool — True where the candidate's shift
            equals the minimum-image shift; False for non-MIC replicated
            copies emitted by selectors that handle PBC image expansion.
    """

    edges: Edges[D]
    is_minimum_image: Array

    @property
    def key_idx(self) -> Index[ParticleId]:
        """Pair-specific: key-side index of shape ``(n,)``. Only meaningful for ``D == 2``."""
        return self.edges.indices[:, 0]

    @property
    def query_idx(self) -> Index[ParticleId]:
        """Pair-specific: query-side index of shape ``(n,)``. Only meaningful for ``D == 2``."""
        return self.edges.indices[:, 1]

class CandidateSelector[D: int](Protocol):
    """Produces a ``CandidateBatch[D]`` from the pipeline context.

    Owns all candidate-set construction, including any PBC image
    replication required when ``max(cutoff/perp_axis) > 0.5``.
    """

    def __call__(self, ctx: PipelineContext) -> CandidateBatch[D]: ...

class CellListCache[Data](NamedTuple):
    """Slot-sorted particle rows with per-cell occupancy and stencil tables.

    Slots ``0 .. n-1`` hold the particles, slots up to ``n_pad`` are padding,
    and slot ``n_pad`` is the inactive sentinel filling empty table entries.
    Inactive rows occupy no cell entry and emit no candidates.

    Attributes:
        rows: Per-slot rows, leaves ``(n_pad + 1, ...)``.
        slot_of_row: Slot of each original particle row, ``(n,)``.
        cells: Slot ids per cell, ``(n_cells + 1, cell_capacity)``, sentinel
            filled; the last row is the always-empty sentinel cell.
        stencil: Neighbor cell ids per cell, ``(n_cells + 1, stencil_width)``,
            invalid entries routed to the sentinel cell.
        bins: Per-system bin counts, ``(n_systems, 3)``.
        block_bounds: Bounds of the active rows in each run of
            ``key_block_size`` slots.
    """

    rows: CellRows[Data]
    slot_of_row: Array
    cells: Array
    stencil: Array
    bins: Array
    block_bounds: BlockBounds

    @property
    def block_size(self) -> int:
        """Slots per spatial block."""
        return self.sentinel_slot // self.block_bounds.lower.shape[1]

    def block_rows(self, blocks: Array) -> CellRows[Data]:
        """Rows of whole blocks, leaves ``(*blocks.shape, block_size, ...)``.

        Block ids must be in bounds; callers mask padded ids themselves.
        """
        size = self.block_size
        n = self.sentinel_slot
        return jax.tree.map(
            lambda x: x[:n].reshape(n // size, size, *x.shape[1:])[blocks],
            self.rows,
        )

    @property
    def sentinel_slot(self) -> int:
        return self.rows.frac.shape[0] - 1

    @property
    def sentinel_cell(self) -> int:
        return self.stencil.shape[0] - 1

    @property
    def max_cells_per_system(self) -> int:
        """Allocated capacity, derived from the stored array shapes."""
        return self.sentinel_cell // self.bins.shape[0]

    def candidates(self, cell: Array) -> Array:
        """Candidate slots of the given cells, ``(..., stencil_width * cell_capacity)``."""
        return self.cells[self.stencil[cell]].reshape(
            *cell.shape,
            self.stencil.shape[1] * self.cells.shape[1],
        )

    def select(self, ctx: PipelineContext) -> Candidates:
        """Look up occupants for ``CellListSelector`` using the stored grid.

        The cache must correspond to ``ctx.keys`` in original row order.
        Candidate ids are mapped back to those rows; the shared selector uses
        original coordinates when computing periodic shifts.
        """
        query = ctx.query_table
        _, cell = assign_cells(
            query.data.positions,
            query.data.system.indices_in(ctx.systems.keys),
            query.data.inclusion.valid_mask,
            self.bins,
            self.max_cells_per_system,
            ctx.systems.data.cell,
        )
        slots = (
            self.candidates(cell) if ctx.keys.size else jnp.zeros((query.size, 0), int)
        )
        row_of_slot = jnp.full(self.sentinel_slot + 1, ctx.keys.size, dtype=int)
        row_of_slot = row_of_slot.at[self.slot_of_row].set(jnp.arange(ctx.keys.size))
        return Candidates(
            Index(ctx.keys.keys, row_of_slot[slots].ravel(), _cls=ctx.keys.cls),
            Index(
                query.keys,
                jnp.repeat(jnp.arange(query.size), slots.shape[1]),
                _cls=query.cls,
            ),
        )

    def candidate_chunks(self, cell: Array, chunk_size: int) -> Array:
        """Group candidate columns into blocks, with a leading chunk axis.

        Columns from every neighboring cell share a block. Entirely unused
        blocks can be skipped without reducing cell occupancy capacity.
        The shape is ``(n_chunks, *cell.shape, stencil_width * chunk_size)``;
        the last block is sentinel-padded.
        """
        if chunk_size < 1:
            raise ValueError("chunk_size must be positive")
        width, capacity = self.stencil.shape[1], self.cells.shape[1]
        indices = self.candidates(cell).reshape(*cell.shape, width, capacity)
        padding = -capacity % chunk_size
        indices = jnp.pad(
            indices,
            ((0, 0),) * (indices.ndim - 1) + ((0, padding),),
            constant_values=self.sentinel_slot,
        )
        n_chunks = (capacity + padding) // chunk_size
        indices = indices.reshape(*cell.shape, width, n_chunks, chunk_size)
        return jnp.moveaxis(indices, -2, 0).reshape(
            n_chunks, *cell.shape, width * chunk_size
        )

    def bin_rows[D](
        self,
        particles: Table[ParticleId, NeighborListPoints],
        systems: Table[SystemId, HasCell[AnyPeriodicity]],
        data: D,
    ) -> CellRows[D]:
        """Bin further particles (e.g. proposal queries) into this table's cells."""
        return cell_rows(particles, systems, self.bins, self.max_cells_per_system, data)

@dataclass
class CellListCacheParameters:
    """Static shapes of a cell table.

    Attributes:
        chunk_size: Query rows processed per local scan step; the
            padded slot count is a multiple of it.
        max_cells_per_system: Capacity for spatial bins per system
            (``prod(bins) <= max_cells_per_system`` is asserted).
        cell_capacity: Maximum active particles per bin (asserted).
        stencil_width: Distinct neighbor cells per cell. 27 is always valid;
            with fewer than three bins along an axis the stencil wraps onto
            duplicates and ``prod(min(bins, 3))`` suffices (asserted).
            Smaller widths shrink every candidate gather.
        key_chunk_size: Optional keys per block: columns of each neighboring cell
            for ``key_layout="cells"``, or particle slots for ``"slots"``.
            Empty blocks skip pair evaluation; ``None`` uses one block.
            This changes traversal, not the occupancy limit.
        key_layout: Read neighboring cells, or traverse all particle slots.
            Slot traversal still applies the exact distance and pair masks.
        max_images_per_pair: Capacity for each pair's periodic-image window.
            Estimated from the cutoff and cell geometry; one in the
            minimum-image regime.
        key_block_size: Slots per spatial block. Slots are ordered by a
            spatial bisection tree at build time and each block keeps
            fractional bounds of its active rows, maintained on updates.
            Local fused evaluations then skip blocks that are provably
            beyond the cutoff of every query and gather only the remaining
            blocks, in batches of ``key_chunk_size`` slots. Exact: skipped
            pairs are outside the cutoff for the evaluated minimum image.
            Culling applies only with ``key_layout="slots"``,
            ``max_images_per_pair == 1`` and at most 64 query rows per
            evaluation, summed over phases and systems
            (``fused.MAX_DENSE_QUERY_ROWS``); otherwise the unculled
            traversal runs and the bounds are maintained unused.
    """

    chunk_size: int = field(static=True, default=32)
    max_cells_per_system: int = field(static=True, default=512)
    cell_capacity: int = field(static=True, default=8)
    stencil_width: int = field(static=True, default=27)
    key_chunk_size: int | None = field(static=True, default=None, kw_only=True)
    key_layout: Literal["cells", "slots"] = field(
        static=True, default="cells", kw_only=True
    )
    max_images_per_pair: int = field(static=True, default=1, kw_only=True)
    key_block_size: int = field(static=True, default=8, kw_only=True)

    def __post_init__(self) -> None:
        if min(self.chunk_size, self.max_cells_per_system, self.cell_capacity) <= 0:
            raise ValueError("Cell table capacities and chunk_size must be positive")
        if self.key_chunk_size is not None and self.key_chunk_size < 1:
            raise ValueError("key_chunk_size must be positive")
        if self.key_layout not in ("cells", "slots"):
            raise ValueError("key_layout must be 'cells' or 'slots'")
        if not 1 <= self.stencil_width <= 27:
            raise ValueError("stencil_width must be between 1 and 27")
        if self.max_images_per_pair < 1:
            raise ValueError("max_images_per_pair must be positive")
        if self.key_block_size < 1:
            raise ValueError("key_block_size must be positive")

    @property
    def slot_multiple(self) -> int:
        """Padded slot counts are multiples of the query chunk and block size."""
        return math.lcm(self.chunk_size, self.key_block_size)

    @classmethod
    @no_jax_tracing
    def estimate(
        cls,
        particles: Table[ParticleId, NeighborListPoints],
        systems: Table[SystemId, HasCell[AnyPeriodicity]],
        cutoffs: Table[SystemId, Array],
        *,
        chunk_size: int = 64,
        occupancy_factor: float = 2.0,
        occupancy_headroom: int = 32,
        key_chunk_size: int | Literal["auto"] | None = None,
        key_layout: Literal["cells", "slots", "auto"] = "cells",
        key_block_size: int = 8,
    ) -> Self:
        """Estimate parameters from a concrete state.

        Sizes ``cell_capacity`` from the exact per-cell occupancy of the
        active particles with multiplicative and additive headroom to absorb
        density growth (e.g. GCMC insertions); violations at runtime fail the
        assertions in [`build_cell_list_cache`][kups.core.neighborlist.cell_list_cache.build_cell_list_cache].

        Args:
            particles: Current (possibly buffered) particle table.
            systems: System table with cells.
            cutoffs: Per-system cutoff radii.
            chunk_size: Query rows per local scan step.
            occupancy_factor: Multiplier on the observed maximum occupancy.
            occupancy_headroom: Additive slack on top of the scaled occupancy.
            key_chunk_size: Keys per block, or ``"auto"`` to round cell occupancy
                up to a multiple of 32. For slot traversal, rounds total occupancy
                plus additive headroom up to a multiple of 128. Later blocks
                remain available as density grows; ``None`` uses one block.
            key_layout: ``"auto"`` selects slots when every system's stencil
                covers all cells and either slot blocks are culled or, for a
                single system, the estimated slot block is smaller than the
                combined cell blocks. Otherwise retains cell traversal.
            key_block_size: Slots per spatial block. Blocks are culled only
                in the minimum-image regime, where ``key_layout="auto"`` then
                prefers slots and ``key_chunk_size="auto"`` sizes gathered
                batches to the blocks near a typical query.

        Returns:
            Static parameters for the table.
        """
        cutoff = Table.broadcast_to(cutoffs, systems).data
        bins, images = _cell_grid(systems.data, cutoff)
        bins_np, images = np.asarray(bins), np.asarray(images)
        max_cells = int(bins_np.prod(axis=-1).max())
        max_occupancy, system_active = _occupancy(
            particles, systems, bins, max_cells=max_cells
        )
        max_occupancy, system_active = int(max_occupancy), np.asarray(system_active)
        active = int(system_active.sum())
        stencil_width = int(np.minimum(bins_np, 3).prod(axis=-1).max())
        cell_chunk_size = max(32, -(-max_occupancy // 32) * 32)
        occupancy = active + occupancy_headroom
        slot_chunk_size = max(128, -(-occupancy // 128) * 128)
        max_images = int(images.prod(axis=-1).max())
        # Slot blocks are culled only in the minimum-image regime.
        culled = max_images == 1
        if key_layout == "auto":
            limits = np.where(np.asarray(systems.data.cell.periodic), 3, 2)
            key_layout = (
                "slots"
                if bool((bins_np <= limits).all())
                and (
                    culled
                    or systems.size == 1
                    and slot_chunk_size < stencil_width * cell_chunk_size
                )
                else "cells"
            )
        if key_chunk_size == "auto":
            key_chunk_size = (
                slot_chunk_size if key_layout == "slots" else cell_chunk_size
            )
            if culled and key_layout == "slots":
                key_chunk_size = _estimate_block_batch(
                    systems, cutoff, system_active, particles.size, key_block_size
                )
        return cls(
            chunk_size=chunk_size,
            key_chunk_size=key_chunk_size,
            key_layout=key_layout,
            max_cells_per_system=max_cells,
            cell_capacity=max(
                1, int(max_occupancy * occupancy_factor) + occupancy_headroom
            ),
            stencil_width=stencil_width,
            max_images_per_pair=max_images,
            key_block_size=key_block_size,
        )