From 2a4f07ece389ab3454afeeeced8480a1c958f8f9 Mon Sep 17 00:00:00 2001
From: Aaron LI <aly@aaronly.me>
Date: Mon, 16 Oct 2017 10:57:51 +0800
Subject: Move python/calc_radial_psd.py to astro/calc_psd.py with some cleanup
---
astro/calc_psd.py | 438 +++++++++++++++++++++++++++++++++++++++++++
python/calc_radial_psd.py | 460 ----------------------------------------------
2 files changed, 438 insertions(+), 460 deletions(-)
create mode 100755 astro/calc_psd.py
delete mode 100755 python/calc_radial_psd.py
diff --git a/astro/calc_psd.py b/astro/calc_psd.py
new file mode 100755
index 0000000..fe02b02
--- /dev/null
+++ b/astro/calc_psd.py
@@ -0,0 +1,438 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2015-2017 Aaron LI
+# MIT License
+#
+
+"""
+Compute the radially averaged power spectral density (i.e., power spectrum)
+of a 2D image.
+
+XXX: If the input image is NOT SQUARE; then are the horizontal frequencies
+ the same as the vertical frequencies ??
+
+Credit
+------
+* 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
+"""
+
+__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))
+ 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()
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()
--
cgit v1.2.2