Skip to content

kups.mcmc.widom

Widom test-particle method.

A ghost move runs the full propose/patch/log-ratio pipeline and discards the resulting state patch; the log acceptance ratio is accumulated into running statistics.

Contents:

  • widom_test: per-system \(\ln\alpha\) for a ghost move
  • GhostProbe: propagator wrapper accumulating the ratio via a lens + update callback
  • WidomStatistics: running-sum accumulator reduced to \(\mu^\mathrm{ex}\), \(K_H\), \(q_\mathrm{st}\) by the post-hoc analyzer.
References

Widom, B. (1963). J. Chem. Phys., 39, 2808. Vlugt, T. J. H. et al. (2008). J. Chem. Theory Comput., 4, 1107.

Energy = Array

Potential energy [energy].

LogAcceptanceRatio = Array

Log Metropolis acceptance ratio \(\ln\alpha\) [dimensionless].

GhostProbe

Bases: Propagator[State]

Propagator running one ghost move and updating a lens-accessed statistic.

Attributes:

Name Type Description
propose_fn / patch_fn / log_probability_ratio_fn

standard MCMC trio (:class:~kups.core.propagator.MCMCPropagator interface); the resulting patch is discarded.
stat_lens Lens[State, Stat]

where in state the accumulator lives.
update_fn Callable[[State, Stat, Array], Stat]

(state, stat, ln_alpha) -> stat.
Source code in src/kups/mcmc/widom.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
@dataclass
class GhostProbe[State, Changes, Move: Patch, Stat](Propagator[State]):
    r"""Propagator running one ghost move and updating a lens-accessed statistic.

    Attributes:
        propose_fn / patch_fn / log_probability_ratio_fn: standard MCMC trio
            (:class:`~kups.core.propagator.MCMCPropagator` interface); the
            resulting patch is discarded.
        stat_lens: where in ``state`` the accumulator lives.
        update_fn: ``(state, stat, ln_alpha) -> stat``.
    """

    propose_fn: ChangesFn[State, Changes] = field(static=True)
    patch_fn: PatchFn[State, Changes, Move] = field(static=True)
    log_probability_ratio_fn: LogProbabilityRatioFn[State, Move] = field(static=True)
    stat_lens: Lens[State, Stat] = field(static=True)
    update_fn: Callable[[State, Stat, Array], Stat] = field(static=True)

    def __call__(self, key: Array, state: State) -> State:
        ln_alpha = widom_test(
            key,
            state,
            self.propose_fn,
            self.patch_fn,
            self.log_probability_ratio_fn,
        )
        current = self.stat_lens.get(state)
        updated = self.update_fn(state, current, ln_alpha.data)
        return self.stat_lens.set(state, updated)

WidomStatistics

Online accumulator for plain Widom insertion sums.

Attributes:

Name Type Description
sum_boltzmann Array

\(\sum \exp(-\beta \Delta U) = \sum W\) [dimensionless].
sum_delta_u_boltzmann Array

\(\sum \Delta U \cdot \exp(-\beta \Delta U)\) [energy], with \(\Delta U\) the ghost insertion (host-guest) energy, not the cell's total potential energy.
n_samples Array

Number of evaluations accumulated.
Source code in src/kups/mcmc/widom.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
@dataclass
class WidomStatistics:
    r"""Online accumulator for plain Widom insertion sums.

    Attributes:
        sum_boltzmann: $\sum \exp(-\beta \Delta U) = \sum W$ [dimensionless].
        sum_delta_u_boltzmann:
            $\sum \Delta U \cdot \exp(-\beta \Delta U)$ [energy], with
            $\Delta U$ the ghost insertion (host-guest) energy, not the
            cell's total potential energy.
        n_samples: Number of evaluations accumulated.
    """

    sum_boltzmann: Array
    sum_delta_u_boltzmann: Array
    n_samples: Array

    @staticmethod
    def zeros(n_systems: int) -> WidomStatistics:
        """Zero-initialize."""
        return WidomStatistics(
            sum_boltzmann=jnp.zeros(n_systems),
            sum_delta_u_boltzmann=jnp.zeros(n_systems),
            n_samples=jnp.zeros(n_systems, dtype=jnp.int32),
        )

    def reset(self) -> WidomStatistics:
        """Zero all fields."""
        return self.zeros(int(self.n_samples.shape[0]))

    def update(self, ln_alpha: LogAcceptanceRatio, delta_u: Energy) -> WidomStatistics:
        r"""Accumulate one ghost insertion.

        Args:
            ln_alpha: Per-system log Metropolis ratio.
            delta_u: Per-system ghost insertion energy. With a zero-move-log
                insertion proposal and a bare Boltzmann log-ratio,
                $\Delta U = -k_BT \ln\alpha$ exactly.
        """
        boltzmann = jnp.exp(ln_alpha)
        return WidomStatistics(
            sum_boltzmann=self.sum_boltzmann + boltzmann,
            sum_delta_u_boltzmann=self.sum_delta_u_boltzmann + delta_u * boltzmann,
            n_samples=self.n_samples + 1,
        )

