def _kconverter(val, allow_unitless=False): if hasattr(val, "unit"): return val.to(littleh / un.Mpc, with_H0(config.COSMO.H0)) if not allow_unitless: raise ValueError("no units supplied!") # Assume it has the 1/Mpc units (default from 21cmFAST) return (val / un.Mpc).to(littleh / un.Mpc, with_H0(config.COSMO.H0))
def test_littleh(): """Test :func:`astropy.cosmology.units.with_H0`.""" H0_70 = 70 * u.km / u.s / u.Mpc h70dist = 70 * u.Mpc / cu.littleh assert_quantity_allclose(h70dist.to(u.Mpc, cu.with_H0(H0_70)), 100 * u.Mpc) # make sure using the default cosmology works cosmodist = default_cosmology.get().H0.value * u.Mpc / cu.littleh assert_quantity_allclose(cosmodist.to(u.Mpc, cu.with_H0()), 100 * u.Mpc) # Now try a luminosity scaling h1lum = 0.49 * u.Lsun * cu.littleh ** -2 assert_quantity_allclose(h1lum.to(u.Lsun, cu.with_H0(H0_70)), 1 * u.Lsun) # And the trickiest one: magnitudes. Using H0=10 here for the round numbers H0_10 = 10 * u.km / u.s / u.Mpc # assume the "true" magnitude M = 12. # Then M - 5*log_10(h) = M + 5 = 17 withlittlehmag = 17 * (u.mag - u.MagUnit(cu.littleh ** 2)) assert_quantity_allclose(withlittlehmag.to(u.mag, cu.with_H0(H0_10)), 12 * u.mag)
def write( self, filename: str | Path, thermal: bool = True, sample: bool = True, prefix: str = None, ) -> Path: """Save sensitivity results to HDF5 file. Returns ------- filename The path to the file that is written. """ out = self._get_all_sensitivity_combos(thermal, sample) prefix = prefix + "_" if prefix else "" if filename is None: filename = Path( f"{prefix}{self.foreground_model}_{self.observation.frequency:.3f}.h5" ) else: filename = Path(filename) logger.info(f"Writing sensitivies to '{filename}'") with h5py.File(filename, "w") as fl: # TODO: We should be careful to try and write everything into this file # i.e. all the parameters etc. for k, v in out.items(): fl[k] = v fl[k.replace("noise", "snr")] = self.delta_squared / v fl["k"] = self.k1d.to("1/Mpc", with_H0(config.COSMO.H0)).value fl["delta_squared"] = self.delta_squared fl.attrs["k_min"] = self.k_min fl.attrs["k_max"] = self.k_max fl.attrs["total_snr"] = self.calculate_significance() fl.attrs["foreground_model"] = self.foreground_model fl.attrs["horizon_buffer"] = self.horizon_buffer fl.attrs["k_unit"] = "1/Mpc" return filename
class PowerSpectrum(Sensitivity): """ A Power Spectrum sensitivity calculator. Parameters ---------- horizon_buffer : float or Quantity A buffer to add to the horizon line in order to excise foreground-contaminated modes. foreground_model : str, {moderate, optimistic} Which approach to take for foreground excision. Moderate uses a defined horizon buffer, while optimistic excludes all k modes inside the primary field of view. k_21 : array or Quantity, optional An array of wavenumbers used to define a cosmological power spectrum, in order to get sample variance. If not a Quantity, will assume k has units of 1/Mpc, though it will convert these units to h/Mpc throughout the class. Default is to use built-in data file from 21cmFAST. delta_21 : array or Quantity, optional An array of Delta^2 power spectrum values used for sample variance. If not a Quantity, will assume units of mK^2. """ horizon_buffer: tp.Wavenumber = attr.ib( default=0.1 * littleh / un.Mpc, validator=( tp.vld_unit(littleh / un.Mpc, with_H0(config.COSMO.H0)), ut.nonnegative, ), converter=_kconverter, ) foreground_model: str = attr.ib( default="moderate", validator=vld.in_(["moderate", "optimistic"]) ) k_21: tp.Wavenumber = attr.ib( _K21_DEFAULT, validator=tp.vld_unit(littleh / un.Mpc, with_H0(config.COSMO.H0)), converter=_kconverter, ) delta_21: tp.Delta = attr.ib(_D21_DEFAULT, validator=(tp.vld_unit(un.mK**2))) @classmethod def from_yaml(cls, yaml_file) -> Sensitivity: """ Construct a PowerSpectrum sensitivity from yaml. YAML spec has p21 as a file which obtains k_21 and delta_21. It is assumed that k in the file is in units of 1/Mpc (true for 21cmFAST). """ data = cls._load_yaml(yaml_file) p21 = data.pop("p21", None) if p21 is not None: data["k_21"] = p21[:, 0] << 1 / un.Mpc data["delta_21"] = p21[:, 1] << un.mK**2 if isinstance(yaml_file, str): obsfile = path.join(path.dirname(yaml_file), data.pop("observation")) else: obsfile = data.pop("observation") data["observation"] = obsfile return super().from_yaml(data) @k_21.validator def _p21k_validator(self, att, val): assert val.ndim == 1 @delta_21.validator def _delta21_validator(self, att, val): assert val.ndim == 1 assert val.shape == self.k_21.shape @cached_property def p21(self): """An interpolation function defining the cosmological power spectrum.""" fnc = interpolate.interp1d( self.k_21.to_value(littleh / un.Mpc), self.delta_21.to_value(un.mK**2), kind="linear", ) return lambda k: fnc(k) * un.mK**2 @cached_property def k_min(self) -> tp.Wavenumber: """Minimum k value to use in estimates.""" return self.k_21.min() @cached_property def k_max(self) -> tp.Wavenumber: """Maximum k value to use in estimates.""" return self.k_21.max() @cached_property def k1d(self) -> tp.Wavenumber: """1D array of wavenumbers for which sensitivities will be generated.""" delta = conv.dk_deta(self.observation.redshift) / self.observation.bandwidth if self.k_max.value < delta.value * self.observation.n_channels: logger.warning( "The maximum k value is being restricted by the theoretical signal " "model. Losing ~" f"{int((delta.value * self.observation.n_channels - self.k_max.value)/delta.value)}" " bins." ) if self.k_min > delta: logger.warning( "The minimum k value is being restricted by the theoretical signal " f"model. Losing ~{int((self.k_min - delta)/delta)} bins between {delta}" f" and {self.k_min}." ) mn = self.k_min else: mn = delta assert delta.unit == self.k_max.unit return np.arange(mn.value, self.k_max.value, delta.value) * delta.unit @cached_property def X2Y(self) -> un.Quantity[un.Mpc**3 / littleh**3 / un.steradian / un.GHz]: """Cosmological scaling factor X^2*Y (eg. Parsons 2012).""" return conv.X2Y(self.observation.redshift) @cached_property def uv_coverage(self) -> np.ndarray: """The UV-coverage of the array, with unused/redundant baselines set to zero.""" grid = self.observation.uv_coverage.copy() size = grid.shape[0] # Cut unnecessary data out of uv coverage: auto-correlations & half of uv # plane (which is not statistically independent for real sky) grid[size // 2, size // 2] = 0.0 grid[:, : size // 2] = 0.0 grid[size // 2 :, size // 2] = 0.0 if self.no_ns_baselines: grid[:, size // 2] = 0.0 return grid def power_normalisation(self, k: tp.Wavenumber) -> float: """Normalisation constant for power spectrum.""" assert k.unit.is_equivalent(littleh / un.Mpc) return ( self.X2Y * self.observation.observatory.beam.b_eff * self.observation.bandwidth * k**3 / (2 * np.pi**2) ).to_value("") def thermal_noise( self, k_par: tp.Wavenumber, k_perp: tp.Wavenumber, trms: tp.Temperature ) -> tp.Delta: """Thermal noise contribution at particular k mode.""" k = np.sqrt(k_par**2 + k_perp**2) scalar = self.power_normalisation(k) return scalar * trms.to("mK") ** 2 def sample_noise(self, k_par: tp.Wavenumber, k_perp: tp.Wavenumber) -> tp.Delta: """Sample variance contribution at a particular k mode.""" k = np.sqrt(k_par**2 + k_perp**2) vals = np.full(k.size, np.inf) * un.mK**2 good_ks = np.logical_and(k >= self.k_min, k <= self.k_max) vals[good_ks] = self.p21(k[good_ks]) return vals @cached_property def _nsamples_2d( self, ) -> dict[str, dict[tp.Wavenumber, un.Quantity[1 / un.mK**4]]]: """Mid-way product specifying thermal and sample variance over the 2D grid.""" # set up blank arrays/dictionaries sense = {"sample": {}, "thermal": {}, "both": {}} # loop over uv_coverage to calculate k_pr nonzero = np.where(self.uv_coverage > 0) for iu, iv in tqdm.tqdm( zip(nonzero[1], nonzero[0]), desc="calculating 2D sensitivity", unit="uv-bins", disable=not config.PROGRESS, total=len(nonzero[1]), ): u, v = self.observation.ugrid[iu], self.observation.ugrid[iv] trms = self.observation.Trms[iv, iu] if np.isinf(trms): continue umag = np.sqrt(u**2 + v**2) k_perp = umag * conv.dk_du(self.observation.redshift) hor = self.horizon_limit(umag) if k_perp not in sense["thermal"]: sense["thermal"][k_perp] = ( np.zeros(len(self.observation.kparallel)) / un.mK**4 ) sense["sample"][k_perp] = ( np.zeros(len(self.observation.kparallel)) / un.mK**4 ) sense["both"][k_perp] = ( np.zeros(len(self.observation.kparallel)) / un.mK**4 ) # Exclude parallel modes dominated by foregrounds kpars = self.observation.kparallel[self.observation.kparallel >= hor] if not len(kpars): continue start = np.where(self.observation.kparallel >= hor)[0][0] n_inds = (self.observation.kparallel.size - 1) // 2 + 1 inds = np.arange(start=start, stop=n_inds) thermal = self.thermal_noise(kpars, k_perp, trms) sample = self.sample_noise(kpars, k_perp) t = 1.0 / thermal**2 s = 1.0 / sample**2 ts = 1.0 / (thermal + sample) ** 2 sense["thermal"][k_perp][inds] += t sense["thermal"][k_perp][-inds] += t sense["sample"][k_perp][inds] += s sense["sample"][k_perp][-inds] += s sense["both"][k_perp][inds] += ts sense["both"][k_perp][-inds] += ts return sense @lru_cache() def calculate_sensitivity_2d( self, thermal: bool = True, sample: bool = True ) -> dict[tp.Wavenumber, tp.Delta]: """ Calculate power spectrum sensitivity for a grid of cylindrical k modes. Parameters ---------- thermal : bool, optional Whether to calculate thermal contribution to the sensitivity sample : bool, optional Whether to calculate sample variance contribution to sensitivity. Returns ------- dict : Keys are cylindrical kperp values and values are arrays aligned with `observation.kparallel`, defining uncertainty in mK^2. """ if thermal and sample: logger.info("Getting Combined Variance") sense = self._nsamples_2d["both"] elif thermal: logger.info("Getting Thermal Variance") sense = self._nsamples_2d["thermal"] elif sample: logger.info("Getting Sample Variance") sense = self._nsamples_2d["sample"] else: raise ValueError("Either thermal or sample must be True") # errors were added in inverse quadrature, now need to invert and take # square root to have error bars; also divide errors by number of indep. fields final_sense = {} for k_perp in sense.keys(): mask = sense[k_perp] > 0 final_sense[k_perp] = np.inf * np.ones(len(mask)) * un.mK**2 final_sense[k_perp][mask] = sense[k_perp][mask] ** -0.5 / np.sqrt( self.observation.n_lst_bins ) return final_sense def calculate_sensitivity_2d_grid( self, kperp_edges: tp.Wavenumber, kpar_edges: tp.Wavenumber, thermal: bool = True, sample: bool = True, ) -> tp.Delta: """Calculate the 2D cylindrical sensitivity on a grid of kperp/kpar. Parameters ---------- kperp_edges The edges of the bins in kperp. kpar_edges The edges of the bins in kpar. """ sense2d_inv = np.zeros((len(kperp_edges) - 1, len(kpar_edges) - 1)) << ( 1 / un.mK**4 ) sense = self.calculate_sensitivity_2d(thermal=thermal, sample=sample) assert np.all(np.diff(kperp_edges) > 0) assert np.all(np.diff(kpar_edges) > 0) if np.sqrt(kpar_edges.min() ** 2 + kperp_edges.min() ** 2) < self.k_min: logger.warning( "The minimum kbin is being restricted by the theoretical model. Some values will be zero." ) if np.sqrt(kpar_edges.max() ** 2 + kperp_edges.max() ** 2) > self.k_max: logger.warning( "The maximum kbin is being restricted by the theoretical model. Some values will be zero." ) for k_perp in tqdm.tqdm( sense.keys(), desc="averaging to 2D grid", unit="kperp-bins", disable=not config.PROGRESS, ): if k_perp < kperp_edges[0] or k_perp >= kperp_edges[-1]: continue # Get the kperp bin it's in. kperp_indx = np.where(k_perp >= kperp_edges)[0][-1] k = np.sqrt(self.observation.kparallel**2 + k_perp**2) good_ks = np.logical_and(self.k_min <= k, k <= self.k_max) kpar_indx = np.digitize(k, kpar_edges) - 1 good_ks &= kpar_indx >= 0 good_ks &= kpar_indx < len(kpar_edges) - 1 sense2d_inv[kperp_indx][kpar_indx[good_ks]] += ( 1.0 / sense[k_perp][good_ks] ** 2 ) # invert errors and take square root again for final answer sense2d = np.ones(sense2d_inv.shape) * un.mK**2 * np.inf mask = sense2d_inv > 0 sense2d[mask] = 1 / np.sqrt(sense2d_inv[mask]) return sense2d def horizon_limit(self, umag: float) -> tp.Wavenumber: """ Calculate a horizon limit, with included buffer, if appropriate. Parameters ---------- umag : float Baseline length (in wavelengths) at which to compute the horizon limit. Returns ------- float : Horizon limit, in h/Mpc. """ horizon = ( conv.dk_deta(self.observation.redshift) * umag / self.observation.frequency ) # calculate horizon limit for baseline of length umag if self.foreground_model in ["moderate", "pessimistic"]: return horizon + self.horizon_buffer elif self.foreground_model in ["optimistic"]: return horizon * np.sin(self.observation.observatory.beam.first_null / 2) def _average_sense_to_1d( self, sense: dict[tp.Wavenumber, tp.Delta], k1d: tp.Wavenumber | None = None ) -> tp.Delta: """Bin 2D sensitivity down to 1D.""" sense1d_inv = np.zeros(len(self.k1d)) / un.mK**4 if k1d is None: k1d = self.k1d for k_perp in tqdm.tqdm( sense.keys(), desc="averaging to 1D", unit="kperp-bins", disable=not config.PROGRESS, ): k = np.sqrt(self.observation.kparallel**2 + k_perp**2) good_ks = np.logical_and(self.k_min <= k, k <= self.k_max) good_ks &= k >= k1d.min() good_ks &= k < k1d.max() sense1d_inv[ut.find_nearest(k1d, k[good_ks])] += ( 1.0 / sense[k_perp][good_ks] ** 2 ) # invert errors and take square root again for final answer sense1d = np.ones(sense1d_inv.shape) * un.mK**2 * np.inf mask = sense1d_inv > 0 sense1d[mask] = 1 / np.sqrt(sense1d_inv[mask]) return sense1d @lru_cache() def calculate_sensitivity_1d( self, thermal: bool = True, sample: bool = True ) -> tp.Delta: """Calculate a 1D sensitivity curve. Parameters ---------- thermal Whether to calculate thermal contribution to the sensitivity sample Whether to calculate sample variance contribution to sensitivity. Returns ------- array : 1D array with units mK^2... the variance of spherical k modes. """ sense = self.calculate_sensitivity_2d(thermal=thermal, sample=sample) return self._average_sense_to_1d(sense) def calculate_sensitivity_1d_binned(self, k: tp.Wavenumber, **kwargs): """Calculate the 1D sensitivity at arbitrary k-bins.""" sense2d = self.calculate_sensitivity_2d(**kwargs) return self._average_sense_to_1d(sense2d, k1d=k) @property def delta_squared(self) -> tp.Delta: """The fiducial 21cm power spectrum evaluated at :attr:`k1d`.""" return self.p21(self.k1d) @lru_cache() def calculate_significance( self, thermal: bool = True, sample: bool = True ) -> float: """ Calculate significance of a detection of the default cosmological power spectrum. Parameters ---------- thermal : bool, optional Whether to calculate thermal contribution to the sensitivity sample : bool, optional Whether to calculate sample variance contribution to sensitivity. Returns ------- float : Significance of detection (in units of sigma). """ mask = np.logical_and(self.k1d >= self.k_min, self.k1d <= self.k_max) sense1d = self.calculate_sensitivity_1d(thermal=thermal, sample=sample) snr = self.delta_squared[mask] / sense1d[mask] return np.sqrt(float(np.dot(snr, snr.T))) def plot_sense_2d(self, sense2d: dict[tp.Wavenumber, tp.Delta]): """Create a colormap plot of the sensitivity un UV bins.""" try: import matplotlib.pyplot as plt except ImportError: # pragma: no cover raise ImportError("matplotlib is required to make plots...") keys = sorted(sense2d.keys()) x = np.array([v.value for v in keys]) x = ( np.repeat(x, len(self.observation.kparallel)) .reshape((len(x), len(self.observation.kparallel))) .T ) y = np.fft.fftshift( np.repeat(self.observation.kparallel.value, x.shape[1]).reshape( (len(self.observation.kparallel), x.shape[1]) ) ) z = np.array([np.fft.fftshift(sense2d[key]) for key in keys]).T plt.pcolormesh(x, y, np.log10(z)) cbar = plt.colorbar() cbar.set_label(r"$\log_{10} \delta \Delta^2$ [mK^2]", fontsize=14) plt.xlabel(r"$k_\perp$ [h/Mpc]", fontsize=14) plt.ylabel(r"$k_{||}$ [h/Mpc]", fontsize=14) def write( self, filename: str | Path, thermal: bool = True, sample: bool = True, prefix: str = None, ) -> Path: """Save sensitivity results to HDF5 file. Returns ------- filename The path to the file that is written. """ out = self._get_all_sensitivity_combos(thermal, sample) prefix = prefix + "_" if prefix else "" if filename is None: filename = Path( f"{prefix}{self.foreground_model}_{self.observation.frequency:.3f}.h5" ) else: filename = Path(filename) logger.info(f"Writing sensitivies to '{filename}'") with h5py.File(filename, "w") as fl: # TODO: We should be careful to try and write everything into this file # i.e. all the parameters etc. for k, v in out.items(): fl[k] = v fl[k.replace("noise", "snr")] = self.delta_squared / v fl["k"] = self.k1d.to("1/Mpc", with_H0(config.COSMO.H0)).value fl["delta_squared"] = self.delta_squared fl.attrs["k_min"] = self.k_min fl.attrs["k_max"] = self.k_max fl.attrs["total_snr"] = self.calculate_significance() fl.attrs["foreground_model"] = self.foreground_model fl.attrs["horizon_buffer"] = self.horizon_buffer fl.attrs["k_unit"] = "1/Mpc" return filename def plot_sense_1d(self, sample: bool = True, thermal: bool = True): """Create a plot of the sensitivity in 1D k-bins.""" try: import matplotlib.pyplot as plt except ImportError: # pragma: no cover raise ImportError("matplotlib is required to make plots...") out = self._get_all_sensitivity_combos(thermal, sample) for key, value in out.items(): plt.plot(self.k1d, value, label=key) plt.xscale("log") plt.yscale("log") plt.xlabel("k [1/Mpc]") plt.ylabel(r"$\Delta^2_N \ [{\rm mK}^2/{\rm Mpc}^3$") plt.legend() plt.title(f"z={conv.f2z(self.observation.frequency):.2f}") return plt.gcf() def _get_all_sensitivity_combos( self, thermal: bool, sample: bool ) -> dict[str, tp.Delta]: result = {} if thermal: result["thermal_noise"] = self.calculate_sensitivity_1d(sample=False) if sample: result["sample_noise"] = self.calculate_sensitivity_1d( thermal=False, sample=True ) if thermal and sample: result["sample+thermal_noise"] = self.calculate_sensitivity_1d( thermal=True, sample=True ) return result
h70dist = 70 * u.Mpc / cu.littleh assert_quantity_allclose(h70dist.to(u.Mpc, cu.with_H0(H0_70)), 100 * u.Mpc) # make sure using the default cosmology works cosmodist = default_cosmology.get().H0.value * u.Mpc / cu.littleh assert_quantity_allclose(cosmodist.to(u.Mpc, cu.with_H0()), 100 * u.Mpc) # Now try a luminosity scaling h1lum = 0.49 * u.Lsun * cu.littleh**-2 assert_quantity_allclose(h1lum.to(u.Lsun, cu.with_H0(H0_70)), 1 * u.Lsun) # And the trickiest one: magnitudes. Using H0=10 here for the round numbers H0_10 = 10 * u.km / u.s / u.Mpc # assume the "true" magnitude M = 12. # Then M - 5*log_10(h) = M + 5 = 17 withlittlehmag = 17 * (u.mag - u.MagUnit(cu.littleh**2)) assert_quantity_allclose(withlittlehmag.to(u.mag, cu.with_H0(H0_10)), 12 * u.mag) @pytest.mark.skipif(not HAS_ASDF, reason="requires ASDF") @pytest.mark.parametrize('equiv', [cu.with_H0()]) def test_equivalencies(tmpdir, equiv): from asdf.tests import helpers tree = {'equiv': equiv} with pytest.warns(AstropyDeprecationWarning, match="`with_H0`"): helpers.assert_roundtrip_tree(tree, tmpdir)