From 51c59cbceab7b1ad71fc94e3073c34ae6c5b7fdb Mon Sep 17 00:00:00 2001 From: DanielaBreitman Date: Mon, 15 Jul 2024 17:42:05 +0200 Subject: [PATCH] main:setting up pytest and github workflows --- .github/workflows/dependabot.yml | 21 ++ .github/workflows/labels.yml | 66 +++++ .github/workflows/release-drafter.yml | 29 +++ src/py21cmfast-tools/__init__.py | 1 + src/py21cmfast-tools/lc2ps.py | 347 ++++++++++++++++++++++++++ tests/__init__.py | 1 + tests/test_ps.py | 27 ++ 7 files changed, 492 insertions(+) create mode 100644 .github/workflows/dependabot.yml create mode 100644 .github/workflows/labels.yml create mode 100644 .github/workflows/release-drafter.yml 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/.github/workflows/dependabot.yml b/.github/workflows/dependabot.yml new file mode 100644 index 0000000..592b5b2 --- /dev/null +++ b/.github/workflows/dependabot.yml @@ -0,0 +1,21 @@ +version: 2 +updates: + - package-ecosystem: github-actions + directory: "/" + schedule: + interval: monthly + - package-ecosystem: pip + directory: "/.github/workflows" + schedule: + interval: monthly + - package-ecosystem: pip + directory: "/docs" + schedule: + interval: monthly + - package-ecosystem: pip + directory: "/" + schedule: + interval: monthly + versioning-strategy: lockfile-only + allow: + - dependency-type: "all" diff --git a/.github/workflows/labels.yml b/.github/workflows/labels.yml new file mode 100644 index 0000000..f7f83aa --- /dev/null +++ b/.github/workflows/labels.yml @@ -0,0 +1,66 @@ +--- +# Labels names are important as they are used by Release Drafter to decide +# regarding where to record them in changelog or if to skip them. +# +# The repository labels will be automatically configured using this file and +# the GitHub Action https://github.com/marketplace/actions/github-labeler. +- name: breaking + description: Breaking Changes + color: bfd4f2 +- name: bug + description: Something isn't working + color: d73a4a +- name: build + description: Build System and Dependencies + color: bfdadc +- name: ci + description: Continuous Integration + color: 4a97d6 +- name: dependencies + description: Pull requests that update a dependency file + color: 0366d6 +- name: documentation + description: Improvements or additions to documentation + color: 0075ca +- name: duplicate + description: This issue or pull request already exists + color: cfd3d7 +- name: enhancement + description: New feature or request + color: a2eeef +- name: github_actions + description: Pull requests that update Github_actions code + color: "000000" +- name: good first issue + description: Good for newcomers + color: 7057ff +- name: help wanted + description: Extra attention is needed + color: 008672 +- name: invalid + description: This doesn't seem right + color: e4e669 +- name: performance + description: Performance + color: "016175" +- name: python + description: Pull requests that update Python code + color: 2b67c6 +- name: question + description: Further information is requested + color: d876e3 +- name: refactoring + description: Refactoring + color: ef67c4 +- name: removal + description: Removals and Deprecations + color: 9ae7ea +- name: style + description: Style + color: c120e5 +- name: testing + description: Testing + color: b1fc6f +- name: wontfix + description: This will not be worked on + color: ffffff diff --git a/.github/workflows/release-drafter.yml b/.github/workflows/release-drafter.yml new file mode 100644 index 0000000..7a04410 --- /dev/null +++ b/.github/workflows/release-drafter.yml @@ -0,0 +1,29 @@ +categories: + - title: ":boom: Breaking Changes" + label: "breaking" + - title: ":rocket: Features" + label: "enhancement" + - title: ":fire: Removals and Deprecations" + label: "removal" + - title: ":beetle: Fixes" + label: "bug" + - title: ":racehorse: Performance" + label: "performance" + - title: ":rotating_light: Testing" + label: "testing" + - title: ":construction_worker: Continuous Integration" + label: "ci" + - title: ":books: Documentation" + label: "documentation" + - title: ":hammer: Refactoring" + label: "refactoring" + - title: ":lipstick: Style" + label: "style" + - title: ":package: Dependencies" + labels: + - "dependencies" + - "build" +template: | + ## Changes + + $CHANGES 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