avoidcorr 0.1.0__tar.gz

avoidcorr-0.1.0/PKG-INFO

@@ -0,0 +1,44 @@
+Metadata-Version: 2.4
+Name: avoidcorr
+Version: 0.1.0
+Summary: Predict avoidance biased 21-cm power spectrum from full k power spectrum.
+Author: satyapan
+Author-email: satyapan.iiserm@gmail.com
+Classifier: Programming Language :: Python :: 2
+Classifier: Programming Language :: Python :: 2.7
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.4
+Classifier: Programming Language :: Python :: 3.5
+Classifier: Programming Language :: Python :: 3.6
+Classifier: Programming Language :: Python :: 3.7
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Description-Content-Type: text/markdown
+
+# Avoidcorr: Predict avoidance power spectrum from full k power spectrum
+
+avoidcorr is a Python package for predicting the avoidance biased 21-cm power spectrum from a full k power spectrum. It can use shape difference degeneracies from a database to estimate biased power spectra corresponding to an input power spectrum and an input arbitrary mask in k space.
+
+# Dependencies
+
+avoidcorr requires the following Python libraries:
+- numpy
+- matplotlib
+- scipy
+- astropy
+- tqdm
+- multiprocessing
+
+# Installation
+avoidcorr can be installed via pip:
+```
+pip install avoidcorr
+```
+
+# Documentation
+A step-by-step guide is presented in the wiki page.

avoidcorr-0.1.0/README.md

@@ -0,0 +1,22 @@
+# Avoidcorr: Predict avoidance power spectrum from full k power spectrum
+
+avoidcorr is a Python package for predicting the avoidance biased 21-cm power spectrum from a full k power spectrum. It can use shape difference degeneracies from a database to estimate biased power spectra corresponding to an input power spectrum and an input arbitrary mask in k space.
+
+# Dependencies
+
+avoidcorr requires the following Python libraries:
+- numpy
+- matplotlib
+- scipy
+- astropy
+- tqdm
+- multiprocessing
+
+# Installation
+avoidcorr can be installed via pip:
+```
+pip install avoidcorr
+```
+
+# Documentation
+A step-by-step guide is presented in the wiki page.

avoidcorr-0.1.0/avoidcorr/__init__.py

@@ -0,0 +1 @@
+from .predictor import *

avoidcorr-0.1.0/avoidcorr/funcs.py

@@ -0,0 +1,55 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from numpy.fft import fftn, fftshift, fftfreq
+from astropy.cosmology import Planck15 as cosmo
+import astropy.constants as const
+
+def legendre_P(ell, x):
+    coeff = np.zeros(ell+1)
+    coeff[ell] = 1.0
+    return np.polynomial.legendre.legval(x, coeff)
+
+def estimate_multipoles(box, boxlen, kb, k_c, los_axis=2, n_multipoles=6):
+    nx, ny, nz = box.shape
+    fx = np.fft.fftfreq(nx, d=boxlen[0]/nx)
+    fy = np.fft.fftfreq(ny, d=boxlen[1]/ny)
+    fz = np.fft.fftfreq(nz, d=boxlen[2]/nz)
+    kx, ky, kz = np.meshgrid(fx*2*np.pi, fy*2*np.pi, fz*2*np.pi, indexing='ij')
+    kx, ky, kz = np.fft.fftshift(kx), np.fft.fftshift(ky), np.fft.fftshift(kz)
+    k = np.sqrt(kx**2 + ky**2 + kz**2)
+    kpar = [kx, ky, kz][los_axis]
+    mu = np.zeros_like(k)
+    m0 = k > 0
+    mu[m0] = np.abs(kpar[m0]) / k[m0]
+    V = boxlen[0]*boxlen[1]*boxlen[2]
+    dk = (box - np.mean(box)) * V/(nx*ny*nz)
+    fk = np.fft.fftshift(np.fft.fftn(dk))
+    pk = np.abs(fk)**2 / V
+    good = (k > 0) & np.isfinite(pk)
+    kvals, muvals, pkvals = k[good].ravel(), mu[good].ravel(), pk[good].ravel()
+    which = np.digitize(kvals, kb) - 1
+    nb = len(kb) - 1
+    ok = (which >= 0) & (which < nb)
+    which, muvals, pkvals = which[ok], muvals[ok], pkvals[ok]
+    nm = np.zeros(nb, int)
+    fac = k_c**3 / (2*np.pi**2)
+    multipoles = []
+    for m in range(n_multipoles):
+        multipoles.append(np.full(nb, np.nan))
+    for i in range(nb):
+        sel = which == i
+        nm[i] = sel.sum()
+        if nm[i] == 0:
+            continue
+        pk_i = pkvals[sel]
+        mu_i = muvals[sel]
+        for j in range(n_multipoles):
+            ell = 2*(j+1)
+            Pl = legendre_P(ell, mu_i)
+            Pkell = (2*ell+1) * np.mean(pk_i * Pl)
+            multipoles[j][i] = fac[i] * Pkell
+    return multipoles
+
+def mu_theta(z, theta=np.pi/2):
+    factor_cosmo = (cosmo.comoving_transverse_distance(z)*cosmo.H(z)/(const.c*(1+z))).decompose().value
+    return np.sin(theta)*factor_cosmo/np.sqrt(1 + (np.sin(theta)*factor_cosmo)**2)

