Merge branch 'improvements-to-bilby_results-cli'

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")
 

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'

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:

cli_bilby/resume.py

+""" 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)

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",