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

1"""Simple counting kernel for testing""" 

2 

3from typing import NamedTuple 

4 

5from .sampling_algorithm import ( 

6 ChainAndKernelState, 

7 ChainState, 

8 SamplingAlgorithm, 

9) 

10 

11__all__ = ["CountingKernelInfo", "CountingKernelState", "counting_kernel"] 

12 

13 

14class CountingKernelState(NamedTuple): 

15 invocation: int 

16 

17 

18class CountingKernelInfo(NamedTuple): 

19 is_accepted: bool 

20 

21 

22def counting_kernel(): 

23 def init_fn(target_log_prob_fn, position): 

24 chain_state = ChainState( 

25 position=position, 

26 log_density=target_log_prob_fn(**position._asdict()), 

27 log_density_grad=(), 

28 ) 

29 kernel_state = CountingKernelState(0) 

30 

31 return ChainAndKernelState(chain_state, kernel_state) 

32 

33 def step_fn(target_log_prob_fn, chain_and_kernel_state, seed): # noqa: ARG001 

34 chain_state, kernel_state = chain_and_kernel_state 

35 

36 new_position = chain_state.position.__class__( 

37 **{k: v + 1.0 for k, v in chain_state.position._asdict().items()} 

38 ) 

39 

40 new_chain_state = ChainState( 

41 position=new_position, 

42 log_density=target_log_prob_fn(**new_position._asdict()), 

43 log_density_grad=(), 

44 ) 

45 new_kernel_state = CountingKernelState(kernel_state.invocation + 1) 

46 

47 return ChainAndKernelState( 

48 new_chain_state, new_kernel_state 

49 ), CountingKernelInfo(True) 

50 

51 return SamplingAlgorithm(init_fn, step_fn)