kups.mcmc.moves ¶

@dataclass
class ExchangeChanges:
    """Combined particle and group changes for a GCMC exchange move.

    Each field pairs an ``Index`` of target slot ids with a ``Buffered``
    holding the new data and an occupation mask.  For insertions the
    occupation is ``True`` (slots become occupied); for deletions it is
    ``False`` (slots are freed).

    Attributes:
        particles: Target particle slots and buffered particle data.
        groups: Target group slots and buffered group data.
    """

    particles: WithIndices[ParticleId, Buffered[ParticleId, ExchangeParticleData]]
    groups: WithIndices[GroupId, Buffered[GroupId, ExchangeGroupData]]

@dataclass
class ExchangeGroupData:
    """Per-group payload for GCMC exchange moves (no identity info).

    Attributes:
        motif: Molecular species of each group.
        system: System assignment of each group.
    """

    motif: Index[MotifId]
    system: Index[SystemId]

@dataclass
class ExchangeMove[State](MonteCarloMove[State, ExchangeChanges]):
    """Grand canonical Monte Carlo (GCMC) insertion/deletion move. Satisfies ``ChangesFn``.

    Randomly proposes either insertion or deletion of a molecular group
    with 50% probability each via ``propose_mixed``.

    Attributes:
        positions: Lens to buffered particle positions.
        groups: Lens to buffered molecular groups.
        motifs: Lens to molecular templates available for insertion.
        unitcell: Lens to unit cell parameters.
        capacity: Lens to capacity constraints.
    """

    positions: View[State, Buffered[ParticleId, HasPositionsGroupSystem]] = field(
        static=True
    )
    groups: View[State, Buffered[GroupId, HasMotifAndSystemIndex]] = field(static=True)
    motifs: View[State, Table[MotifParticleId, IsMotifData]] = field(static=True)
    unitcell: View[State, Table[SystemId, UnitCell]] = field(static=True)
    capacity: View[State, Capacity[int]] = field(static=True)

    def _zero_ratio(self, state: State) -> LogProbabilityRatio:
        n_sys = self.positions(state).data.system.num_labels
        return Table.arange(jnp.zeros((n_sys,)), label=SystemId)

    def _propose_insertion(
        self, key: Array, state: State, /
    ) -> tuple[ExchangeChanges, LogProbabilityRatio]:
        changes = insert_random_motif(
            key,
            self.motifs(state),
            self.positions(state),
            self.groups(state),
            self.unitcell(state),
            self.capacity(state),
        )
        return changes, self._zero_ratio(state)

    def _propose_deletion(
        self, key: Array, state: State, /
    ) -> tuple[ExchangeChanges, LogProbabilityRatio]:
        changes = delete_random_motif(
            key,
            self.motifs(state),
            self.positions(state),
            self.groups(state),
            self.capacity(state),
        )
        return changes, self._zero_ratio(state)

    def __call__(
        self, key: Array, state: State, /
    ) -> tuple[ExchangeChanges, LogProbabilityRatio]:
        chain = key_chain(key)
        changes, log_ratio, _ = propose_mixed(
            next(chain), state, (self._propose_insertion, self._propose_deletion)
        )
        return changes, log_ratio

@dataclass
class ExchangeParticleData:
    """Per-particle payload for GCMC exchange moves (no identity info).

    The ``motif`` field indexes into the motif template
    (``Indexed[MotifParticleId, MotifParticles]``) so the patch function
    can look up masses, charges, labels, etc.

    Attributes:
        new_positions: Proposed coordinates, shape ``(n, 3)``.
        group: Group assignment for each particle.
        system: System assignment for each particle.
        motif: Motif-template index for each particle.
    """

    new_positions: Array
    group: Index[GroupId]
    system: Index[SystemId]
    motif: Index[MotifParticleId]

