Coverage for pesummary/gw/file/formats/bilby.py: 48.8%

121 statements  

« 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 

2 

3import numpy as np 

4from pesummary.core.file.formats.bilby import Bilby as CoreBilby 

5from pesummary.gw.file.formats.base_read import GWSingleAnalysisRead 

6from pesummary.gw.plots.latex_labels import GWlatex_labels 

7from pesummary.utils.utils import logger 

8 

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

10 

11 

12def read_bilby( 

13 path, disable_prior=False, latex_dict=GWlatex_labels, 

14 complex_params=["matched_filter_snr", "optimal_snr"], **kwargs 

15): 

16 """Grab the parameters and samples in a bilby file 

17 

18 Parameters 

19 ---------- 

20 path: str 

21 path to the result file you wish to read in 

22 disable_prior: Bool, optional 

23 if True, do not collect prior samples from the `bilby` result file. 

24 Default False 

25 complex_params: list, optional 

26 list of parameters stored in the bilby result file which are complex 

27 and you wish to store the `amplitude` and `angle` as seperate 

28 posterior distributions 

29 latex_dict: dict, optional 

30 list of latex labels for each parameter 

31 """ 

32 from pesummary.core.file.formats.bilby import ( 

33 read_bilby as _read_bilby 

34 ) 

35 

36 return _read_bilby( 

37 path, disable_prior=disable_prior, latex_dict=latex_dict, 

38 complex_params=complex_params, _bilby_class=Bilby, **kwargs 

39 ) 

40 

41 

42def prior_samples_from_file(path, cls="BBHPriorDict", nsamples=5000, **kwargs): 

43 """Return a dict of prior samples from a `bilby` prior file 

44 

45 Parameters 

46 ---------- 

47 path: str 

48 path to a `bilby` prior file 

49 cls: str, optional 

50 class you wish to read in the prior file 

51 nsamples: int, optional 

52 number of samples to draw from a prior file. Default 5000 

53 """ 

54 from pesummary.core.file.formats.bilby import ( 

55 prior_samples_from_file as _prior_samples_from_file 

56 ) 

57 from bilby.gw import prior 

58 

59 if isinstance(cls, str): 

60 cls = getattr(prior, cls) 

61 return _prior_samples_from_file(path, cls=cls, nsamples=nsamples, **kwargs) 

62 

63 

64class Bilby(GWSingleAnalysisRead): 

65 """PESummary wrapper of `bilby` (https://git.ligo.org/lscsoft/bilby). The 

66 path_to_results_file argument will be passed directly to 

67 `bilby.core.result.read_in_result`. All functions therefore use `bilby` 

68 methods and requires `bilby` to be installed. 

69 

70 Parameters 

71 ---------- 

72 path_to_results_file: str 

73 path to the results file that you wish to read in with `bilby`. 

74 disable_prior: Bool, optional 

75 if True, do not collect prior samples from the `bilby` result file. 

76 Default False 

77 disable_prior_conversion: Bool, optional 

78 if True, disable the conversion module from deriving alternative prior 

79 distributions. Default False 

80 pe_algorithm: str 

81 name of the algorithm used to generate the posterior samples 

82 remove_nan_likelihood_samples: Bool, optional 

83 if True, remove samples which have log_likelihood='nan'. Default True 

84 

85 Attributes 

86 ---------- 

87 parameters: list 

88 list of parameters stored in the result file 

89 converted_parameters: list 

90 list of parameters that have been derived from the sampled distributions 

91 samples: 2d list 

92 list of samples stored in the result file 

93 samples_dict: dict 

94 dictionary of samples stored in the result file keyed by parameters 

95 input_version: str 

96 version of the result file passed. 

97 extra_kwargs: dict 

98 dictionary of kwargs that were extracted from the result file 

99 prior: dict 

100 dictionary of prior samples extracted from the bilby result file. The 

101 analytic priors are evaluated for 5000 samples 

102 injection_parameters: dict 

103 dictionary of injection parameters stored in the result file 

104 converted_parameters: list 

105 list of parameters that have been added 

106 

107 Methods 

108 ------- 

109 to_dat: 

110 save the posterior samples to a .dat file 

111 to_latex_table: 

112 convert the posterior samples to a latex table 

113 generate_latex_macros: 

114 generate a set of latex macros for the stored posterior samples 

115 to_lalinference: 

116 convert the posterior samples to a lalinference result file 

117 generate_all_posterior_samples: 

118 generate all posterior distributions that may be derived from 

119 sampled distributions 

120 """ 

