Coverage for gemlib/mcmc/mcmc_sampler.py: 100%

1"""Higher-order functions to run MCMC""" 

2 

3from functools import partial 

4 

5import jax 

6import jax.numpy as jnp 

7import tensorflow_probability.substrates.jax as tfp 

8 

9from .sampling_algorithm import ( 

10 Position, 

11 SamplingAlgorithm, 

12 SeedType, 

13 TargetDensityFnType, 

14) 

15 

16__all__ = ["mcmc"] 

17 

18 

19def _split_seed(seed, n): 

20 return tfp.random.split_seed(seed, n=n) 

21 

22 

23def _scan(fn, init, xs): 

24 """Scan 

25 

26 This function is equivalent to 

27 

28 ``` 

29 scan :: (c -> a -> (c, b)) -> c -> [a] -> (c, [b]) 

30 ``` 

31 """ 

32 return jax.lax.scan(fn, init, xs) 

33 

34 

35def mcmc( 

36 num_samples: int, 

37 sampling_algorithm: SamplingAlgorithm, 

38 target_log_density_fn: TargetDensityFnType, 

39 initial_position: Position, 

40 seed: SeedType, 

41 kernel_kwargs_fn=lambda _: {}, 

42): 

43 """Runs an MCMC using `sampling_algorithm` 

44 

45 Args: 

46 num_updates: integer giving the number of updates 

47 sampling_algorithm: an instance of `SamplingAlgorithm` 

48 target_log_density_fn: Python callable which takes an argument like 

49 `current_state` and returns its (possibly unnormalized) log-density 

50 under the target distribution. 

51 initial_position: initial state structured tuple 

52 seed: an optional list of two scalar ``int`` tensors. 

53 kernel_kwargs_fn: a callable taking the chain position as an argument, 

54 and returning a dictionary of extra kwargs 

55 

56 

57 Returns: 

58 A tuple containing samples of the Markov chain and information about the 

59 behaviour of the sampler(s) (e.g. whether kernels accepted or rejected, 

60 adaptive covariance matrices, etc). 

61 """ 

62 

63 initial_position = jax.tree.map(lambda x: jnp.asarray(x), initial_position) 

64 initial_state = sampling_algorithm.init( 

65 target_log_density_fn, 

66 initial_position, 

67 **kernel_kwargs_fn(initial_position), 

68 ) 

69 kernel_step_fn = partial(sampling_algorithm.step, target_log_density_fn) 

70 

71 def one_step(state, rng_key): 

72 new_state, info = kernel_step_fn( 

73 state, rng_key, **kernel_kwargs_fn(state) 

74 ) 

75 return new_state, (new_state[0].position, info) 

76 

77 keys = _split_seed(seed, num_samples) 

78 

79 _, trace = _scan(one_step, initial_state, keys) 

80 

81 return trace