add python scripts

1 files changed, 173 insertions, 0 deletions

diff --git a/python/radialPSD2d.py b/python/radialPSD2d.py
new file mode 100755
index 0000000..cc5dd85
--- /dev/null
+++ b/python/radialPSD2d.py
@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+#
+# Aaron LI <aaronly.me@gmail.com>
+# 2015/04/22
+#
+
+"""
+Computes the radially averaged power spectral density (power spectrum).
+"""
+
+
+import numpy as np
+from scipy import fftpack
+
+
+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
+    """
+    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 radially averaged power spectral density (power spectrum)
+    from the provided 2D power spectrum.
+
+    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) )
+        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)
+
+def pol2cart( rho, phi ):
+    """
+    Convert polar coordinates to Cartesian coordinates.
+    """
+    x = rho * np.cos( phi )
+    y = rho * np.sin( phi )
+    return (x, y)
+
+
+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
+
+
+def main():
+    pass
+
+
+if __name__ == "__main__":
+    main()
+