Add astro/ps1d_eorwindow.py: calculate 1D power spectra within EoR window

1 files changed, 216 insertions, 0 deletions

diff --git a/astro/ps1d_eorwindow.py b/astro/ps1d_eorwindow.py
new file mode 100755
index 0000000..bbb0ded
--- /dev/null
+++ b/astro/ps1d_eorwindow.py
@@ -0,0 +1,216 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2017 Weitna LI <weitian@aaronly.me>
+# MIT License
+#
+
+"""
+Average the 2D power spectrum within the EoR window (i.e., excluding the
+foreground contaminated wedge) to derive the 1D spherically averaged
+power spectrum.
+"""
+
+import os
+import argparse
+
+import numpy as np
+
+import matplotlib
+import matplotlib.style
+from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
+from matplotlib.figure import Figure
+
+from eor_window import PS2D
+
+
+# Matplotlib settings
+matplotlib.style.use("ggplot")
+for k, v in [("font.family",       "monospace"),
+             ("image.cmap",        "jet"),
+             ("xtick.major.size",  7.0),
+             ("xtick.major.width", 2.0),
+             ("xtick.minor.size",  4.0),
+             ("xtick.minor.width", 1.5),
+             ("ytick.major.size",  7.0),
+             ("ytick.major.width", 2.0),
+             ("ytick.minor.size",  4.0),
+             ("ytick.minor.width", 1.5)]:
+    matplotlib.rcParams[k] = v
+
+
+class PS1D:
+    """
+    Calculate the 1D spherically averaged power spectrum from 2D PS.
+    """
+    def __init__(self, ps2d):
+        self.ps2d = ps2d
+        self.data = ps2d.ps2d  # shape: [n_k_los, n_k_perp]
+        self.data_err = ps2d.ps2d_err
+
+    @property
+    def k_perp(self):
+        return self.ps2d.k_perp
+
+    @property
+    def k_los(self):
+        return self.ps2d.k_los
+
+    @property
+    def eor_window(self):
+        return self.ps2d.eor_window()
+
+    def calc_ps1d(self, normalize=True):
+        """
+        Calculate the 1D spherically averaged power spectrum by averaging
+        the 2D cylindrical power spectrum.
+
+        Parameters
+        ----------
+        normalize : bool
+            Whether to normalize the 1D power spectrum to obtain the
+            dimensionless power spectrum, i.e.,
+                Δ^2(k) = (k^3 / (2*π^2)) P(k)
+        """
+        eor_window = self.eor_window
+        data = self.data.copy()
+        data_err = self.data_err.copy()
+        data[~eor_window] = np.nan
+        data_err[~eor_window] = np.nan
+
+        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]
+        dk = np.sqrt(dk_perp * dk_los)
+        print("dk = %.6f [Mpc^-1]" % dk)
+        k_max = np.sqrt(k_perp[-1]**2 + k_los[-1]**2)
+        nk = int(k_max / dk) + 1
+        print("number of k points: %d" % nk)
+        ps1d_k = np.arange(nk) * dk
+
+        # PS1D's 3 columns: [k, ps1d, ps1d_err]
+        ps1d = np.zeros(shape=(nk, 3))
+        ps1d[:, 0] = ps1d_k
+
+        print("averaging 2D power spectrum ...")
+        mx, my = np.meshgrid(k_perp, k_los)
+        mk = np.sqrt(mx**2 + my**2)
+        for i, k in enumerate(ps1d_k):
+            ii, jj = (mk <= k).nonzero()
+            mk[ii, jj] = np.inf
+            cells = data[ii, jj]
+            cells = cells[np.isfinite(cells)]
+            if len(cells) > 0:
+                ps1d[i, 1] = np.mean(cells)
+                cells = data_err[ii, jj]
+                cells = cells[np.isfinite(cells)]
+                ps1d[i, 2] = np.sqrt(np.sum((cells/len(cells))**2))
+
+        if normalize:
+            coef = ps1d_k**3 / (2*np.pi**2)
+            ps1d[:, 1] *= coef
+            ps1d[:, 2] *= coef
+            self.ps1d_normalized = True
+        else:
+            self.ps1d_normalized = False
+
+        self.ps1d = ps1d
+        return ps1d
+
+    def save(self, outfile):
+        ps1d = self.ps1d
+        header = [
+            "EoR window:",
+            "  FoV: %f [deg]" % self.ps2d.fov,
+            "  e_ConvWidth: %f" % self.ps2d.e,
+            "  k_perp_min: %f [Mpc^-1]" % self.ps2d.k_perp_min,
+            "  k_perp_max: %f [Mpc^-1]" % self.ps2d.k_perp_max,
+            "  k_los_min: %f [Mpc^-1]" % self.ps2d.k_los_min,
+            "  k_los_max: %f [Mpc^-1]" % self.ps2d.k_los_max,
+            "",
+            "k: wavenumber [Mpc^-1]",
+        ]
+        if self.ps1d_normalized:
+            header += ["ps1d: normalized power [K^2]"]
+        else:
+            header += ["ps1d: power [K^2 Mpc^3]"]
+        header += [
+            "ps1d_err: power errors",
+            "",
+            "k   ps1d   ps1d_err"
+        ]
+        np.savetxt(outfile, ps1d, header="\n".join(header))
+        print("saved 1D power spectrum to file: %s" % outfile)
+
+    def plot(self, ax):
+        ps1d = self.ps1d
+        if self.ps1d_normalized:
+            ylabel = r"$\Delta^2(k)$ [K$^2$]"
+        else:
+            ylabel = r"$P(k)$ [K$^2$ Mpc$^3$]"
+
+        x = ps1d[:, 0]
+        y = ps1d[:, 1]
+        yerr = ps1d[:, 2]
+        ax.errorbar(x[1:], y[1:], yerr=yerr[1:], fmt="none")
+        ax.plot(x[1:], y[1:], marker="o")
+        ax.set(xscale="log", yscale="log",
+               xlabel=r"[Mpc$^{-1}$]", ylabel=ylabel,
+               title="1D Spherically Average Power Spectrum")
+        return ax
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Calculate 1D power spectrum within the EoR window")
+    parser.add_argument("-C", "--clobber", dest="clobber", action="store_true",
+                        help="overwrite the output files if already exist")
+    parser.add_argument("-F", "--fov", dest="fov",
+                        type=float, required=True,
+                        help="instrumental FoV to determine the EoR window; " +
+                        "SKA1-Low has FoV ~ 3.12 / (nu/200MHz) [deg], i.e., " +
+                        "~5.03 @ 124, ~3.95 @ 158, ~3.18 @ 196")
+    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("--no-plot", dest="noplot", action="store_true",
+                        help="do not plot and save the calculated 1D power " +
+                        "power within the EoR window")
+    parser.add_argument("-i", "--infile", dest="infile", required=True,
+                        help="2D power spectrum FITS file")
+    parser.add_argument("-o", "--outfile", dest="outfile", required=True,
+                        help="output TXT file to save the PSD data")
+    args = parser.parse_args()
+
+    if (not args.clobber) and os.path.exists(args.outfile):
+        raise OSError("outfile '%s' already exists" % args.outfile)
+
+    ps2d = PS2D(args.infile, 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)
+    ps1d = PS1D(ps2d)
+    ps1d.calc_ps1d()
+    ps1d.save(args.outfile)
+
+    if not args.noplot:
+        fig = Figure(figsize=(8, 8), dpi=150)
+        FigureCanvas(fig)
+        ax = fig.add_subplot(1, 1, 1)
+        ps1d.plot(ax=ax)
+        fig.tight_layout()
+        plotfile = os.path.splitext(args.outfile)[0] + ".png"
+        fig.savefig(plotfile)
+        print("Plotted 1D power spectrum within EoR window: %s" % plotfile)
+
+
+if __name__ == "__main__":
+    main()