Coverage for gemlib/mcmc/sampling_algorithm.py: 93%

1"""Base MCMC datatypes""" 

2 

3from __future__ import annotations 

4 

5from functools import partial 

6from typing import Any, NamedTuple, Protocol 

7 

8import jax 

9import tensorflow_probability.substrates.jax as tfp 

10 

11split_seed = tfp.random.split_seed 

12 

13 

14__all__ = [ 

15 "ChainState", 

16 "KernelState", 

17 "ChainAndKernelState", 

18 "TargetDensityFnType", 

19 "KernelInfo", 

20 "KernelInitFnType", 

21 "KernelStepFnType", 

22 "Position", 

23 "SamplingAlgorithm", 

24 "SeedType", 

25] 

26 

27 

28# Type aliases 

29Position = NamedTuple 

30KernelInfo = NamedTuple 

31SeedType = tuple[int, int] 

32 

33 

34class ChainState(NamedTuple): 

35 """Represent the state of an MCMC probability space""" 

36 

37 position: Position 

38 log_density: float 

39 log_density_grad: float | None = () 

40 

41 

42class KernelState(NamedTuple): 

43 """Represent the state of a stateful MCMC kernel""" 

44 

45 pass 

46 

47 

48class ChainAndKernelState(NamedTuple): 

49 """A named tuple of ChainState and KernelState""" 

50 

51 chain_state: ChainState #: the chain state named tuple 

52 kernel_state: KernelState #: the kernel state named tuple 

53 

54 

55class TargetDensityFnType(Protocol): 

56 def __call__(self, position: Position) -> float: ... 

57 

58 

59class KernelInitFnType(Protocol): 

60 def __call__( 

61 self, 

62 target_log_density_fn: TargetDensityFnType, 

63 initial_position: Position, 

64 kwargs: dict, 

65 ) -> ChainAndKernelState: ... 

66 

67 

68class KernelStepFnType(Protocol): 

69 def __call__( 

70 self, 

71 target_log_density_fn: TargetDensityFnType, 

72 chain_and_kernel_state: ChainAndKernelState, 

73 seed: SeedType, 

74 **kwargs: dict, 

75 ) -> tuple(ChainAndKernelState, KernelInfo): ... 

76 

77 

78def _maybe_flatten(x: list[Any]): 

79 """Flatten a list if `len(x) <= 1`""" 

80 if len(x) == 0: 

81 return None 

82 if len(x) == 1: 

83 return x[0] 

84 return x 

85 

86 

87def _squeeze(x: list[Any]): 

88 if len(x) == 1: 

89 return x[0] 

90 return x 

91 

92 

93def _maybe_list(x): 

94 if isinstance(x, list): 

95 return x 

96 return [x] 

97 

98 

99def _maybe_tuple(x): 

100 if type(x) is tuple: 

101 return x 

102 return (x,) 

103 

104 

105class KernelInitMonad: 

106 """KernelInitMonad is a Writer monad allowing us to build an initial 

107 state tuple for a Metropolis-within-Gibbs algorithm 

108 """ 

109 

110 def __init__(self, fn: KernelInitFnType): 

111 """The monad 'unit' function""" 

112 self._fn = fn 

113 

114 def __call__( 

115 self, 

116 target_log_density_fn: TargetDensityFnType, 

117 initial_position: Position, 

118 **kwargs: dict, 

119 ) -> ChainAndKernelState: 

120 """Call wrapped function 

121 

122 Args: 

123 log_target_density: a function which when given `position` returns the 

124 corresponding log probability density. 

125 initial_position: the initial position of the MCMC chain. 

126 **kwargs: additional keyword arguments required by the kernel 

127 implementation. 

128 

129 Returns: 

130 an instance of :class:`ChainAndKernelState`. 

131 """ 

132 return self._fn(target_log_density_fn, initial_position, **kwargs) 

133 

134 def then(self, next_kernel_init_fn: KernelInitMonad) -> KernelInitMonad: 

135 """Monad combination, i.e. Haskell fish operator""" 

136 

137 @KernelInitMonad 

138 def compound_init_fn( 

139 target_log_density_fn: TargetDensityFnType, 

140 initial_position: ChainState, 

141 ) -> KernelInitFnType: 

142 _, self_kernel_state = self(target_log_density_fn, initial_position) 

143 next_chain_state, next_kernel_state = next_kernel_init_fn( 

144 target_log_density_fn, initial_position 

145 ) 

146 

147 return ChainAndKernelState( 

148 next_chain_state, 

149 _maybe_list(self_kernel_state) + [next_kernel_state], 

150 ) 

151 

152 return compound_init_fn 

153 

154 def __rshift__(self, next_kernel: KernelInitMonad): 

155 return self.then(next_kernel) 

156 

157 

158class KernelStepMonad: 

159 """StepMonad is a Writer monad that allows us to chain MCMC kernels 

160 together. 

161 """ 

162 

