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

1"""MultiScanKernel calls one_step a number of times on an inner kernel""" 

2 

3from functools import partial 

4from warnings import warn 

5 

6import jax 

7 

8from .sampling_algorithm import ( 

9 ChainAndKernelState, 

10 MultiScanKernelInfo, 

11 MultiScanKernelState, 

12 Position, 

13 SamplingAlgorithm, 

14 TargetDensityFnType, 

15) 

16 

17__all__ = ["multi_scan"] 

18 

19 

20def multi_scan( 

21 num_updates: int, sampling_algorithm: SamplingAlgorithm 

22) -> SamplingAlgorithm: 

23 """Performs multiple applications of a kernel 

24 

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

26 returning the state and info after the last step. 

27 

28 Args: 

29 num_updates: integer giving the number of updates 

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

31 

32 Returns: 

33 An instance of :obj:`SamplingAlgorithm` 

34 """ 

35 warn( 

36 "Use of `multi_scan` is deprecated, and will be removed in future.\ 

37 Instead, please make use of SamplingAlgorithm.__mul__.", 

38 DeprecationWarning, 

39 stacklevel=2, 

40 ) 

41 

42 def init_fn( 

43 target_log_density_fn: TargetDensityFnType, position: Position 

44 ) -> ChainAndKernelState: 

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

46 return ChainAndKernelState(cs, MultiScanKernelState(ks)) 

47 

48 def step_fn( 

49 target_log_density_fn: TargetDensityFnType, 

50 current_state: tuple[Position, MultiScanKernelState], 

51 seed=None, 

52 ) -> tuple[ChainAndKernelState, MultiScanKernelInfo]: 

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

54 step_fn = partial(sampling_algorithm.step, target_log_density_fn) 

55 

56 def body(a): 

57 i, state, info = a 

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

59 return i + 1, state1, info1 

60 

61 def cond(a): 

62 i, state, info = a 

63 return i < num_updates 

64 

65 chain_state, kernel_state = current_state 

66 

67 init_state, init_info = step_fn( 

68 (chain_state, kernel_state.last_results), seed 

69 ) # unrolled first it 

70 

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

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

73 ) 

74 

75 return ( 

76 ChainAndKernelState( 

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

78 ), 

79 MultiScanKernelInfo(last_info), 

80 ) 

81 

82 return SamplingAlgorithm(init_fn, step_fn)