kups.potential.common.graph ¶

@dataclass
class GraphConstructor[
    State,
    Ptch: Patch[Any],
    P: NeighborListPoints,
    S: HasCell[AnyPeriodicity],
    Degree: int,
]:
    """Constructs graphs from a degree-parametrized neighbor list.

    Attributes:
        particles: View extracting ``Indexed[ParticleId, P]``.
        systems: View extracting ``Indexed[SystemId, S]`` (cell).
        neighborlist: View extracting the neighbor list implementation for the state.
        probe: Optional particle probe, which may also supply neighbor lists.
    """

    particles: View[State, Table[ParticleId, P]] = field(static=True)
    systems: View[State, Table[SystemId, S]] = field(static=True)
    neighborlist: View[State, NeighborList[Degree]] = field(static=True)
    probe: Probe[State, Ptch, IsParticleProbe[P]] | None = field(static=True)

    @jit(static_argnames=("old_graph",))
    def __call__(
        self, state: State, patch: Ptch | None, old_graph: bool = False
    ) -> HyperGraph[P, S, Degree]:
        keys = self.particles(state)
        systems = self.systems(state)

        if patch is not None and self.probe is None:
            if old_graph:
                return self(state, None, old_graph=True)
            new_state: State = patch(
                state, systems.set_data(jnp.ones(len(systems), dtype=jnp.bool_))
            )
            return self(new_state, None, old_graph=True)

        if patch is None:
            nnlist = self.neighborlist(state)
            edges = nnlist(keys, systems)
        else:
            assert self.probe is not None, "Expected probe to be set."
            probe = self.probe(state, patch)
            update = probe.particles
            indices = update.indices
            if not old_graph:
                keys = keys.update(indices, update.data)
            if isinstance(probe, IsGraphProbe):
                nnlist = (
                    probe.neighborlist_before if old_graph else probe.neighborlist_after
                )
            else:
                nnlist = self.neighborlist(state)
            edges = nnlist(keys, systems, queried_keys=indices)
        return HyperGraph(keys, systems, edges)

@dataclass
class GraphInputConstructor[
    State,
    Ptch: Patch[Any],
    P: NeighborListPoints,
    S: HasCell[AnyPeriodicity],
    Degree: int,
    Params,
](InputConstructor[State, GraphPotentialInput[Params, P, S, Degree], Ptch]):
    """Construct full or affected graph inputs for the shared sum composers.

    A graph probe restricts proposals to affected interactions, keeping
    parameters from the current state. Without a probe, apply the patch to
    the full state before reading parameters and constructing the graph.
    ``old_input`` selects the current state in either case.
    """

    graph_constructor: GraphConstructor[State, Ptch, P, S, Degree] = field(static=True)
    parameter_view: View[State, Params] = field(static=True)

    def __call__(
        self,
        state: State,
        patch: Ptch | None,
        old_input: bool = False,
    ) -> GraphPotentialInput[Params, P, S, Degree]:
        if patch is not None and self.graph_constructor.probe is None:
            if not old_input:
                systems = self.graph_constructor.systems(state)
                state = patch(
                    state, systems.set_data(jnp.ones(len(systems), dtype=jnp.bool_))
                )
            patch = None
        return GraphPotentialInput(
            self.parameter_view(state),
            self.graph_constructor(state, patch, old_graph=old_input),
        )

@dataclass
class GraphPairEnergy[
    Params,
    Part: HasPositionsAndSystemIndex,
    Feat,
]:
    """Derive edge energies and per-system totals from a pair term.

    Graph edges already encode inclusion, exclusion and periodic-image rules.
    The pair term supplies the features, cutoff and kernel. Directed edges
    contribute half their energy to each system's total. ``fuse_pair_potentials``
    can collect the same term to share geometry and neighbor traversal with
    other pair terms.
    """

    pair: PairEnergy[Params, Part, Feat] = field(static=True)

    @jit
    def edge_energies(
        self,
        inp: GraphPotentialInput[Params, Part, HasCell[AnyPeriodicity], Literal[2]],
    ) -> Array:
        """Evaluate the pair term on each directed edge, ignoring padded edges."""
        keep = jnp.ones(len(inp.graph.edges), dtype=bool)
        return graph_pair_energies(self.pair, inp.parameters, inp.graph, keep, keep)

    def __call__[State](
        self,
        inp: GraphPotentialInput[Params, Part, HasCell[AnyPeriodicity], Literal[2]],
    ) -> WithPatch[Table[SystemId, Energy], IdPatch[State]]:
        system = inp.graph.edge_batch_mask.update_labels(inp.graph.systems.keys).to_cls(
            inp.graph.systems.cls
        )
        energies = system.sum_over(self.edge_energies(inp)) / 2
        return WithPatch(energies, IdPatch[State]())

