kups.potential.common.fused ¶

@dataclass
class FusedFullInput[Params, Part: NeighborListPoints, Feat](
    FusedPotentialInput[Params, Part, Feat]
):
    """Neighbors with the term's shared group masks applied before autodiff."""

    batch: CandidateBatch[Literal[2]]

@dataclass
class FusedInputConstructor[
    State: IsNeighborListState[IsUniversalNeighborlistParams],
    Ptch: Patch[Any],
    P: NeighborListPoints,
    S: HasCell[AnyPeriodicity],
    Params,
    Feat,
](InputConstructor[State, FusedPotentialInput[Params, P, Feat], Ptch]):
    """Construct inputs for the same full/local sum plans used by graph potentials.

    Full inputs contain compacted geometric neighbors selected before autodiff.
    With a probe, restrict to its incident pairs against an initialized cell
    table and let ``LocalSumComposer`` handle old/new weights and the cached
    total. ``old_input`` selects the current state in either case.
    """

    particles: View[State, Table[ParticleId, P]] = field(static=True)
    systems: View[State, Table[SystemId, S]] = field(static=True)
    parameter_view: View[State, Params] = field(static=True)
    pair: PairTerm[Params, P, Feat] = field(static=True)
    probe: Probe[State, Ptch, WithIndices[ParticleId, P]] | None = field(static=True)
    cell_table: View[State, CellListCache[Feat]] = field(static=True)

    def __call__(
        self,
        state: State,
        patch: Ptch | None,
        old_input: bool = False,
    ) -> FusedPotentialInput[Params, P, Feat]:
        if patch is not None and self.probe is None:
            if not old_input:
                systems = self.systems(state)
                state = patch(state, systems.set_data(jnp.ones(systems.size, bool)))
            patch = None
        params = self.parameter_view(state)
        particles = self.particles(state)
        cloud = PointCloud(particles, self.systems(state))
        if patch is None:
            neighbors = AdaptiveNeighborList.from_state(
                state, self.pair.cutoffs(params)
            )
            return FusedFullInput(
                params,
                cloud,
                neighbors.pair_candidates(
                    particles,
                    cloud.systems,
                    inclusion=self.pair.inclusion,
                    exclusion=self.pair.exclusion,
                ),
            )
        assert self.probe is not None
        update = self.probe(state, patch)
        old = particles.subset(update.indices)
        # Padded proposal targets are inactive regardless of their payload.
        data = lens(lambda p: p.inclusion, cls=type(update.data)).apply(
            update.data, lambda index: index.apply_mask(update.indices.valid_mask)
        )
        new = Table(old.keys, data, _cls=old.cls)
        return FusedLocalInput(
            params,
            cloud,
            (old if old_input else new,),
            update.indices,
            self.cell_table(state),
        )

@dataclass
class FusedLocalInput[Params, Part: NeighborListPoints, Feat](
    FusedPotentialInput[Params, Part, Feat]
):
    """Incident pairs against an initialized table, excluding ``removed``.

    One query table evaluates its incident energy. Two tables evaluate new
    minus old, including intra-query pairs once and no old-new cross pairs.
    Query vocabularies may differ from the cloud's; they are matched by key.
    """

    queries: tuple[Table[ParticleId, Part], ...]
    removed: Index[ParticleId]
    cell_table: CellListCache[Feat]

    def __post_init__(self) -> None:
        if len(self.queries) not in (1, 2):
            raise ValueError("Local evaluation requires one query or an old/new pair")