avoidcorr-0.1.0/avoidcorr/predictor.py

@@ -0,0 +1,345 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from numpy.fft import fftn, fftshift, fftfreq
+from astropy.cosmology import Planck15 as cosmo
+from tqdm import tqdm
+import astropy.units as u
+import astropy.constants as const
+from matplotlib.colors import Normalize, LogNorm
+import multiprocessing as mp
+from functools import partial
+from scipy.interpolate import interp1d
+
+from .funcs import *
+from .ps_estimator import *
+
+class Database:
+    """
+    Store a training database of power spectra and multipoles across multiple redshifts.
+
+    Arguments:
+        ps_path (str): Path to the directory where power spectra and multipole files are stored or loaded from.
+        z_vals (array-like): Redshift values corresponding to each snapshot in the database.
+        k_c (numpy.ndarray): k-bin centres. Required when loading pre-computed power spectra (compute_ps_kwargs=None).
+        delta2s_str (str): Filename prefix for the Delta^2 arrays. Default 'delta2s'.
+        multipoles_str (str): Filename prefix for the multipole arrays. Default 'multipoles_list'.
+        box_shape (tuple): Shape of the simulation boxes. Default (128, 128, 128).
+        boxlen (tuple): Physical side lengths of the boxes in Mpc. Default (250, 250, 250).
+        compute_ps_kwargs (dict): If provided, compute power spectra on the fly rather than loading from disk.
+            Required keys: 'ids' (list of simulation IDs), 'simres_path' (path to simulation results).
+            Optional keys: 'n_multipoles' (int, default 6), 'parallel' (bool, default True), 'n_threads' (int, default 12), and any kwargs accepted by PS1D.
+    """
+    def __init__(self, ps_path, z_vals, k_c=None, delta2s_str='delta2s', multipoles_str='multipoles_list', box_shape=(128,128,128), boxlen=(250,250,250), compute_ps_kwargs=None):
+        self.ps_path = ps_path
+        self.z_vals = z_vals
+        self.delta2s_str = delta2s_str
+        self.multipoles_str = multipoles_str
+        self.box_shape = box_shape
+        self.boxlen = boxlen
+        if compute_ps_kwargs is not None:
+            ids = compute_ps_kwargs.get('ids')
+            simres_path = compute_ps_kwargs.get('simres_path')
+            if ids is None or simres_path is None or boxlen is None:
+                raise ValueError("ids, simres_path, and boxlen must be provided in compute_ps_kwargs")
+            n_multipoles = compute_ps_kwargs.get('n_multipoles', 6)
+            parallel = compute_ps_kwargs.get('parallel', True)
+            n_threads = compute_ps_kwargs.get('n_threads', 12)
+            ps_est_kwargs = {k: v for k, v in compute_ps_kwargs.items()
+                            if k not in ('ids', 'simres_path', 'n_multipoles', 'parallel', 'n_threads')}
+            self.compute_ps_allz(ids, simres_path, boxlen, n_multipoles=n_multipoles, parallel=parallel, n_threads=n_threads, **ps_est_kwargs)
+        else:
+            self.delta2s_allz = [np.load(ps_path + f'/{self.delta2s_str}_z{z_val}.npy') for z_val in z_vals]
+            self.multipoles_allz = [np.load(ps_path + f'/{self.multipoles_str}_z{z_val}.npy') for z_val in z_vals]
+            if k_c is None:
+                raise ValueError("k_c must be provided if not computing PS")
+            else:
+                self.k_c = k_c
+        self.N_z = len(z_vals)
+        self.N_k = len(self.k_c)
+        self.filter_ps()
+        self.norm_ps()
+
+    def get_zero_mask(self, delta2s):
+        is_all_zero = np.all(delta2s == 0, axis=1)
+        is_any_nan = np.any(np.isnan(delta2s), axis=1)
+        zero_mask = is_all_zero | is_any_nan
+        return zero_mask
+
+    def filter_ps(self):
+        self.zero_mask_allz = [self.get_zero_mask(self.delta2s_allz[i]) for i in range(self.N_z)]
+        self.delta2s_filt_allz = [self.delta2s_allz[i][~self.zero_mask_allz[i]] for i in range(self.N_z)]
+        self.multipoles_filt_allz = [self.multipoles_allz[i][~self.zero_mask_allz[i]] for i in range(self.N_z)]
+        self.N_filts = [len(self.delta2s_filt_allz[i]) for i in range(self.N_z)]
+
+    def norm_ps(self):
+        self.norm_factors_allz = [np.nansum(self.delta2s_filt_allz[i], axis=1) for i in range(self.N_z)]
+        self.delta2s_norm_allz = [self.delta2s_filt_allz[i] / self.norm_factors_allz[i][:, None] for i in range(self.N_z)]
+        self.multipoles_norm_allz = [self.multipoles_filt_allz[i] / self.norm_factors_allz[i][:, None, None] for i in range(self.N_z)]
+
+    def process_z(self, z_val, ids, simres_path, boxlen, n_multipoles, **kwargs):
+        delta2s = []
+        multipoles = []
+        for id in ids:
+            bt = np.load(simres_path+f'/simres_id{id}_z{z_val}.npy')
+            ps1d = PS1D(bt, boxlen, z_val, **kwargs)
+            multipole = estimate_multipoles(bt, boxlen, ps1d.kb, ps1d.k_c, n_multipoles=n_multipoles)
+            delta2s.append(ps1d.delta2)
+            multipoles.append(multipole)
+        np.save(self.ps_path + f'/{self.delta2s_str}_z{z_val}.npy', np.array(delta2s))
+        np.save(self.ps_path + f'/{self.multipoles_str}_z{z_val}.npy', np.array(multipoles))
+        np.save(self.ps_path + f'/k_c_z{z_val}.npy', ps1d.k_c)
+
+    def compute_ps_allz(self, ids, simres_path, boxlen, n_multipoles=6, parallel=True, n_threads=12, **kwargs):
+        if parallel:
+            fn = partial(self.process_z, ids=ids, simres_path=simres_path, boxlen=boxlen, n_multipoles=n_multipoles, **kwargs)
+            with mp.Pool(n_threads) as pool:
+                pool.map(fn, self.z_vals)
+        else:
+            for z_val in self.z_vals:
+                self.process_z(z_val, ids, simres_path, boxlen, n_multipoles, **kwargs)
+        self.delta2s_allz = [np.load(self.ps_path + f'/{self.delta2s_str}_z{z_val}.npy') for z_val in self.z_vals]
+        self.multipoles_allz = [np.load(self.ps_path + f'/{self.multipoles_str}_z{z_val}.npy') for z_val in self.z_vals]
+        self.k_c = np.load(self.ps_path + f'/k_c_z{self.z_vals[0]}.npy')
+
+    def plot_ps(self, z_idxs=None, normed=True, **kwargs):
+        if z_idxs is None:
+            z_idxs = list(range(self.N_z))
+        z_vals_plot = [self.z_vals[i] for i in z_idxs]
+        norm = Normalize(vmin=min(z_vals_plot), vmax=max(z_vals_plot))
+        colors = plt.cm.viridis(np.linspace(0, 1, len(z_idxs)))
+        if normed:
+            delta2s_plot = self.delta2s_norm_allz
+        else:
+            delta2s_plot = self.delta2s_filt_allz
+        fig, ax = plt.subplots(figsize=(8, 6))
+        for i in range(len(z_idxs)):
+            for j in range(self.N_filts[z_idxs[i]]):
+                ax.plot(self.k_c, delta2s_plot[z_idxs[i]][j], color=colors[i], **kwargs)
+        ax.set_xlabel(r'$k$ [Mpc$^{-1}$]')
+        ax.set_ylabel(r'$\Delta^2(k)$ [mK$^2$]')
+        ax.set_xscale('log')
+        ax.set_yscale('log')
+        sm = plt.cm.ScalarMappable(cmap='viridis', norm=norm)
+        cb = fig.colorbar(ax=ax, mappable=sm, label='Redshift')
+        fig.tight_layout()
+
+
+class BiasPredictor:
+    """
+    Predicts the avoidance biased power spectrum from a full k power spectrum.
+
+    Two modes:
+    - Analytic mode (no database): When multipoles are known: Directly uses known multipole power spectra to predict the biased power spectrum.
+    - Database mode: When multipoles are unknown. Use database of pre-computed power spectra and multipoles to predict biased power spectrum using shape difference degeneracies.
+
+    Arguments:
+        z_val (float): Redshift of the input power spectrum.
+        box_shape (tuple): Shape of the simulation box from which input power spectrum is derived, e.g. (128, 128, 128). Required for estimating allowed k range when not using a database.
+        boxlen (tuple): Physical side lengths of the box in Mpc from which input power spectrum is derived. Required for estimating allowed k range when not using a database.
+        database (Database): Training database instance. Default None.
+        mu_min (float or 'horizon'): Minimum mu value. Default 'horizon': computes the horizon wedge cut at z_val.
+        kpar_range (tuple or list of tuples): One or more (kpar_min, kpar_max) ranges defining rectangular avoidance boxes in (kperp, kpar) space. Default None (no kpar cut).
+        kperp_range (tuple or list of tuples): One or more (kperp_min, kperp_max) ranges. Must have the same number of entries as kpar_range. Default None.
+        combz (bool): If True and using a database, combine all redshifts into one training set. Default False.
+        k_c_inp (numpy.ndarray): k-bin centres of the input power spectrum. Must match the length of the arrays passed to estimate_avoidance_ps or predict. Default None.
+        interpolate (bool): If True and using a database, interpolate between the database k-grid and k_c_inp. Requires k_c_inp. Default False.
+    """
+    def __init__(self, z_val, box_shape=None, boxlen=None, database=None, mu_min='horizon', kpar_range=None, kperp_range=None, combz=False, k_c_inp=None, interpolate=False):
+        self.z_val = z_val
+        self.box_shape = box_shape
+        self.boxlen = boxlen
+        self.database = database
+        self.interpolate = interpolate
+        self.k_c = k_c_inp
+        if mu_min == 'horizon':
+            self.mu_min = mu_theta(self.z_val, theta=np.pi/2)
+        else:
+            self.mu_min = mu_min
+        self.kpar_ranges = self._normalize_range(kpar_range)
+        self.kperp_ranges = self._normalize_range(kperp_range)
+        if len(self.kpar_ranges) != len(self.kperp_ranges):
+            raise ValueError("kpar_range and kperp_range must have the same number of boxes")
+        self.N_boxes = len(self.kpar_ranges)
+        if box_shape is not None and boxlen is not None:
+            self.mask_avoid = self.accessible_kmask(self.k_c)
+            self.k_c_masked = self.k_c[self.mask_avoid]
+        if self.database is not None:
+            self.box_shape = self.database.box_shape
+            self.boxlen = self.database.boxlen
+            self.z_idx = np.argmin(np.abs(np.array(self.database.z_vals) - self.z_val))
+            self.combz = combz
+            self.k_c = self.database.k_c
+            if self.combz:
+                self.delta2s_norm = np.concatenate(self.database.delta2s_norm_allz, axis=0).copy()
+                self.multipoles_norm = np.concatenate(self.database.multipoles_norm_allz, axis=0).copy()
+                self.norm_factors = np.concatenate(self.database.norm_factors_allz, axis=0).copy()
+            else:
+                self.delta2s_norm = self.database.delta2s_norm_allz[self.z_idx].copy()
+                self.multipoles_norm = self.database.multipoles_norm_allz[self.z_idx].copy()
+                self.norm_factors = self.database.norm_factors_allz[self.z_idx].copy()
+            if self.interpolate:
+                if k_c_inp is None:
+                    raise ValueError("k_c_inp must be provided when interpolate=True")
+                self.k_c_inp = k_c_inp
+                self.mask_avoid = self.accessible_kmask(self.k_c_inp)
+                self.k_c_masked = self.k_c_inp[self.mask_avoid]
+            else:
+                self.mask_avoid = self.accessible_kmask(self.k_c)
+                self.k_c_masked = self.k_c[self.mask_avoid]
+        self.mu_bounds()
+
+    @staticmethod
+    def _normalize_range(r):
+        if r is None:
+            return [None]
+        if isinstance(r, (list, tuple)) and len(r) > 0 and isinstance(r[0], (list, tuple)):
+            return list(r)
+        return [r]
+
+    def mu_bounds(self):
+        k = np.asarray(self.k_c, float)
+        self.mu_lo_boxes = []
+        self.mu_hi_boxes = []
+        for kpar_r, kperp_r in zip(self.kpar_ranges, self.kperp_ranges):
+            mu_lo = np.zeros_like(k)
+            mu_hi = np.ones_like(k)
+            if self.mu_min is not None: mu_lo = np.maximum(mu_lo, self.mu_min * np.ones_like(k))
+            if kpar_r is not None:
+                kpar_min, kpar_max = kpar_r
+                if kpar_min is not None: mu_lo = np.maximum(mu_lo, kpar_min / k)
+                if kpar_max is not None: mu_hi = np.minimum(mu_hi, kpar_max / k)
+            if kperp_r is not None:
+                kperp_min, kperp_max = kperp_r
+                if kperp_max is not None: mu_lo = np.maximum(mu_lo, np.sqrt(np.maximum(0.0, 1.0 - (kperp_max / k)**2)))
+                if kperp_min is not None: mu_hi = np.minimum(mu_hi, np.sqrt(np.maximum(0.0, 1.0 - (kperp_min / k)**2)))
+            self.mu_lo_boxes.append(np.clip(mu_lo, 0.0, 1.0))
+            self.mu_hi_boxes.append(np.clip(mu_hi, 0.0, 1.0))
+
+    def coeff_range(self, ell):
+        if ell == 0:
+            return np.ones_like(self.mu_lo_boxes[0])
+        num = np.zeros_like(self.mu_lo_boxes[0])
+        denom = np.zeros_like(self.mu_lo_boxes[0])
+        for mu_lo, mu_hi in zip(self.mu_lo_boxes, self.mu_hi_boxes):
+            valid = mu_hi > mu_lo
+            Ihi = (legendre_P(ell+1, mu_hi) - legendre_P(ell-1, mu_hi)) / (2*ell+1)
+            Ilo = (legendre_P(ell+1, mu_lo) - legendre_P(ell-1, mu_lo)) / (2*ell+1)
+            num += np.where(valid, Ihi - Ilo, 0.0)
+            denom += np.where(valid, mu_hi - mu_lo, 0.0)
+        return num / denom
+
+    def estimate_avoidance_ps(self, ps1d_delta2, multipoles, masked=False):
+        """
+        Predict avoidance biased power spectrum.
+
+        Arguments:
+            ps1d_delta2 (numpy.ndarray): Input full k power spectrum.
+            multipoles (numpy.ndarray or list): Input even multipoles, with shape (n_multipoles, n_k). Multipole j corresponds to ell = 2*(j+1).
+            masked (bool): If True, return only the accessible k values. Default False.
+        Returns:
+            numpy.ndarray: Predicted biased power spectrum. Full array if masked=False, corresponding k values are in self.k_c_masked.
+        """
+        out = ps1d_delta2.copy()
+        for j in range(len(multipoles)):
+            ell = 2*(j+1)
+            aell = self.coeff_range(ell)
+            out += aell * multipoles[j]
+        if masked:
+            out = out[self.mask_avoid]
+            self.k_c_masked = self.k_c[self.mask_avoid]
+        return out
+
+    def accessible_kmask(self, k_c, los_axis=2):
+        Nlos = self.box_shape[los_axis]
+        Llos = self.boxlen[los_axis]
+        kpar_max_grid = np.pi * Nlos / Llos
+        axes = [0, 1, 2]
+        ax_perp = [a for a in axes if a != los_axis]
+        Np0, Np1 = self.box_shape[ax_perp[0]], self.box_shape[ax_perp[1]]
+        Lp0, Lp1 = self.boxlen[ax_perp[0]], self.boxlen[ax_perp[1]]
+        kperp_max_grid = np.sqrt((np.pi*Np0/Lp0)**2 + (np.pi*Np1/Lp1)**2)
+        mask = np.zeros_like(k_c, dtype=bool)
+        for kpar_r, kperp_r in zip(self.kpar_ranges, self.kperp_ranges):
+            mu_lo = np.zeros_like(k_c)
+            mu_hi = np.ones_like(k_c)
+            if self.mu_min is not None: mu_lo = np.maximum(mu_lo, self.mu_min * np.ones_like(k_c))
+            if kpar_r is not None:
+                kpar_min, kpar_max_user = kpar_r
+                if kpar_min is not None: mu_lo = np.maximum(mu_lo, kpar_min / k_c)
+                if kpar_max_user is not None: mu_hi = np.minimum(mu_hi, kpar_max_user / k_c)
+            if kperp_r is not None:
+                kperp_min, kperp_max_user = kperp_r
+                if kperp_max_user is not None: mu_lo = np.maximum(mu_lo, np.sqrt(np.maximum(0.0, 1.0 - (kperp_max_user / k_c)**2)))
+                if kperp_min is not None: mu_hi = np.minimum(mu_hi, np.sqrt(np.maximum(0.0, 1.0 - (kperp_min / k_c)**2)))
+            mu_lo = np.clip(mu_lo, 0.0, 1.0)
+            mu_hi = np.clip(mu_hi, 0.0, 1.0)
+            ok_mu = mu_hi > mu_lo
+            ok_kpar = (mu_lo * k_c) <= kpar_max_grid
+            if kpar_r is not None and kpar_r[0] is not None: ok_kpar &= (kpar_r[0] <= kpar_max_grid)
+            ok_kperp = np.ones_like(k_c, dtype=bool)
+            if kperp_r is not None and kperp_r[0] is not None: ok_kperp &= (kperp_r[0] <= kperp_max_grid)
+            mask |= (ok_mu & ok_kpar & ok_kperp)
+        return mask
+
+    def min_rmse(self, delta2_inp):
+        norm_factor = np.nansum(delta2_inp)
+        if norm_factor == 0:
+            return None, 0
+        target_shape = delta2_inp / norm_factor
+        mse = np.nanmean((self.delta2s_norm - target_shape)**2, axis=1)
+        best_idx = np.argmin(mse)
+        return best_idx, norm_factor
+
+    def interpolate_ps(self, k_old, k_new, delta2, log=True):
+        if log:
+            f = interp1d(np.log(k_old), np.log(np.where(delta2 > 0, delta2, np.nan)),
+                        bounds_error=False, fill_value="extrapolate")
+            return np.exp(f(np.log(k_new)))
+        else:
+            f = interp1d(k_old, delta2, bounds_error=False, fill_value="extrapolate")
+            return f(k_new)
+
+    def predict(self, delta2_inp, plot=True, log_interp=True):
+        """
+        Predict avoidance biased power spectrum using the training database.
+
+        Arguments:
+            delta2_inp (numpy.ndarray): Input full k power spectrum.
+            plot (bool): If True, plot the input and predicted power spectra. Default True.
+            log_interp (bool): If True, interpolate in log-log space when interpolate=True. Default True.
+
+        Returns:
+            numpy.ndarray: Predicted biased power spectrum at the accessible k values (self.k_c_masked).
+        """
+        if self.interpolate:
+            delta2_inp = self.interpolate_ps(self.k_c_inp, self.k_c, delta2_inp, log=log_interp)
+        best_idx, norm_factor = self.min_rmse(delta2_inp)
+        if norm_factor == 0:
+            raise ValueError("Input PS has zero norm.")
+        delta2_norm = delta2_inp/norm_factor
+        delta2_pred = self.estimate_avoidance_ps(delta2_norm, self.multipoles_norm[best_idx])*norm_factor
+        if self.interpolate:
+            delta2_pred = self.interpolate_ps(self.k_c, self.k_c_inp, delta2_pred, log=log_interp)
+        self.delta2_pred = delta2_pred
+        self.delta2_pred_masked = delta2_pred[self.mask_avoid]
+        if plot:
+            self.plot(delta2_inp, self.delta2_pred_masked)
+        return self.delta2_pred_masked
+
+    def plot(self, delta2_inp, delta2_pred_masked, ax=None):
+        """
+        Plot the input full k power spectrum alongside the predicted biased power spectrum.
+
+        Arguments:
+            delta2_inp (numpy.ndarray): Input full-sky power spectrum on the k_c grid.
+            delta2_pred_masked (numpy.ndarray): Predicted biased power spectrum at accessible k values (k_c_masked).
+            ax (matplotlib.axes.Axes): Axes to plot on. If None, a new figure is created.
+        """
+        if ax is None: fig, ax = plt.subplots()
+        ax.plot(self.k_c, delta2_inp, '-o', label='Input', color='red')
+        ax.plot(self.k_c_masked, delta2_pred_masked, '-s', label='Predicted', color='blue')
+        ax.set_xlabel(r'$k$ [Mpc$^{-1}$]')
+        ax.set_ylabel(r'$\Delta^2(k)$ [mK$^2$]')
+        ax.set_xscale('log')
+        ax.set_yscale('log')
+        ax.legend()

