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

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
5import astropy.units as u
7from . import detector
8from . import binary
9from .utils import make_quant
12def Get_SNR_Matrix(
13 source, instrument, var_x, sample_rate_x, var_y, sample_rate_y, **kwargs
14):
15 """Calculates SNR Matrix
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
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
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
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)
56 if "inc" not in locals():
57 inc = None
58 if "integral_consts" not in locals():
59 integral_consts = None
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 )
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
82 sampleSize_x = len(sample_x)
83 sampleSize_y = len(sample_y)
84 SNRMatrix = np.zeros((sampleSize_y, sampleSize_x))
86 for i in range(sampleSize_x):
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])
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)
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 )
149 if switch:
150 return [original_sample_x, original_sample_y, SNRMatrix.T]
151 else:
152 return [sample_x, sample_y, SNRMatrix]
155def Get_Samples(source, instrument, var_x, sample_rate_x, var_y, sample_rate_y):
156 """Gets the x and y-axis samples
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
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
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.
186 """
187 sample_x = []
188 sample_y = []
189 recalculate_strain = False
190 recalculate_noise = "neither"
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.")
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.")
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
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.")
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.")
296 return sample_x, sample_y, recalculate_strain, recalculate_noise
299def Recalculate_Noise(source, instrument):
300 """Recalculate noise curves if something is varied
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")
312 if not isinstance(instrument, detector.GroundBased):
313 if hasattr(instrument, "P_n_f"):
314 del instrument.P_n_f
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
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
331def Calc_Mono_SNR(source, instrument, inc=None):
332 """Calculates the SNR for a monochromatic source
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.
343 """
344 if not hasattr(source, "instrument"):
345 source.instrument = instrument
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()
351 return source.h_gw * np.sqrt(
352 np.max(np.unique(instrument.T_obs.to("s"))) / instrument.S_n_f[indxfgw]
353 )
356def Calc_Chirp_SNR(source, instrument, integral_consts=None):
357 """Calculates the SNR for an evolving source
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
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
372 """
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()
379 if not hasattr(source, "instrument"):
380 source.instrument = instrument
381 if not hasattr(source, "f_T_obs"):
382 source.Check_Freq_Evol()
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]
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
408 if isinstance(instrument, detector.PTA):
409 # Rescaled by 1.5 to make SNR plots match...
410 integral_consts = 4.0 * 1.5
412 elif isinstance(instrument, detector.SpaceBased) or isinstance(
413 instrument, detector.GroundBased
414 ):
415 integral_consts = 16.0 / 5.0
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
428 return np.sqrt(SNRsqrd)