Coverage for pesummary/tests/reweight_test.py: 100.0%

1# Licensed under an MIT style license -- see LICENSE.md 

2 

3import pytest 

4import numpy as np 

5from ..core.reweight import rejection_sampling 

6from ..gw.reweight import uniform_in_comoving_volume_from_uniform_in_volume 

7from ..utils.samples_dict import SamplesDict 

8from .base import gw_parameters 

9 

10n_samples = 100 

11n_params = 10 

12 

13 

14def test_rejection_sampling(): 

15 """Test that the pesummary.core.reweight.rejection_sampling works as 

16 expected 

17 """ 

18 # Check that it works with a numpy array 

19 original_samples = np.random.uniform(0, 10, (n_samples, n_params)) 

20 weights = np.random.uniform(0, 5, n_samples) 

21 new_samples = rejection_sampling(original_samples, weights) 

22 # new_samples should have less samples than what we started with originally 

23 assert len(new_samples) <= n_samples 

24 # Each sample should be in the original posterior table 

25 assert all(new_sample in original_samples for new_sample in new_samples) 

26 # Each sample should be unique 

27 unique = np.unique(new_samples, axis=0) 

28 assert len(unique) == len(new_samples) 

29 

30 # Now check that it works as expected for the 

31 # pesummary.utils.samples_dict.SamplesDict object 

32 original_samples = SamplesDict( 

33 {param: np.random.uniform(0, 10, n_samples) for param in gw_parameters()} 

34 ) 

35 weights = np.random.uniform(0, 5, n_samples) 

36 new_samples = rejection_sampling(original_samples, weights) 

37 assert new_samples.number_of_samples <= original_samples.number_of_samples 

38 assert new_samples.parameters == original_samples.parameters 

39 assert all( 

40 new_sample in original_samples.samples.T for new_sample in 

41 new_samples.samples.T 

42 ) 

43 

44 

45def test_uniform_in_comoving_volume_from_uniform_in_volume(): 

46 """Test that the 

47 pesummary.gw.reweight.uniform_in_comoving_volume_from_uniform_in_volume 

48 function works as expected 

49 """ 

50 original_samples = SamplesDict( 

51 {param: np.random.uniform(0, 10, n_samples) for param in gw_parameters()} 

52 ) 

53 new_samples = uniform_in_comoving_volume_from_uniform_in_volume( 

54 original_samples 

55 ) 

56 assert new_samples.number_of_samples <= original_samples.number_of_samples 

57 assert all( 

58 new_sample in original_samples.samples.T for new_sample in 

59 new_samples.samples.T 

60 ) 

61 # check that if there are no redshift samples it still reweights 

62 original_samples.pop("redshift") 

63 new_samples = uniform_in_comoving_volume_from_uniform_in_volume( 

64 original_samples 

65 ) 

66 assert new_samples.number_of_samples <= original_samples.number_of_samples 

67 assert all( 

68 new_sample in original_samples.samples.T for new_sample in 

69 new_samples.samples.T 

70 ) 

71 # check that if there are no distance samples it still reweights 

72 original_samples = SamplesDict( 

73 {param: np.random.uniform(0, 10, n_samples) for param in gw_parameters()} 

74 ) 

75 original_samples.pop("luminosity_distance") 

76 new_samples = uniform_in_comoving_volume_from_uniform_in_volume( 

77 original_samples 

78 ) 

79 assert new_samples.number_of_samples <= original_samples.number_of_samples 

80 assert all( 

81 new_sample in original_samples.samples.T for new_sample in 

82 new_samples.samples.T 

83 ) 

84 # check that if there are no redshift or distance samples it fails 

85 original_samples.pop("redshift") 

86 with pytest.raises(Exception): 

87 new_samples = uniform_in_comoving_volume_from_uniform_in_volume( 

88 original_samples 

89 )