diff options
Diffstat (limited to 'astro/21cm/get_slice_zfreq.py')
| -rwxr-xr-x | astro/21cm/get_slice_zfreq.py | 142 |
1 files changed, 82 insertions, 60 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__": |