1 files changed, 144 insertions, 0 deletions
diff --git a/calc_pei.py b/calc_pei.py
new file mode 100755
index 0000000..19e71bb
--- /dev/null
+++ b/calc_pei.py
@@ -0,0 +1,144 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+#
+# Aaron LI
+# Created: 2016-04-29
+# Updated: 2016-04-29
+#
+# TODO:
+#   * to calculate the PEI error
+#   * to save a plot with the rectangular and interpolation marked
+#
+
+"""
+Calculate the power excess index (PEI), which is defined the area ratio of
+the lower-left part with respect to the total rectangulr, which is further
+defined by the radial power spectrum and the scale of 0.035R500 and 0.35R500,
+in the logarithmic space.
+
+Reference:
+Zhang, C., et al. 2016, ApJ
+"""
+
+__version__ = "0.1.0"
+__date__    = "2016-04-29"
+
+
+import sys
+import glob
+import argparse
+import json
+from collections import OrderedDict
+
+import numpy as np
+import scipy as sp
+import scipy.interpolate
+import scipy.integrate
+
+from make_r500_regions import get_r500
+
+
+def calc_pei(data, r500, interp_np=101):
+    """
+    Calculate the power excess index (PEI), which is defined the area ratio
+    of the lower-left part with respect to the total rectangulr.
+
+    Arguments:
+      * data: 3-column power spectrum data (frequency, power, power_err)
+      * r500: R500 value in unit of the inverse of the above "frequency"
+      * interp_np: number of data points interpolated to calculate PEI
+    """
+    freqs     = data[:, 0]
+    psd1d     = data[:, 1]
+    psd1d_err = data[:, 2]
+    # frequency values corresponding to 0.35R500 and 0.035R500
+    f1 = 1.0 / (0.350 * r500)
+    f2 = 1.0 / (0.035 * r500)
+    # switch to the logarithmic scale
+    # XXX: how to deal with the errors
+    mask = (freqs > 0.0)
+    x  = np.log10(freqs[mask])
+    y  = np.log10(psd1d[mask])
+    x1 = np.log10(f1)
+    x2 = np.log10(f2)
+    # interpolate the power spectrum
+    f_interp = sp.interpolate.interp1d(x, y, kind="cubic", assume_sorted=True)
+    y1       = f_interp(x1)
+    y2       = f_interp(x2)
+    if interp_np % 2 == 0:
+        # Simpson's rule requires an even number of intervals
+        interp_np += 1
+    x_interp = np.linspace(x1, x2, num=interp_np)
+    y_interp = f_interp(x_interp)
+    # calculate the PEI
+    area_total = abs(x1 - x2) * abs(y1 - y2)
+    area_below = sp.integrate.simps((y_interp-y2), x_interp)
+    pei_value  = area_below / area_total
+    results = {
+            "area_total": area_total,
+            "area_below": area_below,
+            "pei_value":  pei_value,
+            "pei_err":    None,
+    }
+    return results
+
+
+def main():
+    # default arguments
+    default_infile   = "psd.txt"
+    default_outfile  = "pei.json"
+    default_infojson = "../*_INFO.json"
+
+    parser = argparse.ArgumentParser(
+            description="Calculate the power excess index (PEI)",
+            epilog="Version: %s (%s)" % (__version__, __date__))
+    parser.add_argument("-V", "--version", action="version",
+            version="%(prog)s " + "%s (%s)" % (__version__, __date__))
+    parser.add_argument("-j", "--json", dest="json", required=False,
+            help="the *_INFO.json file (default: find %s)" % default_infojson)
+    parser.add_argument("-i", "--infile", dest="infile",
+            required=False, default=default_infile,
+            help="input data of the radial power spectrum " + \
+                 "(default: %s)" % default_infile)
+    parser.add_argument("-o", "--outfile", dest="outfile",
+            required=False, default=default_outfile,
+            help="output json file to save the results " + \
+                 "(default: %s)" % default_outfile)
+    args = parser.parse_args()
+
+    info_json = glob.glob(default_infojson)[0]
+    if args.json:
+        info_json = args.json
+
+    json_str    = open(info_json).read().rstrip().rstrip(",")
+    info        = json.loads(json_str)
+    name        = info["Source Name"]
+    obsid       = int(info["Obs. ID"])
+    r500        = get_r500(info)
+    r500_kpc    = r500["r500_kpc"]
+    r500_pix    = r500["r500_pix"]
+    kpc_per_pix = r500["kpc_per_pix"]
+
+    psd_data = np.loadtxt(args.infile)
+    pei = calc_pei(psd_data, r500=r500_pix)
+
+    pei_data = OrderedDict([
+            ("name",            name),
+            ("obsid",           obsid),
+            ("r500_kpc",        r500_kpc),
+            ("r500_pix",        r500_pix),
+            ("kpc_per_pix",     kpc_per_pix),
+            ("area_total",      pei["area_total"]),
+            ("area_below",      pei["area_below"]),
+            ("pei",             pei["pei_value"]),
+            ("pei_err",         pei["pei_err"]),
+    ])
+    pei_data_json = json.dumps(pei_data, indent=2)
+    print(pei_data_json)
+    open(args.outfile, "w").write(pei_data_json+"\n")
+
+
+if __name__ == "__main__":
+    main()
+
+#  vim: set ts=4 sw=4 tw=0 fenc=utf-8 ft=python: #