radialPSD2d.py: rewrite to use class encapsulation

1 files changed, 243 insertions, 143 deletions

diff --git a/python/radialPSD2d.py b/python/radialPSD2d.py
index cc5dd85..2b5c4d8 100755
--- a/python/radialPSD2d.py
+++ b/python/radialPSD2d.py
@@ -1,171 +1,271 @@
 #!/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
+#
 # Aaron LI <aaronly.me@gmail.com>
-# 2015/04/22
+# Created: 2015-04-22
+# Updated: 2016-04-26
+#
+# Changelog:
+# 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
 #
 
 """
-Computes the radially averaged power spectral density (power spectrum).
+Compute the radially averaged power spectral density (i.e., power spectrum).
 """
 
+__version__ = "0.3.1"
+__date__    = "2016-04-25"
+
+
+import sys
+import os
+import argparse
 
 import numpy as np
 from scipy import fftpack
+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")
 
-def PSD2d( img, normalize=True ):
-    """
-    Computes the 2D power spectrum of the given image.
 
-    Reference:
-    [1] raPsd2d.m by Evan Ruzanski
-        Radially averaged power spectrum of 2D real-valued matrix
-        https://www.mathworks.com/matlabcentral/fileexchange/23636-radially-averaged-power-spectrum-of-2d-real-valued-matrix
+class PSD:
     """
-    img = np.array( img )
-    rows, cols = img.shape
-    ## Compute power spectrum
-    # Perform the Fourier transform and shift the zero-frequency
-    # component to the center of the spectrum.
-    imgf = fftpack.fftshift( fftpack.fft2( img ) )
-    if normalize:
-        norm = rows * cols
-    else:
-        norm = 1.0  # Do not normalize
-    psd2d = ( np.abs( imgf ) / norm ) ** 2
-    return psd2d
-
-
-def radialPSD( psd2d ):
+    Computes the 2D power spectral density and the radially averaged power
+    spectral density (i.e., 1D power spectrum).
     """
-    Computes the radially averaged power spectral density (power spectrum)
-    from the provided 2D 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
 
-    Reference:
-    [1] raPsd2d.m by Evan Ruzanski
-        Radially averaged power spectrum of 2D real-valued matrix
-        https://www.mathworks.com/matlabcentral/fileexchange/23636-radially-averaged-power-spectrum-of-2d-real-valued-matrix
-    """
-    psd2d = np.array( psd2d )
-    rows, cols = psd2d.shape
-    ## Adjust the PSD array size
-    dim_diff = np.abs( rows - cols )
-    dim_max = max( rows, cols )
-    # Pad the PSD array to be sqaure
-    if rows > cols:
-        # pad columns
-        if np.mod( dim_diff, 2 ) == 0:
-            cols_left = np.zeros( (rows, dim_diff/2) )
-            cols_left[:] = np.nan
-            cols_right = np.zeros( (rows, dim_diff/2) )
-            cols_right[:] = np.nan
-            psd2d = np.hstack( (cols_left, psd2d, cols_right) )
-        else:
-            cols_left = np.zeros( (rows, np.floor(dim_diff/2)) )
-            cols_left[:] = np.nan
-            cols_right = np.zeros( (rows, np.floor(dim_diff/2)+1) )
-            cols_right[:] = np.nan
-            psd2d = np.hstack( (cols_left, psd2d, cols_right) )
-    elif rows < cols:
-        # pad rows
-        if np.mod( dim_diff, 2 ) == 0:
-            rows_top = np.zeros( (dim_diff/2, cols) )
-            rows_top[:] = np.nan
-            rows_bottom = np.zeros( (dim_diff/2, cols) )
-            rows_bottom[:] = np.nan
-            psd2d = np.vstack( (rows_top, psd2d, rows_bottom) )
+    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.
+
+        Return:
+            2D power spectral density, which is dimensionless if normalized,
+            otherwise has unit ${pixel_unit}^2.
+        """
+        rows, cols = self.img.shape
+        ## Compute the power spectral density (i.e., power spectrum)
+        imgf = fftpack.fftshift(fftpack.fft2(self.img))
+        if self.normalize:
+            norm = rows * cols * self.pixel[0]**2
         else:
-            rows_top = np.zeros( (np.floor(dim_diff/2), cols) )
-            rows_top[:] = np.nan
-            rows_bottom = np.zeros( (np.floor(dim_diff/2)+1, cols) )
-            rows_bottom[:] = np.nan
-            psd2d = np.vstack( (rows_top, psd2d, rows_bottom) )
-    ## Compute radially average power spectrum
-    px = np.arange( -dim_max/2, dim_max/2 )
-    x, y = np.meshgrid( px, px )
-    rho, phi = cart2pol( x, y )
-    rho = np.around( rho ).astype(int)
-    dim_half = np.floor( dim_max/2 ) + 1
-    radial_psd = np.zeros( dim_half )
-    radial_psd_err = np.zeros( dim_half ) # standard error
-    for r in np.arange( dim_half, dtype=int ):
-        # 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 = fftpack.fftfreq( dim_max, d=1 ) # sample spacing: set to 1 pixel
-    freqs = np.abs( f[:dim_half] )
-    #
-    return (freqs, radial_psd, radial_psd_err)
-
-
-def plotRadialPSD( freqs, radial_psd, radial_psd_err=None ):
-    """
-    Make a plot of the radial 1D PSD with matplotlib.
-    """
-    try:
-        import matplotlib.pyplot as plt
-    except ImportError:
-        import sys
-        print( "Error: matplotlib.pyplot cannot be imported!",
-                file=sys.stderr )
-        sys.exit( 30 )
-    dim_half = radial_psd.size
-    # plot
-    plt.figure()
-    plt.loglog( freqs, radial_psd )
-    plt.title( "Radially averaged power spectrum" )
-    plt.xlabel( "k (/pixel)" )
-    plt.ylabel( "Power" )
-    plt.show()
-
-
-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)
+            norm = 1.0  # Do not normalize
+        self.psd2d = (np.abs(imgf) / norm) ** 2
+        return self.psd2d
 
-def pol2cart( rho, phi ):
-    """
-    Convert polar coordinates to Cartesian coordinates.
-    """
-    x = rho * np.cos( phi )
-    y = rho * np.sin( phi )
-    return (x, y)
+    def calc_radial_psd1d(self, k_geometric=True, k_step=1.2):
+        """
+        Computes the radially averaged power spectral density from the
+        provided 2D power spectral density.
 
