Merge branch 'result-format' into 'master'

Allow different result file formats

Closes #214

See merge request lscsoft/bilby_pipe!395

bilby_pipe/data_analysis.py

@@ -57,6 +57,7 @@ class DataAnalysisInput(Input):
         # Naming arguments
         self.outdir = args.outdir
         self.label = args.label
+        self.result_format = args.result_format
 
         # Data dump file to run on
         self.data_dump_file = args.data_dump_file
@@ -241,6 +242,7 @@ class DataAnalysisInput(Input):
             meta_data=self.meta_data,
             result_class=self.result_class,
             exit_code=CHECKPOINT_EXIT_CODE,
+            save=self.result_format,
             **self.sampler_kwargs,
         )
 

bilby_pipe/job_creation/nodes/analysis_node.py

@@ -78,7 +78,7 @@ class AnalysisNode(Node):
 
     @property
     def result_file(self):
-        return f"{self.inputs.result_directory}/{self.job_name}_result.json"
+        return f"{self.inputs.result_directory}/{self.job_name}_result.{self.inputs.result_format}"
 
     @property
     def slurm_walltime(self):

bilby_pipe/job_creation/nodes/merge_node.py

@@ -17,6 +17,7 @@ class MergeNode(Node):
             self.arguments.append(pn.result_file)
         self.arguments.add("outdir", self.inputs.result_directory)
         self.arguments.add("label", self.label)
+        self.arguments.add("extension", self.inputs.result_format)
         self.arguments.add_flag("merge")
 
         self.process_node()
@@ -37,4 +38,4 @@ class MergeNode(Node):
 
     @property
     def result_file(self):
-        return f"{self.inputs.result_directory}/{self.label}_result.json"
+        return f"{self.inputs.result_directory}/{self.label}_result.{self.inputs.result_format}"

bilby_pipe/main.py

@@ -63,6 +63,7 @@ class MainInput(Input):
         self.transfer_files = args.transfer_files
         self.osg = args.osg
         self.analysis_executable = args.analysis_executable
+        self.result_format = args.result_format
 
         self.webdir = args.webdir
         self.email = args.email

bilby_pipe/parser.py

@@ -739,6 +739,13 @@ def create_parser(top_level=True):
         default=None,
         help="Arguments (in the form of a dictionary) to pass to the summarypages executable",
     )
+    output_parser.add(
+        "--result-format",
+        type=str,
+        default="json",
+        choices=["json", "hdf5", "pickle"],
+        help="Format to save the result file in.",
+    )
 
     prior_parser = parser.add_argument_group(
         title="Prior arguments", description="Specify the prior settings."

bilby_pipe/plot.py

@@ -54,7 +54,15 @@ def _parse_and_load():
 
     logger.info(f"Generating plots for results file {args.result}")
 
-    result = CBCResult.from_json(args.result)
+    extension = os.path.splitext(args.result)[-1][1:]
+    if extension == "json":
+        result = CBCResult.from_json(args.result)
+    elif extension == "hdf5":
+        result = CBCResult.from_hdf5(args.result)
+    elif extension == "pkl":
+        result = CBCResult.from_pickle(args.result)
+    else:
+        raise ValueError(f"Result format {extension} not understood.")
     if "data_dump" in result.meta_data and os.path.exists(
         result.meta_data["data_dump"]
     ):

bilby_pipe/pp_test.py

@@ -50,7 +50,7 @@ def create_parser():
 
 def get_results_filenames(args):
     results_files = []
-    for extension in ["json", "h5", "hdf5"]:
+    for extension in ["json", "h5", "hdf5", "pkl"]:
         glob_string = os.path.join(args.directory, "*result*" + extension)
         results_files += glob.glob(glob_string)
     results_files = [rf for rf in results_files if os.path.isfile(rf)]

tests/main_test.py

@@ -94,6 +94,9 @@ class TestMainInput(unittest.TestCase):
     def test_accounting(self):
         self.assertEqual(self.inputs.accounting, self.args.accounting)
 
+    def test_result_format(self):
+        self.assertEqual(self.inputs.result_format, self.args.result_format)
+
     def test_detectors_single(self):
         # Test the detector set in the ini file
         self.assertEqual(self.inputs.detectors, ["H1"])

tests/test_main_input.ini

@@ -4,6 +4,7 @@ label = label
 detectors = H1
 injection = True
 n-simulation = 1
+result-format = hdf5
 
 create-summary = True
 webdir = path/to/webdir