diff options
Diffstat (limited to 'python')
| -rwxr-xr-x | python/radialPSD2d.py | 386 |
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__": |