diff options
| -rwxr-xr-x | calc_radial_psd.py | 131 |
1 files changed, 69 insertions, 62 deletions
diff --git a/calc_radial_psd.py b/calc_radial_psd.py index 4c1d85b..3f3508a 100755 --- a/calc_radial_psd.py +++ b/calc_radial_psd.py @@ -13,9 +13,11 @@ # # Aaron LI <aaronly.me@gmail.com> # Created: 2015-04-22 -# Updated: 2016-05-01 +# Updated: 2016-05-09 # -# Changelog: +# Change log: +# 2016-05-09: +# * PEP8 fixes # 2016-05-01: # * Adjust plot axis limits # 2016-04-28: @@ -38,16 +40,15 @@ # * Encapsulate the functions within class 'PSD' # * Update docs/comments # +# TODO/FIXME: +# * subtract background with normalization factor (w.r.t particle background) +# considered +# """ Compute the radially averaged power spectral density (i.e., power spectrum). """ -__version__ = "0.5.1" -__date__ = "2016-05-01" - - -import sys import os import argparse @@ -58,6 +59,9 @@ import matplotlib.pyplot as plt from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas from matplotlib.figure import Figure +__version__ = "0.5.2" +__date__ = "2016-05-01" + plt.style.use("ggplot") @@ -75,8 +79,8 @@ class PSD: # 2D power spectral density psd2d = None # 1D (radially averaged) power spectral density - freqs = None - psd1d = None + freqs = None + psd1d = None psd1d_err = None def __init__(self, img, pixel=(1.0, "pixel"), normalize=True): @@ -101,9 +105,9 @@ class PSD: """ if verbose: print("Calculating 2D power spectral density ... ", - end="", flush=True) + end="", flush=True) rows, cols = self.img.shape - ## Compute the power spectral density (i.e., power spectrum) + # 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 @@ -128,37 +132,36 @@ class PSD: """ if verbose: print("Calculating radial (1D) power spectral density ... ", - end="", flush=True) + end="", flush=True) if verbose: print("padding ... ", end="", flush=True) - psd2d = self.pad_square(self.psd2d, value=np.nan) - dim = psd2d.shape[0] + 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) + 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) + rho = np.around(rho).astype(np.int) + radial_psd = np.zeros(dim_half) radial_psd_err = np.zeros(dim_half) if verbose: - print("radially averaging ... ", - end="", flush=True) + 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) + 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.freqs = freqs + self.psd1d = radial_psd self.psd1d_err = radial_psd_err if verbose: print("DONE", end="", flush=True) @@ -193,35 +196,34 @@ class PSD: 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_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)) + 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_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)) + 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_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)) + 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_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)) + mat = np.vstack((rows_top, mat, rows_bottom)) return mat def plot(self, ax=None, fig=None): @@ -237,8 +239,7 @@ class PSD: ymin = np.min(self.psd1d) / 3.0 ymax = np.max(self.psd1d[mask] + self.psd1d_err[mask]) * 1.2 # - eb = ax.errorbar(self.freqs, self.psd1d, yerr=self.psd1d_err, - fmt="none") + 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") @@ -266,7 +267,7 @@ class AstroImage: # background map (e.g., stowed background) bkgmap = None # exposure time of the input image - exposure = None + exposure = None # exposure time of the background map exposure_bkg = None @@ -332,18 +333,21 @@ class AstroImage: if rows > rows_ref: img_fixed = img[:rows_ref, :] else: - img_fixed = np.row_stack((img, + 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, + 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)) + raise ValueError("shape difference exceeds tolerance: " + + "(%d, %d) vs. (%d, %d)" % + (img.shape + ref.shape)) # if self.bkgmap is not None: if verbose: @@ -377,7 +381,7 @@ class AstroImage: with np.errstate(divide="ignore", invalid="ignore"): self.image /= self.expmap # set invalid values to ZERO - self.image[ ~ np.isfinite(self.image) ] = 0.0 + self.image[~ np.isfinite(self.image)] = 0.0 if verbose: print("DONE", flush=True) if cut is not None: @@ -385,7 +389,7 @@ class AstroImage: if verbose: print("Clipping image (%s) ... " % cut, end="", flush=True) threshold = cut * np.max(self.expmap) - self.image[ self.expmap < threshold ] = 0.0 + self.image[self.expmap < threshold] = 0.0 if verbose: print("DONE", flush=True) @@ -395,22 +399,25 @@ def main(): 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__)) + version="%(prog)s " + + "%s (%s)" % (__version__, __date__)) parser.add_argument("-v", "--verbose", dest="verbose", - action="store_true", help="show verbose information") + 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") - parser.add_argument("-i", "--infile", dest="infile", - required=True, help="input image") + 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)") + 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") + 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)") + help="plot the PSD and save " + + "(default: same basename as outfile)") args = parser.parse_args() if args.png is None: @@ -442,7 +449,7 @@ def main(): # Make and save a plot fig = Figure(figsize=(10, 8)) - canvas = FigureCanvas(fig) + FigureCanvas(fig) ax = fig.add_subplot(111) psd.plot(ax=ax, fig=fig) fig.savefig(args.png, format="png", dpi=150) |