163 def __init__(self, fn: KernelStepFnType): 

164 """The monad 'unit' function""" 

165 self._fn = fn # Make private 

166 

167 def __call__( 

168 self, 

169 target_log_density_fn: TargetDensityFnType, 

170 chain_and_kernel_state: ChainAndKernelState, 

171 seed: SeedType, 

172 **kwargs: dict, 

173 ) -> tuple[ChainAndKernelState, KernelInfo]: 

174 """Apply the state transformer computation to a state. 

175 

176 Args: 

177 target_log_prob_fn: a function which when given `position` returns the 

178 corresponding log probability density. 

179 chain_and_kernel_state: an object of type 

180 :class:`ChainAndKernelState` as returned by :meth:`init`. 

181 seed: a random seed/key. 

182 **kwargs: additional keyword arguments required by the kernel 

183 implementation. 

184 

185 Returns: 

186 a tuple of :class:`ChainAndKernelState` and :class:`KernelInfo` 

187 

188 """ 

189 return self._fn( 

190 target_log_density_fn, chain_and_kernel_state, seed, **kwargs 

191 ) 

192 

193 def then(self, next_kernel_fn: KernelStepMonad) -> KernelStepMonad: 

194 """The monad 'bind' operator which allows chaining. 

195 ma >> mb :: ma -> mb -> mc 

196 """ 

197 

198 @KernelStepMonad 

199 def compound_step_kernel( 

200 target_log_density_fn: TargetDensityFnType, 

201 chain_and_kernel_state: ChainAndKernelState, 

202 seed: SeedType, 

203 ) -> tuple[ChainAndKernelState, KernelInfo]: 

204 self_seed, next_seed = split_seed(seed) 

205 

206 chain_state, kernel_state = chain_and_kernel_state 

207 

208 self_kernel_state = _squeeze(kernel_state[:-1]) 

209 next_kernel_state = kernel_state[-1] 

210 

211 (chain_state, self_kernel_state), self_info = self._fn( 

212 target_log_density_fn, 

213 (chain_state, self_kernel_state), 

214 seed=self_seed, 

215 ) 

216 

217 (chain_state, next_kernel_state), next_info = next_kernel_fn( 

218 target_log_density_fn, 

219 (chain_state, next_kernel_state), 

220 seed=next_seed, 

221 ) 

222 

223 return ( 

224 ChainAndKernelState( 

225 chain_state, 

226 _maybe_list(self_kernel_state) + [next_kernel_state], 

227 ), 

228 _maybe_list(self_info) + [next_info], 

229 ) 

230 

231 return compound_step_kernel 

232 

233 def __rshift__(self, next_kernel: KernelStepMonad): 

234 return self.then(next_kernel) 

235 

236 

237class SamplingAlgorithm: 

238 """Represent a sampling algorithm""" 

239 

240 def __init__( 

241 self, 

242 init_fn: KernelInitFnType | KernelInitMonad, 

243 step_fn: KernelStepFnType | KernelStepMonad, 

244 ): 

245 """Create a new sampling algorithm 

246 

247 Args: 

248 init_fn: the kernel initialisation function 

249 step_fn: the kernel step function 

250 """ 

251 if isinstance(init_fn, KernelInitMonad) and isinstance( 

252 step_fn, KernelStepMonad 

253 ): 

254 self._init: KernelInitMonad = init_fn 

255 self._step: KernelStepMonad = step_fn 

256 else: 

257 self._init: KernelInitMonad = KernelInitMonad(init_fn) 

258 self._step: KernelStepMonad = KernelStepMonad(step_fn) 

259 

260 def init( 

261 self, 

262 target_log_density_fn: TargetDensityFnType, 

263 initial_position: Position, 

264 **kwargs: dict, 

265 ) -> ChainAndKernelState: 

266 """Initialize and MCMC chain 

267 

268 Args: 

269 target_log_density_fn: a function which when given `position` returns 

270 the corresponding log probability density. 

271 initial_position: the initial position of the MCMC chain. 

272 **kwargs: additional keyword arguments required by the kernel 

273 implementation. 

274 

275 Returns: 

276 an instance of :class:`ChainAndKernelState`. 

277 """ 

278 return self._init(target_log_density_fn, initial_position, **kwargs) 

279 

280 def step( 

281 self, 

282 target_log_density_fn: TargetDensityFnType, 

283 chain_and_kernel_state: ChainAndKernelState, 

284 seed: SeedType, 

285 **kwargs: dict, 

286 ) -> tuple(ChainAndKernelState, KernelInfo): 

287 """Function to invoke the MCMC kernel 

288 

289 Args: 

290 target_log_density_fn: a function which when given `position` returns 

291 the corresponding log probability density. 

292 chain_and_kernel_state: an object of type 

293 :class:`ChainAndKernelState` as returned by :meth:`init`. 

294 seed: a random seed/key. 

295 **kwargs: additional keyword arguments required by the kernel 

296 implementation. 

297 

298 Returns: 

299 a tuple of :class:`ChainAndKernelState` and :class:`KernelInfo` 

300 """ 

