Add 'calc_centroid.py' to replace 'chandra_xcentroid.sh'

1 files changed, 184 insertions, 0 deletions

diff --git a/scripts/calc_centroid.py b/scripts/calc_centroid.py
new file mode 100755
index 0000000..9a9a67d
--- /dev/null
+++ b/scripts/calc_centroid.py
@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2017 Weitian LI <liweitianux@live.com>
+# MIT license
+
+"""
+Calculate the coordinate of the emission centroid within the image.
+
+The image are smoothed first, and then an iterative procedure with
+two phases is applied to determine the emission centroid.
+"""
+
+import os
+import sys
+import argparse
+import subprocess
+import tempfile
+
+from manifest import get_manifest
+from setup_pfiles import setup_pfiles
+from ds9 import ds9_view
+
+
+def smooth_image(infile, outfile=None,
+                 kernelspec="lib:gaus(2,5,1,10,10)", method="fft",
+                 clobber=False):
+    """
+    Smooth the image by a Gaussian kernel using the ``aconvolve`` tool.
+
+    Parameters
+    ----------
+    infile : str
+        Path to the input image file
+    outfile : str, optional
+        Filename/path of the output smoothed image
+        (default: build in format ``<infile_basename>_aconv.fits``)
+    kernelspec : str, optional
+        Kernel specification for ``aconvolve``
+    method : str, optional
+        Smooth method for ``aconvolve``
+
+    Returns
+    -------
+    outfile : str
+        Filename/path of the smoothed image
+    """
+    clobber = "yes" if clobber else "no"
+    if outfile is None:
+        outfile = os.path.splitext(infile)[0] + "_aconv.fits"
+    subprocess.check_call(["punlearn", "aconvolve"])
+    subprocess.check_call([
+        "aconvolve", "infile=%s" % infile, "outfile=%s" % outfile,
+        "kernelspec=%s" % kernelspec, "method=%s" % method,
+        "clobber=%s" % clobber
+    ])
+    return outfile
+
+
+def get_peak(image, center=None, radius=300):
+    """
+    Get the peak coordinate on the image within the circle if specified.
+
+    Parameters
+    ----------
+    image : str
+        Path to the image file.
+    center : 2-float tuple
+        Central (physical) coordinate of the circle.
+    radius : float
+        Radius (pixel) of the circle.
+
+    Returns
+    -------
+    peak : 2-float tuple
+        (Physical) coordinate of the peak.
+    """
+    subprocess.check_call(["punlearn", "dmstat"])
+    subprocess.check_call([
+        "dmstat", "infile=%s" % image,
+        "centroid=no", "media=no", "sigma=no", "clip=no"
+    ])
+    peak = subprocess.check_output([
+        "pget", "dmstat", "out_max_loc"
+    ]).decode("utf-8").strip()
+    peak = peak.split(",")
+    return (float(peak[0]), float(peak[1]))
+
+
+def get_centroid(image, center, radius=100):
+    """
+    Calculate the centroid on image within the specified circle.
+
+    Parameters
+    ----------
+    image : str
+        Path to the image file.
+    center : 2-float tuple
+        Central (physical) coordinate of the circle.
+    radius : float
+        Radius (pixel) of the circle.
+
+    Returns
+    -------
+    centroid : 2-float tuple
+        (Physical) coordinate of the centroid.
+    """
+    x, y = center
+    with tempfile.NamedTemporaryFile(mode="w+") as fp:
+        fp.file.write("circle(%f,%f,%f)\n" % (x, y, radius))
+        fp.file.flush()
+        subprocess.check_call(["punlearn", "dmstat"])
+        subprocess.check_call([
+            "dmstat", "infile=%s[sky=region(%s)]" % (image, fp.name),
+            "centroid=yes", "media=no", "sigma=no", "clip=no"
+        ])
+    centroid = subprocess.check_output([
+        "pget", "dmstat", "out_cntrd_phys"
+    ]).decode("utf-8").strip()
+    centroid = centroid.split(",")
+    return (float(centroid[0]), float(centroid[1]))
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Calculate the emission centroid within the image")
+    parser.add_argument("-i", "--infile", dest="infile", required=True,
+                        help="input image file (e.g., 0.7-2.0 keV)")
+    parser.add_argument("-o", "--outfile", dest="outfile",
+                        default="centroid.reg",
+                        help="output centroid region file " +
+                        "(default: centroid.reg")
+    parser.add_argument("-R", "--radius1", dest="radius1",
+                        type=float, default=300,
+                        help="circle radius [pixel] for first phase " +
+                        "centroid calculation (default: 300 pixel)")
+    parser.add_argument("-r", "--radius2", dest="radius2",
+                        type=float, default=100,
+                        help="circle radius [pixel] for second phase " +
+                        "calculation to tune centroid (default: 100 pixel)")
+    parser.add_argument("-n", "--niter", dest="niter",
+                        type=int, default=5,
+                        help="iterations for each phase (default: 5)")
+    parser.add_argument("-s", "--start", dest="start",
+                        help="a region file containing a circle/point " +
+                        "that specifies the starting point " +
+                        "(default: using the peak of the image)")
+    parser.add_argument("-V", "--view", dest="view", action="store_true",
+                        help="open DS9 to view output centroid")
+    args = parser.parse_args()
+
+    setup_pfiles(["aconvolve", "dmstat"])
+
+    if args.start:
+        print("Get starting point from region file: %s" % args.start,
+              file=sys.stderr)
+        raise NotImplementedError
+    else:
+        print("Use peak as the starting point ...", file=sys.stderr)
+        center = get_peak(args.infile)
+    print("Starting point: (%f, %f)" % center, file=sys.stderr)
+
+    centroid = center
+    for phase, radius in enumerate([args.radius1, args.radius2]):
+        print("Calculate centroid phase %d (circle radius: %.1f)" %
+              (phase+1, radius), file=sys.stderr)
+        for i in range(args.niter):
+            print("%d..." % (i+1), end="", flush=True, file=sys.stderr)
+            centroid = get_centroid(args.infile, center=centroid,
+                                    radius=radius)
+        print("Done!", file=sys.stderr)
+
+    open(args.outfile, "w").write("point(%f,%f)\n" % centroid).close()
+    print("Saved centroid to file:", args.outfile, file=sys.stderr)
+    if args.view:
+        ds9_view(args.infile, regfile=args.outfile)
+
+    # Add created image to manifest
+    manifest = get_manifest()
+    key = "reg_centroid"
+    manifest.setpath(key, args.outfile)
+
+
+if __name__ == "__main__":
+    main()