kups.potential.classical.blocking ¶

@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 isinstance(self.radii, Array):
            assert (*self.radii.shape, 3) == self.positions.shape

@dataclass
class BlockingSpheresPotentialInput[
    Points: _BlockingParticles,
    MaybeUnitCell: UnitCell | None,
]:
    """Input for blocking spheres energy calculation.

    Attributes:
        parameters: Blocking sphere positions and radii
        particles: Indexed particle data with positions, system, and motif index
        unitcell: Optional unit cell for periodic boundary conditions
        edges: Particle-sphere pairs to check for blocking
    """

    parameters: BlockingSpheresParameters
    particles: Table[ParticleId, Points]
    unitcell: MaybeUnitCell
    edges: Edges[Literal[2]]

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

    Attributes:
        particles_view: Extracts indexed particle data 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, Points]] = field(static=True)
    systems_view: View[State, Table[SystemId, S]] = 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)

        # Build sphere rh as Indexed[ParticleId, _BlockingSpherePoints]
        p = parameters.positions.shape[0]
        sphere_inclusion = Index(
            tuple(InclusionId(lab) for lab in parameters.system.keys),
            parameters.system.indices,
        )
        sphere_exclusion = Index(
            tuple(ExclusionId(-1 - i) for i in range(p)),
            jnp.arange(p),
        )
        spheres = Table.arange(
            _BlockingSpherePoints(
                positions=parameters.positions,
                system=parameters.system,
                inclusion=sphere_inclusion,
                exclusion=sphere_exclusion,
            ),
            label=ParticleId,
        )

        edges = neighborlist(particles, spheres, systems, cutoffs)
        unitcell = systems.data.unitcell
        return Sum(
            Summand(
                BlockingSpheresPotentialInput(parameters, particles, unitcell, 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 systems(self) -> Table[SystemId, HasUnitCell]: ...
    @property
    def blocking_spheres_parameters(self) -> BlockingSpheresParameters: ...
    @property
    def blocking_spheres_neighborlist(self) -> NearestNeighborList: ...