Move fits/fits{cube,image}.py

1 files changed, 522 insertions, 0 deletions

diff --git a/astro/fitsimage.py b/astro/fitsimage.py
new file mode 100755
index 0000000..af49f7e
--- /dev/null
+++ b/astro/fitsimage.py
@@ -0,0 +1,522 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2017-2018 Weitian LI <weitian@aaronly.me>
+# MIT license
+#
+
+"""
+FITS image manipulate tool.
+"""
+
+import sys
+import argparse
+
+import numpy as np
+from astropy.io import fits
+from scipy import ndimage
+
+
+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, pixelsize=None):
+        self.infile = infile
+        with fits.open(infile) as f:
+            self.header = f[0].header.copy(strip=True)
+            self.data = f[0].data
+        self.ndim = self.data.ndim
+        self.shape = self.data.shape
+        if pixelsize is not None:
+            self.pixelsize = pixelsize  # [arcsec]
+
+    @property
+    def bunit(self):
+        return self.header.get("BUNIT")
+
+    @property
+    def Nx(self):
+        """
+        Number of pixels along the X axis (i.e., image width)
+        """
+        return self.shape[-1]
+
+    @property
+    def Ny(self):
+        """
+        Number of pixels along the Y axis (i.e., image height)
+        """
+        return self.shape[-2]
+
+    @property
+    def image(self):
+        """
+        Deal with single-frequency and single-polarized image cube.
+        """
+        if self.ndim == 2:
+            # NAXIS=2: [Y, X]
+            image = self.data[:, :].copy()
+        elif self.ndim == 3 and self.shape[0] == 1:
+            # NAXIS=3: [FREQ=1, Y, X]
+            image = self.data[0, :, :].copy()
+        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, :, :].copy()
+        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 = np.array(value)
+        elif self.ndim == 3:
+            # NAXIS=3: [FREQ=1, Y, X]
+            self.data = np.array(value)[np.newaxis, :, :]
+        else:
+            # NAXIS=4: [STOKES=1, FREQ=1, Y, X]
+            self.data = np.array(value)[np.newaxis, np.newaxis, :, :]
+
+    @property
+    def pixelsize(self):
+        """
+        Image pixel size, in units of [arcsec]
+        """
+        if hasattr(self, "_pixelsize"):
+            return self._pixelsize
+
+        try:
+            return self.header["PixSize"]  # [arcsec]
+        except KeyError:
+            try:
+                return abs(self.header["CDELT1"]) * 3600  # [deg] -> [arcsec]
+            except KeyError:
+                return None
+
+    @pixelsize.setter
+    def pixelsize(self, value):
+        # Unit: [arcsec]
+        oldvalue = self.pixelsize
+        self._pixelsize = value
+        # Update header
+        self.header["PixSize"] = value  # [arcsec]
+        try:
+            self.header["CDELT1"] *= value / oldvalue
+            self.header["CDELT2"] *= value / oldvalue
+        except KeyError:
+            pass
+
+    @property
+    def fov(self):
+        """
+        Image FoV coverage: (fov_x, fov_y)
+        Unit: [deg]
+        """
+        pixelsize = self.pixelsize
+        if pixelsize:
+            return (self.Nx*pixelsize/3600, self.Ny*pixelsize/3600)
+        else:
+            return None
+
+    def zoom(self, newsize, order=1):
+        """
+        Zoom the image to the specified ``newsize``, meanwhile the header
+        information will be updated accordingly to preserve the FoV coverage.
+
+        NOTE
+        ----
+        The image aspect ratio cannot be changed.
+
+        Parameters
+        ----------
+        newsize : (Nx, Ny) or N
+            The size of the zoomed image.
+        order : int, optional
+            The interpolation order, default: 1
+        """
+        try:
+            Nx2, Ny2 = newsize
+        except TypeError:
+            Nx2 = Ny2 = newsize
+        zoom = ((Ny2+0.1)/self.Ny, (Nx2+0.1)/self.Nx)
+        if abs(zoom[0] - zoom[1]) > 1e-3:
+            raise RuntimeError("image aspect ratio cannot be changed")
+
+        pixelsize_old = self.pixelsize
+        self.image = ndimage.zoom(self.image, zoom=zoom, order=order)
+        self.pixelsize = pixelsize_old * (self.Nx / Nx2)
+        return self.image
+
+    def flip(self, direction):
+        if direction == "lr":
+            self.image = np.fliplr(self.image)
+        elif direction == "ud":
+            self.image = np.flipud(self.image)
+        else:
+            raise ValueError("invalid flip direction: %s" % direction)
+        return self.image
+
+    def rotate(self, to):
+        if to == "left":
+            self.image = np.rot90(self.image, k=-1)
+        elif to == "right":
+            self.image = np.rot90(self.image, k=1)
+        elif to == "180":
+            self.image = np.rot90(self.image, k=2)
+        else:
+            raise ValueError("invalid rotate to: %s" % to)
+        return self.image
+
+    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 show_info(filename, abs_=None, center=None):
+    """
+    Show FITS image information.
+    """
+    fimage = FITSImage(filename)
+    print("Image data shape: {0}".format(fimage.shape))
+    print("Image size: %dx%d" % (fimage.Nx, fimage.Ny))
+    print("Data unit: [%s]" % fimage.bunit)
+    pixelsize = fimage.pixelsize
+    if pixelsize:
+        print("Pixel size: %.1f [arcsec]" % pixelsize)
+        print("Field of view: (%.2f, %.2f) [deg]" % fimage.fov)
+    data = fimage.image
+    if abs_:
+        data = np.abs(data)
+    if center:
+        print("Central box size: %d" % center)
+        rows, cols = data.shape
+        rc, cc = rows//2, cols//2
+        cs1, cs2 = center//2, (center+1)//2
+        data = data[(rc-cs1):(rc+cs2), (cc-cs1):(cc+cs2)]
+    min_ = np.nanmin(data)
+    max_ = np.nanmax(data)
+    mean = np.nanmean(data)
+    median = np.nanmedian(data)
+    std = np.nanstd(data)
+    iqr = np.diff(np.nanpercentile(data, q=(25, 75)))
+    mad = np.nanmedian(np.abs(data - median))
+    rms = np.sqrt(np.nanmean(data**2))
+    print("min:    %13.6e" % min_)
+    print("max:    %13.6e" % max_)
+    print("range:  %13.6e  (max - min)" % (max_ - min_))
+    print("mean:   %13.6e" % mean)
+    print("median: %13.6e" % median)
+    print("std:    %13.6e  (standard deviation)" % std)
+    print("iqr:    %13.6e  (interquartile range)" % iqr)
+    print("mad:    %13.6e  (median absolute deviation)" % mad)
+    print("rms:    %13.6e  (root-mean-squared)" % rms)
+
+
+def cmd_info(args):
+    """
+    Sub-command: "info", show FITS image information
+    """
+    for fn in args.files:
+        print(">>> %s <<<" % fn)
+        show_info(fn, abs_=args.abs, center=args.center)
+        print("")
+
+
+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 cmd_div(args):
+    """
+    Sub-command: "div", divide the image by a number or other image(s)
+    """
+    fimage = FITSImage(args.infile)
+    image = fimage.image
+    if args.number:
+        print("Divide by number: %g" % args.number)
+        image /= args.number
+    else:
+        for fn in args.files:
+            print("Divide by another image from: %s" % fn)
+            fimage2 = FITSImage(fn)
+            with np.errstate(divide="warn"):
+                image /= fimage2.image
+
+    if args.fill_value:
+        print("Filling invalid data with: %s" % args.fill_value)
+        image[~np.isfinite(image)] = float(args.fill_value)
+    fimage.image = image
+    fimage.write(args.outfile, clobber=args.clobber)
+    print("Saved FITS image to: %s" % args.outfile)
+
+
+def cmd_zoom(args):
+    """
+    Sub-command: "zoom", zoom the image to a new size with FoV coverage
+    preserved.
+    """
+    fimage = FITSImage(args.infile)
+    print("Image size: %dx%d" % (fimage.Nx, fimage.Ny))
+    pixelsize = fimage.pixelsize
+    if pixelsize is None:
+        raise RuntimeError("--pixelsize required")
+    else:
+        print("Pixel size: %.1f [arcsec]" % pixelsize)
+        print("Field of view: (%.2f, %.2f) [deg]" % fimage.fov)
+
+    print("Zooming image ...")
+    print("Interpolation order: %d" % args.order)
+    print("Zoomed image size: %dx%d" % (args.size, args.size))
+    fimage.zoom(newsize=args.size, order=args.order)
+    print("Zoomed image pixel size: %.1f [arcsec]" % fimage.pixelsize)
+    fimage.write(args.outfile, clobber=args.clobber)
+    print("Saved zoomed FITS image to: %s" % args.outfile)
+
+
+def cmd_flip(args):
+    """
+    Sub-command: "flip", flip the image left-right or up-down.
+    """
+    fimage = FITSImage(args.infile)
+    print("Flipping image ...")
+    direction = "lr" if args.lr else "ud"
+    print("Flip direction: %s" % direction)
+    fimage.flip(direction)
+    fimage.write(args.outfile, clobber=args.clobber)
+    print("Saved flipped FITS image to: %s" % args.outfile)
+
+
+def cmd_rotate(args):
+    """
+    Sub-command: "rotate", rotate the image.
+    """
+    fimage = FITSImage(args.infile)
+    print("Rotating image ...")
+    if args.left:
+        to = "left"
+    elif args.right:
+        to = "right"
+    else:
+        to = "180"
+    print("Rotate to: %s" % to)
+    fimage.rotate(to)
+    fimage.write(args.outfile, clobber=args.clobber)
+    print("Saved rotated 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("files", nargs="+", 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)
+
+    # sub-command: "div"
+    parser_div = subparsers.add_parser(
+        "div", aliases=["divide"],
+        help="divide the image by a number or other image(s)")
+    parser_div.add_argument("-C", "--clobber", dest="clobber",
+                            action="store_true",
+                            help="overwrite existing output file")
+    parser_div.add_argument("-F", "--fill-value", dest="fill_value",
+                            help="value to fill the invalid elements")
+    parser_div.add_argument("-i", "--infile", dest="infile", required=True,
+                            help="input FITS image")
+    parser_div.add_argument("-o", "--outfile", dest="outfile", required=True,
+                            help="output FITS image")
+    exgrp_div = parser_div.add_mutually_exclusive_group(required=True)
+    exgrp_div.add_argument("-n", "--number", dest="number", type=float,
+                           help="number to be divided by")
+    exgrp_div.add_argument("-f", "--files", dest="files", nargs="+",
+                           help="FITS image(s) to be divided by")
+    parser_div.set_defaults(func=cmd_div)
+
+    # sub-command: "zoom"
+    parser_zoom = subparsers.add_parser(
+        "zoom", aliases=["rescale"],
+        help="zoom the image to a new size with FoV coverage preserved")
+    parser_zoom.add_argument("-C", "--clobber", dest="clobber",
+                             action="store_true",
+                             help="overwrite existing output file")
+    parser_zoom.add_argument("--order", type=int, default=1,
+                             help="zoom interpolation order (default: 1)")
+    parser_zoom.add_argument("-s", "--size", type=int, required=True,
+                             help="zoomed image size (number of pixels)")
+    parser_zoom.add_argument("-p", "--pixelsize", type=float,
+                             help="input FITS image pixel size [arcsec] " +
+                             "(default: try to obtain from FITS header)")
+    parser_zoom.add_argument("-i", "--infile", dest="infile", required=True,
+                             help="input FITS image")
+    parser_zoom.add_argument("-o", "--outfile", dest="outfile", required=True,
+                             help="output zoomed FITS image")
+    parser_zoom.set_defaults(func=cmd_zoom)
+
+    # sub-command: "flip"
+    parser_flip = subparsers.add_parser(
+        "flip", help="flip the image left-right or up-down")
+    parser_flip.add_argument("-C", "--clobber", action="store_true",
+                             help="overwrite existing output file")
+    parser_flip.add_argument("-i", "--infile", required=True,
+                             help="input FITS image")
+    parser_flip.add_argument("-o", "--outfile", required=True,
+                             help="output flipped FITS image")
+    exgrp_flip = parser_flip.add_mutually_exclusive_group(required=True)
+    exgrp_flip.add_argument("-l", "--left-right", dest="lr",
+                            action="store_true",
+                            help="flip in the left/right direction")
+    exgrp_flip.add_argument("-u", "--up-down", dest="ud",
+                            action="store_true",
+                            help="flip in the left/right direction")
+    parser_flip.set_defaults(func=cmd_flip)
+
+    # sub-command: "rotate"
+    parser_rot = subparsers.add_parser(
+        "rot", aliases=["rotate"],
+        help="rotate the image")
+    parser_rot.add_argument("-C", "--clobber", action="store_true",
+                            help="overwrite existing output file")
+    parser_rot.add_argument("-i", "--infile", required=True,
+                            help="input FITS image")
+    parser_rot.add_argument("-o", "--outfile", required=True,
+                            help="output rotated FITS image")
+    exgrp_rot = parser_rot.add_mutually_exclusive_group(required=True)
+    exgrp_rot.add_argument("-l", "--left", action="store_true",
+                           help="rotate left")
+    exgrp_rot.add_argument("-r", "--right", action="store_true",
+                           help="rotate right")
+    exgrp_rot.add_argument("-u", "--180", dest="ud",
+                           action="store_true",
+                           help="rotate 180 degree")
+    parser_rot.set_defaults(func=cmd_rotate)
+
+    args = parser.parse_args()
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()