avoidcorr-0.1.0/avoidcorr/ps_estimator.py

@@ -0,0 +1,149 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from numpy.fft import fftn, fftshift, fftfreq
+from astropy.cosmology import Planck15 as cosmo
+import astropy.units as u
+import astropy.constants as const
+import pickle
+from matplotlib.colors import Normalize, LogNorm
+
+class PS1D:
+    """
+    Computes the spherically-averaged 1D dimensionless power spectrum, with optional foreground avoidance masking and/or rectangular cuts in (kperp, kpar) space.
+
+    Arguments:
+        box (numpy.ndarray): 3D brightness temperature box.
+        boxlen (tuple): Physical side lengths of the box in Mpc, e.g. (Lx, Ly, Lz).
+        z_val (float): Redshift of the box, used for the horizon avoidance cut.
+        n_bins (int): Number of k bins. Default 50.
+        avoidance (bool): If True, mask modes inside the foreground wedge. Default False.
+        keep_pk (bool): If True, store the full 3D power spectrum array as self.pk. Default False.
+        log_bins (bool): If True, use logarithmically-spaced k bins. Default False.
+        keep_box (bool): If True, keep the input box in memory after computing. Default False.
+        kperp_range (tuple or list of tuples): One or more (kperp_min, kperp_max) ranges to include.
+            Pass a list of tuples to union multiple rectangular boxes in (kperp, kpar) space.
+        kpar_range (tuple or list of tuples): One or more (kpar_min, kpar_max) ranges to include.
+            Must have the same number of entries as kperp_range.
+        kmin (float): Minimum k for binning (overrides auto-detected minimum).
+        kmax (float): Maximum k for binning (overrides auto-detected maximum).
+        k_edges (array-like): Explicit bin edges. If provided, overrides n_bins, kmin, kmax.
+        apply_window (bool): If True, apply a Blackman-Harris window along the LOS axis. Default False.
+    """
+    def __init__(self, box, boxlen, z_val, n_bins=50, avoidance=False, keep_pk=False, log_bins=False, keep_box=False, kperp_range=None, kpar_range=None, kmin=None, kmax=None, k_edges=None, apply_window=False):
+        self.box, self.boxlen, self.n_bins = box, boxlen, n_bins
+        self.avoidance, self.z_val, self.keep_pk = avoidance, z_val, keep_pk
+        self.log_bins, self.keep_box = log_bins, keep_box
+        self.kmin, self.kmax = kmin, kmax
+        self.k_edges = k_edges
+        self.apply_window = apply_window
+        self.kperp_ranges = self._normalize_range(kperp_range)
+        self.kpar_ranges = self._normalize_range(kpar_range)
+        if len(self.kperp_ranges) != len(self.kpar_ranges):
+            raise ValueError("kperp_range and kpar_range must have the same number of boxes")
+        self.compute_ps()
+
+    @staticmethod
+    def _normalize_range(r):
+        if r is None:
+            return [None]
+        if isinstance(r, (list, tuple)) and len(r) > 0 and isinstance(r[0], (list, tuple)):
+            return list(r)
+        return [r]
+
+    def compute_ps(self):
+        log_bins = self.log_bins
+        box, boxlen, n_bins = self.box, self.boxlen, self.n_bins
+        nx, ny, nz = box.shape
+        fx = fftfreq(nx, d=boxlen[0]/nx)
+        fy = fftfreq(ny, d=boxlen[1]/ny)
+        fz = fftfreq(nz, d=boxlen[2]/nz)
+        kx, ky, kz = np.meshgrid(fx*2*np.pi, fy*2*np.pi, fz*2*np.pi, indexing='ij')
+        kx, ky, kz = fftshift(kx), fftshift(ky), fftshift(kz)
+        kperp, kpar = np.sqrt(kx**2 + ky**2), np.abs(kz)
+        k = np.sqrt(kperp**2 + kpar**2)
+        if self.avoidance:
+            factor_cosmo = (cosmo.comoving_transverse_distance(self.z_val) * cosmo.H(self.z_val) / (const.c * (1 + self.z_val))).decompose().value
+            mask = kpar < factor_cosmo * kperp
+        else:
+            mask = np.zeros_like(k, dtype=bool)
+        has_ranges = any(kp is not None or kz is not None for kp, kz in zip(self.kperp_ranges, self.kpar_ranges))
+        if has_ranges:
+            in_any_box = np.zeros_like(k, dtype=bool)
+            for kperp_r, kpar_r in zip(self.kperp_ranges, self.kpar_ranges):
+                in_box = np.ones_like(k, dtype=bool)
+                if kperp_r is not None:
+                    kp_min, kp_max = kperp_r
+                    if kp_min is not None: in_box &= (kperp >= kp_min)
+                    if kp_max is not None: in_box &= (kperp <= kp_max)
+                if kpar_r is not None:
+                    kz_min, kz_max = kpar_r
+                    if kz_min is not None: in_box &= (kpar >= kz_min)
+                    if kz_max is not None: in_box &= (kpar <= kz_max)
+                in_any_box |= in_box
+            mask |= ~in_any_box
+        V = boxlen[0] * boxlen[1] * boxlen[2]
+        dk = (box - np.mean(box)) * V / (nx * ny * nz)
+        if self.apply_window:
+            freq_window = blackmanharris(nz)
+            dk *= freq_window[None, None, :]
+        fk = fftshift(fftn(dk))
+        pk = np.abs(fk)**2 / V
+        if self.apply_window:
+            window_correction = np.sum(freq_window**2) / nz
+            pk /= window_correction
+        if self.keep_pk: self.pk = pk
+        pk[mask] = np.nan
+        if self.k_edges is not None:
+            kb = np.array(self.k_edges)
+            self.n_bins = len(kb) - 1
+            if log_bins:
+                self.k_c = np.sqrt(kb[1:] * kb[:-1])
+            else:
+                self.k_c = 0.5 * (kb[1:] + kb[:-1])
+        else:
+            if log_bins:
+                kpos = k[(k > 0) & (~np.isnan(pk))]
+                kmin_val = self.kmin if self.kmin is not None else kpos.min()
+                kmax_val = self.kmax if self.kmax is not None else kpos.max()
+                kb = np.logspace(np.log10(kmin_val), np.log10(kmax_val), n_bins + 1)
+                self.k_c = np.sqrt(kb[1:] * kb[:-1])
+            else:
+                kmin_val = self.kmin if self.kmin is not None else k.min()
+                kmax_val = self.kmax if self.kmax is not None else k.max()
+                kb = np.linspace(kmin_val, kmax_val, n_bins + 1)
+                self.k_c = 0.5 * (kb[1:] + kb[:-1])
+        self.kb = kb
+        self.ps1d = np.zeros(len(self.k_c))
+        self.err = np.zeros(len(self.k_c))
+        for i in range(len(self.k_c)):
+            sel = (k >= kb[i]) & (k < kb[i+1]) & (~np.isnan(pk))
+            Nm = np.sum(sel)
+            if Nm > 0:
+                Pk_val = np.nanmean(pk[sel])
+                self.ps1d[i] = Pk_val
+                self.err[i] = Pk_val * np.sqrt(2 / Nm)
+            else:
+                self.ps1d[i], self.err[i] = np.nan, np.nan
+        self.delta2 = self.ps1d * self.k_c**3 / (2 * np.pi**2)
+        self.err_delta2 = self.delta2 * (self.err / self.ps1d)
+        if not self.keep_box: self.box = None
+
+    def plot(self, ax=None, **kwargs):
+        """
+        Plot the dimensionless power spectrum.
+
+        Arguments:
+            ax: Axis to plot on. If None, a new figure is created.
+            **kwargs: Additional keyword arguments passed to ax.errorbar.
+        """
+        if ax is None: fig, ax = plt.subplots()
+        ax.errorbar(self.k_c, self.delta2, yerr=self.err_delta2, fmt='-o', capsize=2, **kwargs)
+        ax.set_xlabel(r'$k$ [Mpc$^{-1}$]')
+        ax.set_ylabel(r'$\Delta^2(k)$ [$\mathrm{mK}^2$]')
+        if self.log_bins:
+            ax.set_xscale('log')
+            ax.set_yscale('log')
+
+    def save(self, filename):
+        with open(filename, 'wb') as outp:
+            pickle.dump(self, outp, pickle.HIGHEST_PROTOCOL)

avoidcorr-0.1.0/pyproject.toml

@@ -0,0 +1,10 @@
+[project]
+name = "avoidcorr"
+version = "0.1.0"
+description = "Predict avoidance biased 21-cm power spectrum from full k power spectrum."
+authors = [{name = "satyapan",email = "satyapan.iiserm@gmail.com"}]
+readme = "README.md"
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"