Coverage for gemlib/mcmc/mwg_step.py: 96%
54 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-15 19:54 +0000
1"""Implementation of Metropolis-within-Gibbs framework"""
3from __future__ import annotations
5from collections import namedtuple
6from collections.abc import Callable
8import tensorflow_probability.substrates.jax as tfp
10from gemlib.mcmc.mcmc_util import is_list_like
11from gemlib.mcmc.sampling_algorithm import (
12 ChainAndKernelState,
13 ChainState,
14 KernelInfo,
15 Position,
16 SamplingAlgorithm,
17 SeedType,
18 TargetDensityFnType,
19)
21split_seed = tfp.random.split_seed
23__all__ = ["MwgStep"]
26def as_list(x):
27 if is_list_like(x):
28 return x
29 return [x]
32def _make_target_type(target_names):
33 if is_list_like(target_names):
34 return namedtuple("_Target", target_names)
35 return lambda x: x # identity
38def _make_position_projector(target_names: list[str]):
39 target_names = as_list(target_names)
41 def fn(position: Position) -> tuple[tuple, dict]:
42 position_dict = position._asdict()
44 for name in target_names:
45 if name not in position_dict:
46 raise ValueError(f"`{name}` is not present in `position`")
48 target_tuple = tuple(position_dict[k] for k in target_names)
49 target_compl_dict = {
50 k: v for k, v in position_dict.items() if k not in target_names
51 }
53 return (
54 target_tuple,
55 target_compl_dict,
56 )
58 return fn
61class MwgStep: # pylint: disable=too-few-public-methods
62 r"""A Metropolis-within-Gibbs step.
64 Given an instance of :class:`SamplingAlgorithm` which operates on
65 :math:`g \subset G` for some global Markov chain state :math:`G`,
66 :class:`MwgStep` returns a new :class:`SamplingAlgorithm` that operates on
67 :math:`G`, but applies the original sampling algorithm to *only* :math:`g`,
68 automatically computing the conditional posterior distribution
69 :math:`\pi(g | g^c)`.
71 More technically, :class:`MwgStep` lifts an instance of
72 :class:`SamplingAlgorithm` into the Metropolis-within-Gibbs monad, making it
73 compatible with the structure of the global state, and composable with other
74 instances of :class:`SamplingAlgorithm` which also operate on the global
75 state.
77 Args:
78 sampling_algorithm: an instance of a sampling algorithms
79 target_names: coordinate name(s) within a global chain position on which
80 the Metropolis-within-Gibbs step is to operate
81 kernel_kwargs_fn: a callable taking the chain position as an argument,
82 and returning a dictionary of extra kwargs to
83 :meth:`SamplingAlgorithm.step`.
85 Returns:
86 An instance of :class:`SamplingAlgorithm`.
88 Note:
89 The structure of the argument supplied to :code:`target_names` determines
90 the structure of the subset of the MCMC chain state forwarded to
91 :code:`sampling_algorithm`. This means there is an important difference
92 between :code:`sampling_algorithm="foo"` and
93 :code:`sampling_algorithm=["foo"]`. The former will forward the
94 :class:`ArrayLike` structure representing the :code:`foo` coordinate to
95 the underlying kernel, whereas the latter will forward a (named) tuple of
96 length 1 _containing_ the :class:`ArrayLike` to the underlying kernel.
97 """
99 def __new__(
100 cls,
101 sampling_algorithm: SamplingAlgorithm,
102 target_names: str | list[str],
103 kernel_kwargs_fn: Callable[[Position], dict] = lambda _: {},
104 ):
105 """Create a new Metropolis-within-Gibbs step"""
107 target_names_list = as_list(target_names)
108 _project_position = _make_position_projector(target_names_list)
110 TargetType = _make_target_type(target_names)
112 def _name_target(target: tuple) -> dict:
113 return dict(target_names, target)
115 def init(
116 target_log_density_fn: TargetDensityFnType,
117 initial_position: Position,
118 ):
119 target, target_compl = _project_position(initial_position)
121 def conditional_tlp(*args):
122 tlp_kwargs = (
123 dict(zip(target_names_list, args, strict=True))
124 | target_compl
125 )
126 return target_log_density_fn(**tlp_kwargs)
128 kernel_state = sampling_algorithm.init(
129 conditional_tlp,
130 TargetType(*target),
131 **kernel_kwargs_fn(initial_position),
132 )
134 chain_state = ChainState(
135 position=initial_position,
136 log_density=kernel_state[0].log_density,
137 log_density_grad=kernel_state[0].log_density_grad,
138 )
140 return chain_state, kernel_state[1]
142 def step(
143 target_log_density_fn: TargetDensityFnType,
144 chain_and_kernel_state: ChainAndKernelState,
145 seed: SeedType,
146 ) -> tuple[ChainAndKernelState, KernelInfo]:
147 chain_state, kernel_state = chain_and_kernel_state
149 # Split global state and generate conditional density
150 target, target_compl = _project_position(chain_state.position)
152 # Calculate the conditional log density
153 def conditional_tlp(*args):
154 tlp_kwargs = (
155 dict(zip(target_names_list, args, strict=True))
156 | target_compl
157 )
158 return target_log_density_fn(**tlp_kwargs)
160 chain_substate = chain_state._replace(position=TargetType(*target))
162 # Invoke the kernel on the target state
163 (new_chain_substate, new_kernel_state), info = (
164 sampling_algorithm.step(
165 conditional_tlp,
166 (chain_substate, kernel_state),
167 seed,
168 **kernel_kwargs_fn(chain_state.position),
169 )
170 )
172 # Stitch the global position back together
173 new_position_dict = dict(
174 zip(
175 target_names_list,
176 as_list(new_chain_substate.position),
177 strict=True,
178 )
179 )
180 new_global_state = new_chain_substate._replace(
181 position=chain_state.position.__class__(
182 **(new_position_dict | target_compl)
183 )
184 )
186 return (new_global_state, new_kernel_state), info
188 return SamplingAlgorithm(init, step)