+        XXX/TODO:
 
-def loadData( filename, ftype="fits" ):
-    """
-    Load data from file into numpy array.
-    """
-    if ftype == "fits":
-        try:
-            from astropy.io import fits
-        except ImportError:
-            import sys
-            print( "Error: astropy.io.fits cannot be imported!",
-                    file=sys.stderr )
-            sys.exit( 20 )
-        ffile = fits.open( filename )
-        data = ffile[0].data.astype( float )
-        ffile.close()
-    else:
-        print( "Error: not implemented yet!",
-                file=sys.stderr )
-        sys.exit( 10 )
-    #
-    return data
+        Arguments:
+          * k_geometric: whether the k (i.e., frequency) varies as
+                         geometric sequences (i.e., k, k*k_step, ...),
+                         otherwise, k varies as (k, k+k_step, ...)
+          * k_step: the step ratio or step length for k
+
+        Return:
+            (freqs, radial_psd, radial_psd_err)
+            freqs: spatial freqencies (unit: ${pixel_unit}^(-1))
+                   if k_geometric=True, frequencies are taken as the
+                   geometric means.
+            radial_psd: radially averaged power spectral density for each
+                        frequency
+            radial_psd_err: standard deviations of each radial_psd
+        """
+        psd2d = self.psd2d.copy()
+        rows, cols = psd2d.shape
+        ## Adjust the PSD array size
+        dim_diff = np.abs(rows - cols)
+        dim_max  = max(rows, cols)
+        # Pad the 2D PSD array to be sqaure
+        if rows > cols:
+            # pad columns
+            if np.mod(dim_diff, 2) == 0:
+                cols_left     = np.zeros((rows, dim_diff/2))
+                cols_left[:]  = np.nan
+                cols_right    = np.zeros((rows, dim_diff/2))
+                cols_right[:] = np.nan
+                psd2d         = np.hstack((cols_left, psd2d, cols_right))
+            else:
+                cols_left     = np.zeros((rows, np.floor(dim_diff/2)))
+                cols_left[:]  = np.nan
+                cols_right    = np.zeros((rows, np.floor(dim_diff/2)+1))
+                cols_right[:] = np.nan
+                psd2d         = np.hstack((cols_left, psd2d, cols_right))
+        elif rows < cols:
+            # pad rows
+            if np.mod(dim_diff, 2) == 0:
+                rows_top       = np.zeros((dim_diff/2, cols))
+                rows_top[:]    = np.nan
+                rows_bottom    = np.zeros((dim_diff/2, cols))
+                rows_bottom[:] = np.nan
+                psd2d          = np.vstack((rows_top, psd2d, rows_bottom))
+            else:
+                rows_top       = np.zeros((np.floor(dim_diff/2), cols))
+                rows_top[:]    = np.nan
+                rows_bottom    = np.zeros((np.floor(dim_diff/2)+1, cols))
+                rows_bottom[:] = np.nan
+                psd2d          = np.vstack((rows_top, psd2d, rows_bottom))
+        ## Compute radially average power spectrum
+        px       = np.arange(-dim_max/2, dim_max/2)
+        x, y     = np.meshgrid(px, px)
+        rho, phi = self.cart2pol(x, y)
+        rho      = np.around(rho).astype(np.int)
+        dim_half = int(np.floor(dim_max/2) + 1)
+        radial_psd     = np.zeros(dim_half)
+        radial_psd_err = np.zeros(dim_half) # standard error
+        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 = fftpack.fftfreq(dim_max, d=1)  # sample spacing: set to 1 pixel
+        freqs = np.abs(f[:dim_half])
+        #
+        self.freqs     = freqs
+        self.psd1d     = radial_psd
+        self.psd1d_err = radial_psd_err
+        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)
+
+    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)
 
 
 def main():
-    pass
+    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("-i", "--infile", dest="infile",
+            required=True, help="input image")
+    parser.add_argument("-o", "--outfile", dest="outfile",
+            required=True, help="output file to store the PSD data")
+    parser.add_argument("-p", "--png", dest="png",
+            help="plot the PSD and save to the given PNG file")
+    parser.add_argument("-v", "--verbose", dest="verbose",
+            action="store_true", help="show verbose information")
+    parser.add_argument("-C", "--clobber", dest="clobber",
+            action="store_true",
+            help="overwrite the output files if already exist")
+    args = parser.parse_args()
+
+    # 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
+    if args.verbose:
+        print("Loading input image ...", file=sys.stderr)
+    with fits.open(args.infile) as ffile:
+        img = ffile[0].data
+        psd = PSD(img, normalize=True)
+
+    # Calculate the power spectral density
+    if args.verbose:
+        print("Calculate 2D power spectral density ...", file=sys.stderr)
+    psd.calc_psd2d()
+    if args.verbose:
+        print("Calculate radially averaged (1D) power spectral density ...",
+                file=sys.stderr)
+    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__":