Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1import numpy as np 

2import scipy.interpolate as interp 

3import scipy.integrate as integrate 

4 

5import astropy.units as u 

6 

7from . import detector 

8from . import binary 

9from .utils import make_quant 

10 

11 

12def Get_SNR_Matrix( 

13 source, instrument, var_x, sample_rate_x, var_y, sample_rate_y, **kwargs 

14): 

15 """Calculates SNR Matrix 

16 

17 Parameters 

18 ---------- 

19 source : object 

20 Instance of a gravitational wave source class 

21 instrument : object 

22 Instance of a gravitational wave detector class 

23 var_x : str 

24 x-axis variable 

25 sample_rate_x : int 

26 Number of samples at which SNRMatrix is calculated corresponding to the x-axis variable 

27 var_y : str 

28 y-axis variable 

29 sample_rate_y : array 

30 samples at which SNRMatrix was calculated corresponding to the y-axis variable 

31 

32 Returns 

33 ------- 

34 sample_x : array 

35 samples at which SNRMatrix was calculated corresponding to the x-axis variable 

36 sample_y : array 

37 samples at which SNRMatrix was calculated corresponding to the y-axis variable 

38 SNRMatrix : array-like 

39 the sample_rate_y X sample_rate_x matrix at which the SNR was calculated corresponding to the particular x and y-axis variable choices 

40 

41 Notes 

42 ----- 

43 Uses the variable given and the data range to sample the space either logrithmically or linearly based on the 

44 selection of variables. Then it computes the SNR for each value. 

45 Returns the variable ranges used to calculate the SNR for each matrix, then returns the SNRs with size of the sample_yXsample_x 

46 

47 """ 

48 for keys, value in kwargs.items(): 

49 if keys == "inc": 

50 inc = value 

51 elif keys == "integral_consts": 

52 integral_consts = value 

53 else: 

54 raise ValueError("%s is not an accepted input option." % keys) 

55 

56 if "inc" not in locals(): 

57 inc = None 

58 if "integral_consts" not in locals(): 

59 integral_consts = None 

60 

61 source.instrument = instrument 

62 # Get Samples for variables 

63 [sample_x, sample_y, recalculate_strain, recalculate_noise] = Get_Samples( 

64 source, instrument, var_x, sample_rate_x, var_y, sample_rate_y 

65 ) 

66 

67 switch = False 

68 if recalculate_noise == "y": 

69 """Calculation runs much faster when it doesn't recalculate 

70 the noise every time.""" 

71 switch = True 

72 recalculate_noise = "x" 

73 original_sample_x = sample_x 

74 original_sample_y = sample_y 

75 original_var_x = var_x 

76 original_var_y = var_y 

77 var_x = original_var_y 

78 var_y = original_var_x 

79 sample_x = original_sample_y 

80 sample_y = original_sample_x 

81 

82 sampleSize_x = len(sample_x) 

83 sampleSize_y = len(sample_y) 

84 SNRMatrix = np.zeros((sampleSize_y, sampleSize_x)) 

85 

86 for i in range(sampleSize_x): 

87 

88 if recalculate_noise in ["x", "both"]: 

89 # Update Attribute (also updates dictionary) 

90 if isinstance(instrument, detector.GroundBased): 

91 var_x_names = var_x.split() 

92 if len(var_x_names) == 2: 

93 updated_dict_x = {var_x_names[0]: {var_x_names[1]: sample_x[i]}} 

94 elif len(var_x_names) == 3: 

95 updated_dict_x = { 

96 var_x_names[0]: {var_x_names[1]: {var_x_names[2]: sample_x[i]}} 

97 } 

98 instrument.Set_Noise_Dict(updated_dict_x) 

99 else: 

100 setattr(instrument, var_x, sample_x[i]) 

101 Recalculate_Noise(source, instrument) 

102 elif recalculate_noise in ["neither"]: 

103 # Update Attribute (also updates dictionary) 

104 setattr(source, var_x, sample_x[i]) 

105 

106 for j in range(sampleSize_y): 

107 if recalculate_noise in ["x", "neither"]: 

108 # Update Attribute (also updates dictionary) 

109 setattr(source, var_y, sample_y[j]) 

110 elif recalculate_noise in ["both"]: 

111 # Update Attribute (also updates dictionary) 

112 if isinstance(instrument, detector.GroundBased): 

113 var_y_names = var_y.split() 

114 if len(var_y_names) == 2: 

115 updated_dict_y = {var_y_names[0]: {var_y_names[1]: sample_y[i]}} 

