Add astro/21cm/tile_slice.py

1 files changed, 93 insertions, 0 deletions

diff --git a/astro/21cm/tile_slice.py b/astro/21cm/tile_slice.py
new file mode 100755
index 0000000..0274524
--- /dev/null
+++ b/astro/21cm/tile_slice.py
@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2017 Weitian LI <weitian@aaronly.me>
+# MIT License
+#
+
+"""
+Tile the given slices to the required FoV size, scale down to the
+wanted size (for faster simulation later).
+"""
+
+import sys
+import argparse
+
+import numpy as np
+from astropy.io import fits
+from astropy.cosmology import FlatLambdaCDM
+import astropy.units as au
+import scipy.ndimage
+
+from z2freq import z2freq, freq2z
+
+
+# Adopted cosmology
+cosmo = FlatLambdaCDM(H0=71.0, Om0=0.27, Ob0=0.046)
+
+
+def main():
+    outfile_default = "{prefix}_f{freq:06.2f}_N{Nside}_fov{fov:.1f}.fits"
+
+    parser = argparse.ArgumentParser(
+        description="Tile slice to match FoV size and scale to required size")
+    parser.add_argument("-C", "--clobber", dest="clobber",
+                        action="store_true",
+                        help="overwrite existing files")
+    parser.add_argument("-F", "--fov", dest="fov", default=5.0, type=float,
+                        help="required FoV [deg] of the output slice " +
+                        "(default: 5.0 deg)")
+    parser.add_argument("-N", "--n-side", dest="Nside", default=600, type=int,
+                        help="required image size of output slice " +
+                        "(default: 500)")
+    parser.add_argument("-i", "--infile", dest="infile", required=True,
+                        help="input slice")
+    parser.add_argument("-o", "--outfile", dest="outfile",
+                        default=outfile_default,
+                        help="output tiled slice (default: %s)" %
+                        outfile_default)
+    parser.add_argument("-p", "--prefix", dest="prefix", required=True,
+                        help="prefix for the output tiled slice")
+    exgrp = parser.add_mutually_exclusive_group(required=True)
+    exgrp.add_argument("-z", "--redshift-c", dest="zc", type=float,
+                       help="central redshift of the selected cube")
+    exgrp.add_argument("-f", "--freq-c", dest="fc", type=float,
+                       help="central frequency [MHz] of the selected cube")
+    args = parser.parse_args()
+
+    if args.zc:
+        zc = args.zc
+        fc = z2freq(zc, print_=False)
+    else:
+        fc = args.fc
+        zc = freq2z(fc, print_=False)
+
+    with fits.open(args.infile) as f:
+        img_in = f[0].data
+        header = f[0].header
+    freq = header["FREQ"]  # [MHz]
+    Lside = header["LSIDE"]  # [Mpc]
+    print("frequency = %.2f [MHz], Lside = %.1f [cMpc]" % (freq, Lside))
+    DM = cosmo.comoving_distance(zc).value  # [Mpc]
+    print("Comoving distance (@z=%.3f/f=%.2fMHz) = %.2f [Mpc]" % (zc, fc, DM))
+    Nside = img_in.shape[0]
+    fov_in = (Lside / DM) * au.rad.to(au.deg)  # [deg]
+    Nup = int(np.ceil(args.fov / fov_in))
+    print("Input FoV: %s [deg], Tiling repeats: %d" % (fov_in, Nup))
+    img_tiled = np.tile(img_in, reps=(Nup, Nup))
+    Nside2 = round(Nside * args.fov / fov_in)
+    img2 = img_tiled[:Nside2, :Nside2]
+    # Rescale to the output size
+    img_out = scipy.ndimage.zoom(img2, zoom=args.Nside/Nside2, order=1)
+    header.add_history(" ".join(sys.argv))
+    hdu = fits.PrimaryHDU(data=img_out, header=header)
+    outfile = args.outfile.format(prefix=args.prefix, freq=freq,
+                                  Nside=args.Nside, fov=args.fov)
+    try:
+        hdu.writeto(outfile, overwrite=args.clobber)
+    except TypeError:
+        hdu.writeto(outfile, clobber=args.clobber)
+    print("Tiled and scaled slice: %s" % outfile)
+
+
+if __name__ == "__main__":
+    main()