Coverage for gemlib/abc/abc_rejection_sampler.py: 100%

1from collections.abc import Callable 

2from typing import TypeVar 

3 

4import jax 

5import tensorflow_probability.substrates.jax as tfp 

6 

7tfd = tfp.distributions 

8Array = jax.Array 

9PyTree = TypeVar("PyTree") 

10 

11__all__ = ["abc_rejection_sampler"] 

12 

13 

14def abc_rejection_sampler( 

15 accept_fn: Callable[[PyTree], Array], 

16 num_samples: int, 

17 model: tfd.JointDistribution, 

18 key: Array, 

19) -> tuple[PyTree, float]: 

20 """Approximate Bayesian Computation rejection sampler 

21 

22 Args: 

23 accept_fn: Python callable which takes a batch of samples from `model` 

24 and returns a boolean array of shape `num_samples` to keep. 

25 Any tolerance/threshold logic is responsibility of `accept_fn`. 

26 num_samples: total number of samples 

27 model: a `tfd.JointDistribution` to draw proposal from 

28 key: a PRNG key 

29 

30 Returns: 

31 A tuple of (posterior, acceptance_rate). 

32 """ 

33 

34 @jax.jit 

35 def fn(key): 

36 samples = model.sample(sample_shape=num_samples, seed=key) 

37 keep = accept_fn(samples) 

38 

39 return samples, keep 

40 

41 samples, accept = fn(key) 

42 

43 return jax.tree.map(lambda x: x[accept], samples), accept.mean()