Coverage for pesummary/gw/file/strain.py: 67.0%
103 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-12-09 22:34 +0000
« prev ^ index » next coverage.py v7.4.4, created at 2024-12-09 22:34 +0000
1# Licensed under an MIT style license -- see LICENSE.md
3import pathlib
4from gwpy.timeseries import TimeSeries
5from pesummary.utils.utils import logger
6from pesummary.utils.dict import Dict
7from pesummary.utils.decorators import docstring_subfunction
9__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
12class StrainDataDict(Dict):
13 """Class to store multiple StrainData objects from different IFOs
15 Parameters
16 ----------
17 data: dict
18 dict keyed by IFO and values StrainData objects
20 Examples
21 --------
22 >>> from pesummary.gw.file.strain import StrainDataDict
23 >>> data = {
24 ... "H1": "./H-H1_LOSC_4_V2-1126257414-4096.gwf",
25 ... "L1": "./L-L1_LOSC_4_V2-1126257414-4096.gwf"
26 ... }
27 >>> channels = {"H1": "H1:LOSC-STRAIN", "L1": "L1:LOSC-STRAIN"}
28 >>> strain = StrainDataDict.read(data, channels=channels)
29 """
30 def __init__(self, *args):
31 super(StrainDataDict, self).__init__(*args, value_class=StrainData)
33 @classmethod
34 def read(cls, data, channels={}):
35 strain_data = {}
36 if not len(channels):
37 _data, _channels = {}, {}
38 for key in data.keys():
39 if ":" in key:
40 try:
41 IFO, _ = key.split(":")
42 _channels[IFO] = key
43 _data[IFO] = data[key]
44 logger.debug(
45 "Found ':' in '{}'. Assuming '{}' is the IFO and "
46 "'{}' is the channel".format(key, IFO, key)
47 )
48 except ValueError:
49 _data[key] = data[key]
50 _channels[key] = None
51 else:
52 _data = data
53 _channels = channels
54 if not len(_data):
55 raise ValueError("Please provide strain data to read")
56 if not all(IFO in _channels.keys() for IFO in _data.keys()):
57 raise ValueError("Please provide a channel for each IFO")
58 for IFO in _data.keys():
59 strain_data[IFO] = StrainData.read(_data[IFO], _channels[IFO], IFO=IFO)
60 return cls(strain_data)
62 @property
63 def detectors(self):
64 return list(self.keys())
67class StrainData(TimeSeries):
68 """Class to extend the gwpy.timeseries.TimeSeries plotting functions to
69 include the pesummary plots
71 Parameters
72 ----------
73 IFO: str, optional
74 IFO for which the strain data corresponds too. This is used to determine
75 the color on plots. Default 'H1'
77 Attributes
78 ----------
79 gwpy: gwpy.timeseries.TimeSeries
80 original gwpy TimeSeries object
81 IFO: str
82 IFO for which the strain data corresponds too
83 strain_dict: dict
84 dictionary of strain data
86 Methods
87 -------
88 plot:
89 Generate a plot based on the stored data
90 """
91 def __new__(cls, *args, IFO="H1", **kwargs):
92 new = super(StrainData, cls).__new__(cls, *args, **kwargs)
93 new.gwpy = TimeSeries(*args, **kwargs)
94 new.IFO = IFO
95 new.strain_dict = {new.IFO: new}
96 return new
98 def __array_finalize__(self, obj):
99 super(StrainData, self).__array_finalize__(obj)
100 try:
101 self.gwpy = getattr(obj, 'gwpy')
102 self.IFO = getattr(obj, 'IFO')
103 self.strain_dict = getattr(obj, 'strain_dict')
104 except (TypeError, AttributeError):
105 pass
107 @classmethod
108 def read(cls, *args, IFO="H1", **kwargs):
109 from pesummary.gw.file.formats.base_read import GWRead
110 if len(args) and isinstance(args[0], str):
111 if GWRead.extension_from_path(args[0]) == "pickle":
112 try:
113 from pesummary.gw.file.formats.bilby import Bilby
114 obj = Bilby._timeseries_from_bilby_pickle(args[0])
115 return StrainDataDict(obj)
116 except Exception as e:
117 pass
118 elif GWRead.extension_from_path(args[0]) == "lcf":
119 from glue.lal import Cache
120 with open(args[0], "r") as f:
121 data = Cache.fromfile(f)
122 args[0] = data
123 if len(args) and isinstance(args[0], pathlib.PosixPath):
124 args = list(args)
125 args[0] = str(args[0])
126 obj = super(StrainData, cls).read(*args, **kwargs)
127 return cls(obj, IFO=IFO)
129 @classmethod
130 def fetch_open_frame(cls, event, **kwargs):
131 """Fetch open frame files for a given event
133 Parameters
134 ----------
135 sampling_rate: int, optional
136 sampling rate of strain data you wish to download. Default 16384
137 format: str, optional
138 format of strain data you wish to download. Default "gwf"
139 duration: int, optional
140 duration of strain data you wish to download. Default 32
141 IFO: str, optional
142 detector strain data you wish to download. Default 'L1'
143 **kwargs: dict, optional
144 all additional kwargs passed to StrainData.read
145 """
146 from ..fetch import fetch_open_strain
147 return fetch_open_strain(event, read_file=True, **kwargs)
149 @classmethod
150 def fetch_open_data(cls, *args, **kwargs):
151 obj = super(StrainData, cls).fetch_open_data(*args, **kwargs)
152 return cls(obj, IFO=args[0])
154 @property
155 def plotting_map(self):
156 return {
157 "td": self._time_domain_plot,
158 "fd": self._frequency_domain_plot,
159 "omegascan": self._omega_scan_plot,
160 "spectrogram": self._spectrogram_plot
161 }
163 @property
164 def available_plots(self):
165 return list(self.plotting_map.keys())
167 @docstring_subfunction([
168 'pesummary.gw.plots.detchar.spectrogram',
169 'pesummary.gw.plots.detchar.omegascan',
170 'pesummary.gw.plots.detchar.time_domain_strain_data',
171 'pesummary.gw.plots.detchar.frequency_domain_strain_data'
172 ])
173 def plot(self, *args, type="td", **kwargs):
174 """Generate a plot displaying the gravitational wave strain data
176 Parameters
177 ----------
178 *args: tuple
179 all arguments are passed to the plotting function
180 type: str
181 name of the plot you wish to make
182 **kwargs: dict
183 all additional kwargs are passed to the plotting function
184 """
185 if type not in self.plotting_map.keys():
186 raise NotImplementedError(
187 "The {} method is not currently implemented. The allowed "
188 "plotting methods are {}".format(
189 type, ", ".join(self.available_plots)
190 )
191 )
192 return self.plotting_map[type](self.strain_dict, *args, **kwargs)
194 def _time_domain_plot(self, *args, **kwargs):
195 """Plot the strain data in the time domain
197 Parameters
198 ----------
199 *args: tuple
200 all args passed to the time_domain_strain_data function
201 **kwargs: dict
202 all kwargs passed to the time_domain_strain_data function
203 """
204 from pesummary.gw.plots.detchar import time_domain_strain_data
206 return time_domain_strain_data(*args, **kwargs)[self.IFO]
208 def _frequency_domain_plot(self, *args, **kwargs):
209 """Plot the strain data in the frequency domain
211 Parameters
212 ----------
213 *args: tuple
214 all args passed to the frequency_domain_strain_data function
215 **kwargs: dict
216 all kwargs passed to the frequency_domain_strain_data function
217 """
218 from pesummary.gw.plots.detchar import frequency_domain_strain_data
220 return frequency_domain_strain_data(*args, **kwargs)[self.IFO]
222 def _omega_scan_plot(self, *args, **kwargs):
223 """Plot an omegascan of the strain data
225 Parameters
226 ----------
227 *args: tuple
228 all args passed to the omegascan function
229 **kwargs: dict
230 all kwargs passed to the omegascan function
231 """
232 from pesummary.gw.plots.detchar import omegascan
234 return omegascan(*args, **kwargs)[self.IFO]
236 def _spectrogram_plot(self, *args, **kwargs):
237 """Plot the spectrogram of the strain data
239 Parameters
240 ----------
241 *args: tuple
242 all args passed to the spectrogram function
243 **kwargs: dict
244 all kwargs passed to the spectrogram function
245 """
246 from pesummary.gw.plots.detchar import spectrogram
248 return spectrogram(*args, **kwargs)[self.IFO]