kups.core.neighborlist.cell_list ¶

class CellListLookup(Protocol):
    """Cached occupant lookup, independent of any per-particle payload type."""

    def select(self, ctx: PipelineContext) -> Candidates: ...

@dataclass
class CellListNeighborList:
    """Neighbor list using spatial hashing with cell lists.

    This is the recommended implementation when the cutoff is much smaller than
    the box size. It divides space into a grid of cells and only checks pairs in
    neighboring cells.

    Honors the cell's per-axis ``periodic`` mask: stencil offsets that cross a
    non-periodic face are routed to an out-of-bounds bin (no key matches), and
    minimum-image shifts are zero on non-periodic axes.

    ``cache`` optionally supplies an initialized ``CellListCache`` for the
    current particles and cutoff. It retains occupant slots and neighboring
    cells between calls; accepted moves must update it along with the particles.
    Rebuild it when the cell, cutoff, or particle row layout changes.
    Without a cache, the same grid algorithm joins particle and query cell hashes.

    Candidate storage is O(N) for well-distributed particles at fixed density
    and cutoff. Uncached lookup sorts cell hashes; cached lookup gathers slots.
    Efficiency improves as cutoff/box ratio decreases.

    Attributes:
        avg_candidates: Capacity for candidate pair storage (from cell list).
        avg_edges: Capacity for final edge array.
        cells: Capacity for cell hash table (grows with box_size³/cutoff³).
        avg_image_candidates: Capacity for image candidate pairs.
        cutoffs: Per-system cutoff distances used by this neighbor list.
        cache: Current cell occupants, or ``None`` to rebuild the hash lookup.

    Algorithm:
        1. Partition space into grid cells of size ~cutoff
        2. Hash each particle to its cell
        3. For each particle, check only neighboring 27 cells (3D)
        4. Filter candidates by actual distance

    When to use:
        - When cutoff/box_size << 1 (cutoff much smaller than box)
        - Typically cutoff/box < 0.3 for good efficiency
        - Nonperiodic outliers remain searchable in boundary bins; many outliers
          can increase the candidate count.

    Example:
        ```python
        # Example: 10 Å cutoff in 50 Å box → cutoff/box = 0.2 -- Good for CellList
        nl = CellListNeighborList.new(state, lens(lambda s: s.nl_params), cutoffs)

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

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

    avg_candidates: Capacity[int]
    avg_edges: Capacity[int]
    cells: Capacity[int]
    avg_image_candidates: Capacity[int]
    cutoffs: Table[SystemId, Array]
    cache: CellListLookup | None = field(default=None, kw_only=True)

    @classmethod
    def new[S](
        cls,
        state: S,
        lens: Lens[S, IsCellListParams],
        cutoffs: Table[SystemId, Array],
    ) -> CellListNeighborList:
        params = lens.get(state)
        return CellListNeighborList(
            avg_candidates=LensCapacity(
                params.avg_candidates, lens.focus(lambda x: x.avg_candidates)
            ),
            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),
            ),
            cells=LensCapacity(params.cells, lens.focus(lambda x: x.cells), base=1),
            cutoffs=cutoffs,
        )

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

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

    @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]]:
        return build_radius_graph(
            self, keys, systems, queries=queries, queried_keys=queried_keys
        )

@dataclass
class CellListSelector:
    """Selector for the cell-list algorithm.

    Joins cell hashes, or reads occupants from an optional initialized cache,
    then replicates per image multiplicity when ``max(cutoff/perp) > 0.5``.
    Set ``promise_unique_cells`` when the caller guarantees at least three cells
    along every periodic axis,
    so each query's stencil cells are already distinct; this skips the
    per-query cell deduplication in the uncached lookup.
    """

    cutoffs: Table[SystemId, Array]
    max_cells: Capacity[int]
    max_candidates: Capacity[int]
    max_image_candidates: Capacity[int]
    promise_unique_cells: bool = field(default=False, static=True)
    cache: CellListLookup | None = field(default=None, kw_only=True)

    def __call__(self, ctx: PipelineContext) -> CandidateBatch[Literal[2]]:
        query = ctx.query_table
        candidates = (
            self.cache.select(ctx)
            if self.cache is not None
            else _cell_list_subselect(
                ctx.keys,
                query,
                ctx.systems,
                cutoffs=self.cutoffs.data,
                max_num_cells=self.max_cells,
                max_num_candidates=self.max_candidates,
                promise_unique_cells=self.promise_unique_cells,
            )
        )
        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,
        )

class IsCellListParams(Protocol):
    """Protocol for parameters required by ``CellListNeighborList``."""

    @property
    def avg_candidates(self) -> int: ...
    @property
    def avg_edges(self) -> int: ...
    @property
    def cells(self) -> int: ...
    @property
    def avg_image_candidates(self) -> int: ...