get_slice_zfreq.py: Rewrite to work with "make_lightcone.py"

2 files changed, 83 insertions, 61 deletions

diff --git a/astro/21cm/get_slice_zfreq.py b/astro/21cm/get_slice_zfreq.py
index cd3dd39..161b714 100755
--- a/astro/21cm/get_slice_zfreq.py
+++ b/astro/21cm/get_slice_zfreq.py
@@ -5,78 +5,118 @@
 #
 
 """
-Get the slice of the specified redshift/frequency (21cm signal) from
-the redshift cube (created by `extract_slice.py` and `fitscube.py`)
-using linear interpolation in redshift.
-
-NOTE:
-The cube slices are ordered in increasing redshifts.
+Get the slices at the specified redshifts/frequencies from the HI
+light-cone cube (created by `make_lightcone.py`), and use linear
+interpolation.
 """
 
-import os
 import sys
 import argparse
+import logging
+from datetime import datetime, timezone
 
 import numpy as np
 from astropy.io import fits
-from astropy.wcs import WCS
+from astropy.cosmology import FlatLambdaCDM
 
 from z2freq import z2freq, freq2z
 
 
-class FITSCube:
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger()
+
+cosmo = FlatLambdaCDM(H0=71, Om0=0.27)
+
+
+class LightCone:
+    """
+    Light-cone cube mimic the observation of HI signal.
+    """
     def __init__(self, infile):
         with fits.open(infile) as f:
             self.data = f[0].data
             self.header = f[0].header
-        self.wcs = WCS(self.header)
+        logger.info("Loaded light-cone cube: %dx%d (cells) * %d (slices)" %
+                    (self.Nside, self.Nside, self.Nslice))
 
     @property
-    def nslice(self):
+    def Nslice(self):
         ns, __, __ = self.data.shape
         return ns
 
     @property
-    def redshifts(self):
-        # Z-axis
-        nslice = self.nslice
-        pix = np.zeros(shape=(nslice, 3), dtype=np.int)
-        pix[:, 2] = np.arange(nslice)
-        world = self.wcs.wcs_pix2world(pix, 0)
-        return world[:, 2]
+    def Nside(self):
+        return self.header["Nside"]
+
+    @property
+    def slices_Dc(self):
+        """
+        The comoving distances of each slice in the light-cone cube.
+        The slices are evenly distributed along the LoS with equal
+        comoving step. [Mpc]
+        """
+        Dc_step = self.header["Dc_step"]
+        Dc_min = self.header["Dc_min"]
+        Dc = np.array([Dc_min + Dc_step*i for i in range(self.Nslice)])
+        return Dc
 
     def get_slice(self, z):
-        redshifts = self.redshifts
-        if z < redshifts.min() or z > redshifts.max():
+        Dc = cosmo.comoving_distance(z).value  # [Mpc]
+        slices_Dc = self.slices_Dc
+        if Dc < slices_Dc.min() or Dc > slices_Dc.max():
             raise ValueError("requested redshift out of range: %.2f" % z)
-        idx2 = (redshifts <= z).sum()
-        idx1 = idx2 - 1
-        z1, slice1 = redshifts[idx1], self.data[idx1, :, :]
-        z2, slice2 = redshifts[idx2], self.data[idx2, :, :]
-        if os.environ.get("DEBUG"):
-            print("DEBUG: redshifts: {0}".format(redshifts), file=sys.stderr)
-            print("DEBUG: z={0}".format(z), file=sys.stderr)
-            print("DEBUG: z1={0}, idx1={1}".format(z1, idx1), file=sys.stderr)
-            print("DEBUG: z2={0}, idx2={1}".format(z2, idx2), file=sys.stderr)
-        return slice1 + (slice2-slice1) * (z-z1) / (z2-z1)
+
+        i2 = (slices_Dc <= Dc).sum()
+        i1 = i2 - 1
+        Dc1, s1 = slices_Dc[i1], self.data[i1, :, :]
+        Dc2, s2 = slices_Dc[i2], self.data[i2, :, :]
+        slope = (s2 - s1) / (Dc2 - Dc1)
+        return s1 + slope * (Dc - Dc1)
+
+    def write_slice(self, outfile, data, z, clobber=False):
+        freq = z2freq(z)
+        Dc = cosmo.comoving_distance(z).value  # [Mpc]
+        header = fits.Header()
+        header["BUNIT"] = (self.header["BUNIT"],
+                           self.header.comments["BUNIT"])
+        header["Lside"] = (self.header["Lside"],
+                           self.header.comments["Lside"])
+        header["Nside"] = (self.header["Nside"],
+                           self.header.comments["Lside"])
+        header["REDSHIFT"] = (z, "redshift of this slice")
+        header["FREQ"] = (freq, "[MHz] observed HI signal frequency")
+        header["Dc"] = (Dc, "[cMpc] comoving distance")
+        header["DATE"] = (datetime.now(timezone.utc).astimezone().isoformat(),
+                          "File creation date")
+        header.add_history(" ".join(sys.argv))
+        hdu = fits.PrimaryHDU(data=data, header=header)
+        try:
+            hdu.writeto(outfile, overwrite=clobber)
+        except TypeError:
+            hdu.writeto(outfile, clobber=clobber)
+        logger.info("Wrote slice to file: %s" % outfile)
 
 
 def main():
