From ea2a7aa7fe5cf28b670c1243f1618e650f7eee25 Mon Sep 17 00:00:00 2001
From: Gregory Ashton <gregory.ashton@ligo.org>
Date: Tue, 1 May 2018 09:17:57 +1000
Subject: [PATCH] Close #17 move output format to h5

Removes use of pickle in favour of h5 format. Uses deepdish
(https://github.com/uchicago-cs/deepdish) over the more standard
h5py for the following reasons

- h5py is built around saving 'arrays' while deepdish around saving
  dictionaries
- provides useful tools to check what data is in the h5 file
- is well written, provides good documentation and no strange error
  messages (personal experience with trying to implement h5py).
---
 peyote/result.py           | 7 +++----
 peyote/sampler.py          | 7 +++++--
 requirements.txt           | 1 +
 tutorials/BasicTutorial.py | 2 +-
 4 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/peyote/result.py b/peyote/result.py
index 372b7cd2..fb12a82e 100644
--- a/peyote/result.py
+++ b/peyote/result.py
@@ -1,6 +1,6 @@
 import logging
 import os
-import pickle
+import deepdish
 
 
 class Result(dict):
@@ -24,7 +24,7 @@ class Result(dict):
                         self.logzerr))
 
     def save_to_file(self, outdir, label):
-        file_name = '{}/{}_results.p'.format(outdir, label)
+        file_name = '{}/{}_result.h5'.format(outdir, label)
         if os.path.isdir(outdir) is False:
             os.makedirs(outdir)
         if os.path.isfile(file_name):
@@ -34,5 +34,4 @@ class Result(dict):
             os.rename(file_name, file_name + '.old')
 
         logging.info("Saving result to {}".format(file_name))
-        with open(file_name, 'wb+') as f:
-            pickle.dump(self, f)
+        deepdish.io.save(file_name, self)
diff --git a/peyote/sampler.py b/peyote/sampler.py
index 6cf63204..d8c426d0 100644
--- a/peyote/sampler.py
+++ b/peyote/sampler.py
@@ -247,7 +247,7 @@ class Pymultinest(Sampler):
 
 
 def run_sampler(likelihood, priors, label='label', outdir='outdir',
-                sampler='nestle', use_ratio=False,
+                sampler='nestle', use_ratio=False, injection_parameters=None,
                 **sampler_kwargs):
     """
     The primary interface to easy parameter estimation
@@ -268,6 +268,9 @@ def run_sampler(likelihood, priors, label='label', outdir='outdir',
     use_ratio: bool (False)
         If True, use the likelihood's loglikelihood_ratio, rather than just
         the loglikelhood.
+    injection_parameters: dict
+        A dictionary of injection parameters used in creating the data (if
+        using simulated data). Appended to the result object and saved.
     **sampler_kwargs:
         All kwargs are passed directly to the samplers `run` functino
 
@@ -289,7 +292,7 @@ def run_sampler(likelihood, priors, label='label', outdir='outdir',
         result = sampler.run_sampler()
         result.noise_logz = likelihood.noise_log_likelihood()
         result.log_bayes_factor = result.logz - result.noise_logz
-        print("")
+        result.injection_parameters = injection_parameters
         result.save_to_file(outdir=outdir, label=label)
         return result, sampler
     else:
diff --git a/requirements.txt b/requirements.txt
index c68adb35..31cb609d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -12,3 +12,4 @@ gwsurrogate
 NRSur7dq2
 chainconsumer
 nestle
+deepdish
diff --git a/tutorials/BasicTutorial.py b/tutorials/BasicTutorial.py
index 927d9177..38288827 100644
--- a/tutorials/BasicTutorial.py
+++ b/tutorials/BasicTutorial.py
@@ -84,6 +84,6 @@ sampling_parameters['luminosity_distance'] = peyote.prior.Uniform(lower=30, uppe
 
 result, sampler = peyote.sampler.run_sampler(
     likelihood, priors=sampling_parameters, label='BasicTutorial',
-    sampler='nestle', verbose=True)
+    sampler='nestle', verbose=True, injection_parameters=injection_parameters)
 sampler.plot_corner()
 print(result)
-- 
GitLab