build(deps): bump nox from 2024.3.2 to 2024.4.15 in /.github/workflows

.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

src/py21cmfast-tools/__init__.py

@@ -1 +1 @@
-from lc2ps import calculate_PS
+from lc2ps import calculate_PS

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,14 +258,14 @@ def generator():
  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
@@ -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]],
+ )