Coverage for src/crispatt/_noise.py: 100%

1import typing 

2 

3import numpy as np 

4 

5 

6class _RandomVariableProtocol(typing.Protocol): 

7 def sample(self, shape, *, rng): ... 

8 

9 

10def _generate_noise_on_lengths( 

11 n_floes: int, 

12 n_members: int, 

13 seed_sequence: np.random.SeedSequence, 

14 noise_distribution: _RandomVariableProtocol, 

15): 

16 # NOTE: querrying a sample of size (n_floes, n_members) internally leads 

17 # (for scipy random variable classes) to querrying a sample of size 

18 # n_floes × n_members that is then reshaped. 

19 # It might be unnecessary---as the draws are IID---but here, we explicitely 

20 # samples n_members times a n_floes sized sample. 

21 noise = np.vstack( 

22 [ 

23 noise_distribution.sample( 

24 n_floes, rng=np.random.default_rng(_child_sequence) 

25 ) 

26 for _child_sequence in seed_sequence.spawn(n_members) 

27 ] 

28 ) 

29 

30 return noise