116 elif len(var_y_names) == 3: 

117 updated_dict_y = { 

118 var_y_names[0]: { 

119 var_y_names[1]: {var_y_names[2]: sample_y[i]} 

120 } 

121 } 

122 instrument.Set_Noise_Dict(updated_dict_y) 

123 else: 

124 setattr(instrument, var_y, sample_y[j]) 

125 Recalculate_Noise(source, instrument) 

126 

127 source.Check_Freq_Evol() 

128 if source.ismono: # Monochromatic Source and not diff EOB SNR 

129 if hasattr(source, "h_gw"): 

130 del source.h_gw 

131 SNRMatrix[j, i] = Calc_Mono_SNR(source, instrument, inc=inc) 

132 else: # Chirping Source 

133 if ( 

134 recalculate_strain == True 

135 ): # If we need to calculate the waveform everytime 

136 # Delete old PhenomD waveform 

137 if hasattr(source, "_phenomD_f"): 

138 del source._phenomD_f 

139 if hasattr(source, "_phenomD_h"): 

140 del source._phenomD_h 

141 if hasattr(source, "f"): 

142 del source.f 

143 if hasattr(source, "h_f"): 

144 del source.h_f 

145 SNRMatrix[j, i] = Calc_Chirp_SNR( 

146 source, instrument, integral_consts=integral_consts 

147 ) 

148 

149 if switch: 

150 return [original_sample_x, original_sample_y, SNRMatrix.T] 

151 else: 

152 return [sample_x, sample_y, SNRMatrix] 

153 

154 

155def Get_Samples(source, instrument, var_x, sample_rate_x, var_y, sample_rate_y): 

156 """Gets the x and y-axis samples 

157 

158 Parameters 

159 ---------- 

160 source : object 

161 Instance of a gravitational wave source class 

162 instrument : object 

163 Instance of a gravitational wave detector class 

164 var_x : str 

165 x-axis variable 

166 sample_rate_x : int 

167 Number of samples at which SNRMatrix is calculated corresponding to the x-axis variable 

168 var_y : str 

169 y-axis variable 

170 sample_rate_y : array 

171 samples at which SNRMatrix was calculated corresponding to the y-axis variable 

172 

173 Returns 

174 ------- 

175 sample_x : array 

176 samples at which SNRMatrix was calculated corresponding to the x-axis variable 

177 sample_y : array 

178 samples at which SNRMatrix was calculated corresponding to the y-axis variable 

179 

180 Notes 

181 ----- 

182 The function uses that to create a 

183 sample space for the variable either in linear space or logspace for M,z,L,A_acc 

184 for everything else. 

185 

186 """ 

187 sample_x = [] 

188 sample_y = [] 

189 recalculate_strain = False 

190 recalculate_noise = "neither" 

191 

192 if var_x in source.var_dict.keys(): 

193 if isinstance(source.var_dict[var_x]["min"], u.Quantity): 

194 var_x_min = source.var_dict[var_x]["min"].value 

195 var_x_max = source.var_dict[var_x]["max"].value 

196 else: 

197 var_x_min = source.var_dict[var_x]["min"] 

198 var_x_max = source.var_dict[var_x]["max"] 

199 elif var_x in instrument.var_dict.keys(): 

200 recalculate_noise = "x" 

201 if isinstance(instrument.var_dict[var_x]["min"], u.Quantity): 

202 var_x_min = instrument.var_dict[var_x]["min"].value 

203 var_x_max = instrument.var_dict[var_x]["max"].value 

204 else: 

205 var_x_min = instrument.var_dict[var_x]["min"] 

206 var_x_max = instrument.var_dict[var_x]["max"] 

207 else: 

208 raise ValueError(var_x + " is not a variable in the source or the instrument.") 

209 

210 if var_y in source.var_dict.keys(): 

211 if isinstance(source.var_dict[var_y]["min"], u.Quantity): 

212 var_y_min = source.var_dict[var_y]["min"].value 

213 var_y_max = source.var_dict[var_y]["max"].value 

214 else: 

215 var_y_min = source.var_dict[var_y]["min"] 

216 var_y_max = source.var_dict[var_y]["max"] 

217 elif var_y in instrument.var_dict.keys(): 

218 if recalculate_noise == "x": 

219 recalculate_noise = "both" 

220 else: 

221 recalculate_noise = "y" 

222 if isinstance(instrument.var_dict[var_y]["min"], u.Quantity): 

223 var_y_min = instrument.var_dict[var_y]["min"].value 

224 var_y_max = instrument.var_dict[var_y]["max"].value 

