Coverage for pesummary/gw/plots/detchar.py: 78.0%
127 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-05-02 08:42 +0000
1# Licensed under an MIT style license -- see LICENSE.md
3from pesummary.core.plots.figure import figure
4from pesummary.utils.utils import logger, _check_latex_install
5from gwpy.plot.colors import GW_OBSERVATORY_COLORS
6import numpy as np
8__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
9_check_latex_install()
12def spectrogram(
13 strain, vmin=1e-23, vmax=1e-19, cmap="viridis", ylim=[40, 2000], **kwargs
14):
15 """Generate a spectrogram from the timeseries
17 Parameters
18 ----------
19 strain: dict
20 dictionary of gw.py timeseries objects containing the strain data for
21 each IFO
22 vmin: float, optional
23 minimum for the colormap
24 vmax: float, optional
25 maximum for the colormap
26 cmap: str, optional
27 cmap for the plot. See `matplotlib.pyplot.colormaps()` for options
28 ylim: list, optional
29 list to give the lower and upper bound of the plot
30 """
31 figs = {}
32 for num, key in enumerate(list(strain.keys())):
33 logger.debug("Generating a spectrogram for {}".format(key))
34 figs[key], ax = figure(figsize=(12, 6), gca=True)
35 try:
36 try:
37 specgram = strain[key].spectrogram(
38 20, fftlength=8, overlap=4
39 ) ** (1 / 2.)
40 except Exception as e:
41 specgram = strain[key].spectrogram(strain[key].duration / 2.)
42 im = ax.pcolormesh(
43 specgram, vmin=vmin, vmax=vmax, norm='log', cmap=cmap
44 )
45 ax.set_ylim(ylim)
46 ax.set_ylabel(r'Frequency [$Hz$]')
47 ax.set_yscale('log')
48 ax.set_xscale('minutes', epoch=strain[key].times[0])
49 cbar = figs[key].colorbar(im)
50 cbar.set_label(r"ASD [strain/$\sqrt{Hz}$]")
51 except Exception as e:
52 logger.info(
53 "Failed to generate an spectrogram for {} because {}".format(key, e)
54 )
55 return figs
58def omegascan(
59 strain, gps, window=4, vmin=0, vmax=25, cmap="viridis", ylim=[40, 2000],
60 **kwargs
61):
62 """Generate an omegascan from the timeseries
64 Parameters
65 ----------
66 strain: dict
67 dictionary of gw.py timeseries objects containing the strain data for
68 each IFO
69 gps: float
70 gps time you wish to center your omegascan around
71 window: float, optional
72 window around gps time to generate omagescan for. Default 4s
73 vmin: float, optional
74 minimum for the colormap
75 vmax: float, optional
76 maximum for the colormap
77 cmap: str, optional
78 cmap for the plot. See `matplotlib.pyplot.colormaps()` for options
79 ylim: list, optional
80 list to give the lower and upper bound of the plot
81 """
82 detectors = list(strain.keys())
83 figs = {}
84 for num, key in enumerate(detectors):
85 logger.debug("Generating an omegascan for {}".format(key))
86 try:
87 try:
88 cropped_data = strain[key].crop(gps - window, gps + window)
89 qtransform = cropped_data.q_transform(
90 gps=gps, outseg=(gps - 0.5 * window, gps + 0.5 * window),
91 logf=True
92 )
93 except Exception as e:
94 qtransform = strain[key].q_transform(gps=gps, logf=True)
95 figs[key], ax = figure(figsize=(12, 6), gca=True)
96 im = ax.pcolormesh(qtransform, vmin=vmin, vmax=vmax, cmap=cmap)
97 ax.set_ylim(ylim)
98 ax.set_ylabel(r'Frequency [$Hz$]')
99 ax.set_xscale('seconds', epoch=gps)
100 ax.set_yscale('log')
101 cbar = figs[key].colorbar(im)
102 cbar.set_label("Signal-to-noise ratio")
103 except Exception as e:
104 logger.info(
105 "Failed to generate an omegascan for {} because {}".format(key, e)
106 )
107 figs[key] = figure(figsize=(12, 6), gca=False)
108 return figs
111def time_domain_strain_data(
112 strain, bandpass_frequencies=[50, 250], notches=[60., 120., 180.],
113 window=None, merger_time=None, template=None, grid=False,
114 xlabel="UTC", UTC_format="%B %d %Y, %H:%M:%S"
115):
116 """Plot the strain data in the time domain. Code based on the GW150914
117 tutorial provided by gwpy:
118 https://gwpy.github.io/docs/latest/examples/signal/gw150914.html
120 Parameters
121 ----------
122 """
123 from gwpy.signal import filter_design
124 import matplotlib.patheffects as pe
126 detectors = list(strain.keys())
127 figs = {}
128 for num, key in enumerate(detectors):
129 if bandpass_frequencies is not None and notches is not None:
130 bp = filter_design.bandpass(*bandpass_frequencies, strain[key].sample_rate)
131 zpks = [
132 filter_design.notch(line, strain[key].sample_rate) for line in
133 notches
134 ]
135 zpk = filter_design.concatenate_zpks(bp, *zpks)
136 hfilt = strain[key].filter(zpk, filtfilt=True)
137 _strain = hfilt.crop(*hfilt.span.contract(1))
138 else:
139 _strain = strain[key]
140 figs[key], ax = figure(figsize=(12, 6), gca=True)
141 if merger_time is not None:
142 x = _strain.times.value - merger_time
143 _xlabel = "Time (seconds) from {} {}"
144 if xlabel == "UTC":
145 from lal import gpstime
147 _merger_time = gpstime.gps_to_str(merger_time, form=UTC_format)
148 xlabel = _xlabel.format(_merger_time, "UTC")
149 else:
150 xlabel = _xlabel.format(merger_time, "GPS")
151 else:
152 x = _strain.times
153 xlabel = "GPS [$s$]"
154 if template is not None:
155 if not isinstance(template[key], dict):
156 template[key] = {"template": template[key]}
157 _x = template[key]["template"].times.value
158 if merger_time is not None:
159 _x -= merger_time
160 ax.plot(
161 _x, template[key]["template"], color='gray', linewidth=3.,
162 path_effects=[pe.Stroke(linewidth=4.5, foreground='k'), pe.Normal()],
163 label="Template"
164 )
165 _bounds = ["upper", "lower", "bound_times"]
166 if all(bound in template[key].keys() for bound in _bounds):
167 _x = template[key]["bound_times"]
168 if merger_time is not None:
169 _x -= merger_time
170 ax.fill_between(
171 _x, template[key]["upper"], template[key]["lower"],
172 color='lightgray', label="Uncertainty"
173 )
174 ax.plot(
175 x, _strain, color=GW_OBSERVATORY_COLORS[key], linewidth=3.,
176 label="Detector data"
177 )
178 if window is not None:
179 ax.set_xlim(*window)
180 ax.set_xlabel(xlabel)
181 ax.grid(visible=grid)
182 ax.legend()
183 return figs
186def frequency_domain_strain_data(
187 strain, window=True, window_kwargs={"roll_off": 0.2}, resolution=1. / 512,
188 fmin=-np.inf, fmax=np.inf, asd={}
189):
190 """Plot the strain data in the frequency domain
192 Parameters
193 ----------
194 strain: dict
195 dictionary of gw.py timeseries objects containing the strain data for
196 each IFO
197 window: Bool, optional
198 if True, apply a window to the data before applying FFT to the data.
199 Default True
200 window_kwargs: dict, optional
201 optional kwargs for the window function
202 resolution: float, optional
203 resolution to downsample the frequency domain data. Default 1./512
204 fmin: float, optional
205 lowest frequency to start plotting the data
206 fmax: float, optional
207 highest frequency to stop plotting the data
208 asd: dict, optional
209 dictionary containing the ASDs for each detector to plot ontop of the
210 detector data
211 """
212 detectors = list(strain.keys())
213 figs = {}
214 if not isinstance(asd, dict):
215 raise ValueError(
216 "Please provide the asd as a dictionary keyed by the detector"
217 )
218 elif not all(ifo in asd.keys() for ifo in detectors):
219 logger.info(
220 ""
221 )
222 for num, key in enumerate(detectors):
223 logger.debug("Plotting strain data in frequency domain")
224 if window:
225 from scipy.signal.windows import tukey
227 if "alpha" in window_kwargs.keys():
228 alpha = window_kwargs["alpha"]
229 elif "roll_off" and "duration" in window_kwargs.keys():
230 alpha = 2 * window_kwargs["roll_off"] / window_kwargs["duration"]
231 elif "roll_off" in window_kwargs.keys():
232 alpha = 2 * window_kwargs["roll_off"] / strain[key].duration.value
233 else:
234 raise ValueError(
235 "Please either provide 'alpha' (the shape parameter of the "
236 "Tukey window) or the 'roll_off' for the Tukey window."
237 )
238 _window = tukey(len(strain[key].value), alpha=alpha)
239 else:
240 _window = None
241 freq = strain[key].average_fft(window=_window)
242 freq = freq.interpolate(resolution)
243 freq = np.absolute(freq) / freq.df.value**0.5
244 figs[key], ax = figure(figsize=(12, 6), gca=True)
245 inds = np.where(
246 (freq.frequencies.value > fmin) & (freq.frequencies.value < fmax)
247 )
248 ax.loglog(
249 freq.frequencies[inds], freq[inds], label=key,
250 color=GW_OBSERVATORY_COLORS[key], alpha=0.2
251 )
252 if key in asd.keys():
253 inds = np.where((asd[key][:, 0] > fmin) & (asd[key][:, 0] < fmax))
254 ax.loglog(
255 asd[key][:, 0][inds], asd[key][:, 1][inds],
256 label="%s ASD" % (key), color=GW_OBSERVATORY_COLORS[key]
257 )
258 ax.set_xlabel(r"Frequency [$Hz$]")
259 ax.set_ylabel(r"Strain [strain/$\sqrt{Hz}$]")
260 ax.legend()
261 figs[key].tight_layout()
262 return figs