Coverage for pesummary/core/file/formats/bilby.py: 87.4%
174 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-05-02 08:42 +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 read_bilby(
74 path, disable_prior=False, complex_params=[], latex_dict=latex_labels,
75 nsamples_for_prior=None, _bilby_class=None, **kwargs
76):
77 """Grab the parameters and samples in a bilby file
79 Parameters
80 ----------
81 path: str
82 path to the result file you wish to read in
83 disable_prior: Bool, optional
84 if True, do not collect prior samples from the `bilby` result file.
85 Default False
86 complex_params: list, optional
87 list of parameters stored in the bilby result file which are complex
88 and you wish to store the `amplitude` and `angle` as seperate
89 posterior distributions
90 latex_dict: dict, optional
91 list of latex labels for each parameter
92 nsamples_for_prior: int, optional
93 number of samples to draw from the analytic priors
94 """
95 if _bilby_class is None:
96 _bilby_class = Bilby
97 bilby_object = _load_bilby(path)
98 posterior = bilby_object.posterior
99 _original_keys = posterior.keys()
100 for key in _original_keys:
101 for param in complex_params:
102 if param in key and any(np.iscomplex(posterior[key])):
103 posterior[key + "_amp"] = abs(posterior[key])
104 posterior[key + "_angle"] = np.angle(posterior[key])
105 posterior[key] = np.real(posterior[key])
106 elif param in key:
107 posterior[key] = np.real(posterior[key])
108 # Drop all non numeric bilby data outputs
109 posterior = posterior.select_dtypes(include=[float, int])
110 parameters = list(posterior.keys())
111 samples = posterior.to_numpy().real
112 injection = bilby_object.injection_parameters
113 if injection is None:
114 injection = {i: j for i, j in zip(
115 parameters, [float("nan")] * len(parameters))}
116 else:
117 for i in parameters:
118 if i not in injection.keys():
119 injection[i] = float("nan")
120 _injection = injection.copy()
121 for key, item in _injection.items():
122 if isinstance(item, (str, np.str_)):
123 injection.pop(key)
125 if all(i for i in (
126 bilby_object.constraint_parameter_keys,
127 bilby_object.search_parameter_keys,
128 bilby_object.fixed_parameter_keys)):
129 for key in (
130 bilby_object.constraint_parameter_keys
131 + bilby_object.search_parameter_keys
132 + bilby_object.fixed_parameter_keys):
133 if key not in latex_dict:
134 label = bilby_object.get_latex_labels_from_parameter_keys(
135 [key])[0]
136 latex_dict[key] = label
137 try:
138 extra_kwargs = _bilby_class.grab_extra_kwargs(bilby_object)
139 except Exception:
140 extra_kwargs = {"sampler": {}, "meta_data": {}}
141 extra_kwargs["sampler"]["nsamples"] = len(samples)
142 extra_kwargs["sampler"]["pe_algorithm"] = "bilby"
143 try:
144 version = bilby_object.version
145 if isinstance(version, (list, np.ndarray)):
146 version = version[0]
147 except Exception as e:
148 version = None
150 data = {
151 "parameters": parameters,
152 "samples": samples.tolist(),
153 "injection": injection,
154 "version": version,
155 "kwargs": extra_kwargs
156 }
157 if bilby_object.meta_data is not None:
158 if "command_line_args" in bilby_object.meta_data.keys():
159 data["config"] = {
160 "config": bilby_object.meta_data["command_line_args"]
161 }
162 if not disable_prior:
163 logger.debug("Drawing prior samples from bilby result file")
164 if nsamples_for_prior is None:
165 nsamples_for_prior = len(samples)
166 prior_samples = Bilby.grab_priors(
167 bilby_object, nsamples=nsamples_for_prior
168 )
169 data["prior"] = {"samples": prior_samples}
170 if len(prior_samples):
171 data["prior"]["analytic"] = prior_samples.analytic
172 else:
173 try:
174 _prior = bilby_object.priors
175 data["prior"] = {
176 "samples": {},
177 "analytic": {key: str(item) for key, item in _prior.items()}
178 }
179 except (AttributeError, KeyError):
180 pass
181 return data
184def to_bilby(
185 parameters, samples, label=None, analytic_priors=None, cls=None,
186 meta_data=None, **kwargs
187):
188 """Convert a set of samples to a bilby object
190 Parameters
191 ----------
192 parameters: list
193 list of parameters
194 samples: 2d list
195 list of samples. Columns correspond to a given parameter
196 label: str, optional
197 The label of the analysis. This is used in the filename if a filename
198 if not specified
199 """
200 from bilby.core.result import Result
201 from bilby.core.prior import Prior, PriorDict
202 from pandas import DataFrame
204 if cls is None:
205 cls = Result
206 if analytic_priors is not None:
207 priors = PriorDict._get_from_json_dict(analytic_priors)
208 search_parameters = priors.keys()
209 else:
210 priors = {param: Prior() for param in parameters}
211 search_parameters = parameters
212 posterior_data_frame = DataFrame(samples, columns=parameters)
213 bilby_object = cls(
214 search_parameter_keys=search_parameters, samples=samples, priors=priors,
215 posterior=posterior_data_frame, label="pesummary_%s" % label,
216 )
217 return bilby_object
220def _write_bilby(
221 parameters, samples, outdir="./", label=None, filename=None, overwrite=False,
222 extension="json", save=True, analytic_priors=None, cls=None,
223 meta_data=None, **kwargs
224):
225 """Write a set of samples to a bilby file
227 Parameters
228 ----------
229 parameters: list
230 list of parameters
231 samples: 2d list
232 list of samples. Columns correspond to a given parameter
233 outdir: str, optional
234 directory to write the dat file
235 label: str, optional
236 The label of the analysis. This is used in the filename if a filename
237 if not specified
238 filename: str, optional
239 The name of the file that you wish to write
240 overwrite: Bool, optional
241 If True, an existing file of the same name will be overwritten
242 extension: str, optional
243 file extension for the bilby result file. Default json.
244 save: Bool, optional
245 if True, save the bilby object to file
246 """
247 bilby_object = to_bilby(
248 parameters, samples, label=None, analytic_priors=None, cls=None,
249 meta_data=None, **kwargs
250 )
251 if save:
252 _filename = os.path.join(outdir, filename)
253 bilby_object.save_to_file(filename=_filename, extension=extension)
254 else:
255 return bilby_object
258def write_bilby(
259 parameters, samples, outdir="./", label=None, filename=None, overwrite=False,
260 extension="json", save=True, analytic_priors=None, cls=None,
261 meta_data=None, labels=None, **kwargs
262):
263 """Write a set of samples to a bilby file
265 Parameters
266 ----------
267 parameters: list
268 list of parameters
269 samples: 2d list
270 list of samples. Columns correspond to a given parameter
271 outdir: str, optional
272 directory to write the dat file
273 label: str, optional
274 The label of the analysis. This is used in the filename if a filename
275 if not specified
276 filename: str, optional
277 The name of the file that you wish to write
278 overwrite: Bool, optional
279 If True, an existing file of the same name will be overwritten
280 extension: str, optional
281 file extension for the bilby result file. Default json.
282 save: Bool, optional
283 if True, save the bilby object to file
284 """
285 from pesummary.io.write import _multi_analysis_write
287 func = _write_bilby
288 if not save:
289 func = to_bilby
290 return _multi_analysis_write(
291 func, parameters, samples, outdir=outdir, label=label,
292 filename=filename, overwrite=overwrite, extension=extension,
293 save=save, analytic_priors=analytic_priors, cls=cls,
294 meta_data=meta_data, file_format="bilby", labels=labels,
295 _return=True, **kwargs
296 )
299def config_from_file(path):
300 """Extract the configuration file stored within a bilby result file
302 Parameters
303 ----------
304 path: str
305 path to the bilby result file you wish to load
306 """
307 bilby_object = _load_bilby(path)
308 return config_from_object(bilby_object)
311def config_from_object(bilby_object):
312 """Extract the configuration file stored within a `bilby.core.result.Result`
313 object (or alike)
315 Parameters
316 ----------
317 bilby_object: bilby.core.result.Result (or alike)
318 a bilby.core.result.Result object (or alike) you wish to extract the
319 configuration file from
320 """
321 config = {}
322 if bilby_object.meta_data is not None:
323 if "command_line_args" in bilby_object.meta_data.keys():
324 config = {
325 "config": bilby_object.meta_data["command_line_args"]
326 }
327 return config
330def prior_samples_from_file(path, cls="PriorDict", nsamples=5000, **kwargs):
331 """Return a dict of prior samples from a `bilby` prior file
333 Parameters
334 ----------
335 path: str
336 path to a `bilby` prior file
337 cls: str, optional
338 class you wish to read in the prior file
339 nsamples: int, optional
340 number of samples to draw from a prior file. Default 5000
341 """
342 from bilby.core import prior
344 if isinstance(cls, str):
345 cls = getattr(prior, cls)
346 _prior = cls(filename=path)
347 samples = _prior.sample(size=nsamples)
348 return _bilby_prior_dict_to_pesummary_samples_dict(samples, prior=_prior)
351def prior_samples_from_bilby_object(bilby_object, nsamples=5000, **kwargs):
352 """Return a dict of prior samples from a `bilby.core.result.Result`
353 object
355 Parameters
356 ----------
357 bilby_object: bilby.core.result.Result
358 a bilby.core.result.Result object you wish to draw prior samples from
359 nsamples: int, optional
360 number of samples to draw from a prior file. Default 5000
361 """
362 samples = bilby_object.priors.sample(size=nsamples)
363 return _bilby_prior_dict_to_pesummary_samples_dict(
364 samples, prior=bilby_object.priors
365 )
368def _bilby_prior_dict_to_pesummary_samples_dict(samples, prior=None):
369 """Return a pesummary.utils.samples_dict.SamplesDict object from a bilby
370 priors dict
371 """
372 from pesummary.utils.samples_dict import SamplesDict
374 _samples = SamplesDict(samples)
375 if prior is not None:
376 analytic = {key: str(item) for key, item in prior.items()}
377 setattr(_samples, "analytic", analytic)
378 return _samples
381class Bilby(SingleAnalysisRead):
382 """PESummary wrapper of `bilby` (https://git.ligo.org/lscsoft/bilby). The
383 path_to_results_file argument will be passed directly to
384 `bilby.core.result.read_in_result`. All functions therefore use `bilby`
385 methods and requires `bilby` to be installed.
387 Parameters
388 ----------
389 path_to_results_file: str
390 path to the results file that you wish to read in with `bilby`.
391 disable_prior: Bool, optional
392 if True, do not collect prior samples from the `bilby` result file.
393 Default False
394 remove_nan_likelihood_samples: Bool, optional
395 if True, remove samples which have log_likelihood='nan'. Default True
397 Attributes
398 ----------
399 parameters: list
400 list of parameters stored in the result file
401 samples: 2d list
402 list of samples stored in the result file
403 samples_dict: dict
404 dictionary of samples stored in the result file keyed by parameters
405 input_version: str
406 version of the result file passed.
407 extra_kwargs: dict
408 dictionary of kwargs that were extracted from the result file
409 injection_parameters: dict
410 dictionary of injection parameters extracted from the result file
411 prior: dict
412 dictionary of prior samples keyed by parameters. The prior functions
413 are evaluated for 5000 samples.
414 pe_algorithm: str
415 name of the algorithm used to generate the posterior samples
417 Methods
418 -------
419 to_dat:
420 save the posterior samples to a .dat file
421 to_latex_table:
422 convert the posterior samples to a latex table
423 generate_latex_macros:
424 generate a set of latex macros for the stored posterior samples
425 """
426 def __init__(self, path_to_results_file, **kwargs):
427 super(Bilby, self).__init__(path_to_results_file, **kwargs)
428 self.load(self._grab_data_from_bilby_file, **kwargs)
430 @staticmethod
431 def grab_priors(bilby_object, nsamples=5000):
432 """Draw samples from the prior functions stored in the bilby file
433 """
434 try:
435 return prior_samples_from_bilby_object(
436 bilby_object, nsamples=nsamples
437 )
438 except Exception as e:
439 logger.info("Failed to draw prior samples because {}".format(e))
440 return {}
442 @staticmethod
443 def grab_extra_kwargs(bilby_object):
444 """Grab any additional information stored in the lalinference file
445 """
446 f = bilby_object
447 kwargs = {"sampler": {
448 conf.log_evidence: np.round(f.log_evidence, 2),
449 conf.log_evidence_error: np.round(f.log_evidence_err, 2),
450 conf.log_bayes_factor: np.round(f.log_bayes_factor, 2),
451 conf.log_noise_evidence: np.round(f.log_noise_evidence, 2)},
452 "meta_data": {}, "other": f.meta_data}
453 return kwargs
455 @staticmethod
456 def _grab_data_from_bilby_file(path, **kwargs):
457 """Load the results file using the `bilby` library
458 """
459 return read_bilby(path, **kwargs)
461 def add_marginalized_parameters_from_config_file(self, config_file):
462 """Search the configuration file and add the marginalized parameters
463 to the list of parameters and samples
465 Parameters
466 ----------
467 config_file: str
468 path to the configuration file
469 """
470 pass