diff options
-rwxr-xr-x | astro/21cm/make_lightcone.py | 26 |
1 files changed, 24 insertions, 2 deletions
diff --git a/astro/21cm/make_lightcone.py b/astro/21cm/make_lightcone.py index 93b29df..1e888d0 100755 --- a/astro/21cm/make_lightcone.py +++ b/astro/21cm/make_lightcone.py @@ -73,6 +73,7 @@ import numpy as np from scipy import interpolate import astropy.io.fits as fits from astropy.cosmology import FlatLambdaCDM +from astropy.wcs import WCS logging.basicConfig(level=logging.INFO) @@ -141,6 +142,13 @@ class Configs: [self.zmin, self.zmax]).value # [Mpc] return (Dc_min, Dc_max) + @property + def Dc_cell(self): + """ + Comoving size of a cell. + """ + return self.Lside / self.Nside + def Dc_to_redshift(self, Dc): """ Calculate the redshift corresponding to the given comoving distance @@ -148,7 +156,7 @@ class Configs: """ if not hasattr(self, "_Dc_interp"): Dc_min, Dc_max = self.Dc_limit - dDc = self.Lside / self.Nside + dDc = self.Dc_cell N = int((Dc_max - Dc_min) / dDc) z_ = np.linspace(self.zmin, self.zmax, num=N) Dc_ = self.cosmo.comoving_distance(z_).value # [Mpc] @@ -228,7 +236,7 @@ class LightCone: The slices evenly distributed along the LoS representing with comoving distances. [Mpc] """ - dDc = self.configs.Lside / self.configs.Nside + dDc = self.configs.Dc_cell Dc_min, Dc_max = self.configs.Dc_limit Dc = np.arange(Dc_min, Dc_max, step=dDc) return Dc @@ -238,6 +246,19 @@ class LightCone: return len(self.slices_Dc) @property + def wcs(self): + Dc_min, __ = self.configs.Dc_limit + w = WCS(naxis=3) + w.wcs.ctype = ["pixel", "pixel", "pixel"] + w.wcs.cunit = ["cMpc", "cMpc", "cMpc"] + w.wcs.crpix = np.array([1.0, 1.0, 1.0]) + w.wcs.crval = np.array([0.0, 0.0, Dc_min]) + w.wcs.cdelt = np.array([self.configs.Dc_cell, + self.configs.Dc_cell, + self.configs.Dc_cell]) + return w + + @property def header(self): dDc = self.configs.Lside / self.configs.Nside Dc_min, Dc_max = self.configs.Dc_limit @@ -254,6 +275,7 @@ class LightCone: 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=None, clobber=None): |