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

1"""Adaptive Hamiltonian Monte Carlo sampling algorithm""" 

2 

3from enum import IntEnum 

4 

5import jax 

6import numpy as np 

7import tensorflow_probability.substrates.jax as tfp 

8from tensorflow_probability.python.internal import unnest 

9 

10from gemlib.mcmc.sampling_algorithm import ( 

11 ChainAndKernelState, 

12 ChainState, 

13 Position, 

14 SamplingAlgorithm, 

15 SeedType, 

16 TargetDensityFnType, 

17) 

18 

19RunningVariance = tfp.experimental.stats.RunningVariance 

20 

21__all__ = [ 

22 "RunningVariance", 

23 "StepSizeAlgorithm", 

24 "make_initial_running_variance", 

25 "adaptive_hmc", 

26] 

27 

28 

29# Base samplers 

30class StepSizeAlgorithm(IntEnum): 

31 SIMPLE = 0 

32 NESTEROV = 1 

33 

34 

35PreconditionedHMC = tfp.experimental.mcmc.PreconditionedHamiltonianMonteCarlo 

36StepSizeKernels = [ 

37 tfp.mcmc.SimpleStepSizeAdaptation, 

38 tfp.mcmc.DualAveragingStepSizeAdaptation, 

39] 

40MassMatrixAdaptation = tfp.experimental.mcmc.DiagonalMassMatrixAdaptation 

41 

42 

43def make_initial_running_variance( 

44 position: Position, 

45 variance: "Position | None" = None, 

46): 

47 def fn(x, v): 

48 return RunningVariance.from_stats( 

49 np.array(1.0, x.dtype), mean=x, variance=v 

50 ) 

51 

52 if variance is None: 

53 variance = jax.tree.map(lambda x: np.ones_like(x), position) 

54 

55 return jax.tree.map(fn, position, variance) 

56 

57 

58def adaptive_hmc( 

59 initial_step_size: float, 

60 num_leapfrog_steps: int, 

61 num_step_size_adaptation_steps: int, 

62 num_mass_matrix_estimation_steps: int, 

63 initial_running_variance: RunningVariance, 

64 step_size_algorithm: StepSizeAlgorithm = StepSizeAlgorithm.NESTEROV, 

65 step_size_kwargs: dict | None = None, 

66) -> SamplingAlgorithm: 

67 """Adaptive Hamiltonian Monte Carlo sampler 

68 

69 Implements a Hamiltonian Monte Carlo sampler with adaptive step size 

70 and mass matrix. 

71 

72 Args: 

73 initial_step_size: the initial step size 

74 num_leapfrog_steps: number of leapfrog steps to perform 

75 num_step_size_adaptation_steps: the number of steps before adaptation 

76 stops 

77 num_mass_matrix_estimation_steps: the number of steps before mass matrix 

78 estimation stops 

79 initial_running_variance: an initial running variance, constructed using 

80 :class:`make_initial_running_variance` 

81 step_size_algorithm: the algorithm used to adapt the step size 

82 step_size_kwargs: arguments to pass to the step size adaptation kernel 

83 name: the name of the kernel 

84 

85 Returns: 

86 a SamplingAlgorithm implementing an adaptive Hamiltonian Monte Carlo 

87 sampler 

88 """ 

89 

90 if step_size_kwargs is None: 

91 step_size_kwargs = {} 

92 

93 def _build_kernel(target_log_prob_fn: TargetDensityFnType): 

94 base_hmc_kernel = PreconditionedHMC( 

95 target_log_prob_fn=target_log_prob_fn, 

96 step_size=initial_step_size, 

97 num_leapfrog_steps=num_leapfrog_steps, 

98 store_parameters_in_results=True, 

99 ) 

100 

101 StepSizeAlg = StepSizeKernels[step_size_algorithm] 

102 step_size_adapt = StepSizeAlg( 

103 inner_kernel=base_hmc_kernel, 

104 num_adaptation_steps=num_step_size_adaptation_steps, 

105 **step_size_kwargs, 

106 ) 

107 

108 kernel = tfp.experimental.mcmc.DiagonalMassMatrixAdaptation( 

109 inner_kernel=step_size_adapt, 

110 initial_running_variance=initial_running_variance, 

111 num_estimation_steps=num_mass_matrix_estimation_steps, 

112 ) 

113 

114 return kernel 

115 

116 def init_fn( 

117 target_log_density_fn: TargetDensityFnType, initial_position: Position 

118 ): 

119 kernel = _build_kernel(target_log_density_fn) 

120 results = kernel.bootstrap_results(initial_position) 

121 

122 # Repack the results data structure into our own 

123 chain_state = ChainState( 

124 position=initial_position, 

125 log_density=unnest.get_innermost(results, "target_log_prob"), 

126 log_density_grad=tuple( 

127 unnest.get_innermost(results, "grads_target_log_prob") 

128 ), 

129 ) 

130 

131 return ChainAndKernelState(chain_state, results) 

132 

133 def step_fn( 

134 target_log_density_fn: TargetDensityFnType, 

135 chain_and_kernel_state: ChainAndKernelState, 

136 seed: SeedType, 

137 ): 

138 chain_state, kernel_state = chain_and_kernel_state 

139 

140 seed = tfp.random.sanitize_seed(seed) 

141 

142 kernel = _build_kernel( 

143 target_log_density_fn, 

144 ) 

145 

146 # If grads have not been forwarded, recompute them 

147 if chain_state.log_density_grad == (): 

148 _, log_density_grad = tfp.math.value_and_gradient( 

149 target_log_density_fn, 

150 *chain_state.position, 

151 ) 

152 else: 

153 log_density_grad = chain_state.log_density_grad 

154 

155 next_results = unnest.replace_innermost( 

156 kernel_state, 

157 target_log_prob=chain_state.log_density, 

158 grads_target_log_prob=list(log_density_grad), 

159 ) 

160 

161 new_position, results = kernel.one_step( 

162 chain_state.position, 

163 next_results, 

164 seed, 

165 ) 

166 

167 info = results 

168 

169 chain_state = ChainState( 

170 position=new_position, 

171 log_density=unnest.get_innermost(results, "target_log_prob"), 

172 log_density_grad=tuple( 

173 unnest.get_innermost(results, "grads_target_log_prob") 

174 ), 

175 ) 

176 

177 return ChainAndKernelState(chain_state, results), info 

178 

179 return SamplingAlgorithm(init_fn, step_fn)