From 2dfe756f0cb655cf635c51f6f503f1fe4752a1d7 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 15 Jul 2024 15:48:56 +0000 Subject: [PATCH 1/2] build(deps): bump nox from 2024.3.2 to 2024.4.15 in /.github/workflows Bumps [nox](https://github.com/wntrblm/nox) from 2024.3.2 to 2024.4.15. - [Release notes](https://github.com/wntrblm/nox/releases) - [Changelog](https://github.com/wntrblm/nox/blob/main/CHANGELOG.md) - [Commits](https://github.com/wntrblm/nox/compare/2024.03.02...2024.04.15) --- updated-dependencies: - dependency-name: nox dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] --- .github/workflows/constraints.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/constraints.txt b/.github/workflows/constraints.txt index 3d46861..8e30a2e 100644 --- a/.github/workflows/constraints.txt +++ b/.github/workflows/constraints.txt @@ -1,4 +1,4 @@ -nox==2024.3.2 +nox==2024.4.15 nox-poetry==1.0.3 pip==24.0 poetry==1.8.2 From 2baab4e19cc5f70447f8403a79b0066a9b6d949f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 15 Jul 2024 15:49:05 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/py21cmfast-tools/__init__.py | 2 +- src/py21cmfast-tools/lc2ps.py | 251 ++++++++++++++++++------------- tests/__init__.py | 2 +- tests/test_ps.py | 63 +++++--- 4 files changed, 192 insertions(+), 126 deletions(-) diff --git a/src/py21cmfast-tools/__init__.py b/src/py21cmfast-tools/__init__.py index 357d5a7..403f7d4 100644 --- a/src/py21cmfast-tools/__init__.py +++ b/src/py21cmfast-tools/__init__.py @@ -1 +1 @@ -from lc2ps import calculate_PS \ No newline at end of file +from lc2ps import calculate_PS diff --git a/src/py21cmfast-tools/lc2ps.py b/src/py21cmfast-tools/lc2ps.py index 787a37a..cd5ffdb 100644 --- a/src/py21cmfast-tools/lc2ps.py +++ b/src/py21cmfast-tools/lc2ps.py @@ -1,37 +1,43 @@ 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 +from powerbox.tools import ( + _magnitude_grid, + get_power, + ignore_zero_ki, + power2delta, + 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 +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 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. @@ -40,7 +46,7 @@ def calculate_PS(lc, 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. + 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. @@ -72,7 +78,7 @@ def calculate_PS(lc, 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. + 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]. @@ -82,7 +88,7 @@ def calculate_PS(lc, 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]. + 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. @@ -95,9 +101,19 @@ def calculate_PS(lc, 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 + 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: @@ -108,7 +124,7 @@ def calculate_PS(lc, if HII_DIM is None: HII_DIM = lc.shape[0] - + for i in chunk_indices: start = i end = i + chunk_size @@ -121,47 +137,56 @@ def calculate_PS(lc, 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 - ) + 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_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) + 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 - + + 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): + + 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. @@ -198,21 +223,22 @@ def generator(): 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 - - ) + 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: @@ -232,7 +258,6 @@ def generator(): if calc_global: out["global_Tb"] = np.array(Tb) out["redshifts"] = np.array(zs) - return out @@ -240,6 +265,7 @@ def generator(): 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 @@ -251,39 +277,40 @@ def log_bin(ps, kperp, kpar, bins=None): 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) + 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) + 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 = 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) + # 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, - ): +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 @@ -309,39 +336,51 @@ def postprocess_ps(ps, kperp, kpar, 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.) + kperp = np.exp((np.log(kperp[1:]) + np.log(kperp[:-1])) / 2.0) else: - kperp = (kperp[1:] + kperp[:-1])/2 + kperp = (kperp[1:] + kperp[:-1]) / 2 kpar = kpar[m] - ps = ps[:,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,:] - + ps = ps[mkperp, :] + # Bin kpar in log - rebinned_ps, kperp, log_kpar, kpar_weights = log_bin(ps, kperp, kpar, bins=kpar_bins) + 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] + 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 + 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 + 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) + 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 + 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]] - + 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 index e4c2f16..0da2541 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1 +1 @@ -"""Tests for the py21cmFAST-tools package.""" \ No newline at end of file +"""Tests for the py21cmFAST-tools package.""" diff --git a/tests/test_ps.py b/tests/test_ps.py index 92ee022..3b45d69 100644 --- a/tests/test_ps.py +++ b/tests/test_ps.py @@ -1,27 +1,54 @@ """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_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.] + zs = [5.0, 6.0, 10.0, 27.0] + + 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) - - 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 + 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, + )