Add astro/fitsimage.py: FITS image manipulation tool

Replace fitsimgadd.py and fitsimgsub.py

1 files changed, 236 insertions, 0 deletions

diff --git a/astro/fits/fitsimage.py b/astro/fits/fitsimage.py
new file mode 100755
index 0000000..74b03cf
--- /dev/null
+++ b/astro/fits/fitsimage.py
@@ -0,0 +1,236 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) Weitian LI <weitian@aaronly.me>
+# MIT license
+#
+
+"""
+FITS image manipulate tool.
+"""
+
+import sys
+import argparse
+
+import numpy as np
+from astropy.io import fits
+
+
+class FITSImage:
+    """
+    FITS image class that deals with plain 2D image (NAXIS=2), but also
+    handles single-frequency single-polarized image cube (NAXIS=3, 4),
+    e.g., created by WSClean.
+    """
+    def __init__(self, infile):
+        with fits.open(infile) as f:
+            self.header = f[0].header
+            self.data = f[0].data
+        self.ndim = self.data.ndim
+        self.shape = self.data.shape
+
+    @property
+    def bunit(self):
+        return self.header.get("BUNIT")
+
+    @property
+    def image(self):
+        """
+        Deal with single-frequency and single-polarized image cube.
+        """
+        if self.ndim == 2:
+            # NAXIS=2: [Y, X]
+            image = self.data[:, :]
+        elif self.ndim == 3 and self.shape[0] == 1:
+            # NAXIS=3: [FREQ=1, Y, X]
+            image = self.data[0, :, :]
+        elif self.ndim == 4 and self.shape[0] == 1 and self.shape[1] == 1:
+            # NAXIS=4: [STOKES=1, FREQ=1, Y, X]
+            image = self.data[0, 0, :, :]
+        else:
+            raise ValueError("invalid data shape: {1}".format(self.shape))
+        return image
+
+    @image.setter
+    def image(self, value):
+        if self.ndim == 2:
+            # NAXIS=2: [Y, X]
+            self.data[:, :] = value
+        elif self.ndim == 3:
+            # NAXIS=3: [FREQ=1, Y, X]
+            self.data[0, :, :] = value
+        else:
+            # NAXIS=4: [STOKES=1, FREQ=1, Y, X]
+            self.data[0, 0, :, :] = value
+
+    def write(self, outfile, clobber=False):
+        self.header.add_history(" ".join(sys.argv))
+        hdu = fits.PrimaryHDU(data=self.data, header=self.header)
+        try:
+            hdu.writeto(outfile, overwrite=clobber)
+        except TypeError:
+            hdu.writeto(outfile, clobber=clobber)
+
+
+def cmd_info(args):
+    """
+    Sub-command: "info", show FITS image information
+    """
+    fimage = FITSImage(args.infile)
+    print("Image data shape: {0}".format(fimage.shape))
+    print("Data unit: %s" % fimage.bunit)
+    image = fimage.image
+    if args.abs:
+        image = np.abs(image)
+    if args.center:
+        print("Central box size: %d" % args.center)
+        rows, cols = image.shape
+        rc, cc = rows//2, cols//2
+        cs1, cs2 = args.center//2, (args.center+1)//2
+        image = image[(rc-cs1):(rc+cs2), (cc-cs1):(cc+cs2)]
+    mean = np.mean(image)
+    median = np.median(image)
+    iqr = np.diff(np.percentile(image, q=(25, 75)))
+    std = np.std(image)
+    rms = np.sqrt(np.mean(image**2))
+    print("mean:   %-12.6e" % mean)
+    print("median: %-12.6e" % median)
+    print("std:    %-12.6e  (standard deviation)" % std)
+    print("iqr:    %-12.6e  (interquartile range)" % iqr)
+    print("rms:    %-12.6e  (root-mean-squared)" % rms)
+
+
+def cmd_add(args):
+    """
+    Sub-command: "add", add the image by a number or other image(s)
+    """
+    fimage = FITSImage(args.infile)
+    image = fimage.image
+    if args.number:
+        print("Add by number: %g" % args.number)
+        image += args.number
+    else:
+        for fn in args.files:
+            print("Add by another image from: %s" % fn)
+            fimage2 = FITSImage(fn)
+            image += fimage2.image
+    fimage.image = image
+    fimage.write(args.outfile, clobber=args.clobber)
+    print("Saved FITS image to: %s" % args.outfile)
+
+
+def cmd_sub(args):
+    """
+    Sub-command: "sub", subtract the image by a number or other image(s)
+    """
+    fimage = FITSImage(args.infile)
+    image = fimage.image
+    if args.number:
+        print("Subtract by number: %g" % args.number)
+        image -= args.number
+    else:
+        for fn in args.files:
+            print("Subtract by another image from: %s" % fn)
+            fimage2 = FITSImage(fn)
+            image -= fimage2.image
+    fimage.image = image
+    fimage.write(args.outfile, clobber=args.clobber)
+    print("Saved FITS image to: %s" % args.outfile)
+
+
+def cmd_mul(args):
+    """
+    Sub-command: "mul", multiply the image by a number or other image(s)
+    """
+    fimage = FITSImage(args.infile)
+    image = fimage.image
+    if args.number:
+        print("Multiply by number: %g" % args.number)
+        image *= args.number
+    else:
+        for fn in args.files:
+            print("Multiply by another image from: %s" % fn)
+            fimage2 = FITSImage(fn)
+            image *= fimage2.image
+    fimage.image = image
+    fimage.write(args.outfile, clobber=args.clobber)
+    print("Saved FITS image to: %s" % args.outfile)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="FITS image manipulation tool")
+    subparsers = parser.add_subparsers(dest="subparser_name",
+                                       title="sub-commands",
+                                       help="additional help")
+
+    # sub-command: "info"
+    parser_info = subparsers.add_parser(
+        "info", aliases=["show"],
+        help="show FITS image info")
+    parser_info.add_argument("-c", "--center", dest="center", type=int,
+                             help="choose central region of specified size")
+    parser_info.add_argument("-a", "--abs", dest="abs", action="store_true",
+                             help="take absolute values of image pixels")
+    parser_info.add_argument("infile", help="FITS image filename")
+    parser_info.set_defaults(func=cmd_info)
+
+    # sub-command: "add"
+    parser_add = subparsers.add_parser(
+        "add",
+        help="add the image by a number or other image(s)")
+    parser_add.add_argument("-C", "--clobber", dest="clobber",
+                            action="store_true",
+                            help="overwrite existing output file")
+    parser_add.add_argument("-i", "--infile", dest="infile", required=True,
+                            help="input FITS image")
+    parser_add.add_argument("-o", "--outfile", dest="outfile", required=True,
+                            help="output FITS image")
+    exgrp_add = parser_add.add_mutually_exclusive_group(required=True)
+    exgrp_add.add_argument("-n", "--number", dest="number", type=float,
+                           help="number to be added by")
+    exgrp_add.add_argument("-f", "--files", dest="files", nargs="+",
+                           help="FITS image(s) to be added by")
+    parser_add.set_defaults(func=cmd_add)
+
+    # sub-command: "sub"
+    parser_sub = subparsers.add_parser(
+        "sub", aliases=["subtract"],
+        help="subtract the image by a number or other image(s)")
+    parser_sub.add_argument("-C", "--clobber", dest="clobber",
+                            action="store_true",
+                            help="overwrite existing output file")
+    parser_sub.add_argument("-i", "--infile", dest="infile", required=True,
+                            help="input FITS image")
+    parser_sub.add_argument("-o", "--outfile", dest="outfile", required=True,
+                            help="output FITS image")
+    exgrp_sub = parser_sub.add_mutually_exclusive_group(required=True)
+    exgrp_sub.add_argument("-n", "--number", dest="number", type=float,
+                           help="number to be subtracted by")
+    exgrp_sub.add_argument("-f", "--files", dest="files", nargs="+",
+                           help="FITS image(s) to be subtracted by")
+    parser_sub.set_defaults(func=cmd_sub)
+
+    # sub-command: "mul"
+    parser_mul = subparsers.add_parser(
+        "mul", aliases=["multiply"],
+        help="multiply the image by a number or other image(s)")
+    parser_mul.add_argument("-C", "--clobber", dest="clobber",
+                            action="store_true",
+                            help="overwrite existing output file")
+    parser_mul.add_argument("-i", "--infile", dest="infile", required=True,
+                            help="input FITS image")
+    parser_mul.add_argument("-o", "--outfile", dest="outfile", required=True,
+                            help="output FITS image")
+    exgrp_mul = parser_mul.add_mutually_exclusive_group(required=True)
+    exgrp_mul.add_argument("-n", "--number", dest="number", type=float,
+                           help="number to be multiplied by")
+    exgrp_mul.add_argument("-f", "--files", dest="files", nargs="+",
+                           help="FITS image(s) to be multiplied by")
+    parser_mul.set_defaults(func=cmd_mul)
+
+    args = parser.parse_args()
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()