121 def __init__(self, path_to_results_file, **kwargs): 

122 super(Bilby, self).__init__(path_to_results_file, **kwargs) 

123 self.load(self._grab_data_from_bilby_file, **kwargs) 

124 

125 @staticmethod 

126 def grab_priors(bilby_object, nsamples=5000): 

127 """Draw samples from the prior functions stored in the bilby file 

128 """ 

129 from pesummary.utils.array import Array 

130 

131 f = bilby_object 

132 try: 

133 samples = f.priors.sample(size=nsamples) 

134 priors = {key: Array(samples[key]) for key in samples} 

135 except Exception as e: 

136 logger.info("Failed to draw prior samples because {}".format(e)) 

137 priors = {} 

138 return priors 

139 

140 @staticmethod 

141 def grab_extra_kwargs(bilby_object): 

142 """Grab any additional information stored in the bilby file 

143 """ 

144 from pesummary.core.file.formats.bilby import config_from_object 

145 import ast 

146 

147 f = bilby_object 

148 kwargs = CoreBilby.grab_extra_kwargs(bilby_object) 

149 try: 

150 kwargs["meta_data"]["f_ref"] = \ 

151 f.meta_data["likelihood"]["waveform_arguments"]["reference_frequency"] 

152 except Exception: 

153 pass 

154 for key, item in f.meta_data["likelihood"].items(): 

155 if not isinstance(item, dict): 

156 try: 

157 if isinstance(item, bool): 

158 kwargs["meta_data"][key] = str(item) 

159 else: 

160 kwargs["meta_data"][key] = item 

161 except Exception: 

162 pass 

163 try: 

164 kwargs["meta_data"]["approximant_flags"] = \ 

165 f.meta_data["likelihood"]["waveform_arguments"] 

166 kwargs["meta_data"]["approximant"] = \ 

167 kwargs["meta_data"]["approximant_flags"].pop("waveform_approximant") 

168 except Exception: 

169 pass 

170 try: 

171 kwargs["meta_data"]["IFOs"] = \ 

172 " ".join(f.meta_data["likelihood"]["interferometers"].keys()) 

173 except Exception: 

174 pass 

175 _config = config_from_object(f) 

176 if len(_config): 

177 if "config" in _config.keys(): 

178 _config = _config["config"] 

179 options = [ 

180 "minimum_frequency", "reference_frequency", 

181 "waveform_approximant", "maximum_frequency", 

182 "waveform_arguments_dict" 

183 ] 

184 pesummary_names = [ 

185 "f_low", "f_ref", "approximant", "f_final", 

186 "approximant_flags" 

187 ] 

188 for opt, name in zip(options, pesummary_names): 

189 if opt in _config.keys(): 

190 try: 

191 _option = ast.literal_eval(_config[opt]) 

192 except ValueError: 

193 try: 

194 import re 

195 _option = re.sub( 

196 r'([A-Za-z/\.0-9\-\+][^\[\],:"}]*)', r'"\g<1>"', _config[opt] 

197 ) 

198 _option = ast.literal_eval(_option) 

199 except Exception: 

200 _option = _config[opt] 

201 if isinstance(_option, dict) and name != "approximant_flags": 

202 _value = np.min([ 

203 ast.literal_eval(value) if not isinstance(value, (float, int)) else value 

204 for key, value in _option.items() if key != "waveform" 

205 ]) 

