Coverage for /home/andrew/Documents/Research/gwent/gwent/snr.py: 86%

# Coverage for /home/andrew/Documents/Research/gwent/gwent/snr.py : 86%

## 199 statements 171 run 28 missing 0 excluded

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

1 import numpy as np

2 import scipy.interpolate as interp

3 import scipy.integrate as integrate

4

5 import astropy.units as u

6

7 from . import detector

8 from . import binary

9 from .utils import make_quant

10

11

12 def 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

155 def 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

299 def 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

331 def 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

356 def 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)