extract_sbpfit.py: new; extract SBP fitting results to CSV format

1 files changed, 92 insertions, 0 deletions

diff --git a/extract_sbpfit.py b/extract_sbpfit.py
new file mode 100755
index 0000000..1305a0d
--- /dev/null
+++ b/extract_sbpfit.py
@@ -0,0 +1,92 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+#
+# Extract the sbpfit results, and output in CSV format.
+#
+# Aaron LI
+# Created: 2016-04-27
+#
+
+import sys
+import os
+import json
+import csv
+import argparse
+from collections import OrderedDict
+
+from configobj import ConfigObj
+
+
+def extract_sbpfit(data, config):
+    """
+    Extract the SBP fitting results (sbpfit) as well as some config
+    information, and return the results as a Python dictionary.
+    """
+    results = OrderedDict([
+            # basic information from config
+            ("name",        config["name"]),
+            ("obsid",       int(config["obsid"])),
+            ("r500_pix",    float(config["r500_pix"])),
+            ("r500_kpc",    float(config["r500_kpc"])),
+            ("unit",        config["unit"]),
+            ("model",       config["model"]),
+            # basic fitting results
+            ("ndata",       data["ndata"]),
+            ("nvarys",      data["nvarys"]),
+            ("nfree",       data["nfree"]),
+            ("nfev",        data["nfev"]),
+            ("chisqr",      data["chisqr"]),
+            ("redchi",      data["redchi"]),
+            ("aic",         data["aic"]),
+            ("bic",         data["bic"]),
+    ])
+    # fitted paramters value and confidence intervals
+    results_params = extract_params(data)
+    results.update(results_params)
+    return results
+
+
+def extract_params(data):
+    """
+    Extract the values and confidence intervals (if present)
+    for each parameter.
+    """
+    results = OrderedDict()
+    for pname, pvalue in data["params"].items():
+        # best value
+        results[pname] = pvalue[0]
+        # confidence intervals (if present)
+        if "ci" in data:
+            for ci_name, ci_value in data["ci"][pname].items():
+                if ci_name == "best":
+                    continue
+                results["%s_%s_L" % (pname, ci_name)] = ci_value[0]
+                results["%s_%s_U" % (pname, ci_name)] = ci_value[1]
+    return results
+
+
+def main():
+    parser = argparse.ArgumentParser(
+            description="Extract SBP fitting results")
+    parser.add_argument("config", nargs="?", default="sbpfit.conf",
+            help="config used for SBP fitting (default: sbpfit.conf)")
+    args = parser.parse_args()
+
+    config = ConfigObj(args.config)
+    modelname = config["model"]
+    sbpfit_outfile = config.get("outfile")
+    sbpfit_outfile = config[modelname].get("outfile", sbpfit_outfile)
+    sbpfit_outfile = os.path.join(os.path.dirname(args.config), sbpfit_outfile)
+    sbpfit_results = json.load(open(sbpfit_outfile),
+                               object_pairs_hook=OrderedDict)
+    results = extract_sbpfit(data=sbpfit_results, config=config)
+
+    # output results
+    csv_writer = csv.writer(sys.stdout)
+    csv_writer.writerow(results.keys())
+    csv_writer.writerow(results.values())
+
+
+if __name__ == "__main__":
+    main()
+