206 _f_start = [ 

207 ast.literal_eval(value) if not isinstance(value, (float, int)) else value 

208 for key, value in _option.items() if key == "waveform" 

209 ] 

210 if len(_f_start): 

211 kwargs["meta_data"]["f_start"] = _f_start[0] 

212 else: 

213 _value = _option 

214 if _value is None and name == "approximant_flags": 

215 _value = {} 

216 kwargs["meta_data"][name] = _value 

217 return kwargs 

218 

219 @property 

220 def calibration_spline_posterior(self): 

221 if not any("recalib_" in i for i in self.parameters): 

222 return super(Bilby, self).calibration_spline_posterior 

223 ifos = np.unique( 

224 [ 

225 param.split('_')[1] for param in self.parameters if 'recalib_' 

226 in param and "non_reweighted" not in param 

227 ] 

228 ) 

229 amp_params, phase_params, log_freqs = {}, {}, {} 

230 for ifo in ifos: 

231 amp_params[ifo], phase_params[ifo] = [], [] 

232 freq_params = np.sort( 

233 [ 

234 param for param in self.parameters if 

235 'recalib_%s_frequency_' % (ifo) in param 

236 and "non_reweighted" not in param 

237 ] 

238 ) 

239 posterior = self.samples_dict 

240 log_freqs[ifo] = np.log( 

241 [posterior[param][0] for param in freq_params] 

242 ) 

243 amp_parameters = np.sort( 

244 [ 

245 param for param in self.parameters if 

246 'recalib_%s_amplitude_' % (ifo) in param 

247 and "non_reweighted" not in param 

248 ] 

249 ) 

250 amplitude = np.array([posterior[param] for param in amp_parameters]) 

251 phase_parameters = np.sort( 

252 [ 

253 param for param in self.parameters if 

254 'recalib_%s_phase_' % (ifo) in param 

255 and "non_reweighted" not in param 

256 ] 

257 ) 

258 phase = np.array([posterior[param] for param in phase_parameters]) 

259 for num, i in enumerate(amplitude): 

260 amp_params[ifo].append(i) 

261 phase_params[ifo].append(phase[num]) 

262 return log_freqs, amp_params, phase_params 

263 

264 @staticmethod 

265 def load_strain_data(path_to_strain_file): 

266 """Load the strain data 

267 

268 Parameters 

269 ---------- 

270 path_to_strain_file: str 

271 path to the strain file that you wish to load 

272 """ 

273 Bilby.load_from_function( 

274 Bilby._timeseries_from_bilby_pickle, path_to_strain_file) 

275 

276 @staticmethod 

277 def _timeseries_from_bilby_pickle(path_to_strain_file): 

278 """Load a bilby pickle file containing the strain data 

279 

280 Parameters 

281 ---------- 

282 path_to_strain_file: str 

283 path to the strain file that you wish to load 

284 """ 

285 import pickle 

286 import gwpy 

287 

288 with open(path_to_strain_file, "rb") as f: 

289 data = pickle.load(f) 

290 

291 strain_data = {} 

292 for i in data.interferometers: 

293 strain_data[i.name] = gwpy.timeseries.TimeSeries( 

294 data=i.strain_data.time_domain_strain, 

295 times=i.strain_data.time_array) 

296 return strain_data 

297 

298 @staticmethod 

299 def _grab_data_from_bilby_file(path, disable_prior=False, **kwargs): 

300 """ 

301 Load the results file using the `bilby` library 

302 

303 Complex matched filter SNRs are stored in the result file. 

304 The amplitude and angle are extracted here. 

305 """ 

306 return read_bilby(path, disable_prior=disable_prior, **kwargs) 

307 

308 def add_marginalized_parameters_from_config_file(self, config_file): 

309 """Search the configuration file and add the marginalized parameters 

310 to the list of parameters and samples 

311 

312 Parameters 

313 ---------- 

314 config_file: str 

315 path to the configuration file 

316 """ 

317 pass