225 else: 

226 var_y_min = instrument.var_dict[var_y]["min"] 

227 var_y_max = instrument.var_dict[var_y]["max"] 

228 else: 

229 raise ValueError(var_y + " is not a variable in the source or the instrument.") 

230 

231 if var_x in ["q", "chi1", "chi2"] or var_y in ["q", "chi1", "chi2"]: 

232 recalculate_strain = True # Must recalculate the waveform at each point 

233 

234 # order of magnitude cut 

235 oom_cut = 2.0 

236 if ( 

237 var_x_min != None and var_x_max != None 

238 ): # If the variable has non-None 'min',and 'max' dictionary attributes 

239 if var_x == "n_p": 

240 instrument.var_dict[var_x]["sampled"] = True 

241 # sample in integer steps 

242 sample_range = var_x_max - var_x_min 

243 if sample_range > 10: 

244 sample_rate = max(2, int(sample_range / 10)) 

245 sample_x = np.arange(var_x_min, var_x_max, sample_rate) 

246 if var_x_max not in sample_x: 

247 sample_x = np.append(sample_x, var_x_max) 

248 else: 

249 sample_x = np.arange(var_x_min, var_x_max + 1) 

250 else: 

251 # Any other variables get sorted to linear if max-min < order of magnitude cut (oom_cut) 

252 # Otherwse the samples are in logspace 

253 if var_x_max <= 0.0 or var_x_min <= 0.0: 

254 sample_x = np.linspace(var_x_min, var_x_max, sample_rate_x) 

255 else: 

256 scale = np.log10(var_x_max) - np.log10(var_x_min) 

257 if scale >= oom_cut: 

258 sample_x = np.logspace( 

259 np.log10(var_x_min), np.log10(var_x_max), sample_rate_x 

260 ) 

261 else: 

262 sample_x = np.linspace(var_x_min, var_x_max, sample_rate_x) 

263 else: 

264 raise ValueError(var_x + " does not have an assigned min and/or max.") 

265 

266 if ( 

267 var_y_min != None and var_y_max != None 

268 ): # If the variable has non-None 'min',and 'max' dictionary attributes 

269 if var_y == "n_p": 

270 instrument.var_dict[var_y]["sampled"] = True 

271 # sample in integer steps 

272 sample_range = var_y_max - var_y_min 

273 if sample_range > 10: 

274 sample_rate = max(2, int(sample_range / 10)) 

275 sample_y = np.arange(var_y_min, var_y_max, sample_rate) 

276 if var_y_max not in sample_y: 

277 sample_y = np.append(sample_y, var_y_max) 

278 else: 

279 sample_y = np.arange(var_y_min, var_y_max + 1) 

280 else: 

281 # Any other variables get sorted to linear if max-min < order of magnitude cut (oom_cut) 

282 # Otherwse the samples are in logspace 

283 if var_y_max <= 0.0 or var_y_min <= 0.0: 

284 sample_y = np.linspace(var_y_min, var_y_max, sample_rate_y) 

285 else: 

286 scale = np.log10(var_y_max) - np.log10(var_y_min) 

287 if scale >= oom_cut: 

288 sample_y = np.logspace( 

289 np.log10(var_y_min), np.log10(var_y_max), sample_rate_y 

290 ) 

291 else: 

292 sample_y = np.linspace(var_y_min, var_y_max, sample_rate_y) 

293 else: 

294 raise ValueError(var_y + " does not have an assigned min and/or max value.") 

295 

296 return sample_x, sample_y, recalculate_strain, recalculate_noise 

297 

298 

299def Recalculate_Noise(source, instrument): 

300 """Recalculate noise curves if something is varied 

301 

302 Parameters 

303 ---------- 

304 source : object 

305 Instance of a gravitational wave source class 

306 instrument : object 

307 Instance of a gravitational wave detector class 

308 """ 

309 if hasattr(instrument, "I_type") or hasattr(instrument, "load_location"): 

310 raise ValueError("Cannot vary a loaded instrument's parameters") 

311 

312 if not isinstance(instrument, detector.GroundBased): 

313 if hasattr(instrument, "P_n_f"): 

314 del instrument.P_n_f 

315 

316 if hasattr(instrument, "fT"): 

317 del instrument.fT 

318 if hasattr(instrument, "S_n_f"): 

319 del instrument.S_n_f 

320 if hasattr(instrument, "h_n_f"): 

321 del instrument.h_n_f 

322 

323 if isinstance(instrument, detector.PTA) and hasattr( 

324 instrument, "_sensitivitycurve" 

325 ): 

