# Licensed under an MIT style license -- see LICENSE.md
import os
import numpy as np
from pesummary.core.file.formats.base_read import SingleAnalysisRead
from pesummary.core.plots.latex_labels import latex_labels
from pesummary import conf
from pesummary.utils.utils import logger
__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
def bilby_version(path=None, data=None):
"""Extract the version of bilby used to create a result file
Parameters
----------
path: str, optional
path to the bilby result file
data: iterable, optional
the contents of the bilby result file, opened with either h5py or
json libraries
"""
if path is None and data is None:
raise ValueError(
"Pass either the path to the result file, or result file object "
"opened with either h5py or json"
)
elif data is None:
import json
import h5py
try:
data = h5py.File(path, "r")
except Exception:
with open(path, "r") as f:
data = json.load(f)
option1 = lambda data: data["version"][()]
option2 = lambda data: data["version"]
for opt in [option1, option2]:
try:
version = opt(data)
if isinstance(version, (list, np.ndarray)):
return version[0]
return version
except Exception:
pass
raise ValueError("Unable to extract version from bilby result file")
def _load_bilby(path):
"""Wrapper for `bilby.core.result.read_in_result`
Parameters
----------
path: str
path to the bilby result file you wish to load
"""
from bilby.core.result import read_in_result
from bilby.core.utils import get_version_information
result_version = bilby_version(path=path)
if isinstance(result_version, bytes):
result_version = result_version.decode("utf-8")
env_version = 'bilby={}'.format(get_version_information())
if result_version != env_version:
logger.warning(
"Result file was written with version {} while your environment "
"has version {}. This may cause problems when reading the result "
"file.".format(result_version, env_version)
)
return read_in_result(filename=path)
[docs]
def read_bilby(
path, disable_prior=False, complex_params=[], latex_dict=latex_labels,
nsamples_for_prior=None, _bilby_class=None, **kwargs
):
"""Grab the parameters and samples in a bilby file
Parameters
----------
path: str
path to the result file you wish to read in
disable_prior: Bool, optional
if True, do not collect prior samples from the `bilby` result file.
Default False
complex_params: list, optional
list of parameters stored in the bilby result file which are complex
and you wish to store the `amplitude` and `angle` as seperate
posterior distributions
latex_dict: dict, optional
list of latex labels for each parameter
nsamples_for_prior: int, optional
number of samples to draw from the analytic priors
"""
if _bilby_class is None:
_bilby_class = Bilby
bilby_object = _load_bilby(path)
posterior = bilby_object.posterior
_original_keys = posterior.keys()
for key in _original_keys:
for param in complex_params:
if param in key and any(np.iscomplex(posterior[key])):
posterior[key + "_amp"] = abs(posterior[key])
posterior[key + "_angle"] = np.angle(posterior[key])
posterior[key] = np.real(posterior[key])
elif param in key:
posterior[key] = np.real(posterior[key])
# Drop all non numeric bilby data outputs
posterior = posterior.select_dtypes(include=[float, int])
parameters = list(posterior.keys())
samples = posterior.to_numpy().real
injection = bilby_object.injection_parameters
if injection is None:
injection = {i: j for i, j in zip(
parameters, [float("nan")] * len(parameters))}
else:
for i in parameters:
if i not in injection.keys():
injection[i] = float("nan")
_injection = injection.copy()
for key, item in _injection.items():
if isinstance(item, (str, np.str_)):
injection.pop(key)
if all(i for i in (
bilby_object.constraint_parameter_keys,
bilby_object.search_parameter_keys,
bilby_object.fixed_parameter_keys)):
for key in (
bilby_object.constraint_parameter_keys
+ bilby_object.search_parameter_keys
+ bilby_object.fixed_parameter_keys):
if key not in latex_dict:
label = bilby_object.get_latex_labels_from_parameter_keys(
[key])[0]
latex_dict[key] = label
try:
extra_kwargs = _bilby_class.grab_extra_kwargs(bilby_object)
except Exception:
extra_kwargs = {"sampler": {}, "meta_data": {}}
extra_kwargs["sampler"]["nsamples"] = len(samples)
extra_kwargs["sampler"]["pe_algorithm"] = "bilby"
try:
version = bilby_object.version
if isinstance(version, (list, np.ndarray)):
version = version[0]
except Exception as e:
version = None
data = {
"parameters": parameters,
"samples": samples.tolist(),
"injection": injection,
"version": version,
"kwargs": extra_kwargs
}
if bilby_object.meta_data is not None:
if "command_line_args" in bilby_object.meta_data.keys():
data["config"] = {
"config": bilby_object.meta_data["command_line_args"]
}
if not disable_prior:
logger.debug("Drawing prior samples from bilby result file")
if nsamples_for_prior is None:
nsamples_for_prior = len(samples)
prior_samples = Bilby.grab_priors(
bilby_object, nsamples=nsamples_for_prior
)
data["prior"] = {"samples": prior_samples}
if len(prior_samples):
data["prior"]["analytic"] = prior_samples.analytic
else:
try:
_prior = bilby_object.priors
data["prior"] = {
"samples": {},
"analytic": {key: str(item) for key, item in _prior.items()}
}
except (AttributeError, KeyError):
pass
return data
def to_bilby(
parameters, samples, label=None, analytic_priors=None, cls=None,
meta_data=None, **kwargs
):
"""Convert a set of samples to a bilby object
Parameters
----------
parameters: list
list of parameters
samples: 2d list
list of samples. Columns correspond to a given parameter
label: str, optional
The label of the analysis. This is used in the filename if a filename
if not specified
"""
from bilby.core.result import Result
from bilby.core.prior import Prior, PriorDict
from pandas import DataFrame
if cls is None:
cls = Result
if analytic_priors is not None:
priors = PriorDict._get_from_json_dict(analytic_priors)
search_parameters = priors.keys()
else:
priors = {param: Prior() for param in parameters}
search_parameters = parameters
posterior_data_frame = DataFrame(samples, columns=parameters)
bilby_object = cls(
search_parameter_keys=search_parameters, samples=samples, priors=priors,
posterior=posterior_data_frame, label="pesummary_%s" % label,
)
return bilby_object
def _write_bilby(
parameters, samples, outdir="./", label=None, filename=None, overwrite=False,
extension="json", save=True, analytic_priors=None, cls=None,
meta_data=None, **kwargs
):
"""Write a set of samples to a bilby file
Parameters
----------
parameters: list
list of parameters
samples: 2d list
list of samples. Columns correspond to a given parameter
outdir: str, optional
directory to write the dat file
label: str, optional
The label of the analysis. This is used in the filename if a filename
if not specified
filename: str, optional
The name of the file that you wish to write
overwrite: Bool, optional
If True, an existing file of the same name will be overwritten
extension: str, optional
file extension for the bilby result file. Default json.
save: Bool, optional
if True, save the bilby object to file
"""
bilby_object = to_bilby(
parameters, samples, label=None, analytic_priors=None, cls=None,
meta_data=None, **kwargs
)
if save:
_filename = os.path.join(outdir, filename)
bilby_object.save_to_file(filename=_filename, extension=extension)
else:
return bilby_object
def write_bilby(
parameters, samples, outdir="./", label=None, filename=None, overwrite=False,
extension="json", save=True, analytic_priors=None, cls=None,
meta_data=None, labels=None, **kwargs
):
"""Write a set of samples to a bilby file
Parameters
----------
parameters: list
list of parameters
samples: 2d list
list of samples. Columns correspond to a given parameter
outdir: str, optional
directory to write the dat file
label: str, optional
The label of the analysis. This is used in the filename if a filename
if not specified
filename: str, optional
The name of the file that you wish to write
overwrite: Bool, optional
If True, an existing file of the same name will be overwritten
extension: str, optional
file extension for the bilby result file. Default json.
save: Bool, optional
if True, save the bilby object to file
"""
from pesummary.io.write import _multi_analysis_write
func = _write_bilby
if not save:
func = to_bilby
return _multi_analysis_write(
func, parameters, samples, outdir=outdir, label=label,
filename=filename, overwrite=overwrite, extension=extension,
save=save, analytic_priors=analytic_priors, cls=cls,
meta_data=meta_data, file_format="bilby", labels=labels,
_return=True, **kwargs
)
def config_from_file(path):
"""Extract the configuration file stored within a bilby result file
Parameters
----------
path: str
path to the bilby result file you wish to load
"""
bilby_object = _load_bilby(path)
return config_from_object(bilby_object)
def config_from_object(bilby_object):
"""Extract the configuration file stored within a `bilby.core.result.Result`
object (or alike)
Parameters
----------
bilby_object: bilby.core.result.Result (or alike)
a bilby.core.result.Result object (or alike) you wish to extract the
configuration file from
"""
config = {}
if bilby_object.meta_data is not None:
if "command_line_args" in bilby_object.meta_data.keys():
config = {
"config": bilby_object.meta_data["command_line_args"]
}
return config
def prior_samples_from_file(path, cls="PriorDict", nsamples=5000, **kwargs):
"""Return a dict of prior samples from a `bilby` prior file
Parameters
----------
path: str
path to a `bilby` prior file
cls: str, optional
class you wish to read in the prior file
nsamples: int, optional
number of samples to draw from a prior file. Default 5000
"""
from bilby.core import prior
if isinstance(cls, str):
cls = getattr(prior, cls)
_prior = cls(filename=path)
samples = _prior.sample(size=nsamples)
return _bilby_prior_dict_to_pesummary_samples_dict(samples, prior=_prior)
def prior_samples_from_bilby_object(bilby_object, nsamples=5000, **kwargs):
"""Return a dict of prior samples from a `bilby.core.result.Result`
object
Parameters
----------
bilby_object: bilby.core.result.Result
a bilby.core.result.Result object you wish to draw prior samples from
nsamples: int, optional
number of samples to draw from a prior file. Default 5000
"""
samples = bilby_object.priors.sample(size=nsamples)
return _bilby_prior_dict_to_pesummary_samples_dict(
samples, prior=bilby_object.priors
)
def _bilby_prior_dict_to_pesummary_samples_dict(samples, prior=None):
"""Return a pesummary.utils.samples_dict.SamplesDict object from a bilby
priors dict
"""
from pesummary.utils.samples_dict import SamplesDict
_samples = SamplesDict(samples)
if prior is not None:
analytic = {key: str(item) for key, item in prior.items()}
setattr(_samples, "analytic", analytic)
return _samples
class Bilby(SingleAnalysisRead):
"""PESummary wrapper of `bilby` (https://git.ligo.org/lscsoft/bilby). The
path_to_results_file argument will be passed directly to
`bilby.core.result.read_in_result`. All functions therefore use `bilby`
methods and requires `bilby` to be installed.
Parameters
----------
path_to_results_file: str
path to the results file that you wish to read in with `bilby`.
disable_prior: Bool, optional
if True, do not collect prior samples from the `bilby` result file.
Default False
remove_nan_likelihood_samples: Bool, optional
if True, remove samples which have log_likelihood='nan'. Default True
Attributes
----------
parameters: list
list of parameters stored in the result file
samples: 2d list
list of samples stored in the result file
samples_dict: dict
dictionary of samples stored in the result file keyed by parameters
input_version: str
version of the result file passed.
extra_kwargs: dict
dictionary of kwargs that were extracted from the result file
injection_parameters: dict
dictionary of injection parameters extracted from the result file
prior: dict
dictionary of prior samples keyed by parameters. The prior functions
are evaluated for 5000 samples.
pe_algorithm: str
name of the algorithm used to generate the posterior samples
Methods
-------
to_dat:
save the posterior samples to a .dat file
to_latex_table:
convert the posterior samples to a latex table
generate_latex_macros:
generate a set of latex macros for the stored posterior samples
"""
def __init__(self, path_to_results_file, **kwargs):
super(Bilby, self).__init__(path_to_results_file, **kwargs)
self.load(self._grab_data_from_bilby_file, **kwargs)
@staticmethod
def grab_priors(bilby_object, nsamples=5000):
"""Draw samples from the prior functions stored in the bilby file
"""
try:
return prior_samples_from_bilby_object(
bilby_object, nsamples=nsamples
)
except Exception as e:
logger.info("Failed to draw prior samples because {}".format(e))
return {}
@staticmethod
def grab_extra_kwargs(bilby_object):
"""Grab any additional information stored in the lalinference file
"""
f = bilby_object
kwargs = {"sampler": {
conf.log_evidence: np.round(f.log_evidence, 2),
conf.log_evidence_error: np.round(f.log_evidence_err, 2),
conf.log_bayes_factor: np.round(f.log_bayes_factor, 2),
conf.log_noise_evidence: np.round(f.log_noise_evidence, 2)},
"meta_data": {}, "other": f.meta_data}
return kwargs
@staticmethod
def _grab_data_from_bilby_file(path, **kwargs):
"""Load the results file using the `bilby` library
"""
return read_bilby(path, **kwargs)
def add_marginalized_parameters_from_config_file(self, config_file):
"""Search the configuration file and add the marginalized parameters
to the list of parameters and samples
Parameters
----------
config_file: str
path to the configuration file
"""
pass