Coverage for pesummary/gw/file/strain.py: 81.6%

1# Licensed under an MIT style license -- see LICENSE.md 

2 

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 

8 

9__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"] 

10 

11 

12class StrainDataDict(Dict): 

13 """Class to store multiple StrainData objects from different IFOs 

14 

15 Parameters 

16 ---------- 

17 data: dict 

18 dict keyed by IFO and values StrainData objects 

19 

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__( 

32 *args, value_class=StrainData, deconstruct_complex_columns=False 

33 ) 

34 

35 @classmethod 

36 def read(cls, data, channels={}): 

37 strain_data = {} 

38 if not len(channels): 

39 _data, _channels = {}, {} 

40 for key in data.keys(): 

41 if ":" in key: 

42 try: 

43 IFO, _ = key.split(":") 

44 _channels[IFO] = key 

45 _data[IFO] = data[key] 

46 logger.debug( 

47 "Found ':' in '{}'. Assuming '{}' is the IFO and " 

48 "'{}' is the channel".format(key, IFO, key) 

49 ) 

50 except ValueError: 

51 _data[key] = data[key] 

52 _channels[key] = None 

53 else: 

54 _data = data 

55 _channels = channels 

56 if not len(_data): 

57 raise ValueError("Please provide strain data to read") 

58 if not all(IFO in _channels.keys() for IFO in _data.keys()): 

59 raise ValueError("Please provide a channel for each IFO") 

60 for IFO in _data.keys(): 

61 strain_data[IFO] = StrainData.read(_data[IFO], _channels[IFO], IFO=IFO) 

62 return cls(strain_data) 

63 

64 @property 

65 def detectors(self): 

66 return list(self.keys()) 

67 

68 

69class StrainData(TimeSeries): 

70 """Class to extend the gwpy.timeseries.TimeSeries plotting functions to 

71 include the pesummary plots 

72 

73 Parameters 

74 ---------- 

75 IFO: str, optional 

76 IFO for which the strain data corresponds too. This is used to determine 

77 the color on plots. Default 'H1' 

78 

79 Attributes 

80 ---------- 

81 gwpy: gwpy.timeseries.TimeSeries 

82 original gwpy TimeSeries object 

83 IFO: str 

84 IFO for which the strain data corresponds too 

85 strain_dict: dict 

86 dictionary of strain data 

87 

88 Methods 

89 ------- 

90 plot: 

91 Generate a plot based on the stored data 

92 """ 

93 def __new__(cls, *args, IFO="H1", **kwargs): 

94 new = super(StrainData, cls).__new__(cls, *args, **kwargs) 

95 new.gwpy = TimeSeries(*args, **kwargs) 

96 new.IFO = IFO 

97 new.strain_dict = {new.IFO: new} 

98 return new 

99 

100 def __array_finalize__(self, obj): 

101 super(StrainData, self).__array_finalize__(obj) 

102 try: 

103 self.gwpy = getattr(obj, 'gwpy') 

104 self.IFO = getattr(obj, 'IFO') 

105 self.strain_dict = getattr(obj, 'strain_dict') 

106 except (TypeError, AttributeError): 

107 pass 

108 

109 @classmethod 

110 def read(cls, *args, IFO="H1", **kwargs): 

111 from pesummary.gw.file.formats.base_read import GWRead 

112 if len(args) and isinstance(args[0], str): 

113 if GWRead.extension_from_path(args[0]) == "pickle": 

114 try: 

115 from pesummary.gw.file.formats.bilby import Bilby 

116 obj = Bilby._timeseries_from_bilby_pickle(args[0]) 

117 return StrainDataDict(obj) 

118 except Exception as e: 

119 pass 

120 elif GWRead.extension_from_path(args[0]) == "lcf": 

121 from glue.lal import Cache 

122 with open(args[0], "r") as f: 

123 data = Cache.fromfile(f) 

124 args[0] = data 

125 if len(args) and isinstance(args[0], pathlib.PosixPath): 

126 args = list(args) 

127 args[0] = str(args[0]) 

128 obj = super(StrainData, cls).read(*args, **kwargs) 

129 return cls(obj, IFO=IFO) 

130 

131 @classmethod 

132 def fetch_open_frame(cls, event, **kwargs): 

133 """Fetch open frame files for a given event 