-    outfile_default = "{prefix}_f{freq:06.2f}_z{z:06.3f}.fits"
+    outfile_pattern = "{prefix}_f{freq:06.2f}_z{z:06.3f}.fits"
+    outfile_prefix = "deltaTb"
 
     parser = argparse.ArgumentParser(
-        description="Get slices at requested redshifts/frequencies")
+        description="Get slices at requested redshifts/frequencies " +
+        "from light-cone cube")
     parser.add_argument("-C", "--clobber", dest="clobber",
                         action="store_true",
                         help="overwrite existing files")
     parser.add_argument("-i", "--infile", dest="infile", required=True,
-                        help="input redshift data cube")
+                        help="input light-cone cube")
     parser.add_argument("-o", "--outfile", dest="outfile",
-                        default=outfile_default,
-                        help="output FITS image slice (default: %s)" %
-                        outfile_default)
-    parser.add_argument("-p", "--prefix", dest="prefix", required=True,
-                        help="prefix for the output FITS image")
+                        default=outfile_pattern,
+                        help="output image slice filename pattern FITS " +
+                        "(default: %s)" % outfile_pattern)
+    parser.add_argument("-p", "--prefix", dest="prefix",
+                        default=outfile_prefix,
+                        help="prefix of output slices (default: %s)" %
+                        outfile_prefix)
     exgrp = parser.add_mutually_exclusive_group(required=True)
     exgrp.add_argument("-z", "--redshifts", dest="redshifts", nargs="+",
                        help="redshifts where to interpolate slices")
@@ -91,30 +131,12 @@ def main():
         freqs = [float(f) for f in args.freqs]
         redshifts = freq2z(freqs, print_=False)
 
-    cube = FITSCube(args.infile)
+    lightcone = LightCone(args.infile)
     for z, f in zip(redshifts, freqs):
         outfile = args.outfile.format(prefix=args.prefix, z=z, freq=f)
-        print("z=%06.3f, freq=%06.2f MHz : %s" % (z, f, outfile))
-        zslice = cube.get_slice(z)
-        header = fits.Header()
-        try:
-            header["BUNIT"] = (cube.header["BUNIT"],
-                               cube.header.comments["BUNIT"])
-        except KeyError:
-            pass
-        try:
-            header["LSIDE"] = (cube.header["LSIDE"],
-                               cube.header.comments["LSIDE"])
-        except KeyError:
-            pass
-        header["REDSHIFT"] = (z, "Slice where interpolated")
-        header["FREQ"] = (f, "21cm signal frequency [MHz]")
-        header.add_history(" ".join(sys.argv))
-        hdu = fits.PrimaryHDU(data=zslice, header=header)
-        try:
-            hdu.writeto(outfile, overwrite=args.clobber)
-        except TypeError:
-            hdu.writeto(outfile, clobber=args.clobber)
+        logger.info("z=%06.3f, freq=%06.2f MHz : %s ..." % (z, f, outfile))
+        data = lightcone.get_slice(z)
+        lightcone.write_slice(outfile, data=data, z=z, clobber=args.clobber)
 
 
 if __name__ == "__main__":
diff --git a/astro/21cm/make_lightcone.py b/astro/21cm/make_lightcone.py
index 1e888d0..a17c217 100755
--- a/astro/21cm/make_lightcone.py
+++ b/astro/21cm/make_lightcone.py
@@ -260,7 +260,7 @@ class LightCone:
 
     @property
     def header(self):
-        dDc = self.configs.Lside / self.configs.Nside
+        dDc = self.configs.Dc_cell
         Dc_min, Dc_max = self.configs.Dc_limit
         header = fits.Header()
         header["BUNIT"] = (self.configs.unit, "Data unit")