301 return self._step( 

302 target_log_density_fn, chain_and_kernel_state, seed, **kwargs 

303 ) 

304 

305 def then(self, next_kernel: SamplingAlgorithm): 

306 """Sequential combinator 

307 

308 Given :code:`b: SamplingAlgorithm`, `c = self.then(b)` is 

309 a new :class:`SamplingAlgorithm` representing the sequential composition 

310 `self` then `b`. 

311 

312 Args: 

313 next_kernel: the next kernel in the kernel sequence 

314 

315 Returns: 

316 A new :class:`SamplingAlgorithm` representing the composite 

317 kernel. 

318 """ 

319 return SamplingAlgorithm( 

320 init_fn=(self._init >> next_kernel._init), 

321 step_fn=(self._step >> next_kernel._step), 

322 ) 

323 

324 def __rshift__(self, next_kernel: SamplingAlgorithm) -> SamplingAlgorithm: 

325 """Syntactic sugar for sequential combinator 

326 

327 For :code:`a: SamplingAlgorithm` and :code:`b: SamplingAlgorithm`, 

328 `a >> b` is equivalent to :meth:`then`. 

329 

330 Args: 

331 next_kernel: the next kernel in the kernel sequence 

332 

333 Returns: 

334 A new :class:`SamplingAlgorithm` representing the composite 

335 kernel. 

336 """ 

337 return self.then(next_kernel) 

338 

339 def __mul__(self, n: int) -> SamplingAlgorithm: 

340 """Performs multiple applications of a kernel 

341 

342 :obj:`sampling_algorithm` is invoked :obj:`n` times 

343 returning the state and info after the last step. 

344 

345 Args: 

346 num_updates: integer giving the number of updates 

347 sampling_algorithm: an instance of :obj:`SamplingAlgorithm` 

348 

349 Returns: 

350 An instance of :obj:`SamplingAlgorithm` 

351 """ 

352 return _repeat_sampling_algorithm(n, self) 

353 

354 def __rmul__(self, n: int) -> SamplingAlgorithm: 

355 """Performs multiple applications of a kernel 

356 

357 :obj:`sampling_algorithm` is invoked :obj:`n` times 

358 returning the state and info after the last step. 

359 

360 Args: 

361 num_updates: integer giving the number of updates 

362 sampling_algorithm: an instance of :obj:`SamplingAlgorithm` 

363 

364 Returns: 

365 An instance of :obj:`SamplingAlgorithm` 

366 """ 

367 

368 return _repeat_sampling_algorithm(n, self) 

369 

370 

371class MultiScanKernelState(NamedTuple): 

372 last_results: NamedTuple 

373 

374 

375class MultiScanKernelInfo(NamedTuple): 

376 last_results: NamedTuple 

377 

378 

379def _repeat_sampling_algorithm( 

380 num_updates: int, sampling_algorithm: SamplingAlgorithm 

381) -> SamplingAlgorithm: 

382 """Performs multiple applications of a kernel 

383 

384 :obj:`sampling_algorithm` is invoked :obj:`num_updates` times 

385 returning the state and info after the last step. 

386 

387 Args: 

388 num_updates: integer giving the number of updates 

389 sampling_algorithm: an instance of :obj:`SamplingAlgorithm` 

390 

391 Returns: 

392 An instance of :obj:`SamplingAlgorithm` 

393 """ 

394 

395 def init_fn(target_log_density_fn, position): 

396 cs, ks = sampling_algorithm.init(target_log_density_fn, position) 

397 return ChainAndKernelState(cs, MultiScanKernelState(ks)) 

398 

399 def step_fn( 

400 target_log_density_fn: TargetDensityFnType, 

401 current_state: tuple[Position, MultiScanKernelState], 

402 seed=None, 

403 ): 

404 seeds = jax.random.split(seed, num=num_updates) 

405 step_fn = partial(sampling_algorithm.step, target_log_density_fn) 

406 

407 def body(a): 

408 i, state, info = a 

409 state1, info1 = step_fn(state, seeds[i]) 

410 return i + 1, state1, info1 

411 

412 def cond(a): 

413 i, state, info = a 

414 return i < num_updates 

415 

416 chain_state, kernel_state = current_state 

417 

418 init_state, init_info = step_fn( 

419 (chain_state, kernel_state.last_results), seed 

420 ) # unrolled first it 

421 

422 _, last_state, last_info = jax.lax.while_loop( 

423 cond, body, init_val=(1, init_state, init_info) 

424 ) 

425 

426 return ( 

427 ChainAndKernelState( 

428 last_state[0], MultiScanKernelState(last_state[1]) 

429 ), 

430 MultiScanKernelInfo(last_info), 

431 ) 

432 

433 return SamplingAlgorithm(init_fn, step_fn)