astro/ps2d.py: Support plotting and saving the 2D power spectrum

1 files changed, 37 insertions, 2 deletions

diff --git a/astro/ps2d.py b/astro/ps2d.py
index 1ae59ad..777d592 100755
--- a/astro/ps2d.py
+++ b/astro/ps2d.py
@@ -42,6 +42,12 @@ from astropy.wcs import WCS
 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")
 
 logging.basicConfig(level=logging.INFO,
                     format="%(asctime)s [%(levelname)s] %(message)s",
@@ -208,6 +214,23 @@ class PS2D:
             hdu.writeto(outfile, clobber=clobber)
         logger.info("Wrote 2D power spectrum to file: %s" % outfile)
 
+    def plot(self, fig, ax):
+        """
+        Plot the calculated 2D power spectrum.
+        """
+        x = self.k_perp
+        y = self.k_los
+        mappable = ax.pcolormesh(x[1:], y[1:], np.log10(self.ps2d[1:, 1:]))
+        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$_{los}$ [Mpc$^{-1}$]",
+               title="2D Power Spectrum")
+        cb = fig.colorbar(mappable, ax=ax, pad=0.01, aspect=30)
+        cb.ax.set_xlabel(r"[%s$^2$ Mpc$^3$]" % self.unit)
+        fig.tight_layout()
+        return (fig, ax)
+
     @property
     def Nx(self):
         """
@@ -397,10 +420,13 @@ def main():
     parser.add_argument("-C", "--clobber", dest="clobber",
                         action="store_true",
                         help="overwrite existing file")
+    parser.add_argument("-P", "--plot", dest="plot",
+                        action="store_true",
+                        help="plot the 2D power spectrum and save")
     parser.add_argument("-p", "--pixelsize", dest="pixelsize", type=float,
-                        help="image cube pixel size [arcsec] (default: " +
+                        help="spatial pixel size [arcsec] (default: " +
                         "obtain from FITS header WCS info)")
-    parser.add_argument("--window", dest="window",
+    parser.add_argument("-w", "--window", dest="window",
                         choices=["nuttall"],
                         help="apply window along frequency axis " +
                         "(default: None)")
@@ -444,6 +470,15 @@ def main():
     ps2d.calc_ps2d()
     ps2d.save(outfile=args.outfile, clobber=args.clobber)
 
+    if args.plot:
+        fig = Figure(figsize=(9, 8))
+        FigureCanvas(fig)
+        ax = fig.add_subplot(1, 1, 1)
+        ps2d.plot(ax=ax, fig=fig)
+        plotfile = os.path.splitext(args.outfile)[0] + ".png"
+        fig.savefig(plotfile, format="png", dpi=150)
+        logger.info("Plotted 2D PSD and saved to image: %s" % plotfile)
+
 
 if __name__ == "__main__":
     main()