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

    @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: ...