calc_radial_psd.py: add class "AstroImage" and support background subtraction & exposure correction

1 files changed, 127 insertions, 34 deletions

diff --git a/python/calc_radial_psd.py b/python/calc_radial_psd.py
index 96bc081..40dba6a 100755
--- a/python/calc_radial_psd.py
+++ b/python/calc_radial_psd.py
@@ -13,6 +13,9 @@
 #
 # Changelog:
 # 2016-04-28:
+#   * 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:
@@ -29,7 +32,7 @@
 Compute the radially averaged power spectral density (i.e., power spectrum).
 """
 
-__version__ = "0.3.2"
+__version__ = "0.4.0"
 __date__    = "2016-04-28"
 
 
@@ -71,7 +74,7 @@ class PSD:
         self.pixel = pixel
         self.normalize = normalize
 
-    def calc_psd2d(self):
+    def calc_psd2d(self, verbose=False):
         """
         Computes the 2D power spectral density of the given image.
         Note that the low frequency components are shifted to the center
@@ -81,6 +84,9 @@ class PSD:
             2D power spectral density, which is dimensionless if normalized,
             otherwise has unit ${pixel_unit}^2.
         """
+        if verbose:
+            print("Calculating 2D power spectral density ... ",
+                    end="", flush=True)
         rows, cols = self.img.shape
         ## Compute the power spectral density (i.e., power spectrum)
         imgf = fftpack.fftshift(fftpack.fft2(self.img))
@@ -89,35 +95,32 @@ class PSD:
         else:
             norm = 1.0  # Do not normalize
         self.psd2d = (np.abs(imgf) / norm) ** 2
+        if verbose:
+            print("DONE", flush=True)
         return self.psd2d
 
-    def calc_radial_psd1d(self, k_geometric=True, k_step=1.2):
+    def calc_radial_psd1d(self, verbose=False):
         """
         Computes the radially averaged power spectral density from the
         provided 2D power spectral density.
 
-        XXX/TODO:
-
-        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
         """
+        if verbose:
+            print("Calculating radial (1D) power spectral density ... ",
+                    end="", flush=True)
         psd2d = self.psd2d.copy()
         rows, cols = psd2d.shape
         ## Adjust the PSD array size
         dim_diff = np.abs(rows - cols)
         dim_max  = max(rows, cols)
+        if verbose:
+            print("padding ... ", end="", flush=True)
         # Pad the 2D PSD array to be sqaure
         if rows > cols:
             # pad columns
@@ -154,7 +157,10 @@ class PSD:
         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
+        radial_psd_err = np.zeros(dim_half)
+        if verbose:
+            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()
@@ -163,12 +169,14 @@ class PSD:
             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
+        f = fftpack.fftfreq(dim_max, d=self.pixel[0])
         freqs = np.abs(f[:dim_half])
         #
         self.freqs     = freqs
         self.psd1d     = radial_psd
         self.psd1d_err = radial_psd_err
+        if verbose:
+            print("DONE", end="", flush=True)
         return (freqs, radial_psd, radial_psd_err)
 
     @staticmethod
@@ -218,26 +226,114 @@ class PSD:
         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, verbose=False):
+        self.load_image(image, verbose=verbose)
+        self.load_expmap(expmap, verbose=verbose)
+        self.load_bkgmap(expmap, verbose=verbose)
+
+    def load_image(self, image, verbose=False):
+        if verbose:
+            print("Loading image ... ", end="", flush=True)
+        with fits.open(image) as imgfits:
+            self.image = imgfits[0].data.astype(np.float)
+            self.exposure = imgfits[0].header["EXPOSURE"]
+        if verbose:
+            print("DONE", flush=True)
+
+    def load_expmap(self, expmap, verbose=False):
+        if expmap:
+            if verbose:
+                print("Loading exposure map ... ", end="", flush=True)
+            with fits.open(expmap) as imgfits:
+                self.expmap = imgfits[0].data.astype(np.float)
+            if verbose:
+                print("DONE", flush=True)
+
+    def load_bkgmap(self, bkgmap, verbose=False):
+        if bkgmap:
+            if verbose:
+                print("Loading background map ... ", end="", flush=True)
+            with fits.open(bkgmap) as imgfits:
+                self.bkgmap = imgfits[0].data.astype(np.float)
+                self.exposure_bkg = imgfits[0].header["EXPOSURE"]
+            if verbose:
+                print("DONE", flush=True)
+
+    def subtract_bkg(self, verbose=False):
+        if verbose:
+            print("Subtracting background ... ", end="", flush=True)
+        self.image -= (self.bkgmap / self.exposure_bkg * self.exposure)
+        if verbose:
+            print("DONE", flush=True)
+
+    def correct_exposure(self, cut=0.015, verbose=False):
+        """
+        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
+        """
+        if verbose:
+            print("Correcting image for exposure ... ", end="", flush=True)
+        self.image /= self.expmap
+        # set invalid values to ZERO
+        self.image[ ~ np.isfinite(self.image) ] = 0.0
+        if verbose:
+            print("DONE", flush=True)
+        if cut is not None:
+            # clip image according the exposure threshold
+            if verbose:
+                print("Clipping image (%s) ... " % cut, end="", flush=True)
+            threshold = cut * np.max(self.expmap)
+            self.image[ self.expmap < threshold ] = 0.0
+            if verbose:
+                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("-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 (default: same basename as outfile)")
     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")
+    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 == "":
+    if args.png is None:
         args.png = os.path.splitext(args.outfile)[0] + ".png"
 
     # Check output files whether already exists
@@ -247,20 +343,17 @@ def main():
         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)
+    image = AstroImage(image=args.infile, expmap=args.expmap,
+                       bkgmap=args.bkgmap, verbose=args.verbose)
+    if args.bkgmap:
+        image.subtract_bkg(verbose=args.verbose)
+    if args.expmap:
+        image.correct_exposure(verbose=args.verbose)
 
     # 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()
+    psd = PSD(img=image.image, normalize=True)
+    psd.calc_psd2d(verbose=args.verbose)
+    freqs, psd1d, psd1d_err = psd.calc_radial_psd1d(verbose=args.verbose)
 
     # Write out PSD results
     psd_data = np.column_stack((freqs, psd1d, psd1d_err))