@dataclass
class FusedNeighborEnergy[State, Params, Part: NeighborListPoints, Feat]:
    """Evaluate shared pair terms over a full graph or local spatial chunks.

    ``pair`` defines the interaction and ``layout`` configures cell storage
    and traversal. A persistent ``cell_table_lens``
    enables acceptance-conditional updates for local proposals. Full calls use
    preselected graph edges; only the compacted interactions are differentiated.

    ``max_queries_per_system`` bounds active rows in each old/new local query
    (asserted). It keeps query-pair storage proportional to the batch size.
    Without a bound, the total query count is used conservatively. Rigid MCMC
    supplies the largest motif size automatically.
    """

    pair: PairTerm[Params, Part, Feat] = field(static=True)
    layout: CellListCacheParameters = field(static=True)
    cell_table_lens: Lens[State, CellListCache[Feat]] | None = field(
        static=True, default=None
    )
    max_queries_per_system: int | None = field(static=True, default=None, kw_only=True)

    def __post_init__(self) -> None:
        if self.max_queries_per_system is not None and self.max_queries_per_system < 1:
            raise ValueError("max_queries_per_system must be positive")

    @jit
    def build_cell_list_cache(
        self,
        parameters: Params,
        cloud: PointCloud[Part, HasCell[AnyPeriodicity]],
    ) -> CellListCache[Feat]:
        return build_cell_list_cache(
            cloud.particles,
            cloud.systems,
            self.pair.cutoffs(parameters),
            self.pair.features(cloud.particles.data),
            self.layout,
        )

    def _energies(
        self,
        inp: FusedLocalInput[Params, Part, Feat],
        rows: CellRows[Feat],
        keys: CellRows[Feat],
        index: Array,
        mask: View[CandidateBatch[Literal[2]], Array],
    ) -> Array:
        batch, ctx = cell_candidates(
            keys,
            rows,
            index,
            inp.cloud.systems,
            self.pair.cutoffs(inp.parameters),
            self.layout.max_images_per_pair,
        )
        multiple_images = self.layout.max_images_per_pair > 1
        pairs = PairBatch.from_candidates(
            batch, ctx, query_lanes=None if multiple_images else index.shape[1]
        )
        pairs = pairs._replace(valid=pairs.valid & mask(batch))
        if multiple_images:
            key_index = batch.key_idx.indices
            query_index = batch.query_idx.indices
            energies = self.pair.evaluate(
                inp.parameters,
                tree_map(itemgetter(query_index), rows.data),
                tree_map(itemgetter(key_index), keys.data),
                pairs,
            )
            return jax.ops.segment_sum(
                energies, query_index, num_segments=rows.frac.shape[0]
            )
        pairs = tree_map(lambda x: x.reshape(*index.shape, *x.shape[1:]), pairs)
        left = tree_map(lambda x: x[:, None], rows.data)
        right = tree_map(lambda x: x[index], keys.data)
        return self.pair.evaluate(inp.parameters, left, right, pairs).sum(-1)

    def _sum_keys(
        self,
        inp: FusedLocalInput[Params, Part, Feat],
        rows: CellRows[Feat],
        excluded: Array,
    ) -> Array:
        table = inp.cell_table

        def consume(index: Array) -> Array:
            return self._energies(
                inp, rows, table.rows, index, lambda b: ~excluded[b.key_idx.indices]
            )

        chunk_size = self.layout.key_chunk_size
        if self.layout.key_layout == "slots":
            n = table.sentinel_slot
            chunk_size = min(chunk_size or n, n) or 1
            index = jnp.minimum(jnp.arange(-(-n // chunk_size) * chunk_size), n)
            return sum_chunks(
                lambda keys: consume(
                    jnp.broadcast_to(keys, (rows.frac.shape[0], chunk_size))
                ),
                index,
                chunk_size,
                active=lambda keys: jnp.any(
                    table.rows.cell[keys] != table.sentinel_cell
                ),
            )
        if chunk_size is None or table.cells.shape[1] <= chunk_size:
            return consume(table.candidates(rows.cell))
        return sum_chunks(
            lambda block: consume(block[0]),
            table.candidate_chunks(rows.cell, chunk_size),
            1,
            active=lambda block: jnp.any(block != table.sentinel_slot),
        )

    def _sum_environment(
        self,
        inp: FusedLocalInput[Params, Part, Feat],
        rows: CellRows[Feat],
        excluded: Array,
        weights: Array,
        query_chunk_size: int,
    ) -> Array:
        def consume(chunk: tuple[CellRows[Feat], Array]) -> Array:
            rows, weights = chunk
            energies = self._sum_keys(inp, rows, excluded)
            return jax.ops.segment_sum(
                energies * weights,
                rows.system.indices,
                num_segments=inp.cloud.systems.size,
            )

        n = rows.frac.shape[0]
        chunk_size = (
            min(n, self.layout.chunk_size, query_chunk_size)
            if n
            else self.layout.chunk_size
        )
        padding = -n % chunk_size
        if padding:
            rows = rows.pad(padding, sentinel_cell=inp.cell_table.sentinel_cell)
            weights = jnp.pad(weights, (0, padding))
        return sum_chunks(
            consume,
            (rows, weights),
            chunk_size,
            active=lambda chunk: jnp.any(chunk[0].inclusion.valid_mask),
        )

    def _full(self, inp: FusedFullInput[Params, Part, Feat]) -> Array:
        graph = HyperGraph(inp.cloud.particles, inp.cloud.systems, inp.batch.edges)
        particles = graph.particles[graph.edges.indices]
        inclusion = exclusion = jnp.ones(len(graph.edges), dtype=bool)
        if not self.pair.inclusion:
            inclusion = (
                particles.inclusion.indices[:, 0] == particles.inclusion.indices[:, 1]
            )
        if not self.pair.exclusion:
            exclusion = (
                particles.exclusion.indices[:, 0] != particles.exclusion.indices[:, 1]
            ) | ~inp.batch.is_minimum_image
        energies = graph_pair_energies(
            self.pair, inp.parameters, graph, inclusion, exclusion
        )
        system = graph.edge_batch_mask.update_labels(graph.systems.keys).to_cls(
            graph.systems.cls
        )
        return system.sum_over(energies).data / 2

    def _local(
        self,
        inp: FusedLocalInput[Params, Part, Feat],
    ) -> tuple[Array, CellRows[Feat], Array]:
        table = inp.cell_table
        queries = inp.queries
        parts = [
            table.bin_rows(q, inp.cloud.systems, self.pair.features(q.data))
            for q in queries
        ]
        rows = CellRows.concatenate(*parts)
        # The same phase labels determine energy signs and exclude old-new pairs.
        phase = jnp.concatenate(
            [jnp.full(q.size, i, dtype=int) for i, q in enumerate(queries)]
        )
        weights = jnp.where((len(queries) == 2) & (phase == 0), -1, 1)
        removed_rows = inp.removed.indices_in(inp.cloud.particles.keys)
        slots = table.slot_of_row.at[removed_rows].get(
            mode="fill",
            fill_value=table.sentinel_slot,
        )
        removed = jnp.zeros(table.sentinel_slot + 1, bool).at[slots].set(True)
        n = rows.frac.shape[0]
        environment = self._sum_environment(
            inp,
            rows,
            removed,
            weights,
            # Separate halves let CPU skip empty work; GPU favors larger chunks.
            n
            if jax.default_backend() in {"gpu", "cuda", "rocm"}
            else max(q.size for q in queries),
        )
        # Keep old and new query pairs separate: no old-new cross interactions.
        n_systems = inp.cloud.systems.size
        system = Index(
            tuple(range(len(queries) * n_systems)),
            rows.system.indices + phase * n_systems,
        ).apply_mask(rows.inclusion.valid_mask)
        width = min(max(q.size for q in queries), self.max_queries_per_system or n)
        runtime_assert(
            (system.counts.data <= width).all(),
            "max_queries_per_system exceeded in fused local evaluation.",
        )
        query_index = system.where_rectangular(system, width)
        query_index = jnp.where(system.valid_mask[:, None], query_index, n)
        pairs = self._energies(
            inp,
            rows,
            rows,
            query_index,
            lambda b: (b.key_idx.indices != b.query_idx.indices) | ~b.is_minimum_image,
        )
        values = environment + jax.ops.segment_sum(
            pairs * weights / 2,
            rows.system.indices,
            num_segments=inp.cloud.systems.size,
        )
        return values, parts[-1], slots

    @jit
    def __call__(
        self,
        inp: FusedPotentialInput[Params, Part, Feat],
    ) -> WithPatch[Table[SystemId, Energy], Patch[State]]:
        keys = inp.cloud.systems.keys
        if isinstance(inp, FusedFullInput):
            return WithPatch(Table(keys, self._full(inp)), IdPatch[State]())
        if not isinstance(inp, FusedLocalInput):
            raise TypeError("Construct a FusedFullInput or FusedLocalInput")
        patch: Patch[State] = IdPatch[State]()
        values, rows, slots = self._local(inp)
        if self.cell_table_lens is not None:
            patch = CellListCacheUpdatePatch(
                slots,
                inp.queries[-1].data.system,
                rows,
                self.cell_table_lens,
            )
        return WithPatch(Table(keys, values), patch)

@dataclass
class FusedPotentialCache[Feat]:
    """Persistent neighbor table, combined energy, and the layout used to build it."""

    table: CellListCache[Feat]
    energy: KahanSummand[PotentialOut[EmptyType, EmptyType]]
    layout: CellListCacheParameters = field(static=True)

    @staticmethod
    def create[Params, Part: NeighborListPoints, Data](
        pair: PairTerm[Params, Part, Data],
        parameters: Params,
        cloud: PointCloud[Part, HasCell[AnyPeriodicity]],
        layout: CellListCacheParameters | None = None,
    ) -> FusedPotentialCache[Data]:
        """Build a shared table and empty energy cache before creating a state."""
        if layout is None:
            backend = jax.default_backend()
            on_cpu = backend == "cpu"
            on_gpu = backend in {"gpu", "cuda", "rocm"}
            layout = CellListCacheParameters.estimate(
                cloud.particles,
                cloud.systems,
                pair.cutoffs(parameters),
                key_chunk_size="auto" if on_cpu or on_gpu else None,
                key_layout="auto" if on_cpu else "cells",
            )
            if on_gpu:
                layout = bind(layout, lambda x: x.chunk_size).set(
                    _gpu_query_chunk_size(layout, cloud.particles.size)
                )
        table = FusedNeighborEnergy(pair, layout).build_cell_list_cache(
            parameters, cloud
        )
        energy = KahanSummand.init(
            PotentialOut(
                cloud.systems.set_data(jnp.zeros(cloud.systems.size)), EMPTY, EMPTY
            )
        )
        return FusedPotentialCache(table, energy, layout)

@dataclass
class FusedPotentialInput[Params, Part: NeighborListPoints, Feat]:
    """Shared parameters and geometry of initialized full/local pair inputs."""

    parameters: Params
    cloud: PointCloud[Part, HasCell[AnyPeriodicity]]