Add astro/eor_window.py: Calc total power within the defined EoR window

2 files changed, 289 insertions, 3 deletions

diff --git a/astro/eor_window.py b/astro/eor_window.py
new file mode 100755
index 0000000..39f645a
--- /dev/null
+++ b/astro/eor_window.py
@@ -0,0 +1,287 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) Weitian LI <weitian@aaronly.me>
+# MIT license
+#
+
+"""
+Calculate the total power within the EoR window on the 2D power spectrum.
+
+The adopted EoR window definition is from [thyagarajan2013],Eq.(26),Fig.(11).
+
+.. [thyagarajan2013]
+   Thyagarajan et al. 2013, ApJ, 776, 6
+"""
+
+import argparse
+from functools import lru_cache
+
+import numpy as np
+from astropy.io import fits
+from astropy.cosmology import FlatLambdaCDM
+import astropy.constants as ac
+
+import matplotlib.pyplot as plt
+from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
+from matplotlib.figure import Figure
+
+
+plt.style.use("ggplot")
+
+# HI line frequency
+freq21cm = 1420.405751  # [MHz]
+# Adopted cosmology
+H0 = 71.0  # [km/s/Mpc]
+OmegaM0 = 0.27
+cosmo = FlatLambdaCDM(H0=H0, Om0=OmegaM0)
+
+
+@lru_cache()
+def freq2z(freq):
+    z = freq21cm / freq - 1.0
+    return z
+
+
+class PS2D:
+    """
+    2D cylindrically averaged power spectrum; calculated by ``ps2d.py``.
+
+    Attributes
+    ----------
+    ps2d : 2D `~numpy.ndarray`
+        Shape: [n_k_los, n_k_perp]
+    """
+    def __init__(self, infile):
+        self.infile = infile
+        with fits.open(infile) as f:
+            self.header = f[0].header
+            self.ps2d = f[0].data[0, :, :]  # errors ignored
+        self.freqc = self.header["Freq_C"]
+        self.freqmin = self.header["Freq_Min"]
+        self.freqmax = self.header["Freq_Max"]
+        self.bandwidth = self.freqmax - self.freqmin  # [MHz]
+        self.zc = self.header["Z_C"]
+        self.pixelsize = self.header["PixSize"]
+        self.unit = self.header["BUNIT"]
+
+    @property
+    def k_perp(self):
+        dk = self.header["CDELT1P"]
+        nk = self.header["NAXIS1"]
+        return np.arange(nk) * dk
+
+    @property
+    def k_los(self):
+        dk = self.header["CDELT2P"]
+        nk = self.header["NAXIS2"]
+        return np.arange(nk) * dk
+
+    @property
+    def k_perp_min(self):
+        return self.k_perp[1]  # ignore the first 0
+
+    @property
+    def k_perp_max(self):
+        return self.k_perp[-1]
+
+    @property
+    def k_los_min(self):
+        return self.k_los[1]  # ignore the first 0
+
+    @property
+    def k_los_max(self):
+        return self.k_los[-1]
+
+    def sum_power(self, window):
+        """
+        Sum the power within the defined window.
+
+        NOTE: The cylindrical average should be accounted for.
+        """
+        k_perp = self.k_perp
+        k_los = self.k_los
+        dk_perp = k_perp[1] - k_perp[0]
+        dk_los = k_los[1] - k_los[0]
+        volume = np.zeros_like(self.ps2d)
+        volume[0, :] = 2*np.pi * k_perp * dk_perp * dk_los
+        for i in range(1, len(k_los)):
+            # The extra "2" to account for the average on +k_los and -k_los
+            volume[i, :] = 2*np.pi * k_perp * dk_perp * dk_los * 2
+
+        power = np.sum(self.ps2d * window * volume)
+        return power
+
+    def eor_window(self, fov, e,
+                   k_perp_min=None, k_perp_max=None,
+                   k_los_min=None, k_los_max=None):
+        """
+        Determine the EoR window region.
+
+        Parameters
+        ----------
+        fov : float
+            instrumental field of view (FoV)
+            Unit: [deg]
+        e : float
+            Thyagarajan proposed characteristic convolution width factor,
+            generally 0-3
+
+        Returns
+        -------
+        window : 2D bool `~numpy.ndarray`
+            2D array mask of the same size of the power spectrum indicating
+            the defined EoR window region.
+        header : fits.Header
+            FITS header with the keywords recording the EoR window variables
+        """
+        if k_perp_min is None:
+            k_perp_min = self.k_perp_min
+        if k_perp_max is None:
+            k_perp_max = self.k_perp_max
+        if k_los_min is None:
+            k_los_min = self.k_los_min
+        if k_los_max is None:
+            k_los_max = self.k_los_max
+
+        window = np.ones_like(self.ps2d, dtype=bool)
+        k_perp = self.k_perp
+        k_los = self.k_los
+        k_wedge = self.wedge_edge(k_perp, fov=fov, e=e)
+        window[k_los < k_los_min, :] = False
+        window[k_los > k_los_max, :] = False
+        window[:, k_perp < k_perp_min] = False
+        window[:, k_perp > k_perp_max] = False
+        for i, k in enumerate(k_wedge):
+            window[k_los < k, i] = False
+
+        header = self.eor_window_header(fov=fov, e=e,
+                                        k_perp_min=k_perp_min,
+                                        k_perp_max=k_perp_max,
+                                        k_los_min=k_los_min,
+                                        k_los_max=k_los_max)
+        return (window, header)
+
+    def eor_window_header(self, fov, e, k_perp_min, k_perp_max,
+                          k_los_min, k_los_max):
+        header = self.header.copy(strip=True)
+        header["FoV"] = (fov, "[deg] Field of view to determine EoR window")
+        header["e_ConvW"] = (e, "characteristic convolution width")
+        header["kper_min"] = (k_perp_min, "[Mpc^-1] minimum k_perp")
+        header["kper_max"] = (k_perp_max, "[Mpc^-1] maximum k_perp")
+        header["klos_min"] = (k_los_min, "[Mpc^-1] minimum k_los")
+        header["klos_max"] = (k_los_max, "[Mpc^-1] maximum k_los")
+        return header
+
+    def wedge_edge(self, k_perp, fov, e):
+        """
+        The boundary/edge between the EoR window (top-left) and the
+        foreground wedge (bottom-right).
+        """
+        Hz = cosmo.H(self.zc).value  # [km/s/Mpc]
+        Dc = cosmo.comoving_distance(self.zc).value  # [Mpc]
+        c = ac.c.to("km/s").value  # [km/s]
+        coef = Hz * Dc / (c * (1+self.zc))
+        term1 = np.sin(np.deg2rad(fov)) * k_perp  # [Mpc^-1]
+        term2 = ((2*np.pi * e * freq21cm / self.bandwidth) /
+                 ((1 + self.zc) * Dc))  # [Mpc^-1]
+        k_los = coef * (term1 + term2)
+        return k_los
+
+    def plot(self, ax, fov, e,
+             k_perp_min=None, k_perp_max=None,
+             k_los_min=None, k_los_max=None,
+             colormap="jet"):
+        """
+        Plot the 2D power spectrum with EoR window marked on.
+        """
+        if k_perp_min is None:
+            k_perp_min = self.k_perp_min
+        if k_perp_max is None:
+            k_perp_max = self.k_perp_max
+        if k_los_min is None:
+            k_los_min = self.k_los_min
+        if k_los_max is None:
+            k_los_max = self.k_los_max
+
+        x = self.k_perp
+        y = self.k_los
+        y_wedge = self.wedge_edge(x, fov=fov, e=e)
+        title = "EoR Window (fov=%.1f[deg], e=%.1f)" % (fov, e)
+
+        # data
+        mappable = ax.pcolormesh(x[1:], y[1:],
+                                 np.log10(self.ps2d[1:, 1:]),
+                                 cmap=colormap)
+        # EoR window
+        ax.axvline(x=k_perp_min, color="black", linewidth=2, linestyle="--")
+        ax.axvline(x=k_perp_max, color="black", linewidth=2, linestyle="--")
+        ax.axhline(y=k_los_min, color="black", linewidth=2, linestyle="--")
+        ax.axhline(y=k_los_max, color="black", linewidth=2, linestyle="--")
+        ax.plot(x, y_wedge, color="black", linewidth=2, linestyle="--")
+        #
+        ax.set(xscale="log", yscale="log",
+               xlim=(x[1], x[-1]), ylim=(y[1], y[-1]),
+               xlabel=r"k$_{\perp}$ [Mpc$^{-1}$]",
+               ylabel=r"k$_{||}$ [Mpc$^{-1}$]",
+               title=title)
+        cb = ax.figure.colorbar(mappable, ax=ax, pad=0.01, aspect=30)
+        cb.ax.set_xlabel(r"[%s$^2$ Mpc$^3$]" % self.unit)
+        return ax
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Determine EoR window region and calculate total power")
+    parser.add_argument("-F", "--fov", dest="fov",
+                        type=float, required=True,
+                        help="instrumental FoV to determine the EoR window")
+    parser.add_argument("-e", "--conv-width", dest="conv_width",
+                        type=float, default=3.0,
+                        help="characteristic convolution width (default: 3.0)")
+    parser.add_argument("-p", "--k-perp-min", dest="k_perp_min", type=float,
+                        help="minimum k wavenumber perpendicular to LoS; " +
+                        "unit: [Mpc^-1]")
+    parser.add_argument("-P", "--k-perp-max", dest="k_perp_max", type=float,
+                        help="maximum k wavenumber perpendicular to LoS")
+    parser.add_argument("-l", "--k-los-min", dest="k_los_min", type=float,
+                        help="minimum k wavenumber along LoS")
+    parser.add_argument("-L", "--k-los-max", dest="k_los_max", type=float,
+                        help="maximum k wavenumber along LoS")
+    parser.add_argument("--save-window", dest="save_window",
+                        help="save the determined EoR window into FITS " +
+                        "file with the provided filename")
+    parser.add_argument("--plot", dest="plot",
+                        help="plot the 2D power spectrum with the " +
+                        "determined EoR window marked, and save into " +
+                        "the specified file")
+    parser.add_argument("infile", help="2D power spectrum FITS file")
+    args = parser.parse_args()
+
+    ps2d = PS2D(args.infile)
+    window, header = ps2d.eor_window(fov=args.fov, e=args.conv_width,
+                                     k_perp_min=args.k_perp_min,
+                                     k_perp_max=args.k_perp_max,
+                                     k_los_min=args.k_los_min,
+                                     k_los_max=args.k_los_max)
+    power = ps2d.sum_power(window)
+    print("Total power within EoR window: %g [%s]" % (power, ps2d.unit))
+
+    if args.save_window:
+        hdu = fits.PrimaryHDU(data=window.astype(np.int16), header=header)
+        hdu.writeto(args.save_window)
+        print("Saved EoR window into file: %s" % args.save_window)
+
+    if args.plot:
+        fig = Figure(figsize=(8, 8), dpi=150)
+        FigureCanvas(fig)
+        ax = fig.add_subplot(1, 1, 1)
+        ps2d.plot(ax=ax, fov=args.fov, e=args.conv_width,
+                  k_perp_min=args.k_perp_min, k_perp_max=args.k_perp_max,
+                  k_los_min=args.k_los_min, k_los_max=args.k_los_max)
+        fig.tight_layout()
+        fig.savefig(args.plot)
+        print("Plotted 2D PSD with EoR window and saved to: %s" % args.plot)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/astro/ps2d.py b/astro/ps2d.py
index 718b59a..0bd0ac4 100755
--- a/astro/ps2d.py
+++ b/astro/ps2d.py
@@ -339,9 +339,8 @@ class PS2D:
         Reference: Ref.[liu2014].Eq.(A9)
         """
         dfreq = self.dfreq  # [MHz]
-        c = ac.c.si.value  # [m/s]
-        Ez = cosmo.efunc(self.zc)
-        Hz = Ez * H0 * 1000.0  # [m/s/Mpc]
+        c = ac.c.to("km/s").value  # [km/s]
+        Hz = cosmo.H(self.zc).value  # [km/s/Mpc]
         d_z = c * (1+self.zc)**2 * dfreq / Hz / freq21cm
         return d_z