diff --git a/bilby/core/sampler/dynesty.py b/bilby/core/sampler/dynesty.py
index 64efb33dbf1efaae05fcdaf0b3f800f39919ff43..87540c340a2dc261ccc9472fb40eb5343a962229 100644
--- a/bilby/core/sampler/dynesty.py
+++ b/bilby/core/sampler/dynesty.py
@@ -438,6 +438,7 @@ class Dynesty(NestedSampler):
current_state['posterior'] = resample_equal(
np.array(current_state['physical_samples']), weights)
+ current_state['search_parameter_keys'] = self.search_parameter_keys
except ValueError:
logger.debug("Unable to create posterior")
diff --git a/bilby/core/utils.py b/bilby/core/utils.py
index d14c642e928044524ad7105a2774dc092939c0e4..a406c5c6818cb7edaf7f4d3a445a619c732b61be 100644
--- a/bilby/core/utils.py
+++ b/bilby/core/utils.py
@@ -1006,3 +1006,10 @@ def decode_astropy_quantity(dct):
class IllegalDurationAndSamplingFrequencyException(Exception):
pass
+
+
+class tcolors:
+ KEY = '\033[93m'
+ VALUE = '\033[91m'
+ HIGHLIGHT = '\033[95m'
+ END = '\033[0m'
diff --git a/cli_bilby/bilby_result.py b/cli_bilby/bilby_result.py
index ae3a6545b05fcc29aa52d93130a839f76d950053..237a9622e19389a8203f6ec4548324abc3879c51 100644
--- a/cli_bilby/bilby_result.py
+++ b/cli_bilby/bilby_result.py
@@ -25,13 +25,18 @@ import argparse
import pandas as pd
import bilby
+from bilby.core.utils import tcolors
def setup_command_line_args():
parser = argparse.ArgumentParser(
description="Helper tool for bilby result files")
- parser.add_argument("-r", "--results", nargs='+', required=True,
- help="List of results files.")
+ parser.add_argument(
+ "results", nargs='?',
+ help="List of results files.")
+ parser.add_argument(
+ "-r", "--results", nargs='+', dest="option_results", default=list(),
+ help="List of results files (alternative to passing results as a positional argument).")
parser.add_argument("-c", "--convert", type=str, choices=['json', 'hdf5'],
help="Convert all results.", default=False)
parser.add_argument("-m", "--merge", action='store_true',
@@ -49,7 +54,16 @@ def setup_command_line_args():
parser.add_argument("--ipython", action='store_true',
help=("For each result given, drops the user into an "
"IPython shell with the result loaded in"))
- args, _ = parser.parse_known_args()
+ args = parser.parse_args()
+
+ if args.results is None:
+ args.results = []
+ if isinstance(args.results, str):
+ args.results = [args.results]
+ args.results += args.option_results
+
+ if len(args.results) == 0:
+ raise ValueError("You have not passed any results to bilby_result")
return args
@@ -81,6 +95,18 @@ def drop_to_ipython(results_list):
IPython.embed(header=message)
+def print_matches(results_list, args):
+ for r in results_list:
+ print("\nResult file: {}/{}".format(r.outdir, r.label))
+ for key in args.print:
+ for attr in r.__dict__:
+ if key in attr:
+ print_line = [
+ " ", tcolors.KEY, attr, ":", tcolors.VALUE,
+ str(getattr(r, attr)), tcolors.END]
+ print(" ".join(print_line))
+
+
def main():
args = setup_command_line_args()
results_list = read_in_results(args.results)
@@ -88,10 +114,7 @@ def main():
for r in results_list:
r.save_to_file(extension=args.convert, outdir=args.outdir)
if args.print is not None:
- for r in results_list:
- print("\nResult file: {}/{}".format(r.outdir, r.label))
- for key in args.print:
- print(" {}: {}".format(key, getattr(r, key, 'None')))
+ print_matches(results_list, args)
if args.call is not None:
for r in results_list:
for call in args.call:
diff --git a/cli_bilby/resume.py b/cli_bilby/resume.py
new file mode 100644
index 0000000000000000000000000000000000000000..e808c5febd896d9100105d2d6d0dc6e314341e63
--- /dev/null
+++ b/cli_bilby/resume.py
@@ -0,0 +1,98 @@
+""" A command line interface for converting resume files into results files """
+import argparse
+import os
+import pickle
+
+import pandas as pd
+import bilby as bilby
+
+
+def setup_command_line_args():
+ parser = argparse.ArgumentParser(
+ description=__doc__)
+ parser.add_argument(
+ "resume_files", nargs='+', help="List of resume files")
+ parser.add_argument(
+ "-f", '--format', default="json", help="Output format, defaults to json",
+ choices=["json", "hdf5", "dat"])
+ args = parser.parse_args()
+ return args
+
+
+def check_file(resume_file):
+ """ Verify the file exists and is a resume file """
+ if "resume.pickle" not in resume_file:
+ raise ValueError("File {} is not a resume file".format(resume_file))
+ if os.path.isfile(resume_file) is False:
+ raise ValueError("No file {}".format(resume_file))
+
+
+def get_outdir_and_label(resume_file):
+ """ Infer the appropriate outdir and label from the resume file name """
+ label = os.path.basename(resume_file).replace("_resume.pickle", "")
+ outdir = os.path.dirname(resume_file)
+ return outdir, label
+
+
+def read_in_pickle_file(resume_file):
+ """ Read in the pickle file
+
+ Parameters
+ ----------
+ resume_file: str
+ Input resume file path
+
+ Returns
+ -------
+ df: pandas.DataFrame
+ A data frame of the posterior
+
+ """
+ with open(resume_file, "rb") as file:
+ data = pickle.load(file)
+
+ if "posterior" in data:
+ posterior = data["posterior"]
+ else:
+ raise ValueError("Resume file has no posterior, unable to convert")
+
+ if "search_parameter_keys" in data:
+ search_parameter_keys = data["search_parameter_keys"]
+ else:
+ search_parameter_keys = ["x{}".format(i) for i in range(posterior.shape[1])]
+
+ df = pd.DataFrame(posterior, columns=search_parameter_keys)
+ return df
+
+
+def convert_df_to_posterior_samples(df, resume_file):
+ filename = resume_file.replace("pickle", "dat")
+ filename = filename.replace("resume", "preresult")
+ df.to_csv(filename, index=False, header=True, sep=' ')
+
+
+def convert_df_to_preresult(df, format, resume_file):
+ outdir, label = get_outdir_and_label(resume_file)
+ result = bilby.core.result.Result(
+ label=label, outdir=outdir, search_parameter_keys=list(df.keys()))
+ result.posterior = df
+ result.priors = dict()
+ filename = bilby.core.result.result_file_name(outdir, label, format)
+ filename = filename.replace("result.{}".format(format), "preresult.{}".format(format))
+ result.save_to_file(filename=filename, extension=format)
+
+
+def convert_resume(resume_file, args):
+ check_file(resume_file)
+ print("Converting file {} to {}".format(resume_file, args.format))
+ df = read_in_pickle_file(resume_file)
+ if args.format == "dat":
+ convert_df_to_posterior_samples(df, resume_file)
+ elif args.format in ["json", "hdf5"]:
+ convert_df_to_preresult(df, args.format, resume_file)
+
+
+def main():
+ args = setup_command_line_args()
+ for resume_file in args.resume_files:
+ convert_resume(resume_file, args)
diff --git a/setup.py b/setup.py
index 63ae9e7352b5c305cd4809c4305d94e7ef1390e6..d31ba1a55f0c77f07c07a8adaa7728ef483631ab 100644
--- a/setup.py
+++ b/setup.py
@@ -72,7 +72,7 @@ setup(name='bilby',
packages=['bilby', 'bilby.core', 'bilby.core.sampler',
'bilby.gw', 'bilby.gw.detector', 'bilby.gw.sampler',
'bilby.hyper', 'cli_bilby'],
- package_dir={'bilby': 'bilby'},
+ package_dir={'bilby': 'bilby', 'cli_bilby': 'cli_bilby'},
package_data={'bilby.gw': ['prior_files/*'],
'bilby.gw.detector': ['noise_curves/*.txt', 'detectors/*'],
'bilby': [version_file]},
@@ -87,7 +87,8 @@ setup(name='bilby',
'scipy'],
entry_points={'console_scripts':
['bilby_plot=cli_bilby.plot_multiple_posteriors:main',
- 'bilby_result=cli_bilby.bilby_result:main']
+ 'bilby_result=cli_bilby.bilby_result:main',
+ 'bilby_convert_resume=cli_bilby.resume:main']
},
classifiers=[
"Programming Language :: Python :: 2.7",