reset()

Zero all fields.

Source code in src/kups/mcmc/widom.py 
113
114
115
def reset(self) -> WidomStatistics:
    """Zero all fields."""
    return self.zeros(int(self.n_samples.shape[0]))

update(ln_alpha, delta_u)

Accumulate one ghost insertion.

Parameters:

Name Type Description Default
ln_alpha LogAcceptanceRatio

Per-system log Metropolis ratio.

 required
delta_u Energy

Per-system ghost insertion energy. With a zero-move-log insertion proposal and a bare Boltzmann log-ratio, \(\Delta U = -k_BT \ln\alpha\) exactly.

 required
Source code in src/kups/mcmc/widom.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def update(self, ln_alpha: LogAcceptanceRatio, delta_u: Energy) -> WidomStatistics:
    r"""Accumulate one ghost insertion.

    Args:
        ln_alpha: Per-system log Metropolis ratio.
        delta_u: Per-system ghost insertion energy. With a zero-move-log
            insertion proposal and a bare Boltzmann log-ratio,
            $\Delta U = -k_BT \ln\alpha$ exactly.
    """
    boltzmann = jnp.exp(ln_alpha)
    return WidomStatistics(
        sum_boltzmann=self.sum_boltzmann + boltzmann,
        sum_delta_u_boltzmann=self.sum_delta_u_boltzmann + delta_u * boltzmann,
        n_samples=self.n_samples + 1,
    )

zeros(n_systems) staticmethod

Zero-initialize.

Source code in src/kups/mcmc/widom.py 
104
105
106
107
108
109
110
111
@staticmethod
def zeros(n_systems: int) -> WidomStatistics:
    """Zero-initialize."""
    return WidomStatistics(
        sum_boltzmann=jnp.zeros(n_systems),
        sum_delta_u_boltzmann=jnp.zeros(n_systems),
        n_samples=jnp.zeros(n_systems, dtype=jnp.int32),
    )

widom_test(key, state, propose_fn, patch_fn, log_probability_ratio_fn)

Evaluate per-system \(\ln\alpha\) for a ghost move without modifying state.

Runs the full MCMC proposal \(\to\) patch \(\to\) log-ratio pipeline and intentionally discards the resulting state patch. The physical state is untouched --- this is the Widom test-particle method applied as a reusable subroutine. The returned value is raw \(\ln\alpha\), not clamped by \(\min(1, \cdot)\); callers decide how to consume it:

  • Excess chemical potential: average \(\exp\ln\alpha\), take \(-k_BT \ln\langle\cdot\rangle\).
  • Henry coefficient: same average evaluated at \(N = 0\).

Parameters:

Name Type Description Default
key Array

JAX PRNG key.

 required
state State

Current simulation state. Not modified.

 required
propose_fn ChangesFn[State, Changes]

Move proposal (e.g. insertion or deletion).

 required
patch_fn PatchFn[State, Changes, Move]

Converts proposal to a state patch.

 required
log_probability_ratio_fn LogProbabilityRatioFn[State, Move]

Evaluates the acceptance log-ratio against the proposed patch.

 required

Returns:

Type Description
Table[SystemId, LogAcceptanceRatio]

Per-system log acceptance ratio as Table[SystemId, Array].
Source code in src/kups/mcmc/widom.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def widom_test[State, Changes, Move: Patch](
    key: Array,
    state: State,
    propose_fn: ChangesFn[State, Changes],
    patch_fn: PatchFn[State, Changes, Move],
    log_probability_ratio_fn: LogProbabilityRatioFn[State, Move],
) -> Table[SystemId, LogAcceptanceRatio]:
    r"""Evaluate per-system $\ln\alpha$ for a ghost move without modifying state.

    Runs the full MCMC proposal $\to$ patch $\to$ log-ratio pipeline and
    intentionally discards the resulting state patch. The physical state is
    untouched --- this is the Widom test-particle method applied as a reusable
    subroutine. The returned value is **raw** $\ln\alpha$, not clamped by
    $\min(1, \cdot)$; callers decide how to consume it:

    - Excess chemical potential: average $\exp\ln\alpha$, take $-k_BT \ln\langle\cdot\rangle$.
    - Henry coefficient: same average evaluated at $N = 0$.

    Args:
        key: JAX PRNG key.
        state: Current simulation state. Not modified.
        propose_fn: Move proposal (e.g. insertion or deletion).
        patch_fn: Converts proposal to a state patch.
        log_probability_ratio_fn: Evaluates the acceptance log-ratio against
            the proposed patch.

    Returns:
        Per-system log acceptance ratio as ``Table[SystemId, Array]``.
    """
    chain = key_chain(key)
    changes, move_lr = propose_fn(next(chain), state)
    patch = patch_fn(next(chain), state, changes)
    result = log_probability_ratio_fn(state, patch)
    # result.patch is intentionally discarded --- state is NOT modified.
    return move_lr + result.data