kups.mcmc.probability ¶

@dataclass
class BoltzmannLogProbabilityRatio[State, Move: Patch](
    LogProbabilityRatioFn[State, Move]
):
    """Boltzmann acceptance criterion for canonical (NVT) ensemble.

    Computes the log probability ratio based on energy differences:

    $$
    \\log p_{\\text{ratio}} = -\\frac{\\Delta U}{k_B T} = \\frac{U_{\\text{old}} - U_{\\text{new}}}{k_B T}
    $$

    This corresponds to the Metropolis criterion for constant N, V, T simulations.

    Type Parameters:
        State: Simulation state type
        Move: Patch type for state updates

    Attributes:
        temperature: Lens to system temperature [K]
        potential: Cached potential for efficient energy evaluations

    Example:
        ```python
        boltzmann = BoltzmannLogProbabilityRatio(
            temperature=lambda s: s.temperature,
            potential=cached_lj_potential
        )

        # Evaluate acceptance for a proposed move
        result = boltzmann(state, move_patch)
        accept_prob = jnp.exp(result.log_probability_ratio)
        ```
    """

    temperature: View[State, Table[SystemId, Array]] = field(static=True)
    potential: CachedPotential[State, Any, Any, Move] = field(static=True)

    @jit
    def __call__(
        self, state: State, patch: Move
    ) -> WithPatch[LogProbabilityRatio, Patch[State]]:
        old_energies = self.potential.cached_value(state).total_energies
        potential_out = self.potential(state, patch)
        new_energies = potential_out.data.total_energies
        temperature = self.temperature(state)

        log_ratio = (old_energies - new_energies) / (temperature * BOLTZMANN_CONSTANT)
        out_patch = potential_out.patch
        return WithPatch(log_ratio, out_patch)

class IsBoltzmannState(Protocol):
    """State protocol for Boltzmann acceptance probability computation."""

    @property
    def systems(self) -> Table[SystemId, IsBoltzmannSystems]: ...

class IsFugacityState(Protocol):
    """State protocol for fugacity acceptance probability computation."""

    @property
    def groups(self) -> Table[GroupId, HasMotifAndSystemIndex]: ...
    @property
    def systems(self) -> Table[SystemId, IsFugacitySystems]: ...

class IsMuVTState(Protocol):
    """State protocol for μVT (grand canonical) acceptance probability computation."""

    @property
    def groups(self) -> Table[GroupId, HasMotifAndSystemIndex]: ...
    @property
    def systems(self) -> Table[SystemId, MuVTSystems]: ...

@dataclass
class LogFugacityRatio[State, Move: Patch](LogProbabilityRatioFn[State, Move]):
    """Chemical potential contribution for grand canonical (μVT) moves.

    Computes the log probability ratio from particle number changes in GCMC:

    $$
    \\log p_{\\text{ratio}} = \\sum_i (N_i^{\\text{new}} - N_i^{\\text{old}}) \\log(f_i V) + \\log\\frac{N_i^{\\text{old}}!}{N_i^{\\text{new}}!}
    $$

    where $f_i$ is the fugacity of species $i$, $V$ is the volume, and $N_i$ is
    the particle count. This implements the acceptance criterion for insertion/deletion
    moves in the grand canonical ensemble.

    Type Parameters:
        State: Simulation state type
        Move: Patch type for state updates

    Attributes:
        log_activity: Lens to log(fugacity/k_B T) per species, shape `(n_systems, n_species)`
        volume: Lens to system volumes, shape `(n_systems,)`
        particle_counts: Lens to current particle counts, shape `(n_systems, n_species)`
        patched_particle_counts: Probe for particle counts after applying patch

    Note:
        Activity is fugacity divided by k_B T: $a = f/(k_B T)$.
        The log_activity should be computed from equation of state (e.g., Peng-Robinson).
    """

    log_activity: View[State, Table[SystemId, Array]] = field(static=True)
    volume: View[State, Table[SystemId, Array]] = field(static=True)
    particle_counts: View[State, Table[SystemId, Array]] = field(static=True)
    patched_particle_counts: Probe[State, Move, Table[SystemId, Array]] = field(
        static=True
    )

    @jit
    def __call__(
        self, state: State, patch: Move
    ) -> WithPatch[LogProbabilityRatio, Patch[State]]:
        log_activity = self.log_activity(state)
        volume = self.volume(state)
        counts = self.particle_counts(state)
        patched_counts = self.patched_particle_counts(state, patch)
        log_activity, volume, counts, patched_counts = Table.broadcast(
            log_activity, volume, counts, patched_counts
        )

        @Table.transform
        def inner(la: Array, vol: Array, c: Array, pc: Array) -> Array:
            ln_fV = la + jnp.log(vol[:, None])
            ln_fact_ratio = log_factorial_ratio(c, pc)
            ln_ratio = (pc - c) * ln_fV + ln_fact_ratio
            ln_ratio = jnp.where(pc == c, 0.0, ln_ratio)
            ln_ratio = ln_ratio.sum(axis=-1)
            return ln_ratio

        return WithPatch(inner(log_activity, volume, counts, patched_counts), IdPatch())

@dataclass
class MuVTLogProbabilityRatio[State, Move: Patch](LogProbabilityRatioFn[State, Move]):
    """Combined acceptance criterion for grand canonical (μVT) ensemble.

    Combines chemical potential and Boltzmann factors for GCMC acceptance:

    $$
    \\log p_{\\text{ratio}} = \\log p_{\\text{fugacity}} + \\log p_{\\text{Boltzmann}}
    $$

    This is the full acceptance criterion for grand canonical Monte Carlo simulations,
    accounting for both particle number fluctuations (fugacity) and energy changes
    (Boltzmann factor).

    Type Parameters:
        State: Simulation state type
        Move: Patch type for state updates

    Attributes:
        log_fugacity_ratio: Chemical potential term for particle insertions/deletions
        boltzmann_log_likelihood_ratio: Energy term for configuration changes

    Example:
        ```python
        muvt = MuVTLogProbabilityRatio(
            log_fugacity_ratio=LogFugacityRatio(...),
            boltzmann_log_likelihood_ratio=BoltzmannLogProbabilityRatio(...)
        )

        # Evaluate GCMC acceptance
        result = muvt(state, exchange_move_patch)
        alpha = jnp.minimum(1.0, jnp.exp(result.log_probability_ratio))
        ```
    """

    log_fugacity_ratio: LogFugacityRatio[State, Move] = field(static=True)
    boltzmann_log_likelihood_ratio: BoltzmannLogProbabilityRatio[State, Move] = field(
        static=True
    )

    @jit
    def __call__(
        self, state: State, patch: Move
    ) -> WithPatch[LogProbabilityRatio, Patch[State]]:
        pot_result = self.boltzmann_log_likelihood_ratio(state, patch)
        fug_result = self.log_fugacity_ratio(state, patch)
        result_ratio = fug_result.data + pot_result.data
        return WithPatch(
            result_ratio, ComposedPatch((fug_result.patch, pot_result.patch))
        )