Merge pull request #4 from fact-project/more_output

Add option to store more features from the unfolding to json

Author: maxnoe

fact_funfolding/io.py

@@ -1,10 +1,11 @@
 import json
 import astropy.units as u
+import numpy as np
 
 from fact.analysis.statistics import POINT_SOURCE_FLUX_UNIT
 
 
-def save_spectrum(outputfile, bins, flux, flux_err, counts, counts_err, g=None, bg=None, **metadata):
+def save_spectrum(outputfile, bins, flux, flux_err, counts, counts_err, g=None, bg=None, add_features=None, **metadata):
     bins = bins.to(u.GeV)
     bin_centers = 0.5 * (bins[:-1] + bins[1:])
 
@@ -39,16 +40,34 @@ def save_spectrum(outputfile, bins, flux, flux_err, counts, counts_err, g=None,
         **metadata
     }
 
-    if g is not None:
-        data['g'] = g.tolist()
+    add_if_not_none(data, g, 'g')
+    add_if_not_none(data, bg, 'bg')
 
-    if bg is not None:
-        data['bg'] = bg.tolist()
+    if add_features is not None:
+        for key in add_features:
+            add_if_not_none(data, add_features[key], key)
+
+    if metadata is not None:
+        for key in metadata:
+            add_if_not_none(data, metadata[key], key)
 
     with open(outputfile, 'w') as f:
         json.dump(data, f, indent=4)
 
 
+def add_if_not_none(result_dict, value, key):
+    if value is not None:
+        if isinstance(value, u.quantity.Quantity):
+            result_dict[key] = {
+                'value': value.value if value.isscalar else value.value.tolist(),
+                'unit': str(value.unit)
+            }
+        elif isinstance(value, np.ndarray):
+            result_dict[key] = value.tolist() if len(value) > 1 else value
+        else:
+            result_dict[key] = value
+
+
 def read_spectrum(inputfile):
     with open(inputfile, 'r') as f:
         data = json.load(f)

fact_funfolding/scripts/unfold_observations.py

@@ -167,6 +167,11 @@ def main(
         vec_f_est = vec_f_est[1:]
         sigma_vec_f = sigma_vec_f[:, 1:]
 
+    additional_features_to_save = dict()
+    additional_features_to_save['a_eff'] = a_eff
+    additional_features_to_save['a_eff_low'] = a_eff_low
+    additional_features_to_save['a_eff_high'] = a_eff_high
+
     save_spectrum(
         output_file,
         bins_true,
@@ -178,6 +183,7 @@ def main(
         bg=vec_g_bg,
         tau=config.tau,
         label=config.label,
+        add_features=additional_features_to_save,
 
     )
 

fact_funfolding/scripts/unfold_simulations.py

@@ -162,6 +162,11 @@ def main(
         vec_f_est = vec_f_est[1:]
         sigma_vec_f = sigma_vec_f[:, 1:]
 
+    additional_features_to_save = dict()
+    additional_features_to_save['a_eff'] = a_eff
+    additional_features_to_save['a_eff_low'] = a_eff_low
+    additional_features_to_save['a_eff_high'] = a_eff_high
+
     save_spectrum(
         output_file,
         bins_true,
@@ -171,6 +176,7 @@ def main(
         counts_err=sigma_vec_f,
         tau=config.tau,
         label=config.label,
+        add_features=additional_features_to_save,
     )