kups.potential.classical.blocking ¶

class BlockingSpheresConfig(Protocol):
    """Protocol for a single blocking sphere description (center and radius)."""

    @property
    def center(self) -> tuple[float, float, float]: ...
    @property
    def radius(self) -> float: ...

@dataclass
class BlockingSpheresParameters:
    """Parameters defining blocking sphere positions and radii.

    Attributes:
        radii: Sphere radii, shape `(n_spheres,)`
        positions: Sphere centers, shape `(n_spheres, 3)`
        system: System assignment per sphere
        motif: Motif assignment per sphere
    """

    radii: Array
    positions: Array
    system: Index[SystemId]
    motif: Index[MotifId]

    def __post_init__(self) -> None:
        if not isinstance(self.radii, Array):
            return
        assert (*self.radii.shape, 3) == self.positions.shape, (
            f"Positions shape {self.positions.shape} must match radii shape {self.radii.shape} with last dimension 3"
        )
        assert self.radii.shape == self.system.shape == self.motif.shape, (
            f"Radii, system, and motif must have the same shape: {self.radii.shape}, {self.system.shape}, {self.motif.shape}"
        )

    @staticmethod
    def from_data(
        data: Sequence[Sequence[Sequence[BlockingSpheresConfig]]],
    ) -> BlockingSpheresParameters:
        """Build parameters from a nested sequence of sphere configurations.

        Args:
            data: Nested sequence indexed as ``data[system_idx][motif_idx][sphere_idx]``,
                yielding the sphere configurations for each (system, motif) pair.

        Returns:
            BlockingSpheresParameters with flattened arrays of radii, positions,
            system assignments, and motif assignments.
        """
        radii_list: list[float] = []
        positions_list: list[tuple[float, float, float]] = []
        system_list: list[SystemId] = []
        motif_list: list[MotifId] = []
        for sys_idx, sys_spheres in enumerate(data):
            for motif_idx, motif_spheres in enumerate(sys_spheres):
                for sphere in motif_spheres:
                    radii_list.append(sphere.radius)
                    positions_list.append(sphere.center)
                    system_list.append(SystemId(sys_idx))
                    motif_list.append(MotifId(motif_idx))
        radii = jnp.array(radii_list)
        positions = jnp.array(positions_list).reshape(-1, 3)
        system = Index.new(system_list, label=SystemId).populate_max_count()
        motif = Index.new(motif_list, label=MotifId).populate_max_count()
        return BlockingSpheresParameters(radii, positions, system, motif)

@dataclass
class BlockingSpheresPotentialInput:
    """Input for blocking spheres energy calculation.

    Attributes:
        parameters: Blocking sphere positions and radii
        particles: Indexed particle data with positions, system, and group index
        groups: Indexed group data providing the motif index for each group
        cell: Per-system cell used for periodic boundary wrapping
        edges: Particle-sphere pairs to check for blocking
    """

    parameters: BlockingSpheresParameters
    particles: Table[ParticleId, _BlockingParticles]
    groups: Table[GroupId, HasMotifIndex]
    cell: Table[SystemId, Cell[AnyPeriodicity]]
    edges: Edges[Literal[2]]

@dataclass
class BlockingSpheresSumComposer[State, Ptch: Patch[Any]](
    SumComposer[State, BlockingSpheresPotentialInput, Ptch]
):
    """Composer for blocking spheres potential in energy summation.

    Attributes:
        particles_view: Extracts indexed particle data from state
        groups_view: Extracts indexed group data (motif assignments) from state
        systems_view: Extracts indexed systems from state
        parameters_view: Extracts blocking sphere parameters from state
        neighborlist_view: Extracts a factory that binds cutoffs to a neighbor list
        probe: Probe providing a IsBlockingSpheresProbe; without it a patched
            state is evaluated over every particle
    """

    particles_view: View[State, Table[ParticleId, _BlockingParticles]] = field(
        static=True
    )
    groups_view: View[State, Table[GroupId, HasMotifIndex]] = field(static=True)
    systems_view: View[State, Table[SystemId, HasCell[AnyPeriodicity]]] = field(
        static=True
    )
    parameters_view: View[State, BlockingSpheresParameters] = field(static=True)
    neighborlist_view: View[State, BlockingSpheresNeighborListFactory] = field(
        static=True
    )
    probe: Probe[State, Ptch, IsBlockingSpheresProbe] | None = field(static=True)

    def __call__(
        self, state: State, patch: Ptch | None
    ) -> Sum[BlockingSpheresPotentialInput]:  # type: ignore[reportReturnType]
        changed: Index[ParticleId] | None = None
        if patch is not None:
            # The patch is always applied; a probe only narrows the sphere query
            # to the particles it touches, otherwise everything is re-evaluated.
            if self.probe is not None:
                changed = self.probe(state, patch).particles.indices
            systems = self.systems_view(state)
            state = patch(
                state, systems.set_data(jnp.ones(len(systems), dtype=jnp.bool_))
            )

        particles = self.particles_view(state)
        if changed is not None:
            particles = Table.arange(particles[changed], label=ParticleId)
            queried = particles.data.group.valid_mask & changed.valid_mask
        else:
            queried = particles.data.group.valid_mask
        systems = self.systems_view(state)
        parameters = self.parameters_view(state)
        neighborlist_factory = self.neighborlist_view(state)

        # Build cutoffs: remap sphere system indices into systems index space
        seg_ids = parameters.system.indices_in(tuple(systems.keys))
        max_radii = jax.ops.segment_max(parameters.radii, seg_ids, len(systems.keys))
        cutoffs = Table(systems.keys, max_radii)

        nnlist_particles = particles.map_data(
            lambda p: _BlockingSpherePoints(
                positions=p.positions,
                system=(sys := p.system.apply_mask(queried)),
                inclusion=sys.to_cls(InclusionId),
                exclusion=Index.arange(len(sys), label=ExclusionId),
            )
        )

        # Build sphere rh as Indexed[ParticleId, _BlockingSpherePoints]
        p = parameters.positions.shape[0]
        sphere_inclusion = parameters.system.to_cls(InclusionId)
        sphere_exclusion = Index.new(
            tuple(ExclusionId(len(particles) + i) for i in range(p))
        )
        spheres = Table.arange(
            _BlockingSpherePoints(
                positions=parameters.positions,
                system=parameters.system,
                inclusion=sphere_inclusion,
                exclusion=sphere_exclusion,
            ),
            label=ParticleId,
        )

        neighborlist = neighborlist_factory(cutoffs)
        edges = neighborlist(nnlist_particles, systems, queries=spheres)
        cell = systems.map_data(lambda s: s.cell)
        groups = self.groups_view(state)
        return Sum(
            Summand(
                BlockingSpheresPotentialInput(
                    parameters, particles, groups, cell, edges
                )
            )
        )

class IsBlockingSpheresProbe(Protocol):
    """Probe result for blocking spheres incremental updates.

    Exposes the particles a patch touches, so the sphere query can be
    restricted to them instead of sweeping the whole configuration.
    """

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