Add astro/21cm/21cmfast_lightcone.py: convert 21cmFAST lighttravel cubes

1 files changed, 192 insertions, 0 deletions

diff --git a/astro/21cm/21cmfast_lightcone.py b/astro/21cm/21cmfast_lightcone.py
new file mode 100755
index 0000000..b7fe7bc
--- /dev/null
+++ b/astro/21cm/21cmfast_lightcone.py
@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2017 Weitian LI <weitian@aaronly.me>
+# MIT License
+#
+
+"""
+Convert the 21cmFAST simulated and generated lighttravel cubes into the
+FITS format with proper headers that are consistent with the lightcone
+created by ``make_lightcone.py``, and therefore ``get_slice_zfreq.py``
+can be used to extract needed slices.
+
+The 21cmFAST lighttravel cube is created by the tool
+``redshift_interpolate_boxes.c`` shipped with 21cmFAST.
+
+For example, Mesinger et al. (2016) released their simulations at:
+http://homepage.sns.it/mesinger/EOS.html
+where the *faint galaxies* model is recommended by the authors.
+
+The lighttravel cubes have filenames:
+``delta_T_v3_no_halos__zstart*_zend*_FLIPBOXES0_1024_1600Mpc_lighttravel``
+data type: float32, little endidian, C ordering
+cube dimension: 1024x1024 (XY-spatial), 1024 (LoS/z)
+length: 1600 comoving Mpc
+unit: mK
+
+.. 21cmFAST: https://github.com/andreimesinger/21cmFAST
+
+.. Mesinger et al. 2016, MNRAS, 459, 2342
+"""
+
+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
+import astropy.units as au
+
+
+logging.basicConfig(level=logging.INFO,
+                    format="[%(levelname)s:%(lineno)d] %(message)s")
+logger = logging.getLogger()
+
+
+class LightCone:
+    """
+    Light-cone cube mimic the observation of HI signal.
+
+    cube shape:
+      * [X, Y, LoS] (axes_swapped = False)
+      * [LoS, Y, Y] (axes_swapped = True)
+    """
+    def __init__(self, cube, zmin, zmax, unit, Lside, cosmo,
+                 axes_swapped=False):
+        self.cube = cube
+        self.zmin = zmin
+        self.zmax = zmax
+        self.unit = unit
+        self.Lside = Lside  # [cMpc] simulation side length
+        self.cosmo = cosmo
+        self.Dc_min = cosmo.comoving_distance(zmin).value  # [cMpc]
+        self.Dc_max = cosmo.comoving_distance(zmax).value  # [cMpc]
+        self.axes_swapped = axes_swapped
+
+    def swap_axes(self):
+        self.cube = np.swapaxes(self.cube, 0, 2)
+        self.axes_swapped = not self.axes_swapped
+        logger.info("Axes swapped: %s" % self.axes_swapped)
+        if self.axes_swapped:
+            logger.info("Cube axes: [LoS, Y, X]")
+        else:
+            logger.info("Cube axes: [X, Y, LoS]")
+
+    @property
+    def Nside(self):
+        return self.cube.shape[1]
+
+    @property
+    def Nslice(self):
+        if self.axes_swapped:
+            return self.cube.shape[0]
+        else:
+            return self.cube.shape[2]
+
+    @property
+    def Dc_cell(self):
+        return self.Lside / self.Nslice
+
+    @property
+    def wcs(self):
+        w = WCS(naxis=3)
+        w.wcs.ctype = ["pixel", "pixel", "pixel"]
+        w.wcs.cunit = ["Mpc", "Mpc", "Mpc"]  # comoving
+        w.wcs.crpix = np.array([1.0, 1.0, 1.0])
+        w.wcs.crval = np.array([0.0, 0.0, self.Dc_min])
+        w.wcs.cdelt = np.array([self.Dc_cell, self.Dc_cell, self.Dc_cell])
+        return w
+
+    @property
+    def header(self):
+        dDc = self.Dc_cell
+        header = fits.Header()
+        header["BUNIT"] = (str(self.unit), "Data unit")
+        header["zmin"] = (self.zmin, "HI simulation minimum redshift")
+        header["zmax"] = (self.zmax, "HI simulation maximum redshift")
+        header["Dc_min"] = (self.Dc_min, "[cMpc] comoving distance at zmin")
+        header["Dc_max"] = (self.Dc_max, "[cMpc] comoving distance at zmax")
+        header["Dc_step"] = (dDc, "[cMpc] comoving distance between slices")
+        header["Lside"] = (self.Lside, "[cMpc] Simulation side length")
+        header["Nside"] = (self.Nside, "Number of cells at each side")
+        header["DATE"] = (datetime.now(timezone.utc).astimezone().isoformat(),
+                          "File creation date")
+        header.add_history(" ".join(sys.argv))
+        header.extend(self.wcs.to_header(), update=True)
+        return header
+
+    def write(self, outfile, clobber=False):
+        hdu = fits.PrimaryHDU(data=self.cube, header=self.header)
+        logger.info("Created FITS object, writing to disk ...")
+        try:
+            hdu.writeto(outfile, overwrite=clobber)
+        except TypeError:
+            hdu.writeto(outfile, clobber=clobber)
+        logger.info("Wrote light-cone cube to: %s" % outfile)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="convert 21cmFAST lighttravel cube to FITS lightcone cube")
+    parser.add_argument("-C", "--clobber", dest="clobber",
+                        action="store_true", help="overwrite existing file")
+    parser.add_argument("--data-type", dest="data_type", default="float32",
+                        help="input cube data type (default: float32)")
+    parser.add_argument("--unit", dest="unit", default="mK",
+                        help="input cube data unit (default: mK)")
+    parser.add_argument("--unit-out", dest="unit_out", default="K",
+                        help="output data unit (default: K)")
+    parser.add_argument("--side-length", dest="side_length",
+                        type=float, default=1600.0,
+                        help="input cube simulation side length [cMpc]")
+    parser.add_argument("--dimension", dest="dimension",
+                        nargs=3, type=int, default=[1024, 1024, 1024],
+                        help="input cube dimension (C ordering assumed)")
+    parser.add_argument("--H0", dest="H0", type=float, default=67.8,
+                        help="simulation adopted H0 (default: 67.8)")
+    parser.add_argument("--omega-m0", dest="Om0", type=float, default=0.308,
+                        help="simulation adopted OmegaM0 (default: 0.308)")
+    parser.add_argument("--omega-b0", dest="Ob0", type=float, default=0.0484,
+                        help="simulation adopted Omegab0 (default: 0.0484)")
+    parser.add_argument("-z", "--z-min", dest="zmin", type=float,
+                        required=True, help="minimum/beginning redshift")
+    parser.add_argument("-Z", "--z-max", dest="zmax", type=float,
+                        required=True, help="maximum/end redshift")
+    parser.add_argument("infile", help="input 21cmFAST lighttravel cube")
+    parser.add_argument("outfile", help="output FITS lightcone cube")
+    args = parser.parse_args()
+
+    if os.path.exists(args.outfile):
+        if args.clobber:
+            os.remove(args.outfile)
+            logger.warning("Removed existing output file: %s" % args.outfile)
+        else:
+            raise OSError("output file already exists: %s" % args.outfile)
+
+    cosmo = FlatLambdaCDM(H0=args.H0, Om0=args.Om0, Ob0=args.Ob0)
+    logger.info("Cosmology: {0}".format(cosmo))
+    unit_in = au.Unit(args.unit)
+    unit_out = au.Unit(args.unit_out)
+    logger.info("Unit: %s (input) -> %s (output)" % (unit_in, unit_out))
+
+    cube_in = np.fromfile(args.infile, dtype=np.dtype(args.data_type))
+    cube_in = cube_in.reshape(args.dimension)
+    logger.info("Loaded lighttravel cube from file: %s" % args.infile)
+    logger.info("Data type: {0}; dimension: {1}".format(
+        cube_in.dtype, cube_in.shape))
+
+    cube_in *= unit_in.to(unit_out)
+    logger.info("Converted unit to: %s" % unit_out)
+    lightcone = LightCone(cube_in, zmin=args.zmin, zmax=args.zmax,
+                          unit=unit_out, Lside=args.side_length,
+                          cosmo=cosmo, axes_swapped=False)
+    lightcone.swap_axes()
+    lightcone.write(args.outfile)
+
+
+if __name__ == "__main__":
+    main()