kups.application.simulations.mcmc_rigid ¶

class Config(BaseModel):
    """Top-level configuration for rigid-body MCMC simulations."""

    adsorbates: tuple[AdsorbateConfig, ...]
    hosts: tuple[HostConfig, ...]
    run: RunConfig
    lj: LJConfig
    ewald: EwaldConfig
    max_num_adsorbates: int
    compute_stress: bool = False

class EwaldConfig(BaseModel):
    """Ewald summation configuration."""

    real_cutoff: float
    precision: float

class LJConfig(BaseModel):
    """Lennard-Jones potential configuration."""

    cutoff: float
    parameters: dict[str, tuple[float | None, float | None]]
    tail_correction: bool
    mixing_rule: MixingRule

@dataclass
class MCMCState:
    """Full state for a rigid-body grand-canonical MCMC simulation.

    Holds buffered particle/group arrays, motif templates, system
    thermodynamic data, neighbor lists, potential parameters with
    caches, and per-move adaptive step-size schedulers.
    """

    particles: Buffered[ParticleId, MCMCParticles]
    groups: Buffered[GroupId, MCMCGroup]
    motifs: Table[MotifParticleId, MotifParticles]
    systems: Table[SystemId, MCMCSystems]
    neighborlist_params: UniversalNeighborlistParameters
    lj_parameters: WithCache[
        GlobalTailCorrectedLennardJonesParameters, PotentialOut[EmptyType, EmptyType]
    ]
    ewald_parameters: WithCache[EwaldParameters, EwaldCache[EmptyType, EmptyType]]
    translation_params: Table[SystemId, ParameterSchedulerState]
    rotation_params: Table[SystemId, ParameterSchedulerState]
    reinsertion_params: Table[SystemId, ParameterSchedulerState]
    exchange_params: Table[SystemId, ParameterSchedulerState]

    @property
    def max_cutoff(self) -> Table[SystemId, Array]:
        """Per-system maximum cutoff across LJ and Ewald potentials."""
        return Table(
            self.systems.keys,
            jnp.maximum(
                self.lj_parameters.data.cutoff.data,
                self.ewald_parameters.data.cutoff.data,
            ),
        )

    @property
    def move_capacity(self) -> Capacity[int]:
        """Maximum number of particles per motif (move buffer size)."""
        motif_size = self.motifs.data.motif.max_count
        assert motif_size is not None
        return FixedCapacity(motif_size * len(self.systems))

    @property
    def neighborlist(self) -> NearestNeighborList:
        return DenseNearestNeighborList.from_state(self)

    @property
    def guest_only(self) -> MCMCState:
        return bind(self, lambda x: x.particles.data).apply(MCMCParticles.guest_only)

@dataclass
class MCMCStateUpdate:
    """Proposed MCMC state change with pre-computed neighbor lists.

    Stores the proposed particle and group modifications together with
    neighbor-list edges computed *before* and *after* the change, so
    that energy differences can be evaluated without rebuilding the
    full neighbor list.

    Calling an instance applies the update conditionally on ``accept``.
    """

    _particles: WithIndices[ParticleId, Buffered[ParticleId, MCMCParticles]]
    groups: WithIndices[GroupId, Buffered[GroupId, MCMCGroup]]
    edges_after: Edges[Literal[2]]
    edges_before: Edges[Literal[2]]

    @staticmethod
    def from_changes(
        key: Array,
        state: MCMCState,
        proposal: ExchangeChanges,
    ) -> MCMCStateUpdate:
        """Build an update from exchange changes."""
        p_data = proposal.particles.data.data
        g_data = proposal.groups.data.data

        motif_data = state.motifs[p_data.motif]
        new_particles = Buffered(
            proposal.particles.data.keys,
            MCMCParticles(
                positions=p_data.new_positions,
                masses=motif_data.masses,
                atomic_numbers=motif_data.atomic_numbers,
                charges=motif_data.charges,
                labels=motif_data.labels,
                system=p_data.system,
                group=p_data.group,
                motif=p_data.motif,
            ),
            system_view,
        )
        new_groups = Buffered(
            proposal.groups.data.keys,
            MCMCGroup(g_data.system, g_data.motif),
            system_view,
        )
        particle_changes = WithIndices(proposal.particles.indices, new_particles)
        group_changes = WithIndices(proposal.groups.indices, new_groups)

        result = neighborlist_changes(
            state.neighborlist,
            state.particles,
            particle_changes,
            state.systems,
            state.max_cutoff,
            compaction=1.0,
        )
        return MCMCStateUpdate(
            _particles=particle_changes,
            groups=group_changes,
            edges_after=result.added,
            edges_before=result.removed,
        )

    def __call__(self, state: MCMCState, accept: Accept) -> MCMCState:
        """Apply the update to ``state``, conditional on ``accept``."""
        acc = Table.broadcast_to(accept, state.systems)
        new_groups = state.groups.update_if(
            acc, self.groups.indices, self.groups.data.data
        )
        new_particles = state.particles.update_if(
            acc, self._particles.indices, self._particles.data.data
        )
        return bind(state, lambda x: (x.particles, x.groups)).set(
            (new_particles, new_groups)
        )

    @property
    def particles(self):
        """Proposed particle data (without the buffer wrapper)."""
        return self._particles.map_data(lambda x: x.data)

    @property
    def neighborlist_before(self):
        """Neighbor list for the *current* (pre-move) configuration."""
        return RefineMaskNeighborList(self.edges_before)

    @property
    def neighborlist_after(self):
        """Neighbor list for the *proposed* (post-move) configuration."""
        return RefineMaskNeighborList(self.edges_after)