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__(*args, value_class=StrainData) 

32 

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) 

61 

62 @property 

63 def detectors(self): 

64 return list(self.keys()) 

65 

66 

67class StrainData(TimeSeries): 

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

69 include the pesummary plots 

70 

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' 

76 

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 

85 

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 

97 

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 

106 

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) 

128 

129 @classmethod 

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

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

132 

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) 

148 

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]) 

153 

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 } 

162 

163 @property 

164 def available_plots(self): 

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

166 

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 

175 

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) 

193 

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

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

196 

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 

205 

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

207 

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

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

210 

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 

219 

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

221 

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

223 """Plot an omegascan of the strain data 

224 

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 

233 

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

235 

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

237 """Plot the spectrogram of the strain data 

238 

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 

247 

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