Move python/calc_radial_psd.py to astro/calc_psd.py with some cleanup

1 files changed, 0 insertions, 460 deletions

diff --git a/python/calc_radial_psd.py b/python/calc_radial_psd.py
deleted file mode 100755
index b7d1743..0000000
--- a/python/calc_radial_psd.py
+++ /dev/null
@@ -1,460 +0,0 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-#
-# Credit:
-# [1]  Radially averaged power spectrum of 2D real-valued matrix
-#      Evan Ruzanski
-#      'raPsd2d.m'
-#      https://www.mathworks.com/matlabcentral/fileexchange/23636-radially-averaged-power-spectrum-of-2d-real-valued-matrix
-#
-# XXX:
-# * If the input image is NOT SQUARE; then are the horizontal frequencies
-#   the same as the vertical frequencies ??
-#
-# Aaron LI <aaronly.me@gmail.com>
-# Created: 2015-04-22
-# Updated: 2016-04-28
-#
-# Changelog:
-# 2016-04-28:
-#   * Fix wrong meshgrid with respect to the shift zero-frequency component
-#   * Use "numpy.fft" instead of "scipy.fftpack"
-#   * Split method "pad_square()" from "calc_radial_psd()"
-#   * Hide numpy warning when dividing by zero
-#   * Add method "AstroImage.fix_shapes()"
-#   * Add support for background subtraction and exposure correction
-#   * Show verbose information during calculation
-#   * Add class "AstroImage"
-#   * Set default value for 'args.png'
-#   * Rename from 'radialPSD2d.py' to 'calc_radial_psd.py'
-# 2016-04-26:
-#   * Adjust plot function
-#   * Update normalize argument; Add pixel argument
-# 2016-04-25:
-#   * Update plot function
-#   * Add command line scripting support
-#   * Encapsulate the functions within class 'PSD'
-#   * Update docs/comments
-#
-
-"""
-Compute the radially averaged power spectral density (i.e., power spectrum).
-"""
-
-__version__ = "0.5.0"
-__date__    = "2016-04-28"
-
-
-import sys
-import os
-import argparse
-
-import numpy as np
-from astropy.io import fits
-
-import matplotlib.pyplot as plt
-from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
-from matplotlib.figure import Figure
-
-plt.style.use("ggplot")
-
-
-class PSD:
-    """
-    Computes the 2D power spectral density and the radially averaged power
-    spectral density (i.e., 1D power spectrum).
-    """
-    # 2D image data
-    img = None
-    # value and unit of 1 pixel for the input image
-    pixel = (None, None)
-    # whether to normalize the power spectral density by image size
-    normalize = True
-    # 2D power spectral density
-    psd2d = None
-    # 1D (radially averaged) power spectral density
-    freqs     = None
-    psd1d     = None
-    psd1d_err = None
-
-    def __init__(self, img, pixel=(1.0, "pixel"), normalize=True):
-        self.img = img.astype(np.float)
-        self.pixel = pixel
-        self.normalize = normalize
-
-    def calc_psd2d(self):
-        """
-        Computes the 2D power spectral density of the given image.
-        Note that the low frequency components are shifted to the center
-        of the FFT'ed image.
-
-        NOTE:
-        The zero-frequency component is shifted to position of index (0-based)
-            (ceil((n-1) / 2), ceil((m-1) / 2)),
-        where (n, m) are the number of rows and columns of the image/psd2d.
-
-        Return:
-            2D power spectral density, which is dimensionless if normalized,
-            otherwise has unit ${pixel_unit}^2.
-        """
-        print("Calculating 2D power spectral density ... ", end="", flush=True)
-        rows, cols = self.img.shape
-        # Compute the power spectral density (i.e., power spectrum)
-        imgf = np.fft.fftshift(np.fft.fft2(self.img))
-        if self.normalize:
-            norm = rows * cols * self.pixel[0]**2
-        else:
-            norm = 1.0  # Do not normalize
-        self.psd2d = (np.abs(imgf) / norm) ** 2
-        print("DONE", flush=True)
-        return self.psd2d
-
-    def calc_radial_psd1d(self):
-        """
-        Computes the radially averaged power spectral density from the
-        provided 2D power spectral density.
-
-        Return:
-            (freqs, radial_psd, radial_psd_err)
-            freqs: spatial freqencies (unit: ${pixel_unit}^(-1))
-            radial_psd: radially averaged power spectral density for each
-                        frequency
-            radial_psd_err: standard deviations of each radial_psd
-        """
-        print("Calculating radial (1D) power spectral density ... ",
-              end="", flush=True)
-        print("padding ... ", end="", flush=True)
-        psd2d    = self.pad_square(self.psd2d, value=np.nan)
-        dim      = psd2d.shape[0]
-        dim_half = (dim+1) // 2
-        # NOTE:
-        # The zero-frequency component is shifted to position of index
-        # (0-based): (ceil((n-1) / 2), ceil((m-1) / 2))
-        px       = np.arange(dim_half-dim, dim_half)
-        x, y     = np.meshgrid(px, px)
-        rho, phi = self.cart2pol(x, y)
-        rho      = np.around(rho).astype(np.int)
-        radial_psd     = np.zeros(dim_half)
-        radial_psd_err = np.zeros(dim_half)
-        print("radially averaging ... ", end="", flush=True)
-        for r in range(dim_half):
-            # Get the indices of the elements satisfying rho[i,j]==r
-            ii, jj = (rho == r).nonzero()
-            # Calculate the mean value at a given radii
-            data              = psd2d[ii, jj]
-            radial_psd[r]     = np.nanmean(data)
-            radial_psd_err[r] = np.nanstd(data)
-        # Calculate frequencies
-        f = np.fft.fftfreq(dim, d=self.pixel[0])
-        freqs = np.abs(f[:dim_half])
-        #
-        self.freqs     = freqs
-        self.psd1d     = radial_psd
-        self.psd1d_err = radial_psd_err
-        print("DONE", flush=True)
-        return (freqs, radial_psd, radial_psd_err)
-
-    @staticmethod
-    def cart2pol(x, y):
-        """
-        Convert Cartesian coordinates to polar coordinates.
-        """
-        rho = np.sqrt(x**2 + y**2)
-        phi = np.arctan2(y, x)
-        return (rho, phi)
-
-    @staticmethod
-    def pol2cart(rho, phi):
-        """
-        Convert polar coordinates to Cartesian coordinates.
-        """
-        x = rho * np.cos(phi)
-        y = rho * np.sin(phi)
-        return (x, y)
-
-    @staticmethod
-    def pad_square(data, value=np.nan):
-        """
-        Symmetrically pad the supplied data matrix to make it square.
-        The padding rows are equally added to the top and bottom,
-        as well as the columns to the left and right sides.
-        The padded rows/columns are filled with the specified value.
-        """
-        mat = data.copy()
-        rows, cols = mat.shape
-        dim_diff = abs(rows - cols)
-        dim_max  = max(rows, cols)
-        if rows > cols:
-            # pad columns
-            if dim_diff // 2 == 0:
-                cols_left     = np.zeros((rows, dim_diff/2))
-                cols_left[:]  = value
-                cols_right    = np.zeros((rows, dim_diff/2))
-                cols_right[:] = value
-                mat           = np.hstack((cols_left, mat, cols_right))
-            else:
-                cols_left     = np.zeros((rows, np.floor(dim_diff/2)))
-                cols_left[:]  = value
-                cols_right    = np.zeros((rows, np.floor(dim_diff/2)+1))
-                cols_right[:] = value
-                mat           = np.hstack((cols_left, mat, cols_right))
-        elif rows < cols:
-            # pad rows
-            if dim_diff // 2 == 0:
-                rows_top       = np.zeros((dim_diff/2, cols))
-                rows_top[:]    = value
-                rows_bottom    = np.zeros((dim_diff/2, cols))
-                rows_bottom[:] = value
-                mat            = np.vstack((rows_top, mat, rows_bottom))
-            else:
-                rows_top       = np.zeros((np.floor(dim_diff/2), cols))
-                rows_top[:]    = value
-                rows_bottom    = np.zeros((np.floor(dim_diff/2)+1, cols))
-                rows_bottom[:] = value
-                mat            = np.vstack((rows_top, mat, rows_bottom))
-        return mat
-
-    def plot(self, ax=None, fig=None):
-        """
-        Make a plot of the radial (1D) PSD with matplotlib.
-        """
-        if ax is None:
-            fig, ax = plt.subplots(1, 1)
-        #
-        xmin = self.freqs[1] / 1.2  # ignore the first 0
-        xmax = self.freqs[-1]
-        ymin = np.nanmin(self.psd1d) / 10.0
-        ymax = np.nanmax(self.psd1d + self.psd1d_err)
-        #
-        eb = ax.errorbar(self.freqs, self.psd1d, yerr=self.psd1d_err,
-                         fmt="none")
-        ax.plot(self.freqs, self.psd1d, "ko")
-        ax.set_xscale("log")
-        ax.set_yscale("log")
-        ax.set_xlim(xmin, xmax)
-        ax.set_ylim(ymin, ymax)
-        ax.set_title("Radially Averaged Power Spectral Density")
-        ax.set_xlabel(r"k (%s$^{-1}$)" % self.pixel[1])
-        if self.normalize:
-            ax.set_ylabel("Power")
-        else:
-            ax.set_ylabel(r"Power (%s$^2$)" % self.pixel[1])
-        fig.tight_layout()
-        return (fig, ax)
-
-
-class AstroImage:
-    """
-    Manipulate the astronimcal counts image, as well as the corresponding
-    exposure map and background map.
-    """
-    # input counts image
-    image = None
-    # exposure map with respect to the input counts image
-    expmap = None
-    # background map (e.g., stowed background)
-    bkgmap = None
-    # exposure time of the input image
-    exposure     = None
-    # exposure time of the background map
-    exposure_bkg = None
-
-    def __init__(self, image, expmap=None, bkgmap=None):
-        self.load_image(image)
-        self.load_expmap(expmap)
-        self.load_bkgmap(bkgmap)
-
-    @staticmethod
-    def open_image(infile):
-        """
-        Open the slice image and return its header and 2D image data.
-
-        NOTE
-        ----
-        The input slice image may have following dimensions:
-        * NAXIS=2: [Y, X]
-        * NAXIS=3: [FREQ=1, Y, X]
-        * NAXIS=4: [STOKES=1, FREQ=1, Y, X]
-
-        NOTE
-        ----
-        Only open slice image that has only ONE frequency and ONE Stokes
-        parameter.
-
-        Returns
-        -------
-        header : `~astropy.io.fits.Header`
-        image : 2D `~numpy.ndarray`
-            The 2D [Y, X] image part of the slice image.
-        """
-        with fits.open(infile) as f:
-            header = f[0].header
-            data = f[0].data
-        if data.ndim == 2:
-            # NAXIS=2: [Y, X]
-            image = data
-        elif data.ndim == 3 and data.shape[0] == 1:
-            # NAXIS=3: [FREQ=1, Y, X]
-            image = data[0, :, :]
-        elif data.ndim == 4 and data.shape[0] == 1 and data.shape[1] == 1:
-            # NAXIS=4: [STOKES=1, FREQ=1, Y, X]
-            image = data[0, 0, :, :]
-        else:
-            raise ValueError("Slice '{0}' has invalid dimensions: {1}".format(
-                infile, data.shape))
-        return (header, image)
-
-    def load_image(self, image):
-        print("Loading image ... ", end="", flush=True)
-        self.header, self.image = self.open_image(image)
-        self.exposure = self.header.get("EXPOSURE")
-        print("DONE", flush=True)
-
-    def load_expmap(self, expmap):
-        if expmap:
-            print("Loading exposure map ... ", end="", flush=True)
-            __, self.expmap = self.open_image(expmap)
-            print("DONE", flush=True)
-
-    def load_bkgmap(self, bkgmap):
-        if bkgmap:
-            print("Loading background map ... ", end="", flush=True)
-            header, self.bkgmap = self.open_image(bkgmap)
-            self.exposure_bkg = header.get("EXPOSURE")
-            print("DONE", flush=True)
-
-    def fix_shapes(self, tolerance=2):
-        """
-        Fix the shapes of self.expmap and self.bkgmap to make them have
-        the same shape as the self.image.
-
-        NOTE:
-        * if the image is bigger than the reference image, then its
-          columns on the right and rows on the botton are clipped;
-        * if the image is smaller than the reference image, then padding
-          columns on the right and rows on the botton are added.
-        * Original images are REPLACED!
-
-        Arguments:
-          * tolerance: allow absolute difference between images
-        """
-        def _fix_shape(img, ref, tol=tolerance):
-            if img.shape == ref.shape:
-                print("SKIPPED", flush=True)
-                return img
-            elif np.allclose(img.shape, ref.shape, atol=tol):
-                print(img.shape, "->", ref.shape, flush=True)
-                rows, cols = img.shape
-                rows_ref, cols_ref = ref.shape
-                # rows
-                if rows > rows_ref:
-                    img_fixed = img[:rows_ref, :]
-                else:
-                    img_fixed = np.row_stack((img,
-                        np.zeros((rows_ref-rows, cols), dtype=img.dtype)))
-                # columns
-                if cols > cols_ref:
-                    img_fixed = img_fixed[:, :cols_ref]
-                else:
-                    img_fixed = np.column_stack((img_fixed,
-                        np.zeros((rows_ref, cols_ref-cols), dtype=img.dtype)))
-                return img_fixed
-            else:
-                raise ValueError("shape difference exceeds tolerance: " + \
-                        "(%d, %d) vs. (%d, %d)" % (img.shape + ref.shape))
-        #
-        if self.bkgmap is not None:
-            print("Fixing shape for bkgmap ... ", end="", flush=True)
-            self.bkgmap = _fix_shape(self.bkgmap, self.image)
-        if self.expmap is not None:
-            print("Fixing shape for expmap ... ", end="", flush=True)
-            self.expmap = _fix_shape(self.expmap, self.image)
-
-    def subtract_bkg(self):
-        print("Subtracting background ... ", end="", flush=True)
-        self.image -= (self.bkgmap / self.exposure_bkg * self.exposure)
-        print("DONE", flush=True)
-
-    def correct_exposure(self, cut=0.015):
-        """
-        Correct the image for exposure by dividing by the expmap to
-        create the exposure-corrected image.
-
-        Arguments:
-          * cut: the threshold percentage with respect to the maximum
-                 exposure map value; and those pixels with lower values
-                 than this threshold will be excluded/clipped (set to ZERO)
-                 if set to None, then skip clipping image
-        """
-        print("Correcting image for exposure ... ", end="", flush=True)
-        with np.errstate(divide="ignore", invalid="ignore"):
-            self.image /= self.expmap
-        # set invalid values to ZERO
-        self.image[ ~ np.isfinite(self.image) ] = 0.0
-        print("DONE", flush=True)
-        if cut is not None:
-            # clip image according the exposure threshold
-            print("Clipping image (%s) ... " % cut, end="", flush=True)
-            threshold = cut * np.max(self.expmap)
-            self.image[ self.expmap < threshold ] = 0.0
-            print("DONE", flush=True)
-
-
-def main():
-    parser = argparse.ArgumentParser(
-            description="Compute the radially averaged power spectral density",
-            epilog="Version: %s (%s)" % (__version__, __date__))
-    parser.add_argument("-V", "--version", action="version",
-            version="%(prog)s " + "%s (%s)" % (__version__, __date__))
-    parser.add_argument("-C", "--clobber", dest="clobber",
-            action="store_true",
-            help="overwrite the output files if already exist")
-    parser.add_argument("-i", "--infile", dest="infile",
-            required=True, help="input image")
-    parser.add_argument("-b", "--bkgmap", dest="bkgmap", default=None,
-            help="background map (for background subtraction)")
-    parser.add_argument("-e", "--expmap", dest="expmap", default=None,
-            help="exposure map (for exposure correction)")
-    parser.add_argument("-o", "--outfile", dest="outfile",
-            required=True, help="output file to store the PSD data")
-    parser.add_argument("-p", "--png", dest="png", default=None,
-            help="plot the PSD and save (default: same basename as outfile)")
-    args = parser.parse_args()
-
-    if args.png is None:
-        args.png = os.path.splitext(args.outfile)[0] + ".png"
-
-    # Check output files whether already exists
-    if (not args.clobber) and os.path.exists(args.outfile):
-        raise ValueError("outfile '%s' already exists" % args.outfile)
-    if (not args.clobber) and os.path.exists(args.png):
-        raise ValueError("output png '%s' already exists" % args.png)
-
-    # Load image data
-    image = AstroImage(image=args.infile, expmap=args.expmap,
-                       bkgmap=args.bkgmap)
-    image.fix_shapes()
-    if args.bkgmap:
-        image.subtract_bkg()
-    if args.expmap:
-        image.correct_exposure()
-
-    # Calculate the power spectral density
-    psd = PSD(img=image.image, normalize=True)
-    psd.calc_psd2d()
-    freqs, psd1d, psd1d_err = psd.calc_radial_psd1d()
-
-    # Write out PSD results
-    psd_data = np.column_stack((freqs, psd1d, psd1d_err))
-    np.savetxt(args.outfile, psd_data, header="freqs  psd1d  psd1d_err")
-
-    # Make and save a plot
-    fig = Figure(figsize=(10, 8))
-    canvas = FigureCanvas(fig)
-    ax = fig.add_subplot(111)
-    psd.plot(ax=ax, fig=fig)
-    fig.savefig(args.png, format="png", dpi=150)
-
-
-if __name__ == "__main__":
-    main()