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

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

2 

3from pycbc.psd import aLIGOZeroDetHighPower 

4from pycbc import pnutils 

5from pesummary.gw.waveform import fd_waveform 

6from pesummary.gw.pycbc import optimal_snr, compute_the_overlap 

7import numpy as np 

8 

9__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"] 

10 

11 

12class TestPyCBCModule(object): 

13 """Class to test pesummary.gw.pycbc module 

14 """ 

15 def setup_method(self): 

16 """Setup the testing class 

17 """ 

18 self.n_samples = 20 

19 self.approx = "IMRPhenomPv2" 

20 self.mass_1 = np.random.uniform(20, 100, self.n_samples) 

21 self.mass_2 = np.random.uniform(5, self.mass_1, self.n_samples) 

22 self.a_1 = np.random.uniform(0, 1, self.n_samples) 

23 self.a_2 = np.random.uniform(0, 1, self.n_samples) 

24 self.tilt_1 = np.arccos(np.random.uniform(-1, 1, self.n_samples)) 

25 self.tilt_2 = np.arccos(np.random.uniform(-1, 1, self.n_samples)) 

26 self.phi_12 = np.random.uniform(0, 1, self.n_samples) 

27 self.theta_jn = np.arccos(np.random.uniform(-1, 1, self.n_samples)) 

28 self.phi_jl = np.random.uniform(0, 1, self.n_samples) 

29 self.f_low = [20.] * self.n_samples 

30 self.f_final = [1024.] * self.n_samples 

31 self.phase = np.random.uniform(0, 1, self.n_samples) 

32 self.distance = np.random.uniform(100, 500, self.n_samples) 

33 self.ra = np.random.uniform(0, np.pi, self.n_samples) 

34 self.dec = np.arccos(np.random.uniform(-1, 1, self.n_samples)) 

35 self.psi_l = np.random.uniform(0, 1, self.n_samples) 

36 self.time = [10000.] * self.n_samples 

37 self.spin_1z = self.a_1 * np.cos(self.tilt_1) 

38 self.spin_2z = self.a_2 * np.cos(self.tilt_2) 

39 

40 def _make_psd(self, index): 

41 """Make a PSD for testing 

42 

43 Parameters 

44 ---------- 

45 index: int 

46 index of mass_1/mass_2/... array to use when computing the PSD 

47 """ 

48 duration = pnutils.get_imr_duration( 

49 self.mass_1[index], self.mass_2[index], self.spin_1z[index], 

50 self.spin_2z[index], self.f_low[index], "IMRPhenomD" 

51 ) 

52 t_len = 2**np.ceil(np.log2(duration) + 1) 

53 df = 1./t_len 

54 flen = int(self.f_final[index] / df) + 1 

55 aLIGOpsd = aLIGOZeroDetHighPower(flen, df, self.f_low[index]) 

56 psd = aLIGOpsd 

57 return psd, df 

58 

59 def _make_waveform(self, index, df): 

60 """Make a waveform for testing 

61 

62 Parameters 

63 ---------- 

64 index: int 

65 index of mass_1/mass_2/... array to use for waveform generation 

66 df: float 

67 difference in frequency samples to use when generating waveform. 

68 This should match the difference in frequency samples used for the 

69 PSD 

70 """ 

71 wvfs = fd_waveform( 

72 dict( 

73 theta_jn=self.theta_jn[index], 

74 phi_jl=self.phi_jl[index], phase=self.phase[index], 

75 mass_1=self.mass_1[index], mass_2=self.mass_2[index], 

76 tilt_1=self.tilt_1[index], tilt_2=self.tilt_2[index], 

77 phi_12=self.phi_12[index], a_1=self.a_1[index], 

78 a_2=self.a_2[index], luminosity_distance=self.distance[index] 

79 ), self.approx, df, self.f_low[index], self.f_final[index] / 2, 

80 f_ref=self.f_low[index], pycbc=True 

81 ) 

82 hp, hc = wvfs["h_plus"], wvfs["h_cross"] 

83 return hp, hc 

84 

85 def test_optimal_snr(self): 

86 """Test the pesummary.gw.pycbc.optimal_snr function 

87 """ 

88 from pycbc.filter import sigma 

89 

90 for i in range(self.n_samples): 

91 psd, df = self._make_psd(i) 

92 hp, hc = self._make_waveform(i, df) 

93 pycbc = sigma( 

94 hp, psd, low_frequency_cutoff=self.f_low[i], 

95 high_frequency_cutoff=self.f_final[i] / 2 

96 ) 

97 np.testing.assert_almost_equal(optimal_snr( 

98 hp, psd, low_frequency_cutoff=self.f_low[i], 

99 high_frequency_cutoff=self.f_final[i] / 2 

100 ), pycbc) 

101 

102 def test_compute_the_overlap(self): 

103 """Test the pesummary.gw.pycbc.compute_the_overlap function 

104 """ 

105 from pycbc.filter import overlap_cplx 

106 

107 for i in range(self.n_samples): 

108 psd, df = self._make_psd(i) 

109 hp, hc = self._make_waveform(i, df) 

110 np.testing.assert_almost_equal(compute_the_overlap( 

111 hp, hp, psd, low_frequency_cutoff=self.f_low[i] 

112 ), 1.0) 

113 np.testing.assert_almost_equal(compute_the_overlap( 

114 hc, hc, psd, low_frequency_cutoff=self.f_low[i] 

115 ), 1.0) 

116 pycbc = overlap_cplx( 

117 hp, hc, psd=psd, low_frequency_cutoff=self.f_low[i], 

118 normalized=True 

119 ) 

120 np.testing.assert_almost_equal(compute_the_overlap( 

121 hp, hc, psd, low_frequency_cutoff=self.f_low[i] 

122 ), pycbc)