From 006b0f6c4cc52c2d72933b82456422af1edbf402 Mon Sep 17 00:00:00 2001 From: DanielaBreitman Date: Mon, 15 Jul 2024 17:27:50 +0200 Subject: [PATCH] Finally some code --- src/py21cmfast-tools/__init__.py | 1 + src/py21cmfast-tools/lc2ps.py | 347 +++++++++++++++++++++++++++++++ tests/__init__.py | 1 + tests/test_ps.py | 27 +++ 4 files changed, 376 insertions(+) create mode 100644 src/py21cmfast-tools/__init__.py create mode 100644 src/py21cmfast-tools/lc2ps.py create mode 100644 tests/__init__.py create mode 100644 tests/test_ps.py diff --git a/src/py21cmfast-tools/__init__.py b/src/py21cmfast-tools/__init__.py new file mode 100644 index 0000000..357d5a7 --- /dev/null +++ b/src/py21cmfast-tools/__init__.py @@ -0,0 +1 @@ +from lc2ps import calculate_PS \ No newline at end of file diff --git a/src/py21cmfast-tools/lc2ps.py b/src/py21cmfast-tools/lc2ps.py new file mode 100644 index 0000000..787a37a --- /dev/null +++ b/src/py21cmfast-tools/lc2ps.py @@ -0,0 +1,347 @@ +import numpy as np +import powerbox as pb +from powerbox.tools import get_power, ignore_zero_ki, ignore_zero_absk, power2delta, _magnitude_grid, regular_angular_generator + + +def calculate_PS(lc, + lc_redshifts, + L, + HII_DIM = None, + zs = None, + chunk_size = None, + chunk_skip = 37, + calc_2d = True, + nbins = 50, + k_weights = ignore_zero_ki, + postprocess = True, + kpar_bins = None, + log_bins = True, + crop = None, + calc_1d = False, + nbins_1d = 14, + calc_global = False, + mu = None, + bin_ave=True, + interp=None, + prefactor_fnc=power2delta, + interp_points_generator=None, + ): + + """ Calculate 1D and 2D PS + Parameters + ---------- + lc : np.ndarray + The lightcone whose power spectrum we want to calculate. + The lightcone should be a 3D array with shape [HII_DIM, HII_DIM, len(lc_redshifts)]. + lc_redshifts : np.ndarray + The redshifts of the lightcone. + L : float + The side length of the box in cMpc. + HII_DIM : int, optional + The size of the lightcone in pixels. If None, inferred from the shape of the lightcone. + zs : np.ndarray, optional + The redshifts at which to calculate the power spectrum. + If None, the lightcone is broken up into chunks using arguments chunk_skip and chunk_size. + chunk_size : int, optional + The size of the chunks to break the lightcone into. + If None, the chunk is assumed to be a cube i.e. chunk_size = HII_DIM. + chunk_skip : int, optional + The number of lightcone slices to skip between chunks. Default is 37. + calc_2d : bool, optional + If True, calculate the 2D power spectrum. + nbins : int, optional + The number of bins to use for the kperp axis of the 2D PS. + k_weights : callable, optional + A function that takes a frequency tuple and returns a boolean mask for the k values to ignore. + See powerbox.tools.ignore_zero_ki for an example and powerbox.tools.get_power documentation for more details. + Default is powerbox.tools.ignore_zero_ki, which excludes the power any k_i = 0 mode. + Typically, only the central zero mode |k| = 0 is excluded, in which case use powerbox.tools.ignore_zero_absk. + postprocess : bool, optional + If True, postprocess the 2D PS. This step involves cropping out empty bins and/or log binning the kpar axis. + kpar_bins : int or np.ndarray, optional + Affects only the postprocessing step. The number of bins or the bin edges to use for binning the kpar axis. + If None, produces 16 bins. + log_bins : bool, optional + Affects only the postprocessing step. If True, log bin the kpar axis. + crop : list, optional + Affects only the postprocessing step. The crop range for the (log-binned) PS. If None, crops out only the empty bins. + calc_1d : bool, optional + If True, calculate the 1D power spectrum. + nbins_1d : int, optional + The number of bins on which to calculate 1D PS. + calc_global : bool, optional + If True, calculate the global brightness temperature. + mu : float, optional + The minimum value of :math:`\cos(\theta), \theta = \arctan (k_\perp/k_\parallel)` for all calculated PS. + If None, all modes are included. + bin_ave : bool, optional + If True, return the center value of each kperp and kpar bin i.e. len(kperp) = PS_2D.shape[0]. + If False, return the left edge of each bin i.e. len(kperp) = PS_2D.shape[0] + 1. + interp : str, optional + If True, use linear interpolation to calculate the PS at the points specified by interp_points_generator. + Note that this significantly slows down the calculation. + prefactor_fnc : callable, optional + A function that takes a frequency tuple and returns the prefactor to multiply the PS with. + Default is powerbox.tools.power2delta, which converts the power P [mK^2 Mpc^{-3}] to the dimensionless power \delta^2 [mK^2]. + interp_points_generator : callable, optional + A function that generates the points at which to interpolate the PS. + See powerbox.tools.get_power documentation for more details. + """ + # Split the lightcone into chunks for each redshift bin + if zs is None: + if chunk_size is None: + chunk_size = HII_DIM + n_slices = lc.shape[-1] + chunk_indices = list(range(0, n_slices - chunk_size, chunk_skip)) + else: + if chunk_size is None: + chunk_size=HII_DIM + chunk_indices = np.array(np.max([np.zeros_like(zs),np.array([np.argmin(abs(lc_redshifts - z)) for z in zs]) - chunk_size//2], axis = 0), dtype = np.int32) + zs = [] # all redshifts that will be computed + lc_PS_2D = [] + clean_lc_PS_2D = [] + if calc_global: + Tb = [] + if calc_1d: + lc_PS_1D = [] + out = {} + + if HII_DIM is None: + HII_DIM = lc.shape[0] + + for i in chunk_indices: + start = i + end = i + chunk_size + if end > len(lc_redshifts): + shift_it_back_by_a_few_bins = end - len(lc_redshifts) + start -= shift_it_back_by_a_few_bins + end = len(lc_redshifts) + chunk = lc[..., start:end] + zs.append(lc_redshifts[(start + end) // 2]) + if calc_global: + Tb.append(np.mean(chunk)) + if calc_2d: + PS_2D, kperp, Nmodes,kpar = get_power(chunk, + (L,L,L*chunk.shape[-1]/HII_DIM), + res_ndim = 2, + bin_ave = bin_ave, + bins = nbins, + log_bins = log_bins, + nthreads=1, + k_weights=k_weights, + prefactor_fnc=prefactor_fnc, + interpolation_method=interp, + return_sumweights=True + ) + if postprocess: + clean_PS_2D, clean_kperp, clean_kpar, clean_Nmodes = postprocess_ps(PS_2D, kperp, kpar, + log_bins = log_bins, + kpar_bins = kpar_bins, + crop = crop.copy() if crop is not None else crop, + kperp_modes = Nmodes, + return_modes = True, + ) + clean_lc_PS_2D.append(clean_PS_2D) + + lc_PS_2D.append(PS_2D) + + if calc_1d: + if mu is not None: + if interp is None: + def mask_fnc(freq, absk): + kz_mesh = np.zeros((len(freq[0]), len(freq[1]), len(freq[2]))) + kz = freq[2] + for i in range(len(kz)): + kz_mesh[:,:,i] = kz[i] + phi = np.arccos(kz_mesh/absk) + mu_mesh = abs(np.cos(phi)) + kmag = _magnitude_grid([c for i, c in enumerate(freq) if i < 2]) + return np.logical_and(mu_mesh > mu, ignore_zero_ki(freq, kmag)) + k_weights1d=mask_fnc + + if interp is not None: + k_weights1d = ignore_zero_ki + def above_mu_min_angular_generator(bins, dims2avg, angular_resolution=0.1, mu=mu): + r""" + Returns a set of spherical coordinates above a certain :math:`\\mu` value. + + Parameters + ---------- + bins : array-like + 1D array of radii at which we want to spherically average the field. + dims2avg : int + The number of dimensions to average over. + angular_resolution : float, optional + The angular resolution in radians for the sample points for the interpolation. + mu : float, optional + The minimum value of :math:`\\cos(\theta), \theta = \arctan (k_\\perp/k_\\parallel)` + for the sample points generated for the interpolation. + + Returns + ------- + r_n : array-like + 1D array of radii. + phi_n : array-like + 2D array of azimuthal angles with shape (ndim-1, N), where N is the number of points. + phi_n[0,:] :math:`\\in [0,2*\\pi]`, and phi_n[1:,:] :math:`\\in [0,\\pi]`. + """ + + def generator(): + r_n, phi_n = regular_angular_generator( + bins, dims2avg, angular_resolution=angular_resolution + ) + # sine because the phi_n are wrt x-axis and we need them wrt z-axis. + if len(phi_n) == 1: + mask = np.sin(phi_n[0, :]) >= mu + else: + mask = np.all(np.sin(phi_n[1:, :]) >= mu, axis=0) + return r_n[mask], phi_n[:, mask] + + return generator() + interp_points_generator = above_mu_min_angular_generator + else: + k_weights1d = ignore_zero_ki + PS_1D, k, Nmodes_1D = get_power(chunk, + (L,L,L*chunk.shape[-1]/HII_DIM), + bin_ave = bin_ave, + bins = nbins_1d, + log_bins = log_bins, + k_weights=k_weights1d, + prefactor_fnc=power2delta, + interpolation_method=interp, + interp_points_generator=interp_points_generator, + return_sumweights=True + + ) + lc_PS_1D.append(PS_1D) + + if calc_1d: + out["k"] = k + out["PS_1D"] = np.array(lc_PS_1D) + out["Nmodes_1D"] = Nmodes_1D + out["mu"] = mu + if calc_2d: + out["full_kperp"] = kperp + out["full_kpar"] = kpar + out["full_PS_2D"] = np.array(lc_PS_2D) + out["final_PS_2D"] = np.array(clean_lc_PS_2D) + out["final_kpar"] = clean_kpar + out["final_kperp"] = clean_kperp + out["full_Nmodes"] = Nmodes + out["final_Nmodes"] = clean_Nmodes + if calc_global: + out["global_Tb"] = np.array(Tb) + out["redshifts"] = np.array(zs) + + + return out + + +def log_bin(ps, kperp, kpar, bins=None): + """ + Log bin a 2D PS along the kpar axis and crop out empty bins in both axes. + Parameters + ---------- + ps : np.ndarray + The 2D power spectrum of shape [len(kperp), len(kpar)]. + kperp : np.ndarray + Values of kperp. + kpar : np.ndarray + Values of kpar. + bins : np.ndarray or int, optional + The number of bins or the bin edges to use for binning the kpar axis. + If None, produces 16 bins logarithmically spaced between the minimum and maximum `kpar` supplied. + + """ + if bins is None: + bins = np.logspace(np.log10(kpar[0]), np.log10(kpar[-1]),17) + elif isinstance(bins, int): + bins = np.logspace(np.log10(kpar[0]), np.log10(kpar[-1]),bins+1) + elif isinstance(bins, np.ndarray) or isinstance(bins, list): + bins = np.array(bins) + else: + raise ValueError("Bins should be np.ndarray or int") + modes = np.zeros(len(bins)-1) + new_ps = np.zeros((len(kperp), len(bins)-1)) + for i in range(len(bins)-1): + m = np.logical_and(kpar > bins[i], kpar < bins[i+1]) + new_ps[:,i] = np.nanmean(ps[:,m], axis = 1) + modes[i] = np.sum(m) + #print('kpar modes', modes) + bin_centers = np.exp((np.log(bins[1:]) + np.log(bins[:-1])) / 2) + return new_ps, kperp, bin_centers, modes + + + + +def postprocess_ps(ps, kperp, kpar, + kpar_bins = None, + log_bins=True, + crop=None, + kperp_modes=None, + return_modes=False, + ): + """ + Postprocess a cylindrical PS by cropping out empty bins and log binning the kpar axis. + + Parameters + ---------- + ps : np.ndarray + The 2D power spectrum of shape [len(kperp), len(kpar)]. + kperp : np.ndarray + Values of kperp. + kpar : np.ndarray + Values of kpar. + kpar_bins : np.ndarray or int, optional + The number of bins or the bin edges to use for binning the kpar axis. + If None, produces 16 bins logarithmically spaced between the minimum and maximum `kpar` supplied. + log_bins : bool, optional + If True, log bin the kpar axis. + crop : list, optional + The crop range for the log-binned PS. If None, crops out all empty bins. + kperp_modes : np.ndarray, optional + The number of modes in each kperp bin. + return_modes : bool, optional + If True, return a grid with the number of modes in each bin. + Requires kperp_modes to be supplied. + """ + kpar = kpar[0] + m = kpar > 1e-10 + if ps.shape[0] < len(kperp): + if log_bins: + kperp = np.exp((np.log(kperp[1:]) + np.log(kperp[:-1]))/2.) + else: + kperp = (kperp[1:] + kperp[:-1])/2 + kpar = kpar[m] + ps = ps[:,m] + mkperp = ~np.isnan(kperp) + if kperp_modes is not None: + kperp_modes = kperp_modes[mkperp] + kperp = kperp[mkperp] + ps = ps[mkperp,:] + + # Bin kpar in log + rebinned_ps, kperp, log_kpar, kpar_weights = log_bin(ps, kperp, kpar, bins=kpar_bins) + if crop is None: + crop = [0, rebinned_ps.shape[0]+1, 0, rebinned_ps.shape[1]+1] + # Find last bin that is NaN and cut out all bins before + try: + lastnan_perp = np.where(np.isnan(np.nanmean(rebinned_ps, axis = 1)))[0][-1]+1 + crop[0] = crop[0] + lastnan_perp + except: + pass + try: + lastnan_par = np.where(np.isnan(np.nanmean(rebinned_ps, axis = 0)))[0][-1]+1 + crop[2] = crop[2] + lastnan_par + except: + pass + if kperp_modes is not None: + final_kperp_modes = kperp_modes[crop[0]:crop[1]] + kpar_grid, kperp_grid = np.meshgrid(kpar_weights[crop[2]:crop[3]], final_kperp_modes) + + Ngrid = np.sqrt(kperp_grid**2 + kpar_grid**2) + if return_modes: + return rebinned_ps[crop[0]:crop[1]][:,crop[2]:crop[3]], kperp[crop[0]:crop[1]], log_kpar[crop[2]:crop[3]], Ngrid + else: + return rebinned_ps[crop[0]:crop[1]][:,crop[2]:crop[3]], kperp[crop[0]:crop[1]], log_kpar[crop[2]:crop[3]] + diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e4c2f16 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Tests for the py21cmFAST-tools package.""" \ No newline at end of file diff --git a/tests/test_ps.py b/tests/test_ps.py new file mode 100644 index 0000000..92ee022 --- /dev/null +++ b/tests/test_ps.py @@ -0,0 +1,27 @@ +"""Test cases for the __main__ module.""" + +import shutil + +import numpy as np +import pytest +from typeguard import suppress_type_checks + +from py21cmfast_tools import calculate_PS + + +def test_calculate_PS(): + test_lc = np.random.random(100*100*1000).reshape((100, 100, 1000)) + test_redshifts = np.logspace(np.log10(5), np.log10(30), 1000) + zs = [5.,6.,10.,27.] + + out = calculate_PS(test_lc, test_redshifts, HII_DIM=100, L=200, zs=zs, + calc_1D = True, calc_global=True) + + out = calculate_PS(test_lc, test_redshifts, HII_DIM=100, L=200, zs=zs, + calc_1D = True, calc_global=True, interp=True) + + out = calculate_PS(test_lc, test_redshifts, HII_DIM=100, L=200, zs=zs, + calc_1D = True, calc_global=True, mu = 0.5) + + out = calculate_PS(test_lc, test_redshifts, HII_DIM=100, L=200, zs=zs, + calc_1D = True, calc_global=True, interp=True, mu = 0.5) \ No newline at end of file