@dataclass
class GroupRotationMove[State](MonteCarloMove[State, ParticlePositionChanges]):
    """Rigid body rotation of molecular groups. Satisfies ``ChangesFn``.

    Attributes:
        particles: Lens to particle positions.
        groups: Lens to groups eligible for moves.
        systems: Lens to indexed systems with unit cell data.
        step_width: Lens to rotation magnitude (0=no rotation, 1=full).
        capacity: Lens to capacity constraints.
    """

    particles: View[State, BatchedPositions] = field(static=True)
    groups: View[State, Table[GroupId, HasSystemIndex]] = field(static=True)
    systems: View[State, Table[SystemId, HasUnitCell]] = field(static=True)
    step_width: View[State, Table[SystemId, Array]] = field(static=True)
    capacity: View[State, Capacity[int]] = field(static=True)

    def __call__(
        self, key: Array, state: State, /
    ) -> tuple[ParticlePositionChanges, LogProbabilityRatio]:
        n_sys = self.particles(state).data.system.num_labels
        changes = propose_group_rotation(
            key,
            self.particles(state),
            self.groups(state),
            self.systems(state),
            self.step_width(state),
            self.capacity(state),
        )
        return changes, Table.arange(jnp.zeros((n_sys,)), label=SystemId)

@dataclass
class GroupTranslationMove[State](MonteCarloMove[State, ParticlePositionChanges]):
    """Rigid body translation of molecular groups. Satisfies ``ChangesFn``.

    Attributes:
        particles: Lens to particle positions.
        groups: Lens to groups eligible for moves.
        systems: Lens to indexed systems with unit cells.
        step_width: Lens to maximum translation magnitude.
        capacity: Lens to capacity constraints.
        distribution: Symmetric distribution for displacements (default: normal).
    """

    particles: View[State, BatchedPositions] = field(static=True)
    groups: View[State, Table[GroupId, HasSystemIndex]] = field(static=True)
    systems: View[State, Table[SystemId, HasUnitCell]] = field(static=True)
    step_width: View[State, Table[SystemId, Array]] = field(static=True)
    capacity: View[State, Capacity[int]] = field(static=True)
    distribution: SymmetricTranslationDistribution = field(
        static=True, default=jax.random.normal
    )

    def __call__(
        self, key: Array, state: State, /
    ) -> tuple[ParticlePositionChanges, LogProbabilityRatio]:
        n_sys = self.particles(state).data.system.num_labels
        changes = propose_group_translation(
            key,
            self.particles(state),
            self.groups(state),
            self.systems(state),
            self.step_width(state),
            self.capacity(state),
            self.distribution,
        )
        return changes, Table.arange(jnp.zeros((n_sys,)), label=SystemId)

class HasGroupSystemIndex(HasGroupIndex, HasSystemIndex, Protocol):
    """Combined group and system index for particles."""

    ...

class HasPositionsGroupSystem(HasPositions, HasGroupIndex, HasSystemIndex, Protocol):
    """Combined positions, group and system index for particles."""

    ...

class IsDisplacementState(
    IsTranslationState, IsRotationState, IsReinsertionState, Protocol
):
    """State protocol for the combined translation/rotation/reinsertion move."""

    ...

class IsExchangeState(IsMCMCMoveState, Protocol):
    """State protocol for particle exchange (grand canonical) MCMC moves."""

    @property
    def particles(self) -> Buffered[ParticleId, HasPositionsGroupSystem]: ...
    @property
    def groups(self) -> Buffered[GroupId, HasMotifAndSystemIndex]: ...
    @property
    def motifs(self) -> Table[MotifParticleId, IsMotifData]: ...
    @property
    def exchange_params(self) -> Table[SystemId, ParameterSchedulerState]: ...

class IsGCMCState(IsDisplacementState, Protocol):
    """State protocol for the combined displacement + exchange move."""

    @property
    def particles(self) -> Buffered[ParticleId, IsExchangeParticles]: ...
    @property
    def groups(self) -> Buffered[GroupId, HasMotifAndSystemIndex]: ...
    @property
    def motifs(self) -> Table[MotifParticleId, IsMotifData]: ...
    @property
    def exchange_params(self) -> Table[SystemId, ParameterSchedulerState]: ...

class IsMCMCMoveState(Protocol):
    """Base state protocol for MCMC moves with particles, groups, systems, and capacity."""

    @property
    def particles(self) -> BatchedPositions: ...
    @property
    def groups(self) -> Table[GroupId, HasSystemIndex]: ...
    @property
    def systems(self) -> Table[SystemId, HasUnitCell]: ...
    @property
    def move_capacity(self) -> Capacity[int]: ...

