From fa8208c993f5d74e67f506ed963b33bf4342401d Mon Sep 17 00:00:00 2001
From: Aaron LI <aly@aaronly.me>
Date: Sun, 11 Jun 2017 20:07:11 +0800
Subject: Add astro/21cm/get_slice_zfreq.py
---
astro/21cm/get_slice_zfreq.py | 112 ++++++++++++++++++++++++++++++++++++++++++
1 file changed, 112 insertions(+)
create mode 100755 astro/21cm/get_slice_zfreq.py
(limited to 'astro')
diff --git a/astro/21cm/get_slice_zfreq.py b/astro/21cm/get_slice_zfreq.py
new file mode 100755
index 0000000..c67d513
--- /dev/null
+++ b/astro/21cm/get_slice_zfreq.py
@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2017 Weitian LI <weitian@aaronly.me>
+# MIT License
+#
+
+"""
+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.
+"""
+
+import os
+import sys
+import argparse
+
+import numpy as np
+from astropy.io import fits
+from astropy.wcs import WCS
+
+from z2freq import z2freq, freq2z
+
+
+class FITSCube:
+ 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)
+
+ @property
+ 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 get_slice(self, z):
+ redshifts = self.redshifts
+ if z < redshifts.min() or z > redshifts.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)
+
+
+def main():
+ outfile_default = "{prefix}_z{z:05.2f}_f{freq:06.2f}.fits"
+
+ parser = argparse.ArgumentParser(
+ description="Get slices at requested redshifts/frequencies")
+ 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")
+ 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")
+ exgrp = parser.add_mutually_exclusive_group(required=True)
+ exgrp.add_argument("-z", "--redshifts", dest="redshifts", nargs="+",
+ help="redshifts where to interpolate slices")
+ exgrp.add_argument("-f", "--freqs", dest="freqs", nargs="+",
+ help="21cm frequencies [MHz] to interpolate slices")
+ args = parser.parse_args()
+
+ if args.redshifts:
+ redshifts = [float(z) for z in args.redshifts]
+ freqs = z2freq(redshifts, print_=False)
+ else:
+ freqs = [float(f) for f in args.freqs]
+ redshifts = freq2z(freqs, print_=False)
+
+ cube = FITSCube(args.infile)
+ for z, f in zip(redshifts, freqs):
+ outfile = args.outfile.format(prefix=args.prefix, z=z, freq=f)
+ print("z=%05.2f, freq=%06.2f MHz : %s" % (z, f, outfile))
+ zslice = cube.get_slice(z)
+ header = fits.Header()
+ header["BUNIT"] = cube.header.get("BUNIT")
+ 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)
+
+
+if __name__ == "__main__":
+ main()
--
cgit v1.2.2