kups.core.neighborlist.adaptive ¶

@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 NeighborListCandidate:
    """A backing neighbor list paired with its cost estimator.

    Attributes:
        neighborlist: The candidate implementation.
        cost: Estimates the implementation's cost for a call's counts.
    """

    neighborlist: NeighborList[Literal[2]]
    cost: NeighborListCost = field(static=True)

class NeighborListCost(Protocol):
    """Estimates the relative runtime cost of a neighbor list for one call.

    Lower is cheaper; :class:`AdaptiveNeighborList` dispatches to the minimum.
    Return ``math.inf`` to mark an implementation invalid for the given shape.
    """

    def __call__(self, num_particles: int, num_systems: int) -> float: ...