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):
        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]]]):
        """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 = []
        positions = []
        system = []
        motif = []
        for sys_idx, sys_spheres in enumerate(data):
            for motif_idx, motif_spheres in enumerate(sys_spheres):
                for sphere in motif_spheres:
                    radii.append(sphere.radius)
                    positions.append(sphere.center)
                    system.append(SystemId(sys_idx))
                    motif.append(MotifId(motif_idx))
        radii = jnp.array(radii)
        positions = jnp.array(positions).reshape(-1, 3)
        system = Index.new(system, label=SystemId).populate_max_count()
        motif = Index.new(motif, 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]
    edges: Edges[Literal[2]]

@dataclass
class BlockingSpheresSumComposer[State, Ptch: Patch](
    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 neighbor list instance from state
        probe: Probe providing a IsBlockingSpheresProbe
    """

    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]] = field(static=True)
    parameters_view: View[State, BlockingSpheresParameters] = field(static=True)
    neighborlist_view: View[State, NearestNeighborList] = field(static=True)
    probe: Probe[State, Ptch, IsBlockingSpheresProbe] | None = field(static=True)

    def __call__(self, state: State, patch: Ptch | None):  # type: ignore[reportReturnType]
        particles = self.particles_view(state)
        systems = self.systems_view(state)
        parameters = self.parameters_view(state)
        neighborlist = self.neighborlist_view(state)

        if patch is not None and self.probe is not None:
            n_sys = particles.data.system.num_labels
            patched_state = patch(
                state, systems.set_data(jnp.ones((n_sys,), dtype=jnp.bool_))
            )
            probe_result = self.probe(state, patch)
            neighborlist = probe_result.neighborlist
            particles = self.particles_view(patched_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
        nnlist_particles = particles.map_data(
            lambda p: _BlockingSpherePoints(
                positions=p.positions,
                system=(sys := p.system.apply_mask(p.group.valid_mask)),
                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)
        # Let's just pick negative exclusion IDs for spheres to avoid any possible overlap with particle exclusion IDs
        sphere_exclusion = Index.new(tuple(ExclusionId(-1 - i) for i in range(p)))
        spheres = Table.arange(
            _BlockingSpherePoints(
                positions=parameters.positions,
                system=parameters.system,
                inclusion=sphere_inclusion,
                exclusion=sphere_exclusion,
            ),
            label=ParticleId,
        )

        edges = neighborlist(nnlist_particles, spheres, systems, cutoffs)
        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.

    Bundles changed particle indices with the updated neighbor list,
    enabling efficient re-evaluation when only a subset of particles move.
    """

    @property
    def changed_particle_idx(self) -> Array: ...
    @property
    def neighborlist(self) -> NearestNeighborList: ...

class IsBlockingSpheresState(Protocol):
    """Protocol for states providing all inputs for the blocking spheres potential."""

    @property
    def particles(self) -> Table[ParticleId, _BlockingParticles]: ...
    @property
    def groups(self) -> Table[GroupId, HasMotifIndex]: ...
    @property
    def systems(self) -> Table[SystemId, HasCell]: ...
    @property
    def blocking_spheres_parameters(self) -> BlockingSpheresParameters: ...
    @property
    def blocking_spheres_neighborlist(self) -> NearestNeighborList: ...