134 

135 Parameters 

136 ---------- 

137 sampling_rate: int, optional 

138 sampling rate of strain data you wish to download. Default 16384 

139 format: str, optional 

140 format of strain data you wish to download. Default "gwf" 

141 duration: int, optional 

142 duration of strain data you wish to download. Default 32 

143 IFO: str, optional 

144 detector strain data you wish to download. Default 'L1' 

145 **kwargs: dict, optional 

146 all additional kwargs passed to StrainData.read 

147 """ 

148 from ..fetch import fetch_open_strain 

149 return fetch_open_strain(event, read_file=True, **kwargs) 

150 

151 @classmethod 

152 def fetch_open_data(cls, *args, **kwargs): 

153 obj = super(StrainData, cls).fetch_open_data(*args, **kwargs) 

154 return cls(obj, IFO=args[0]) 

155 

156 @property 

157 def plotting_map(self): 

158 return { 

159 "td": self._time_domain_plot, 

160 "fd": self._frequency_domain_plot, 

161 "omegascan": self._omega_scan_plot, 

162 "spectrogram": self._spectrogram_plot 

163 } 

164 

165 @property 

166 def available_plots(self): 

167 return list(self.plotting_map.keys()) 

168 

169 @docstring_subfunction([ 

170 'pesummary.gw.plots.detchar.spectrogram', 

171 'pesummary.gw.plots.detchar.omegascan', 

172 'pesummary.gw.plots.detchar.time_domain_strain_data', 

173 'pesummary.gw.plots.detchar.frequency_domain_strain_data' 

174 ]) 

175 def plot(self, *args, type="td", **kwargs): 

176 """Generate a plot displaying the gravitational wave strain data 

177 

178 Parameters 

179 ---------- 

180 *args: tuple 

181 all arguments are passed to the plotting function 

182 type: str 

183 name of the plot you wish to make 

184 **kwargs: dict 

185 all additional kwargs are passed to the plotting function 

186 """ 

187 if type not in self.plotting_map.keys(): 

188 raise NotImplementedError( 

189 "The {} method is not currently implemented. The allowed " 

190 "plotting methods are {}".format( 

191 type, ", ".join(self.available_plots) 

192 ) 

193 ) 

194 return self.plotting_map[type](self.strain_dict, *args, **kwargs) 

195 

196 def _time_domain_plot(self, *args, **kwargs): 

197 """Plot the strain data in the time domain 

198 

199 Parameters 

200 ---------- 

201 *args: tuple 

202 all args passed to the time_domain_strain_data function 

203 **kwargs: dict 

204 all kwargs passed to the time_domain_strain_data function 

205 """ 

206 from pesummary.gw.plots.detchar import time_domain_strain_data 

207 

208 return time_domain_strain_data(*args, **kwargs)[self.IFO] 

209 

210 def _frequency_domain_plot(self, *args, **kwargs): 

211 """Plot the strain data in the frequency domain 

212 

213 Parameters 

214 ---------- 

215 *args: tuple 

216 all args passed to the frequency_domain_strain_data function 

217 **kwargs: dict 

218 all kwargs passed to the frequency_domain_strain_data function 

219 """ 

220 from pesummary.gw.plots.detchar import frequency_domain_strain_data 

221 

222 return frequency_domain_strain_data(*args, **kwargs)[self.IFO] 

223 

224 def _omega_scan_plot(self, *args, **kwargs): 

225 """Plot an omegascan of the strain data 

226 

227 Parameters 

228 ---------- 

229 *args: tuple 

230 all args passed to the omegascan function 

231 **kwargs: dict 

232 all kwargs passed to the omegascan function 

233 """ 

234 from pesummary.gw.plots.detchar import omegascan 

235 

236 return omegascan(*args, **kwargs)[self.IFO] 

237 

238 def _spectrogram_plot(self, *args, **kwargs): 

239 """Plot the spectrogram of the strain data 

240 

241 Parameters 

242 ---------- 

243 *args: tuple 

244 all args passed to the spectrogram function 

245 **kwargs: dict 

246 all kwargs passed to the spectrogram function 

247 """ 

248 from pesummary.gw.plots.detchar import spectrogram 

249 

250 return spectrogram(*args, **kwargs)[self.IFO]