class GraphPotentialInput(NamedTuple, Generic[Params, Part, Sys, Degree]):
    """Input bundle for graph-based potential energy functions."""

    parameters: Params
    graph: HyperGraph[Part, Sys, Degree]

@dataclass
class HyperGraph(PointCloud[Part, Sys], Generic[Part, Sys, Degree]):
    """Point cloud with edges representing particle interactions.

    Generic in ``Part`` (particle data), ``Sys`` (system data), and
    ``Degree`` (number of particles per edge: 2=pairs, 3=triplets).

    Attributes:
        particles: Inherited -- indexed particle data.
        systems: Inherited -- indexed system data.
        edges: Edge connectivity with ``Index[ParticleId]`` indices and periodic shifts.
    """

    edges: Edges[Degree]

    @property
    def edge_offsets(self) -> Array:
        return self.edges.shifts

    @property
    def edge_shifts(self) -> Array:
        return self.edges.difference_vectors(self.particles, self.systems)

    @property
    def edge_batch_mask(self) -> Index[SystemId]:
        return self.particles[self.edges.indices[:, 0]].system

    @overload
    def sorted_by_system(
        self,
        sort_edges: bool = ...,
        *,
        return_sort_order: Literal[False] = ...,
    ) -> HyperGraph[Part, Sys, Degree]: ...
    @overload
    def sorted_by_system(
        self,
        sort_edges: bool = ...,
        *,
        return_sort_order: Literal[True],
    ) -> tuple[HyperGraph[Part, Sys, Degree], Array]: ...
    def sorted_by_system(
        self,
        sort_edges: bool = False,
        *,
        return_sort_order: bool = False,
    ) -> HyperGraph[Part, Sys, Degree] | tuple[HyperGraph[Part, Sys, Degree], Array]:
        """Sort particles by system index and remap edges accordingly.

        Args:
            sort_edges: If True, also sort edges by the system of their first
                particle.
            return_sort_order: If True, also return the sort permutation.
        """
        n = len(self.particles)
        sort_order = jnp.argsort(self.particles.data.system.raw_indices, stable=True)
        sorted_data = bind(self.particles.data).at((sort_order,)).get()
        sorted_particles = Table(
            self.particles.keys, sorted_data, _cls=self.particles.cls
        )
        # Each particle's new row, so edges follow their endpoints.
        inverse_order = Table(
            self.particles.keys,
            Index(
                self.particles.keys,
                jnp.zeros_like(sort_order).at[sort_order].set(jnp.arange(n)),
                max_count=self.edges.indices.max_count,
                _cls=self.particles.cls,
            ),
        )
        remapped_indices = inverse_order[self.edges.indices]
        sorted_edges = Edges(indices=remapped_indices, shifts=self.edges.shifts)
        if sort_edges:
            edge_systems = sorted_particles[remapped_indices[:, 0]].system
            edge_order = jnp.argsort(edge_systems.raw_indices, stable=True)
            sorted_edges = Edges(
                indices=remapped_indices[edge_order],
                shifts=sorted_edges.shifts[edge_order],
            )
        result = HyperGraph(sorted_particles, self.systems, sorted_edges)
        if return_sort_order:
            return result, sort_order
        return result

class IsGraphInput(Protocol):
    """A potential input exposing a ``graph`` of particles and systems."""

    @property
    def graph(
        self,
    ) -> PointCloud[HasPositionsAndSystemIndex, HasCell[AnyPeriodicity]]: ...

@runtime_checkable
class IsGraphProbe[P: NeighborListPoints, Degree: int](IsParticleProbe[P], Protocol):
    """Particle probe supplying neighbor lists for incremental graph construction."""

    @property
    def neighborlist_after(self) -> NeighborList[Degree]: ...
    @property
    def neighborlist_before(self) -> NeighborList[Degree]: ...

class IsParticleProbe[P: NeighborListPoints](Protocol):
    """Probe result identifying changed particles."""

    @property
    def particles(self) -> WithIndices[ParticleId, P]: ...

@dataclass
class PointCloud(Generic[Part, Sys]):
    """Indexed particles and systems, the base for all graph representations.

    Generic in ``Part`` (particle data with positions and system assignment)
    and ``Sys`` (system data with cell).

    Attributes:
        particles: Indexed particle data with positions and system assignment.
        systems: Indexed system data with cell information.
    """

    particles: Table[ParticleId, Part]
    systems: Table[SystemId, Sys]

    @property
    def batch_size(self) -> int:
        return self.particles.data.system.num_labels