Coverage for pesummary/core/file/formats/bilby.py: 85.7%
182 statements
« prev ^ index » next coverage.py v7.4.4, created at 2025-11-05 13:38 +0000
1# Licensed under an MIT style license -- see LICENSE.md
3import os
4import numpy as np
5from pesummary.core.file.formats.base_read import SingleAnalysisRead
6from pesummary.core.plots.latex_labels import latex_labels
7from pesummary import conf
8from pesummary.utils.utils import logger
10__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
13def bilby_version(path=None, data=None):
14 """Extract the version of bilby used to create a result file
16 Parameters
17 ----------
18 path: str, optional
19 path to the bilby result file
20 data: iterable, optional
21 the contents of the bilby result file, opened with either h5py or
22 json libraries
23 """
24 if path is None and data is None:
25 raise ValueError(
26 "Pass either the path to the result file, or result file object "
27 "opened with either h5py or json"
28 )
29 elif data is None:
30 import json
31 import h5py
32 try:
33 data = h5py.File(path, "r")
34 except Exception:
35 with open(path, "r") as f:
36 data = json.load(f)
37 option1 = lambda data: data["version"][()]
38 option2 = lambda data: data["version"]
39 for opt in [option1, option2]:
40 try:
41 version = opt(data)
42 if isinstance(version, (list, np.ndarray)):
43 return version[0]
44 return version
45 except Exception:
46 pass
47 raise ValueError("Unable to extract version from bilby result file")
50def _load_bilby(path):
51 """Wrapper for `bilby.core.result.read_in_result`
53 Parameters
54 ----------
55 path: str
56 path to the bilby result file you wish to load
57 """
58 from bilby.core.result import read_in_result
59 from bilby.core.utils import get_version_information
60 result_version = bilby_version(path=path)
61 if isinstance(result_version, bytes):
62 result_version = result_version.decode("utf-8")
63 env_version = 'bilby={}'.format(get_version_information())
64 if result_version != env_version:
65 logger.warning(
66 "Result file was written with version {} while your environment "
67 "has version {}. This may cause problems when reading the result "
68 "file.".format(result_version, env_version)
69 )
70 return read_in_result(filename=path)
73def _check_bilby_attributes(result, attribute):
74 """Check the bilby attributes are the same format across different versions.
75 Currently this checks that an empty list is returned if the attribute is
76 None or if the length is 0. This can be deprecated after a few bilby releases.
78 Parameters
79 ----------
80 result: bilby.core.result.Result
81 bilby result file to check
82 attribute: str
83 attribute you wish to check
84 """
85 _attr = getattr(result, attribute)
86 if (_attr is None) or (not len(_attr)):
87 _attr = []
88 return _attr
91def read_bilby(
92 path, disable_prior=False, complex_params=[], latex_dict=latex_labels,
93 nsamples_for_prior=None, _bilby_class=None, **kwargs
94):
95 """Grab the parameters and samples in a bilby file
97 Parameters
98 ----------
99 path: str
100 path to the result file you wish to read in
101 disable_prior: Bool, optional
102 if True, do not collect prior samples from the `bilby` result file.
103 Default False
104 complex_params: list, optional
105 list of parameters stored in the bilby result file which are complex
106 and you wish to store the `amplitude` and `angle` as seperate
107 posterior distributions
108 latex_dict: dict, optional
109 list of latex labels for each parameter
110 nsamples_for_prior: int, optional
111 number of samples to draw from the analytic priors
112 """
113 if _bilby_class is None:
114 _bilby_class = Bilby
115 bilby_object = _load_bilby(path)
116 posterior = bilby_object.posterior
117 _original_keys = posterior.keys()
118 for key in _original_keys:
119 for param in complex_params:
120 if param in key and any(np.iscomplex(posterior[key])):
121 posterior[key + "_abs"] = abs(posterior[key])
122 posterior[key + "_angle"] = np.angle(posterior[key])
123 posterior[key] = np.real(posterior[key])
124 elif param in key:
125 posterior[key] = np.real(posterior[key])
126 # Drop all non numeric bilby data outputs
127 posterior = posterior.select_dtypes(include=[float, int])
128 parameters = list(posterior.keys())
129 samples = posterior.to_numpy().real
130 injection = bilby_object.injection_parameters
131 if injection is None:
132 injection = {i: j for i, j in zip(
133 parameters, [float("nan")] * len(parameters))}
134 else:
135 for i in parameters:
136 if i not in injection.keys():
137 injection[i] = float("nan")
138 _injection = injection.copy()
139 for key, item in _injection.items():
140 if isinstance(item, (str, np.str_)):
141 injection.pop(key)
143 _fixed_keys = _check_bilby_attributes(
144 bilby_object, "fixed_parameter_keys"
145 )
146 _constraint_keys = _check_bilby_attributes(
147 bilby_object, "constraint_parameter_keys"
148 )
149 _search_keys = _check_bilby_attributes(
150 bilby_object, "search_parameter_keys"
151 )
152 if all(i for i in (_constraint_keys, _search_keys, _fixed_keys)):
153 for key in (_constraint_keys + _search_keys + _fixed_keys):
154 if key not in latex_dict:
155 label = bilby_object.get_latex_labels_from_parameter_keys(
156 [key])[0]
157 latex_dict[key] = label
158 try:
159 extra_kwargs = _bilby_class.grab_extra_kwargs(bilby_object)
160 except Exception:
161 extra_kwargs = {"sampler": {}, "meta_data": {}}
162 extra_kwargs["sampler"]["nsamples"] = len(samples)
163 extra_kwargs["sampler"]["pe_algorithm"] = "bilby"
164 try:
165 version = bilby_object.version
166 if isinstance(version, (list, np.ndarray)):
167 version = version[0]
168 except Exception as e:
169 version = None
171 data = {
172 "parameters": parameters,
173 "samples": samples.tolist(),
174 "injection": injection,
175 "version": version,
176 "kwargs": extra_kwargs
177 }
178 if bilby_object.meta_data is not None:
179 if "command_line_args" in bilby_object.meta_data.keys():
180 data["config"] = {
181 "config": bilby_object.meta_data["command_line_args"]
182 }
183 if not disable_prior:
184 logger.debug("Drawing prior samples from bilby result file")
185 if nsamples_for_prior is None:
186 nsamples_for_prior = len(samples)
187 prior_samples = Bilby.grab_priors(
188 bilby_object, nsamples=nsamples_for_prior
189 )
190 data["prior"] = {"samples": prior_samples}
191 if len(prior_samples):
192 data["prior"]["analytic"] = prior_samples.analytic
193 else:
194 try:
195 _prior = bilby_object.priors
196 data["prior"] = {
197 "samples": {},
198 "analytic": {key: str(item) for key, item in _prior.items()}
199 }
200 except (AttributeError, KeyError):
201 pass
202 return data
205def to_bilby(
206 parameters, samples, label=None, analytic_priors=None, cls=None,
207 meta_data=None, **kwargs
208):
209 """Convert a set of samples to a bilby object
211 Parameters
212 ----------
213 parameters: list
214 list of parameters
215 samples: 2d list
216 list of samples. Columns correspond to a given parameter
217 label: str, optional
218 The label of the analysis. This is used in the filename if a filename
219 if not specified
220 """
221 from bilby.core.result import Result
222 from bilby.core.prior import Prior, PriorDict
223 from pandas import DataFrame
225 if cls is None:
226 cls = Result
227 if analytic_priors is not None:
228 priors = PriorDict._get_from_json_dict(analytic_priors)
229 search_parameters = priors.keys()
230 else:
231 priors = {param: Prior() for param in parameters}
232 search_parameters = parameters
233 posterior_data_frame = DataFrame(samples, columns=parameters)
234 bilby_object = cls(
235 search_parameter_keys=search_parameters, samples=samples, priors=priors,
236 posterior=posterior_data_frame, label="pesummary_%s" % label,
237 )
238 return bilby_object
241def _write_bilby(
242 parameters, samples, outdir="./", label=None, filename=None, overwrite=False,
243 extension="json", save=True, analytic_priors=None, cls=None,
244 meta_data=None, **kwargs
245):
246 """Write a set of samples to a bilby file
248 Parameters
249 ----------
250 parameters: list
251 list of parameters
252 samples: 2d list
253 list of samples. Columns correspond to a given parameter
254 outdir: str, optional
255 directory to write the dat file
256 label: str, optional
257 The label of the analysis. This is used in the filename if a filename
258 if not specified
259 filename: str, optional
260 The name of the file that you wish to write
261 overwrite: Bool, optional
262 If True, an existing file of the same name will be overwritten
263 extension: str, optional
264 file extension for the bilby result file. Default json.
265 save: Bool, optional
266 if True, save the bilby object to file
267 """
268 bilby_object = to_bilby(
269 parameters, samples, label=None, analytic_priors=None, cls=None,
270 meta_data=None, **kwargs
271 )
272 if save:
273 _filename = os.path.join(outdir, filename)
274 bilby_object.save_to_file(filename=_filename, extension=extension)
275 else:
276 return bilby_object
279def write_bilby(
280 parameters, samples, outdir="./", label=None, filename=None, overwrite=False,
281 extension="json", save=True, analytic_priors=None, cls=None,
282 meta_data=None, labels=None, **kwargs
283):
284 """Write a set of samples to a bilby file
286 Parameters
287 ----------
288 parameters: list
289 list of parameters
290 samples: 2d list
291 list of samples. Columns correspond to a given parameter
292 outdir: str, optional
293 directory to write the dat file
294 label: str, optional
295 The label of the analysis. This is used in the filename if a filename
296 if not specified
297 filename: str, optional
298 The name of the file that you wish to write
299 overwrite: Bool, optional
300 If True, an existing file of the same name will be overwritten
301 extension: str, optional
302 file extension for the bilby result file. Default json.
303 save: Bool, optional
304 if True, save the bilby object to file
305 """
306 from pesummary.io.write import _multi_analysis_write
308 func = _write_bilby
309 if not save:
310 func = to_bilby
311 return _multi_analysis_write(
312 func, parameters, samples, outdir=outdir, label=label,
313 filename=filename, overwrite=overwrite, extension=extension,
314 save=save, analytic_priors=analytic_priors, cls=cls,
315 meta_data=meta_data, file_format="bilby", labels=labels,
316 _return=True, **kwargs
317 )
320def config_from_file(path):
321 """Extract the configuration file stored within a bilby result file
323 Parameters
324 ----------
325 path: str
326 path to the bilby result file you wish to load
327 """
328 bilby_object = _load_bilby(path)
329 return config_from_object(bilby_object)
332def config_from_object(bilby_object):
333 """Extract the configuration file stored within a `bilby.core.result.Result`
334 object (or alike)
336 Parameters
337 ----------
338 bilby_object: bilby.core.result.Result (or alike)
339 a bilby.core.result.Result object (or alike) you wish to extract the
340 configuration file from
341 """
342 config = {}
343 if bilby_object.meta_data is not None:
344 if "command_line_args" in bilby_object.meta_data.keys():
345 config = {
346 "config": bilby_object.meta_data["command_line_args"]
347 }
348 return config
351def prior_samples_from_file(path, cls="PriorDict", nsamples=5000, **kwargs):
352 """Return a dict of prior samples from a `bilby` prior file
354 Parameters
355 ----------
356 path: str
357 path to a `bilby` prior file
358 cls: str, optional
359 class you wish to read in the prior file
360 nsamples: int, optional
361 number of samples to draw from a prior file. Default 5000
362 """
363 from bilby.core import prior
365 if isinstance(cls, str):
366 cls = getattr(prior, cls)
367 _prior = cls(filename=path)
368 samples = _prior.sample(size=nsamples)
369 return _bilby_prior_dict_to_pesummary_samples_dict(samples, prior=_prior)
372def prior_samples_from_bilby_object(bilby_object, nsamples=5000, **kwargs):
373 """Return a dict of prior samples from a `bilby.core.result.Result`
374 object
376 Parameters
377 ----------
378 bilby_object: bilby.core.result.Result
379 a bilby.core.result.Result object you wish to draw prior samples from
380 nsamples: int, optional
381 number of samples to draw from a prior file. Default 5000
382 """
383 samples = bilby_object.priors.sample(size=nsamples)
384 return _bilby_prior_dict_to_pesummary_samples_dict(
385 samples, prior=bilby_object.priors
386 )
389def _bilby_prior_dict_to_pesummary_samples_dict(samples, prior=None):
390 """Return a pesummary.utils.samples_dict.SamplesDict object from a bilby
391 priors dict
392 """
393 from pesummary.utils.samples_dict import SamplesDict
395 _samples = SamplesDict(samples)
396 if prior is not None:
397 analytic = {key: str(item) for key, item in prior.items()}
398 setattr(_samples, "analytic", analytic)
399 return _samples
402class Bilby(SingleAnalysisRead):
403 """PESummary wrapper of `bilby` (https://git.ligo.org/lscsoft/bilby). The
404 path_to_results_file argument will be passed directly to
405 `bilby.core.result.read_in_result`. All functions therefore use `bilby`
406 methods and requires `bilby` to be installed.
408 Parameters
409 ----------
410 path_to_results_file: str
411 path to the results file that you wish to read in with `bilby`.
412 disable_prior: Bool, optional
413 if True, do not collect prior samples from the `bilby` result file.
414 Default False
415 remove_nan_likelihood_samples: Bool, optional
416 if True, remove samples which have log_likelihood='nan'. Default True
418 Attributes
419 ----------
420 parameters: list
421 list of parameters stored in the result file
422 samples: 2d list
423 list of samples stored in the result file
424 samples_dict: dict
425 dictionary of samples stored in the result file keyed by parameters
426 input_version: str
427 version of the result file passed.
428 extra_kwargs: dict
429 dictionary of kwargs that were extracted from the result file
430 injection_parameters: dict
431 dictionary of injection parameters extracted from the result file
432 prior: dict
433 dictionary of prior samples keyed by parameters. The prior functions
434 are evaluated for 5000 samples.
435 pe_algorithm: str
436 name of the algorithm used to generate the posterior samples
438 Methods
439 -------
440 to_dat:
441 save the posterior samples to a .dat file
442 to_latex_table:
443 convert the posterior samples to a latex table
444 generate_latex_macros:
445 generate a set of latex macros for the stored posterior samples
446 """
447 def __init__(self, path_to_results_file, **kwargs):
448 super(Bilby, self).__init__(path_to_results_file, **kwargs)
449 self.load(self._grab_data_from_bilby_file, **kwargs)
451 @staticmethod
452 def grab_priors(bilby_object, nsamples=5000):
453 """Draw samples from the prior functions stored in the bilby file
454 """
455 try:
456 return prior_samples_from_bilby_object(
457 bilby_object, nsamples=nsamples
458 )
459 except Exception as e:
460 logger.info("Failed to draw prior samples because {}".format(e))
461 return {}
463 @staticmethod
464 def grab_extra_kwargs(bilby_object):
465 """Grab any additional information stored in the lalinference file
466 """
467 f = bilby_object
468 kwargs = {"sampler": {
469 conf.log_evidence: np.round(f.log_evidence, 2),
470 conf.log_evidence_error: np.round(f.log_evidence_err, 2),
471 conf.log_bayes_factor: np.round(f.log_bayes_factor, 2),
472 conf.log_noise_evidence: np.round(f.log_noise_evidence, 2)},
473 "meta_data": {}, "other": f.meta_data}
474 return kwargs
476 @staticmethod
477 def _grab_data_from_bilby_file(path, **kwargs):
478 """Load the results file using the `bilby` library
479 """
480 return read_bilby(path, **kwargs)
482 def add_marginalized_parameters_from_config_file(self, config_file):
483 """Search the configuration file and add the marginalized parameters
484 to the list of parameters and samples
486 Parameters
487 ----------
488 config_file: str
489 path to the configuration file
490 """
491 pass