326 del instrument._sensitivitycurve 

327 if hasattr(source, "instrument"): 

328 source.instrument = instrument 

329 

330 

331def Calc_Mono_SNR(source, instrument, inc=None): 

332 """Calculates the SNR for a monochromatic source 

333 

334 Parameters 

335 ---------- 

336 source : object 

337 Instance of a gravitational wave source class 

338 instrument : object 

339 Instance of a gravitational wave detector class 

340 inc : None,float,int, optional 

341 The inclination of the monochromatic source in radians. 

342 

343 """ 

344 if not hasattr(source, "instrument"): 

345 source.instrument = instrument 

346 

347 # Assumes mass and frequency in source class are in the source frame 

348 source.h_gw = binary.Get_Mono_Strain(source, inc=inc, frame="source") 

349 indxfgw = np.abs(instrument.fT - source.f_gw).argmin() 

350 

351 return source.h_gw * np.sqrt( 

352 np.max(np.unique(instrument.T_obs.to("s"))) / instrument.S_n_f[indxfgw] 

353 ) 

354 

355 

356def Calc_Chirp_SNR(source, instrument, integral_consts=None): 

357 """Calculates the SNR for an evolving source 

358 

359 Parameters 

360 ---------- 

361 source : object 

362 Instance of a gravitational wave source class 

363 instrument : object 

364 Instance of a gravitational wave detector class 

365 

366 Notes 

367 ----- 

368 Uses an interpolated method to align waveform and instrument noise, then integrates 

369 over the overlapping region. See eqn 18 from Robson,Cornish,and Liu 2018 <https://arxiv.org/abs/1803.01944> 

370 Values outside of the sensitivity curve are arbitrarily set to 1e30 so the SNR is effectively 0 

371 

372 """ 

373 

374 # Previously, it was designed to integrate from initial observed frequency f(t_init) to f(t_init-T_obs) 

375 # Does not work unless t_init is randomly sampled, which we don't do 

376 # indxfgw_start = np.abs(source.f-source.f_init).argmin() 

377 # indxfgw_end = np.abs(source.f-source.f_T_obs).argmin() 

378 

379 if not hasattr(source, "instrument"): 

380 source.instrument = instrument 

381 if not hasattr(source, "f_T_obs"): 

382 source.Check_Freq_Evol() 

383 

384 # Only want to integrate from observed frequency (f(T_obs_before_merger)) till merger 

385 indxfgw_start = np.abs(source.f - source.f_T_obs).argmin() 

386 indxfgw_end = len(source.f) 

387 if indxfgw_start >= len(source.f) - 1: 

388 # If the SMBH has already merged set the SNR to ~0 

389 return 1e-30 

390 else: 

391 f_cut = source.f[indxfgw_start:indxfgw_end] 

392 h_cut = source.h_f[indxfgw_start:indxfgw_end] 

393 

394 ################################# 

395 # Interpolate the Strain Noise Spectral Density to only the frequencies the 

396 # strain runs over 

397 # Set Noise to 1e30 outside of signal frequencies 

398 S_n_f_interp_old = interp.interp1d( 

399 np.log10(instrument.fT.value), 

400 np.log10(instrument.S_n_f.value), 

401 kind="cubic", 

402 fill_value=30.0, 

403 bounds_error=False, 

404 ) 

405 S_n_f_interp_new = S_n_f_interp_old(np.log10(f_cut.value)) 

406 S_n_f_interp = 10 ** S_n_f_interp_new 

407 

408 if isinstance(instrument, detector.PTA): 

409 # Rescaled by 1.5 to make SNR plots match... 

410 integral_consts = 4.0 * 1.5 

411 

412 elif isinstance(instrument, detector.SpaceBased) or isinstance( 

413 instrument, detector.GroundBased 

414 ): 

415 integral_consts = 16.0 / 5.0 

416 

417 # CALCULATE SNR FOR BOTH NOISE CURVES 

418 denom = S_n_f_interp # Sky Averaged Noise Spectral Density 

419 numer = h_cut ** 2 

420 integrand = numer / denom 

421 if isinstance(integrand, u.Quantity) and isinstance(f_cut, u.Quantity): 

422 SNRsqrd = integral_consts * np.trapz( 

423 integrand.value, f_cut.value, axis=0 

424 ) # SNR**2 

425 else: 

426 SNRsqrd = integral_consts * np.trapz(integrand, f_cut, axis=0) # SNR**2 

427 

428 return np.sqrt(SNRsqrd)