@runtime_checkable
class IsMotifData(Protocol):
    """Motif template data: positions indexed by motif assignment.

    Attributes:
        positions: Template particle positions, shape ``(n_template_particles, 3)``.
        motif: Per-particle motif assignment.
    """

    @property
    def positions(self) -> Array: ...
    @property
    def motif(self) -> Index[MotifId]: ...

class IsReinsertionState(IsMCMCMoveState, Protocol):
    """State protocol for particle reinsertion MCMC moves."""

    @property
    def reinsertion_params(self) -> Table[SystemId, ParameterSchedulerState]: ...

class IsRotationState(IsMCMCMoveState, Protocol):
    """State protocol for group rotation MCMC moves."""

    @property
    def rotation_params(self) -> Table[SystemId, ParameterSchedulerState]: ...

class IsTranslationState(IsMCMCMoveState, Protocol):
    """State protocol for group translation MCMC moves."""

    @property
    def translation_params(self) -> Table[SystemId, ParameterSchedulerState]: ...

class MonteCarloMove[State, Changes](ChangesFn[State, Changes], abc.ABC):
    """Base class for Monte Carlo move proposals that satisfy ``ChangesFn``."""

    @abc.abstractmethod
    def __call__(
        self, key: Array, state: State, /
    ) -> tuple[Changes, LogProbabilityRatio]: ...

@dataclass
class ParticlePositionChanges:
    """Description of particle position updates.

    Attributes:
        particle_ids: Indices of particles to update
        new_positions: New position coordinates for specified particles
    """

    particle_ids: Index[ParticleId]
    new_positions: Array

@dataclass
class ParticleTranslationMove[State](MonteCarloMove[State, ParticlePositionChanges]):
    """Single particle translation move. Satisfies ``ChangesFn``.

    Attributes:
        positions: Lens to particle positions in state.
        systems: Lens to indexed systems with unit cells.
        step_width: Lens to maximum displacement magnitude (tunable).
        distribution: Symmetric distribution for displacements (default: normal).
    """

    positions: View[State, BatchedPositions] = field(static=True)
    systems: View[State, Table[SystemId, HasUnitCell]] = field(static=True)
    step_width: View[State, Table[SystemId, Array]] = field(static=True)
    distribution: SymmetricTranslationDistribution = field(
        static=True, default=jax.random.normal
    )

    def __call__(
        self, key: Array, state: State, /
    ) -> tuple[ParticlePositionChanges, LogProbabilityRatio]:
        particles = self.positions(state)
        n_sys = particles.data.system.num_labels
        changes = propose_particle_translation(
            key,
            particles,
            self.systems(state),
            self.step_width(state),
            self.distribution,
        )
        return changes, Table.arange(jnp.zeros((n_sys,)), label=SystemId)

@dataclass
class ReinsertionMove[State](MonteCarloMove[State, ParticlePositionChanges]):
    """Random reinsertion of molecular groups. Satisfies ``ChangesFn``.

    Attributes:
        positions: Lens to particle positions.
        groups: Lens to groups eligible for moves.
        systems: Lens to indexed systems with unit cell data.
        capacity: Lens to capacity constraints.
    """

    positions: View[State, BatchedPositions] = field(static=True)
    groups: View[State, Table[GroupId, HasSystemIndex]] = field(static=True)
    systems: View[State, Table[SystemId, HasUnitCell]] = field(static=True)
    capacity: View[State, Capacity[int]] = field(static=True)

    def __call__(
        self, key: Array, state: State, /
    ) -> tuple[ParticlePositionChanges, LogProbabilityRatio]:
        n_sys = self.positions(state).data.system.num_labels
        changes = propose_reinsertion(
            key,
            self.positions(state),
            self.groups(state),
            self.systems(state),
            self.capacity(state),
        )
        return changes, Table.arange(jnp.zeros((n_sys,)), label=SystemId)

class SymmetricTranslationDistribution(Protocol):
    """Protocol for symmetric translation distributions."""

    def __call__(
        self, key: Array, shape: tuple[int, ...], dtype: DTypeLike = float
    ) -> Array: ...

@dataclass
class SystemIndex:
    """Index metadata for systems."""

    system: Index[SystemId]