Coverage for pybmc/rng.py: 75%

1"""Package-wide random-number generation. 

2 

3All stochastic pybmc routines (the MCMC samplers in 

4:mod:`pybmc.inference_utils` and the posterior-predictive draws in 

5:mod:`pybmc.sampling_utils`) obtain their generator through 

6:func:`get_rng`, so the whole pipeline is driven by one seeded state: 

7 

8- With ``seed=None`` a function draws from the shared package-wide 

9 generator, which is seeded with `DEFAULT_SEED` at import time. A fresh 

10 session that performs the same sequence of calls is therefore fully 

11 reproducible, training included. 

12- With an explicit ``seed`` a function uses an independent generator 

13 seeded with that value, so a single call is reproducible in isolation 

14 regardless of what ran before it. 

15 

16Use `set_seed` to re-seed the shared generator mid-session (e.g. at the 

17top of a script or between repetitions of an experiment). 

18""" 

19 

20import numpy as np 

21 

22#: Seed for the shared package-wide generator (and the default for the 

23#: per-call reproducible posterior-predictive draws). 

24DEFAULT_SEED = 142858 

25 

26_global_rng = np.random.default_rng(DEFAULT_SEED) 

27 

28 

29def get_rng(seed=None): 

30 """ 

31 Returns the generator to use for a stochastic routine. 

32 

33 Args: 

34 seed (int | numpy.random.Generator | None): If None, the shared 

35 package-wide generator (seeded with `DEFAULT_SEED` at import, 

36 or the last `set_seed` call). If an integer, a fresh 

37 independent generator seeded with it. A ready-made 

38 `numpy.random.Generator` is returned unchanged. 

39 

40 Returns: 

41 numpy.random.Generator: The generator to draw from. 

42 """ 

43 if seed is None: 

44 return _global_rng 

45 if isinstance(seed, np.random.Generator): 

46 return seed 

47 return np.random.default_rng(seed) 

48 

49 

50def set_seed(seed=DEFAULT_SEED): 

51 """ 

52 Re-seeds the shared package-wide generator. 

53 

54 Args: 

55 seed (int, optional): New seed (default: `DEFAULT_SEED`). 

56 

57 Returns: 

58 numpy.random.Generator: The freshly seeded shared generator. 

59 """ 

60 global _global_rng 

61 _global_rng = np.random.default_rng(seed) 

62 return _global_rng