From bb3348e008ee01663c498d1ec3f5f50ffb0bfe0f Mon Sep 17 00:00:00 2001
From: Yonatan Sklansky <yonatan.sklansky@gmail.com>
Date: Wed, 5 Mar 2025 15:31:24 -0600
Subject: [PATCH 01/37] Add working bubble model for reionization

---
 requirements.txt           |   6 +-
 tests/test_astrophysics.py |   7 +
 zeus21/__init__.py         |   1 +
 zeus21/maps.py             | 149 +++++++++++++++++---
 zeus21/reionization.py     | 279 +++++++++++++++++++++++++++++++++++++
 zeus21/sfrd.py             |  51 +++++--
 zeus21/z21_utilities.py    |  51 +++++++
 7 files changed, 515 insertions(+), 29 deletions(-)
 create mode 100644 zeus21/reionization.py
 create mode 100644 zeus21/z21_utilities.py

diff --git a/requirements.txt b/requirements.txt
index f04ef0a..b967da2 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,4 +4,8 @@ scipy
 sphinx
 myst_parser
 numexpr
-astropy
\ No newline at end of file
+astropy
+tqdm
+matplotlib
+powerbox
+pyfftw
\ No newline at end of file
diff --git a/tests/test_astrophysics.py b/tests/test_astrophysics.py
index f20c976..4c4ee02 100644
--- a/tests/test_astrophysics.py
+++ b/tests/test_astrophysics.py
@@ -15,6 +15,7 @@
 
 from zeus21.sfrd import *
 from zeus21.correlations import *
+from zeus21.reionization import *
 
 UserParams = zeus21.User_Parameters()
 
@@ -41,6 +42,9 @@
 CosmoParams_21cmfast = zeus21.Cosmo_Parameters(UserParams, CosmoParams_input_21cmfast, ClassyCosmo_21cmfast)
 AstroParams_21cmfast = zeus21.Astro_Parameters(UserParams,CosmoParams_21cmfast, astromodel = 1)
 
+#and for bubbles:
+BMF_class = zeus21.BMF(UserParams, Coeffs, HMFintclass, CosmoParams, AstroParams)
+
 
 ztest = 20.
 iztest = min(range(len(Coeffs.zintegral)), key=lambda i: np.abs(Coeffs.zintegral[i]-ztest))
@@ -128,6 +132,9 @@ def test_background():
 
     assert( (Coeffs.gamma_index2D >= 0.0).all()) #effective biases have to be larger than 0 in reasonable models, since galaxies live in haloes that are more clustered than average matter (in other words, SFRD grows monotonically with density)
 
+    assert( ((BMF_class.ion_frac >= 0.0)*(BMF_class.ion_frac <= 1.0) ).all())
+    assert( (BMF_class.BMF >= 0.0).all())
+
 
 
 
diff --git a/zeus21/__init__.py b/zeus21/__init__.py
index 656a475..7a9caee 100644
--- a/zeus21/__init__.py
+++ b/zeus21/__init__.py
@@ -6,6 +6,7 @@
 from .xrays import Xray_class
 from .UVLFs import UVLF_binned
 from .maps import CoevalMaps
+from .reionization import BMF
 
 import warnings
 warnings.filterwarnings("ignore", category=UserWarning) #to silence unnecessary warning in mcfit
diff --git a/zeus21/maps.py b/zeus21/maps.py
index e8bcd4c..d007909 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -2,18 +2,22 @@
 
 Make maps! For fun and science
 
-Author: Julian B. Muñoz
-UT Austin - August 2024
+Authors: Julian B. Muñoz, Yonatan Sklansky, Emilie Thelie
+UT Austin - February 2025
 
 """
 
 from . import cosmology
 from . import constants
+from . import z21_utilities
 
 import numpy as np
 import powerbox as pbox
 from scipy.interpolate import interp1d
+from scipy.interpolate import InterpolatedUnivariateSpline as spline
 from pyfftw import empty_aligned as empty
+from tqdm import trange
+import time
 
 
 class CoevalMaps:
@@ -81,7 +85,7 @@ def __init__(self, T21_coefficients, Power_Spectrum, z, Lbox=600, Nbox=200, KIND
 
             powerratio = powerratioint(pb.k())
             T21lin_k = powerratio * deltak
-            self.T21maplin= self.T21global + powerboxCtoR(pb,mapkin = T21lin_k)
+            self.T21maplin= self.T21global + z21_utilities.powerboxCtoR(pb,mapkin = T21lin_k)
 
             #now make a nonlinear correction, built as \sum_R [e^(gR dR) - gR dR]. Uncorrelatd with all dR so just a separate field!
             #NOTE: its not guaranteed to work, excess power can be negative in some cases! Not for each component xa, Tk, but yes for T21
@@ -108,18 +112,129 @@ def __init__(self, T21_coefficients, Power_Spectrum, z, Lbox=600, Nbox=200, KIND
             print('ERROR, KIND not implemented yet!')
 
 
-
-
-def powerboxCtoR(pbobject,mapkin = None):
-    'Function to convert a complex field to real 3D (eg density, T21...) on the powerbox notation'
-    'Takes a powerbox object pbobject, and a map in k space (mapkin), or otherwise assumes its pbobject.delta_k() (tho in that case it should be delta_x() so...'
-
-    realmap = empty((pbobject.N,) * pbobject.dim, dtype='complex128')
-    if (mapkin is None):
-        realmap[...] = pbobject.delta_k()
+def make_ion_fields(CosmoParams, CoeffStructure, ClassyCosmo, CorrFClass, BMF, input_z, boxlength=300., ncells=300, seed=1234, r_precision=1., timer=True, logd = False, barrier = None, spherize=False):
+    """
+    Generates a 3D map of ionized fields and ionized fraction of hydrogen.
+    
+    Uses a density threshold barrier determined from a converged bubble mass function. With default parameters, the code takes about 20 minutes on laptop to run.
+
+    Parameters
+    ----------
+    CosmoParams: zeus21.Cosmo_Parameters class
+        Stores cosmology.
+    CoeffStructure: zeus21.get_T21_coefficients class
+        Stores sfrd and 21cm coefficients.
+    ClassyCosmo: zeus21.runclass class
+        Sets up Class cosmology.
+    CorrFClass: zeus21.Correlations class
+        Calculates correlation functions.
+    input_z: 1D np.array
+        The redshifts at which to compute output maps. Narrowed down later to select available redshifts from CoeffStructure.zintegral.
+    boxlength: float
+        Comoving physical side length of the box. Default is 300 cMpc.
+    ncells: int
+        Number of cells on a side. Default is 300 cells.
+    seed: int
+        Sets the predetermined generation of maps. Default is 1234.
+    r_precision: float
+        Allows to change the steps of the radii for faster computation. Default (and max) is 1, lower values make the computation faster at the cost of accuracy.
+    timer: bool
+        Whether to print the time elapsed along the process. Default is True.
+    logd: bool
+        Whether to use lognormal (True) or Gaussian (False) density fields. Default is False.
+    barrier: function
+        Input density barrier to be used as the threshold for map generation. Takes z index //\\Change//\\ (relative to CoeffStructure.zintegral) as input and returns np.array of shape of radii. Default is None.
+    spherize: bool
+        Whether to flag spheres around ionized cells (True) instead of only central pixel flagging (False). Default is False. Central pixel flagging is generally more consistent with the bubble mass function than spherizing.
+
+    Outputs
+    ----------
+    ion_fields: 4D np.array
+        Resultant 3D ionized field maps at each redshift. The first dimension is redshifts, and the three other dimensions are spatial.
+    ion_frac: 1D np.array
+        Ionized fraction at each redshift in ion_fields.
+    """
+    
+    #Measure time elapsed from start
+    start_time = time.time()
+    
+    if timer:
+        z21_utilities.print_timer(start_time, 'Making density field...')
+
+    #selecting redshifts and radii from available redshifts
+    zlist = CoeffStructure.zintegral
+    z_idx = z21_utilities.find_nearest_idx(zlist, input_z)
+    z = zlist[z_idx]
+
+    Rs = CoeffStructure.Rtabsmoo
+    r_idx = np.linspace(0,len(Rs)-1,int(len(Rs)*r_precision),dtype=int)
+    r = Rs[r_idx]
+    dx = boxlength/ncells
+
+    #Generating matter power spectrum at z=5
+    klist = CorrFClass._klistCF
+    pk_matter = np.zeros_like(klist)
+    for i, k in enumerate(klist):
+        pk_matter[i] = ClassyCosmo.pk(k, z[0])
+    pk_spl = spline(np.log(klist), np.log(pk_matter))
+
+    #generating density map
+    if logd:
+        pb = pbox.LogNormalPowerBox(N=ncells, dim=3, pk=(lambda k: np.exp(pk_spl(np.log(k)))), boxlength=boxlength, seed=seed)
     else:
-        realmap[...] = mapkin
-    realmap[...] = pbobject.V * pbox.dft.ifft(realmap, L=pbobject.boxlength, a=pbobject.fourier_a, b=pbobject.fourier_b)[0]
-    realmap = np.real(realmap)
-
-    return realmap
\ No newline at end of file
+        pb = pbox.PowerBox(N=ncells, dim=3, pk=(lambda k: np.exp(pk_spl(np.log(k)))), boxlength=boxlength, seed=seed)
+    density_field = pb.delta_x()
+
+    if timer:
+        z21_utilities.print_timer(start_time, 'Creating smoothing function...')
+
+    #comment
+    klistfftx = np.fft.fftfreq(density_field.shape[0],dx)*2*np.pi
+    klist3Dfft = np.sqrt(np.sum(np.meshgrid(klistfftx**2, klistfftx**2, klistfftx**2, indexing='ij'), axis=0))
+    density_fft = np.fft.fftn(density_field)
+
+    if timer:
+        z21_utilities.print_timer(start_time, 'Smoothing...')
+
+    #comment
+    smooth_density_fields = np.array([z21_utilities.tophat_smooth(rr, klist3Dfft, density_fft) for rr in r])
+
+    if timer:
+        z21_utilities.print_timer(start_time, 'Creating ionized field...')
+
+    if barrier is None:
+        barrier = BMF.barrier
+    ion_fields = []
+    ion_frac = np.zeros(len(z))
+    for i in trange(len(z)):
+        curr_z_idx = z_idx[i]
+        ion_field = ionize(CosmoParams, zlist, Rs, curr_z_idx, smooth_density_fields, barrier, r_idx, klist3Dfft, spherize)
+        ion_fields.append(ion_field)
+        ion_frac[i] = np.sum(ion_field)/ncells**3
+    
+    ion_fields = np.array(ion_fields)
+
+    print('Done!')
+    print('\nTotal time:')
+    z21_utilities.print_timer(start_time)
+
+    return ion_fields, ion_frac
+
+#look over this again
+def ionize(CosmoParams, zlist, Rs, curr_z_idx, smooth_density_fields, barrier, r_idx, klist3Dfft, spherize):
+    Dg0 = CosmoParams.growthint(zlist[0])
+    Dg = CosmoParams.growthint(zlist[curr_z_idx])
+    if not spherize:
+        ion_field = np.any(smooth_density_fields > (Dg0/Dg)*barrier(curr_z_idx)[r_idx][:, None, None, None], axis=0)
+    else:
+        ion_field_Rs = []
+        for j in range(len(r_idx)):
+            curr_R = r_idx[j]
+            ion_field_oneR = smooth_density_fields[j] > (Dg0/Dg)*barrier(curr_z_idx)[curr_R]
+            ionk = np.fft.fftn(ion_field_oneR)
+            cutoff = 1/(4/3*np.pi*Rs[curr_R]**3)/2*(1+barrier(curr_z_idx)[curr_R]) #comment
+            ion_spheres = z21_utilities.tophat_smooth(Rs[curr_R]/(1+barrier(curr_z_idx)[curr_R])**(1/3), klist3Dfft, ionk) > cutoff
+            ion_field_Rs.append(ion_spheres)
+                
+        ion_field = np.any(ion_field_Rs, axis=0)
+    return ion_field
\ No newline at end of file
diff --git a/zeus21/reionization.py b/zeus21/reionization.py
new file mode 100644
index 0000000..90917c3
--- /dev/null
+++ b/zeus21/reionization.py
@@ -0,0 +1,279 @@
+"""
+
+Models reionization using an analogy of a halo mass function to ionized bubbles 
+See Sklansky et al. (in prep)
+
+Authors: Yonatan Sklansky, Emilie Thelie
+UT Austin - January 2025
+
+"""
+
+from . import z21_utilities
+from . import cosmology
+from . import constants
+import numpy as np
+import astropy.units as u
+from scipy.integrate import cumulative_trapezoid
+from tqdm import trange
+
+
+class BMF:
+    """
+    Computes the bubble mass function (BMF). 
+
+    
+    """
+    
+    def __init__(self, UserParams, CoeffStructure, HMFintclass, CosmoParams, AstroParams, R_linear_sigma_fit_input=10, FLAG_converge=True, max_iter=10, ZMAX_REION = 30):
+
+        self.ZMAX_REION = ZMAX_REION #max redshift up to which we calculate reionization observables
+        self.zlist = CoeffStructure.zintegral
+        self.Rs = CoeffStructure.Rtabsmoo
+
+        
+        self.gamma = CoeffStructure.gamma_niondot_II_index2D
+        self.gamma2 = CoeffStructure.gamma2_niondot_II_index2D
+        self.sigma = CoeffStructure.sigmaofRtab
+
+        self.Hz = cosmology.Hubinvyr(CosmoParams, self.zlist)
+        #self.Hzconv = (u.km/u.s/u.Mpc).to(1/u.yr) #multiply this times Hubble parameter to convert from km/s/Mpc to 1/yr
+        self.trec0 = 1/(constants.alphaB * cosmology.n_H(CosmoParams,0) * AstroParams._clumping) #seconds
+        self.trec = self.trec0/(1+self.zlist)**3/constants.yrTos #years
+        self.niondot_avg = CoeffStructure.niondot_avg_II
+
+        #self.xHI = CoeffStructure.xHI_avg ----> need to figure out what's going on with this
+        self.ion_frac = np.fmin(1, [self.calc_Q(CosmoParams, i) for i in np.arange(len(self.zlist))])
+        self.ion_frac_initial = np.copy(self.ion_frac)
+
+        zr_mesh = np.meshgrid(np.arange(len(self.Rs)), np.arange(len(self.zlist)))
+        self.nion_norm = self.nion_normalization(zr_mesh[1], zr_mesh[0])
+        self.barrier = self.compute_barrier(self.ion_frac)
+        self.barrier_initial = np.copy(self.barrier)
+
+        self.R_linear_sigma_fit_idx = z21_utilities.find_nearest_idx(self.Rs, R_linear_sigma_fit_input)
+        self.R_linear_sigma_fit = self.Rs[self.R_linear_sigma_fit_idx]
+
+        self.BMF = np.array([self.VRdn_dR(HMFintclass, np.arange(len(self.Rs)), i) for i in range(len(self.zlist))]) #initial bubble mass function
+        self.BMF_initial = np.copy(self.BMF)
+        self.ion_frac = np.nan_to_num([np.trapz(self.BMF[i], np.log(self.Rs)) for i in range(len(self.zlist))]) #ion_frac by integrating the BMF
+
+        if FLAG_converge:
+            self.converge_BMF(HMFintclass, self.ion_frac, max_iter=max_iter)
+        #two functions: compute BMF and iterate
+        
+
+
+    def compute_barrier(self, ion_frac):
+        """
+        Computes the density barrier threshold for ionization.
+        
+        Using the analytic model from Sklansky et al. (in prep), if the total number of ionized photons produced in an overdensity exceeds the sum of the number of hydrogens present and total number of recombinations occurred, then the overdensity is ionized. The density required to ionized is recorded.
+
+        Parameters
+        ----------
+        ion_frac: 1D np.array
+            The ionized fractions to be used to compute the number of recombinations. 
+
+        Output
+        ----------
+        barrier: 2D np.array
+            The resultant density threshold array. First dimension is each redshift, second dimension is each radius scale.
+        """
+        barrier = np.zeros((len(self.zlist), len(self.Rs)))
+        ds_array = np.linspace(-1, 2, 21)
+       
+        
+        for ir in range(len(self.Rs)):
+            #Compute nion_values and nrec_values for this 'ir'
+            nion_values = self.nion_delta_r_int(ds_array, ir)  #Shape (nd, nz)
+            nrec_values = self.nrec(ds_array, ion_frac)             #Shape (nd, nz)
+        
+            total_values = np.log10(nion_values / (1 + nrec_values))   #taking difference in logspace to find zero-crossing 
+        
+            #Loop over redshift indices
+            for i in range(len(self.zlist)):
+                y_values = total_values[:, i]  #Shape (nd,)
+        
+                # Find zero crossings
+                sign_change = np.diff(np.sign(y_values))
+                idx = np.where(sign_change)[0]
+                if idx.size > 0:
+                    #Linear interpolation to find zero crossings
+                    x0 = ds_array[idx]
+                    x1 = ds_array[idx + 1]
+                    y0 = y_values[idx]
+                    y1 = y_values[idx + 1]
+                    x_intersect = x0 - y0 * (x1 - x0) / (y1 - y0)
+                    barrier[i, ir] = x_intersect[0]  #Assuming we take the first crossing
+                else:
+                    barrier[i, ir] = np.nan #Never crosses
+
+        barrier[self.zlist > self.ZMAX_REION] = 100 #sets density to an unreachable barrier, as if reionization isn't happening
+        return barrier
+
+
+
+    #normalizing the nion/sfrd model
+    def nion_normalization(self, z, R):
+        return 1/np.sqrt(1-2*self.gamma2[z, R]*self.sigma[z, R]**2)*np.exp(self.gamma[z, R]**2 * self.sigma[z, R]**2 / (2-4*self.gamma2[z, R]*self.sigma[z, R]**2))
+
+    def nrec(self, d_array, ion_frac):
+        """
+        Vectorized computation of nrec over an array of overdensities d_array.
+
+        Parameters
+        ----------
+        d_array: 1D np.array
+            A list of sample overdensity values to evaluate nrec over.
+        ion_frac: 1D np.array
+            The ionized fraction over all redshifts.
+
+        Output
+        ----------
+        nrecs: 2D np.array
+            The total number of recombinations at each overdensity for a certain ionized fraction history at each redshift. The first dimension is densities, the second dimension is redshifts.
+        """
+        
+        #reverse the inputs to make the integral easier to compute
+        z_rev = self.zlist[::-1]
+        Hz_rev = self.Hz[::-1]
+        trec_rev = self.trec[::-1]
+        ion_frac_rev = ion_frac[::-1]
+    
+        denom = -1 / (1 + z_rev) / Hz_rev / trec_rev
+        integrand_base = denom * ion_frac_rev 
+    
+        # Broadcast over d_array
+        integrand = integrand_base[np.newaxis, :] * (1 + d_array[:, np.newaxis])
+        cumulative_integral = np.concatenate(
+            (np.zeros((len(d_array), 1)), cumulative_trapezoid(integrand, x=z_rev, axis=1)), axis=1
+        )
+        nrecs = cumulative_integral[:, ::-1]  # Reverse to match self.zlist order
+        return nrecs
+    
+    def niondot_delta_r(self, d_array, ir):
+        """
+        Compute niondot over an array of overdensities d_array for a given ir.
+
+        Parameters
+        ----------
+        d_array: 1D np.array
+            A list of sample overdensity values to evaluate niondot over.
+        ir: int
+            Index corresponding to a certain radius value from self.Rs.
+
+        Output
+        ----------
+        niondot: 2D np.array
+            The rates of ionizing photon production. The first dimension is densities, the second dimension is redshifts.
+        """
+        
+        d_array = d_array[:, np.newaxis]
+    
+        gamma_ir = self.gamma[:, ir]   
+        gamma2_ir = self.gamma2[:, ir]  
+        nion_norm_ir = self.nion_norm[:, ir] 
+    
+        exp_term = np.exp(gamma_ir[np.newaxis, :] * d_array + gamma2_ir[np.newaxis, :] * d_array**2)
+        niondot = (self.niondot_avg[np.newaxis, :] / nion_norm_ir[np.newaxis, :]) * exp_term
+        
+        return niondot
+    
+    def nion_delta_r_int(self, d_array, ir):
+        """
+        Vectorized computation of nion over an array of overdensities d_array for a given ir.
+
+        Parameters
+        ----------
+        d_array: 1D np.array
+            A list of sample overdensity values to evaluate niondot over.
+        ir: int
+            Index corresponding to a certain radius value from self.Rs.
+
+        Output
+        ----------
+        nion: 2D np.array
+            The total number of ionizing photons produced since z=50. The first dimension is densities, the second dimension is redshifts.
+        """
+        
+        #reverse the inputs to make the integral easier to compute
+        z_rev = self.zlist[::-1]
+        Hz_rev = self.Hz[::-1]
+    
+        niondot_values = self.niondot_delta_r(d_array, ir)
+    
+        integrand = -1 / (1 + z_rev) / Hz_rev * niondot_values[:, ::-1]
+        cumulative_integral = np.concatenate(
+            (np.zeros((len(d_array), 1)), cumulative_trapezoid(integrand, x=z_rev, axis=1)), axis=1
+        )
+        nion = cumulative_integral[:, ::-1]  # Reverse back to match self.zlist order
+        return nion
+
+    #calculating ionized fraction, put outside the class
+    def calc_Q(self, CosmoParams, iz):
+        z = self.zlist[iz]
+        dtdz = 1/cosmology.Hubinvyr(CosmoParams, self.zlist[iz:])/(1 + self.zlist[iz:])
+        tau0 = self.trec0 * np.sqrt(CosmoParams.OmegaM) * cosmology.Hubinvyr(CosmoParams,0) / constants.yrTos
+        exp = np.exp(2/3/tau0 * (np.power(1 + z, 3/2) - np.power(1 + self.zlist[iz:], 3/2))) #switched order around to be correct
+        
+        integrand = dtdz * self.niondot_avg[iz:] * exp
+    
+        return np.trapz(integrand, x = self.zlist[iz:])
+
+    #computing linear barrier
+    def B_1(self, HMFintclass, ir, iz):
+        sigmax = HMFintclass.sigmaR_int(self.Rs[self.R_linear_sigma_fit_idx+1], self.zlist[iz])
+        sigmin = HMFintclass.sigmaR_int(self.Rs[self.R_linear_sigma_fit_idx-1], self.zlist[iz])
+        barriermax = self.barrier[iz, self.R_linear_sigma_fit_idx+1]
+        barriermin = self.barrier[iz, self.R_linear_sigma_fit_idx-1]
+        return (barriermax - barriermin)/(sigmax**2 - sigmin**2)
+        
+    def B_0(self, HMFintclass, ir, iz):
+        sigmin = HMFintclass.sigmaR_int(self.Rs[self.R_linear_sigma_fit_idx-1], self.zlist[iz])
+        barriermin = self.barrier[iz, self.R_linear_sigma_fit_idx-1]
+        return barriermin - sigmin**2 * self.B_1(HMFintclass, ir, iz)
+    
+    def B(self, HMFintclass, ir, iz): 
+        sig = HMFintclass.sigmaR_int(self.Rs[ir], self.zlist[iz])
+        return self.B_0(HMFintclass, ir, iz) + self.B_1(HMFintclass, ir, iz)*sig**2
+
+    
+    #computing other terms in the BMF
+    def dsigma_dR(self, HMFintclass, ir, iz):
+        R = self.Rs[ir]
+        z = self.zlist[iz]
+        sigma = HMFintclass.sigmaR_int(R, z)
+        return sigma/R*np.gradient(np.log(sigma), np.log(R))
+    
+    def dlogsigma_dlogR(self, HMFintclass, ir, iz):
+        R = self.Rs[ir]
+        z = self.zlist[iz]
+        return self.dsigma_dR(HMFintclass, ir, iz) * R/HMFintclass.sigmaR_int(R,z)
+    
+    def VRdn_dR(self, HMFintclass, ir, iz):
+        R = self.Rs[ir]
+        z = self.zlist[iz]
+        sig = HMFintclass.sigmaR_int(self.Rs[ir], self.zlist[iz])
+        return np.sqrt(2/np.pi) * np.abs(self.dlogsigma_dlogR(HMFintclass, ir, iz)) * np.abs(self.B_0(HMFintclass, ir, iz))/sig * np.exp(-self.B(HMFintclass, ir, iz)**2/2/sig**2)
+    
+    def Rdn_dR(self, ir, iz):
+        R = self.Rs[ir]
+        z = self.zlist[iz]
+        return self.VRdn_dR(HMFintclass, ir, iz)*3/(4*np.pi*R**3)
+
+    def converge_BMF(self, HMFintclass, ion_frac_input, max_iter):
+        self.ion_frac = ion_frac_input
+        for j in trange(max_iter):
+            ion_frac_prev = np.copy(self.ion_frac)
+            
+            self.barrier = self.compute_barrier(self.ion_frac)
+            self.BMF = np.array([self.VRdn_dR(HMFintclass, np.arange(len(self.Rs)), i) for i in range(len(self.zlist))])
+            self.ion_frac = np.nan_to_num([np.trapz(self.BMF[i], np.log(self.Rs)) for i in range(len(self.zlist))])
+
+            if np.allclose(ion_frac_prev, self.ion_frac):
+                print(f'SUCCESS: BMF converged in {j} iterations.')
+                return 
+            
+        print(f"WARNING: BMF didn't converge within {max_iter} iterations.")
+            
+        
\ No newline at end of file
diff --git a/zeus21/sfrd.py b/zeus21/sfrd.py
index daad9af..ddd5f58 100644
--- a/zeus21/sfrd.py
+++ b/zeus21/sfrd.py
@@ -81,17 +81,19 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
         
         self.SFRDbar2D = np.zeros((self.Nzintegral, Cosmo_Parameters.NRs)) #SFR at z=zprime when averaged over a radius R (so up to a higher z)
         
-        self.gamma_index2D = np.zeros_like(self.SFRDbar2D) #index of SFR ~ exp(\gamma delta)
-        self.gamma_II_index2D = np.zeros_like(self.SFRDbar2D) #index of SFR ~ exp(\gamma delta)
-        self.gamma_III_index2D = np.zeros_like(self.SFRDbar2D) #index of SFR ~ exp(\gamma delta)
+        self.gamma_index2D = np.zeros_like(self.SFRDbar2D) #index of SFR ~ exp(\gamma delta + \gamma_2 delta^2)
+        self.gamma_II_index2D = np.zeros_like(self.SFRDbar2D) #index of SFR ~ exp(\gamma delta + \gamma_2 delta^2)
+        self.gamma2_II_index2D = np.zeros_like(self.SFRDbar2D) #index of SFR ~ exp(\gamma delta + \gamma_2 delta^2)
+        self.gamma_III_index2D = np.zeros_like(self.SFRDbar2D) #index of SFR ~ exp(\gamma delta + \gamma_2 delta^2)
+        self.gamma2_III_index2D = np.zeros_like(self.SFRDbar2D) #index of SFR ~ exp(\gamma \delta + \gamma_2 \delta^2)
         
         self.niondot_avg = np.zeros_like(self.zintegral) #\dot nion at each z (int d(SFRD)/dM *fesc(M) dM)/rhobaryon
-        self.gamma_Niondot_index2D = np.zeros_like(self.SFRDbar2D) #index of SFR ~ exp(\gamma delta)
+        self.gamma_Niondot_index2D = np.zeros_like(self.SFRDbar2D) #index of SFR ~ exp(\gamma delta + \gamma_2 delta^2)
 
 
         
 #        ###HAC: OLD, TO BE DELETED, added SFRD variables for pop II and III stars
-#        self.gamma_index2D_old = np.zeros_like(self.SFRDbar2D) #index of SFR ~ exp(\gamma delta)
+#        self.gamma_index2D_old = np.zeros_like(self.SFRDbar2D) #index of SFR ~ exp(\gamma delta + \gamma_2 delta^2)
 
         #EoR coeffs
         self.sigmaofRtab = np.array([HMF_interpolator.sigmaR_int(self.Rtabsmoo, zz) for zz in self.zintegral]) #to be used in correlations.py, in get_bubbles()
@@ -112,7 +114,7 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
 
         #and EPS factors
         Nsigmad = 1.0 #how many sigmas we explore
-        Nds = 2 #how many deltas
+        Nds = 3 #how many deltas
         deltatab_norm = np.linspace(-Nsigmad,Nsigmad,Nds)
 
         #initialize Xrays
@@ -225,8 +227,9 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
             print("ERROR: Need to set FLAG_EMULATE_21CMFAST at True or False in the self.gamma_index2D calculation.")
 
         ########
-        # Compute SFRD quantities
+        # Compute SFRD and niondot quantities
         SFRD_II_dR = np.trapz(integrand_II, HMF_interpolator.logtabMh, axis = 2)
+        niondot_II_dR = np.trapz(integrand_II*fesctab_II[None, None, :, None], HMF_interpolator.logtabMh, axis = 2)
 
         ###
         if Astro_Parameters.USE_POPIII == True:
@@ -236,19 +239,41 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
                 integrand_III = PS_HMF_corr * SFR_III(Astro_Parameters, Cosmo_Parameters, ClassCosmo, HMF_interpolator, mArray, J21LW_interp, zGreaterArray, zGreaterArray, ClassCosmo.pars['v_avg']) * mArray
 
             SFRD_III_dR = np.trapz(integrand_III, HMF_interpolator.logtabMh, axis = 2)
+            niondot_III_dR = np.trapz(integrand_III*fesctab_III[None, None, :, None], HMF_interpolator.logtabMh, axis = 2)
         else:
             SFRD_III_dR = np.zeros_like(SFRD_II_dR)
             
-            
         #compute gammas
-        self.gamma_II_index2D = np.log(SFRD_II_dR[:,:,-1]/SFRD_II_dR[:,:,0]) / (deltaArray[:,:,0,-1] - deltaArray[:,:,0,0])
+        midpoint = deltaArray.shape[-1]//2 #midpoint of deltaArray at delta = 0
+
+        self.gamma_II_index2D = np.log(SFRD_II_dR[:,:,midpoint+1]/SFRD_II_dR[:,:,midpoint-1]) / (deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1])
         self.gamma_II_index2D[np.isnan(self.gamma_II_index2D)] = 0.0
 
+        self.gamma_niondot_II_index2D = np.log(niondot_II_dR[:,:,midpoint+1]/niondot_II_dR[:,:,midpoint-1]) / (deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1])
+        self.gamma_niondot_II_index2D[np.isnan(self.gamma_niondot_II_index2D)] = 0.0
+
+        #compute second-order derivative gammas by computing two first-order derivatives
+        der1_II =  np.log(niondot_II_dR[:,:,midpoint]/niondot_II_dR[:,:,midpoint-1])/(deltaArray[:,:,0,midpoint] - deltaArray[:,:,0,midpoint-1]) #ln(y2/y1)/(x2-x1)
+        der2_II =  np.log(niondot_II_dR[:,:,midpoint+1]/niondot_II_dR[:,:,midpoint])/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint]) #ln(y3/y2)/(x3-x2)
+        self.gamma2_niondot_II_index2D = (der2_II - der1_II)/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1]) #second derivative: (der2-der1)/((x3-x1)/2)
+        self.gamma2_niondot_II_index2D[np.isnan(self.gamma2_niondot_II_index2D)] = 0.0
+
         if Astro_Parameters.USE_POPIII == True:
-            self.gamma_III_index2D = np.log(SFRD_III_dR[:,:,-1]/SFRD_III_dR[:,:,0]) / (deltaArray[:,:,0,-1] - deltaArray[:,:,0,0])
+            self.gamma_III_index2D = np.log(SFRD_III_dR[:,:,midpoint+1]/SFRD_III_dR[:,:,midpoint-1]) / (deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1])
             self.gamma_III_index2D[np.isnan(self.gamma_III_index2D)] = 0.0
+
+            self.gamma_niondot_III_index2D = np.log(niondot_III_dR[:,:,midpoint+1]/niondot_III_dR[:,:,midpoint-1]) / (deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1])
+            self.gamma_niondot_III_index2D[np.isnan(self.gamma_niondot_III_index2D)] = 0.0
+
+            der1_III =  np.log(niondot_III_dR[:,:,midpoint]/niondot_III_dR[:,:,midpoint-1])/(deltaArray[:,:,0,midpoint] - deltaArray[:,:,0,midpoint-1]) #ln(y2/y1)/(x2-x1)
+            der2_III =  np.log(niondot_III_dR[:,:,midpoint+1]/niondot_III_dR[:,:,midpoint])/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint]) #ln(y3/y2)/(x3-x2)
+            self.gamma2_niondot_III_index2D = (der2_III - der1_III)/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1]) #second derivative: (der2-der1)/((x3-x1)/2)
+            self.gamma2_niondot_III_index2D[np.isnan(self.gamma2_niondot_III_index2D)] = 0.0
+
         else:
             self.gamma_III_index2D = np.zeros_like(self.gamma_II_index2D)
+            self.gamma_niondot_III_index2D = np.zeros_like(self.gamma_niondot_II_index2D)
+            self.gamma2_niondot_III_index2D = np.zeros_like(self.gamma2_niondot_II_index2D)
            
         #####################################################################################################
         ### STEP 3: Computing lambdas in velocity anisotropies
@@ -547,6 +572,11 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
         else:
             self._coeff_Ja_xa_0 = 8.0*np.pi*(constants.wavelengthLyA/1e7)**2 * constants.widthLyA * constants.Tstar_21/(9.0*constants.A10_21*self.T_CMB) #units of (cm^2 s Hz sr), convert from Ja to xa. should give 1.81e11/(1+self.zintegral) for Tcmb_0=2.725 K
 
+        #collision coefficient fh(1-x_e)
+        #self.xc_HH = Cosmo_Parameters.f_H * (1.0 - self.xe_avg) * cosmology.n_baryon(Cosmo_Parameters, self.zintegral) * np.interp(self.Tk_avg, constants.Tk_HH, constants.k_HH) / constants.A10_21 * constants.Tstar_21 / cosmology.Tcmb(ClassCosmo, self.zintegral)
+        #self.xc_He = self.xe_avg * cosmology.n_baryon(Cosmo_Parameters, self.zintegral) * np.interp(self.Tk_avg, constants.Tk_He, constants.k_He) / constants.A10_21 * constants.Tstar_21 / cosmology.Tcmb(ClassCosmo, self.zintegral) #xe
+        #self.xc_avg = self.xc_HH + self.xc_He
+
         if(User_Parameters.FLAG_WF_ITERATIVE==True): #iteratively find Tcolor and Ts. Could initialize one to zero, but this should converge faster
             _invTs_tryfirst = 1.0/self.T_CMB
             self._invTs_avg = 1.0/self.Tk_avg
@@ -611,7 +641,6 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
                 self._Q0iteration = self._nion - self._nrecombinations
             self.Qion_avg = self._Q0iteration
 
-
         #common to both methods.
         self.Qion_avg = np.fmin(1.0, self.Qion_avg)
         self._xHII_avg = self.Qion_avg + (1.0 - self.Qion_avg) * self.xe_avg #accounts for partial ionization, small effect
diff --git a/zeus21/z21_utilities.py b/zeus21/z21_utilities.py
new file mode 100644
index 0000000..e05f231
--- /dev/null
+++ b/zeus21/z21_utilities.py
@@ -0,0 +1,51 @@
+"""
+Helper functions to be used across zeus21
+
+Authors: Yonatan Sklansky, Emilie Thelie
+UT Austin - February 2025
+
+"""
+
+import numpy as np
+import time
+
+def powerboxCtoR(pbobject,mapkin = None):
+    'Function to convert a complex field to real 3D (eg density, T21...) on the powerbox notation'
+    'Takes a powerbox object pbobject, and a map in k space (mapkin), or otherwise assumes its pbobject.delta_k() (tho in that case it should be delta_x() so...'
+
+    realmap = empty((pbobject.N,) * pbobject.dim, dtype='complex128')
+    if (mapkin is None):
+        realmap[...] = pbobject.delta_k()
+    else:
+        realmap[...] = mapkin
+    realmap[...] = pbobject.V * pbox.dft.ifft(realmap, L=pbobject.boxlength, a=pbobject.fourier_a, b=pbobject.fourier_b)[0]
+    realmap = np.real(realmap)
+
+    return realmap
+
+def tophat_smooth(rr, ks, dk):
+    x = ks * rr + 1e-5
+    win_k = 3/(x**3) * (np.sin(x) - x*np.cos(x))
+    deltakfilt = dk * win_k
+    return np.real(np.fft.ifftn(deltakfilt))
+
+def find_nearest_idx(array, values):
+    array = np.atleast_1d(array)
+    values = np.atleast_1d(values)
+    idx = []
+    for i in range(len(values)):
+        idx.append((np.abs(array - values[i])).argmin())
+    return np.unique(idx)
+
+def print_timer(start_time, text=""):
+    elapsed_time = time.time() - start_time
+    mins = int(elapsed_time//60)
+    secs = int(elapsed_time - mins*60)
+    print(f"\n{mins}min {secs}s\n")
+    print(text)
+
+def v2r(v):
+    return (3/4/np.pi * v)**(1/3)
+    
+def r2v(r):
+    return 4/3 * np.pi * r**3
\ No newline at end of file

From 723066ad4e4c4044f11abbf828c3582208126c7c Mon Sep 17 00:00:00 2001
From: Yonatan Sklansky <yonatan.sklansky@gmail.com>
Date: Mon, 10 Mar 2025 16:15:08 -0500
Subject: [PATCH 02/37] Fixed map function

---
 zeus21/maps.py | 36 ++++++++++++++++++++++++++++++------
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index d007909..b602866 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -112,7 +112,7 @@ def __init__(self, T21_coefficients, Power_Spectrum, z, Lbox=600, Nbox=200, KIND
             print('ERROR, KIND not implemented yet!')
 
 
-def make_ion_fields(CosmoParams, CoeffStructure, ClassyCosmo, CorrFClass, BMF, input_z, boxlength=300., ncells=300, seed=1234, r_precision=1., timer=True, logd = False, barrier = None, spherize=False):
+def make_ion_fields(CosmoParams, CoeffStructure, ClassyCosmo, CorrFClass, BMF, input_z, boxlength=300., ncells=300, seed=1234, r_precision=1., timer=True, logd = False, barrier = None, spherize=False, FLAG_return_densities = 0):
     """
     Generates a 3D map of ionized fields and ionized fraction of hydrogen.
     
@@ -128,6 +128,8 @@ def make_ion_fields(CosmoParams, CoeffStructure, ClassyCosmo, CorrFClass, BMF, i
         Sets up Class cosmology.
     CorrFClass: zeus21.Correlations class
         Calculates correlation functions.
+    BMF: zeus21.reionization class
+        Computes bubble mass functions and barriers.
     input_z: 1D np.array
         The redshifts at which to compute output maps. Narrowed down later to select available redshifts from CoeffStructure.zintegral.
     boxlength: float
@@ -146,6 +148,11 @@ def make_ion_fields(CosmoParams, CoeffStructure, ClassyCosmo, CorrFClass, BMF, i
         Input density barrier to be used as the threshold for map generation. Takes z index //\\Change//\\ (relative to CoeffStructure.zintegral) as input and returns np.array of shape of radii. Default is None.
     spherize: bool
         Whether to flag spheres around ionized cells (True) instead of only central pixel flagging (False). Default is False. Central pixel flagging is generally more consistent with the bubble mass function than spherizing.
+    FLAG_return_densities: int
+        Options: (0, 1, 2). Default is 0.
+            0: returns only the ionized fields and ionized fractions.
+            1: returns only the ionized fields, ionized fractions, and density field at the last redshift.
+            2: returns the ionized fields, ionized fractions, density field at the last redshift, and the smoothed density fields at each scale at the last redshift.
 
     Outputs
     ----------
@@ -153,6 +160,12 @@ def make_ion_fields(CosmoParams, CoeffStructure, ClassyCosmo, CorrFClass, BMF, i
         Resultant 3D ionized field maps at each redshift. The first dimension is redshifts, and the three other dimensions are spatial.
     ion_frac: 1D np.array
         Ionized fraction at each redshift in ion_fields.
+    Optional:
+        density_field: 3D np.array
+            The field of densities at the last redshift.
+        smoothed_density_fields: 4D np.array
+            The field of densities at the last redshift, smoothed over each radius scale. The first dimension is the smoothing scales, the other three are spatial.
+    
     """
     
     #Measure time elapsed from start
@@ -218,22 +231,33 @@ def make_ion_fields(CosmoParams, CoeffStructure, ClassyCosmo, CorrFClass, BMF, i
     print('\nTotal time:')
     z21_utilities.print_timer(start_time)
 
-    return ion_fields, ion_frac
+    if FLAG_return_densities == 0:
+        return ion_fields, ion_frac
+
+    elif FLAG_return_densities == 1:
+        return ion_fields, ion_frac, density_field
+
+    elif FLAG_return_densities == 2:
+        return ion_fields, ion_frac, density_field, smooth_density_fields
+
+    else:
+        print('WARNING: FLAG_return_densities is not set to (0, 1, or 2). Defaulting to 0.')
+        return ion_fields, ion_frac
 
 #look over this again
 def ionize(CosmoParams, zlist, Rs, curr_z_idx, smooth_density_fields, barrier, r_idx, klist3Dfft, spherize):
     Dg0 = CosmoParams.growthint(zlist[0])
     Dg = CosmoParams.growthint(zlist[curr_z_idx])
     if not spherize:
-        ion_field = np.any(smooth_density_fields > (Dg0/Dg)*barrier(curr_z_idx)[r_idx][:, None, None, None], axis=0)
+        ion_field = np.any(smooth_density_fields > (Dg0/Dg)*barrier[curr_z_idx, r_idx][:, None, None, None], axis=0)
     else:
         ion_field_Rs = []
         for j in range(len(r_idx)):
             curr_R = r_idx[j]
-            ion_field_oneR = smooth_density_fields[j] > (Dg0/Dg)*barrier(curr_z_idx)[curr_R]
+            ion_field_oneR = smooth_density_fields[j] > (Dg0/Dg)*barrier[curr_z_idx, curr_R]
             ionk = np.fft.fftn(ion_field_oneR)
-            cutoff = 1/(4/3*np.pi*Rs[curr_R]**3)/2*(1+barrier(curr_z_idx)[curr_R]) #comment
-            ion_spheres = z21_utilities.tophat_smooth(Rs[curr_R]/(1+barrier(curr_z_idx)[curr_R])**(1/3), klist3Dfft, ionk) > cutoff
+            cutoff = 1/(4/3*np.pi*Rs[curr_R]**3)/2*(1+barrier[curr_z_idx, curr_R]) #comment
+            ion_spheres = z21_utilities.tophat_smooth(Rs[curr_R]/(1+barrier[curr_z_idx, curr_R])**(1/3), klist3Dfft, ionk) > cutoff
             ion_field_Rs.append(ion_spheres)
                 
         ion_field = np.any(ion_field_Rs, axis=0)

From db1c39762851ce1e8ebe5e48e124ba62af77c662 Mon Sep 17 00:00:00 2001
From: Yonatan Sklansky <yonatan.sklansky@gmail.com>
Date: Tue, 11 Mar 2025 13:31:02 -0500
Subject: [PATCH 03/37] Added gamma2 for SFRD_II and _III

---
 zeus21/sfrd.py | 25 ++++++++++++++++++-------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/zeus21/sfrd.py b/zeus21/sfrd.py
index ddd5f58..101f538 100644
--- a/zeus21/sfrd.py
+++ b/zeus21/sfrd.py
@@ -252,10 +252,15 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
         self.gamma_niondot_II_index2D = np.log(niondot_II_dR[:,:,midpoint+1]/niondot_II_dR[:,:,midpoint-1]) / (deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1])
         self.gamma_niondot_II_index2D[np.isnan(self.gamma_niondot_II_index2D)] = 0.0
 
-        #compute second-order derivative gammas by computing two first-order derivatives
-        der1_II =  np.log(niondot_II_dR[:,:,midpoint]/niondot_II_dR[:,:,midpoint-1])/(deltaArray[:,:,0,midpoint] - deltaArray[:,:,0,midpoint-1]) #ln(y2/y1)/(x2-x1)
-        der2_II =  np.log(niondot_II_dR[:,:,midpoint+1]/niondot_II_dR[:,:,midpoint])/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint]) #ln(y3/y2)/(x3-x2)
-        self.gamma2_niondot_II_index2D = (der2_II - der1_II)/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1]) #second derivative: (der2-der1)/((x3-x1)/2)
+        #compute second-order derivative gammas by computing two first-order derivatives #TODO: functionalize derivatives
+        der1_II =  np.log(SFRD_II_dR[:,:,midpoint]/SFRD_II_dR[:,:,midpoint-1])/(deltaArray[:,:,0,midpoint] - deltaArray[:,:,0,midpoint-1]) #ln(y2/y1)/(x2-x1)
+        der2_II =  np.log(SFRD_II_dR[:,:,midpoint+1]/SFRD_II_dR[:,:,midpoint])/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint]) #ln(y3/y2)/(x3-x2)
+        self.gamma2_II_index2D = (der2_II - der1_II)/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1]) #second derivative: (der2-der1)/((x3-x1)/2)
+        self.gamma2_II_index2D[np.isnan(self.gamma2_II_index2D)] = 0.0
+        
+        der1_niondot_II =  np.log(niondot_II_dR[:,:,midpoint]/niondot_II_dR[:,:,midpoint-1])/(deltaArray[:,:,0,midpoint] - deltaArray[:,:,0,midpoint-1]) #ln(y2/y1)/(x2-x1)
+        der2_niondot_II =  np.log(niondot_II_dR[:,:,midpoint+1]/niondot_II_dR[:,:,midpoint])/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint]) #ln(y3/y2)/(x3-x2)
+        self.gamma2_niondot_II_index2D = (der2_niondot_II - der1_niondot_II)/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1]) #second derivative: (der2-der1)/((x3-x1)/2)
         self.gamma2_niondot_II_index2D[np.isnan(self.gamma2_niondot_II_index2D)] = 0.0
 
         if Astro_Parameters.USE_POPIII == True:
@@ -265,13 +270,19 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
             self.gamma_niondot_III_index2D = np.log(niondot_III_dR[:,:,midpoint+1]/niondot_III_dR[:,:,midpoint-1]) / (deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1])
             self.gamma_niondot_III_index2D[np.isnan(self.gamma_niondot_III_index2D)] = 0.0
 
-            der1_III =  np.log(niondot_III_dR[:,:,midpoint]/niondot_III_dR[:,:,midpoint-1])/(deltaArray[:,:,0,midpoint] - deltaArray[:,:,0,midpoint-1]) #ln(y2/y1)/(x2-x1)
-            der2_III =  np.log(niondot_III_dR[:,:,midpoint+1]/niondot_III_dR[:,:,midpoint])/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint]) #ln(y3/y2)/(x3-x2)
-            self.gamma2_niondot_III_index2D = (der2_III - der1_III)/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1]) #second derivative: (der2-der1)/((x3-x1)/2)
+            der1_III =  np.log(SFRD_III_dR[:,:,midpoint]/SFRD_III_dR[:,:,midpoint-1])/(deltaArray[:,:,0,midpoint] - deltaArray[:,:,0,midpoint-1]) #ln(y2/y1)/(x2-x1)
+            der2_III =  np.log(SFRD_III_dR[:,:,midpoint+1]/SFRD_III_dR[:,:,midpoint])/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint]) #ln(y3/y2)/(x3-x2)
+            self.gamma2_III_index2D = (der2_III - der1_III)/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1]) #second derivative: (der2-der1)/((x3-x1)/2)
+            self.gamma2_III_index2D[np.isnan(self.gamma2_III_index2D)] = 0.0
+
+            der1_niondot_III =  np.log(niondot_III_dR[:,:,midpoint]/niondot_III_dR[:,:,midpoint-1])/(deltaArray[:,:,0,midpoint] - deltaArray[:,:,0,midpoint-1]) #ln(y2/y1)/(x2-x1)
+            der2_niondot_III =  np.log(niondot_III_dR[:,:,midpoint+1]/niondot_III_dR[:,:,midpoint])/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint]) #ln(y3/y2)/(x3-x2)
+            self.gamma2_niondot_III_index2D = (der2_niondot_III - der1_niondot_III)/(deltaArray[:,:,0,midpoint+1] - deltaArray[:,:,0,midpoint-1]) #second derivative: (der2-der1)/((x3-x1)/2)
             self.gamma2_niondot_III_index2D[np.isnan(self.gamma2_niondot_III_index2D)] = 0.0
 
         else:
             self.gamma_III_index2D = np.zeros_like(self.gamma_II_index2D)
+            self.gamma2_III_index2D = np.zeros_like(self.gamma2_II_index2D)
             self.gamma_niondot_III_index2D = np.zeros_like(self.gamma_niondot_II_index2D)
             self.gamma2_niondot_III_index2D = np.zeros_like(self.gamma2_niondot_II_index2D)
            

From c8b3b2be3f267b9d417cc53887ae93a1b3bd8d44 Mon Sep 17 00:00:00 2001
From: Yonatan Sklansky <yonatan.sklansky@gmail.com>
Date: Thu, 3 Apr 2025 17:04:38 -0500
Subject: [PATCH 04/37] Changes to BMF, setup.py, z21utilities, and added
 ionized maps class

---
 setup.py                |   3 +
 zeus21/maps.py          | 318 +++++++++++++++++++++++++---------------
 zeus21/reionization.py  |  77 +++++-----
 zeus21/sfrd.py          |   2 +-
 zeus21/z21_utilities.py |   5 +-
 5 files changed, 247 insertions(+), 158 deletions(-)

diff --git a/setup.py b/setup.py
index 898653c..e3a4154 100755
--- a/setup.py
+++ b/setup.py
@@ -20,5 +20,8 @@
            "classy",
            "numexpr",
            "astropy",
+           "powerbox",
+           "pyfftw",
+           "tqdm"
        ],
 )
diff --git a/zeus21/maps.py b/zeus21/maps.py
index b602866..e94ccf2 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -112,9 +112,9 @@ def __init__(self, T21_coefficients, Power_Spectrum, z, Lbox=600, Nbox=200, KIND
             print('ERROR, KIND not implemented yet!')
 
 
-def make_ion_fields(CosmoParams, CoeffStructure, ClassyCosmo, CorrFClass, BMF, input_z, boxlength=300., ncells=300, seed=1234, r_precision=1., timer=True, logd = False, barrier = None, spherize=False, FLAG_return_densities = 0):
+class reionization_maps:
     """
-    Generates a 3D map of ionized fields and ionized fraction of hydrogen.
+    Generates 3D maps of the reionization fields.
     
     Uses a density threshold barrier determined from a converged bubble mass function. With default parameters, the code takes about 20 minutes on laptop to run.
 
@@ -132,7 +132,7 @@ def make_ion_fields(CosmoParams, CoeffStructure, ClassyCosmo, CorrFClass, BMF, i
         Computes bubble mass functions and barriers.
     input_z: 1D np.array
         The redshifts at which to compute output maps. Narrowed down later to select available redshifts from CoeffStructure.zintegral.
-    boxlength: float
+    input_boxlength: float
         Comoving physical side length of the box. Default is 300 cMpc.
     ncells: int
         Number of cells on a side. Default is 300 cells.
@@ -140,125 +140,207 @@ def make_ion_fields(CosmoParams, CoeffStructure, ClassyCosmo, CorrFClass, BMF, i
         Sets the predetermined generation of maps. Default is 1234.
     r_precision: float
         Allows to change the steps of the radii for faster computation. Default (and max) is 1, lower values make the computation faster at the cost of accuracy.
-    timer: bool
-        Whether to print the time elapsed along the process. Default is True.
-    logd: bool
-        Whether to use lognormal (True) or Gaussian (False) density fields. Default is False.
     barrier: function
-        Input density barrier to be used as the threshold for map generation. Takes z index //\\Change//\\ (relative to CoeffStructure.zintegral) as input and returns np.array of shape of radii. Default is None.
-    spherize: bool
-        Whether to flag spheres around ionized cells (True) instead of only central pixel flagging (False). Default is False. Central pixel flagging is generally more consistent with the bubble mass function than spherizing.
-    FLAG_return_densities: int
-        Options: (0, 1, 2). Default is 0.
-            0: returns only the ionized fields and ionized fractions.
-            1: returns only the ionized fields, ionized fractions, and density field at the last redshift.
-            2: returns the ionized fields, ionized fractions, density field at the last redshift, and the smoothed density fields at each scale at the last redshift.
-
-    Outputs
+        Input density barrier to be used as the threshold for map generation. Takes z index //\\Change//\\ (relative to CoeffStructure.zintegral) as input and returns np.array of shape. Default is None.
+    PRINT_TIMER: bool
+        Whether to print the time elapsed along the process. Default is True.
+    LOGNORMAL_DENSITY: bool
+        Whether to use lognormal (True) or Gaussian (False) density fields. Default is False.of radii. Default is False.
+    COMPUTE_DENSITY_AT_ALLZ: bool
+        Whether to output the density field at all redshifts. If False, only the density at the lower input redshift is computed. If True, the computation time and memory usage dramatically increases. Default is False.
+    SPHERIZE: bool
+        Whether to flag spheres around ionized cells (True) instead of only central pixel flagging (False). Default is False. Central pixel flagging is generally more consistent with the bubble mass function than spherizing. Default is False.
+    COMPUTE_MASSWEIGHTED_IONFRAC: bool
+        Whether to compute the mass weighted ionized fraction. If True, COMPUTE_DENSITY_AT_ALLZ will be forced to True, thus increasing dramatically computation time. Default is False.
+    lowres_massweighting: int
+        Compute the mass-weighted ionized fraction more efficiently by using lower resolution density and ionized fields. Has to be >=1 and an integer. Default is 1.
+
+    Attributes
     ----------
-    ion_fields: 4D np.array
-        Resultant 3D ionized field maps at each redshift. The first dimension is redshifts, and the three other dimensions are spatial.
+    dx: float
+        Cell resolution of a side of the boxes.
+    z: 1D np.array
+        Redshifts at which the output maps are computed. Selected to be the closest to the input redshifts from the available ones in zeus21.
+    r: 1D np.array
+        Radii at which the density field is smoothed. Selected using r_precision from the available ones in zeus21.
+    z_of_density: float
+        Redshift at which the density is computed.
+    density: 3D np.array
+        Overdensity field at the lowest redshift asked by the user.
+    density_allz: 4D np.array
+        Overdensity field at all the redshifts asked by the user. First dimension correponds to redshifts. Only computed if COMPUTE_DENSITY_AT_ALLZ is True.
+    ion_field_allz: 4D np.array
+        Ionized fraction field at all the redshifts asked by the user. First dimension correponds to redshifts.
     ion_frac: 1D np.array
-        Ionized fraction at each redshift in ion_fields.
-    Optional:
-        density_field: 3D np.array
-            The field of densities at the last redshift.
-        smoothed_density_fields: 4D np.array
-            The field of densities at the last redshift, smoothed over each radius scale. The first dimension is the smoothing scales, the other three are spatial.
-    
+        Volume weighted ionized fraction at all the redshifts asked by the user.
+    ion_frac_massweighted: 1D np.array
+        Mass weighted ionized fraction at all the redshifts asked by the user. Only computed if COMPUTE_MASSWEIGHTED_IONFRAC is True.
     """
     
-    #Measure time elapsed from start
-    start_time = time.time()
+    def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, input_z, 
+                 input_boxlength=300., ncells=300, seed=1234, r_precision=1., barrier=None, 
+                 PRINT_TIMER=True, ENFORCE_BMF_SCALE=True, 
+                 LOGNORMAL_DENSITY=False, COMPUTE_DENSITY_AT_ALLZ=False, SPHERIZE=False, 
+                 COMPUTE_MASSWEIGHTED_IONFRAC=False, lowres_massweighting=1):
+        #Measure time elapsed from start
+        self._start_time = time.time()
+        
+        ### boxes parameters
+        self.input_z = input_z
+        self.ENFORCE_BMF_SCALE = ENFORCE_BMF_SCALE
+        self.ncells = ncells
+        self.boxlength = input_boxlength
+        self.dx = self.boxlength/self.ncells
+        if self.ENFORCE_BMF_SCALE:
+            self.dx = CoeffStructure.Rtabsmoo[0] * (4*np.pi/3)**(1/3) #change to constant
+            self.boxlength = self.dx * self.ncells
+            print(f"WARNING: with ENFORCE_BMF_SCALE=True, boxlength is changed to match BMF scale. Current value: {self.boxlength:.2f} cMpc.")
+        self.seed = seed
+
+        ### FLAGS
+        self.PRINT_TIMER = PRINT_TIMER
+        self.LOGNORMAL_DENSITY = LOGNORMAL_DENSITY
+        self.COMPUTE_DENSITY_AT_ALLZ = COMPUTE_DENSITY_AT_ALLZ
+        self.SPHERIZE = SPHERIZE
+        self.COMPUTE_MASSWEIGHTED_IONFRAC = COMPUTE_MASSWEIGHTED_IONFRAC
+        if self.COMPUTE_MASSWEIGHTED_IONFRAC:
+            self.COMPUTE_DENSITY_AT_ALLZ = True
+
+        ### selecting redshifts and radii from available redshifts
+        # redshifts
+        self._z_idx = z21_utilities.find_nearest_idx(CoeffStructure.zintegral, self.input_z)
+        self.z = CoeffStructure.zintegral[self._z_idx]
+        # radii
+        self.r_precision = r_precision
+        self._r_idx = np.linspace(0, len(CoeffStructure.Rtabsmoo)-1, int(len(CoeffStructure.Rtabsmoo)*self.r_precision), dtype=int)
+        self.r = CoeffStructure.Rtabsmoo[self._r_idx]
+
+        ### generating the density field at the closest redshift to the lower one inputed
+        self.z_of_density = self.z[0]
+        self.density = self.generate_density(ClassyCosmo, CorrFClass)
+        self.density_allz = np.empty((len(self.z), self.ncells, self.ncells, self.ncells))
+        if self.COMPUTE_DENSITY_AT_ALLZ:
+            self.generate_density_allz(CosmoParams)
+
+        ### smoothing the density field
+        self._k = self.compute_k()
+        self.density_smoothed_allr = self.smooth_density()
+
+        ### generating the ionized field, and computing the ionized fraction
+        self.barrier = barrier
+        if self.barrier is None:
+            self.barrier = BMF.barrier
+        self.ion_field_allz, self.ion_frac = self.generate_xHII(CosmoParams, CoeffStructure, BMF)
+
+        ### computing the mass weighted ionized fraction
+        self.lowres_massweighting = lowres_massweighting
+        self.ion_frac_massweighted = np.empty(len(self.z))
+        if self.COMPUTE_MASSWEIGHTED_IONFRAC:
+            self.compute_massweighted_xHII(CosmoParams, self.lowres_massweighting)
+        
+        if self.PRINT_TIMER:
+            print_timer(self._start_time, text_before="Total computation time: ")
+        
+
+    def generate_density(self, ClassyCosmo, CorrFClass):
+        if self.PRINT_TIMER:
+            start_time = time.time()
+            print("Generating density field...")
+        #Generating matter power spectrum at the lowest redshift
+        klist = CorrFClass._klistCF
+        pk_matter = np.zeros_like(klist)
+        for i, k in enumerate(klist):
+            pk_matter[i] = ClassyCosmo.pk(k, self.z_of_density)
+        pk_spl = spline(np.log(klist), np.log(pk_matter))
     
-    if timer:
-        z21_utilities.print_timer(start_time, 'Making density field...')
-
-    #selecting redshifts and radii from available redshifts
-    zlist = CoeffStructure.zintegral
-    z_idx = z21_utilities.find_nearest_idx(zlist, input_z)
-    z = zlist[z_idx]
-
-    Rs = CoeffStructure.Rtabsmoo
-    r_idx = np.linspace(0,len(Rs)-1,int(len(Rs)*r_precision),dtype=int)
-    r = Rs[r_idx]
-    dx = boxlength/ncells
-
-    #Generating matter power spectrum at z=5
-    klist = CorrFClass._klistCF
-    pk_matter = np.zeros_like(klist)
-    for i, k in enumerate(klist):
-        pk_matter[i] = ClassyCosmo.pk(k, z[0])
-    pk_spl = spline(np.log(klist), np.log(pk_matter))
-
-    #generating density map
-    if logd:
-        pb = pbox.LogNormalPowerBox(N=ncells, dim=3, pk=(lambda k: np.exp(pk_spl(np.log(k)))), boxlength=boxlength, seed=seed)
-    else:
-        pb = pbox.PowerBox(N=ncells, dim=3, pk=(lambda k: np.exp(pk_spl(np.log(k)))), boxlength=boxlength, seed=seed)
-    density_field = pb.delta_x()
-
-    if timer:
-        z21_utilities.print_timer(start_time, 'Creating smoothing function...')
-
-    #comment
-    klistfftx = np.fft.fftfreq(density_field.shape[0],dx)*2*np.pi
-    klist3Dfft = np.sqrt(np.sum(np.meshgrid(klistfftx**2, klistfftx**2, klistfftx**2, indexing='ij'), axis=0))
-    density_fft = np.fft.fftn(density_field)
-
-    if timer:
-        z21_utilities.print_timer(start_time, 'Smoothing...')
-
-    #comment
-    smooth_density_fields = np.array([z21_utilities.tophat_smooth(rr, klist3Dfft, density_fft) for rr in r])
-
-    if timer:
-        z21_utilities.print_timer(start_time, 'Creating ionized field...')
-
-    if barrier is None:
-        barrier = BMF.barrier
-    ion_fields = []
-    ion_frac = np.zeros(len(z))
-    for i in trange(len(z)):
-        curr_z_idx = z_idx[i]
-        ion_field = ionize(CosmoParams, zlist, Rs, curr_z_idx, smooth_density_fields, barrier, r_idx, klist3Dfft, spherize)
-        ion_fields.append(ion_field)
-        ion_frac[i] = np.sum(ion_field)/ncells**3
+        #generating density map
+        if self.LOGNORMAL_DENSITY:
+            pb = pbox.LogNormalPowerBox(N=self.ncells, dim=3, pk=(lambda k: np.exp(pk_spl(np.log(k)))), boxlength=self.boxlength, seed=self.seed)
+        else:
+            pb = pbox.PowerBox(N=self.ncells, dim=3, pk=(lambda k: np.exp(pk_spl(np.log(k)))), boxlength=self.boxlength, seed=self.seed)
+        density_field = pb.delta_x()
+        if self.PRINT_TIMER:
+            print_timer(start_time, text_before="    done in ")
+        return density_field
+
+    def generate_density_allz(self, CosmoParams):
+        if self.PRINT_TIMER:
+            start_time = time.time()
+            print('Evolving density field...')
+        Dg = CosmoParams.growthint(self.z)
+        growthfactor_ratio = (Dg/Dg[0])[:, np.newaxis, np.newaxis, np.newaxis]
+        density_lastz = np.copy(self.density)
+        self.density_allz = density_lastz[np.newaxis]*growthfactor_ratio
+        if self.PRINT_TIMER:
+            print_timer(start_time, text_before="    done in ")
+        return self.density_allz
+
+    def compute_k(self):
+        klistfftx = np.fft.fftfreq(self.ncells,self.dx)*2*np.pi
+        k = np.sqrt(np.sum(np.meshgrid(klistfftx**2, klistfftx**2, klistfftx**2, indexing='ij'), axis=0))
+        return k
+
+    def smooth_density(self):
+        if self.PRINT_TIMER:
+            start_time = time.time()
+            print("Smoothing density field...")
+        density_fft = np.fft.fftn(self.density)
+        density_smoothed_allr = np.array([z21_utilities.tophat_smooth(rr, self._k, density_fft) for rr in self.r])
+        if self.PRINT_TIMER:
+            print_timer(start_time, text_before="    done in ")
+        return density_smoothed_allr
+
+    def generate_xHII(self, CosmoParams, CoeffStructure, BMF):
+        if self.PRINT_TIMER:
+            start_time = time.time()
+            print("Generating ionized field...")
+        ion_field_allz = np.zeros((len(self.z),self.ncells,self.ncells,self.ncells))
+        ion_frac = np.zeros(len(self.z))
+        for i in trange(len(self.z)):
+            curr_z_idx = self._z_idx[i]
+            ion_field = self.ionize(CosmoParams, CoeffStructure, curr_z_idx)
+            ion_field_allz[i] = ion_field
+            ion_frac[i] = np.sum(ion_field)/(self.ncells**3)
+        if self.PRINT_TIMER:
+            print_timer(start_time, text_before="    done in ")
+        return ion_field_allz, ion_frac
+
+    def ionize(self,CosmoParams, CoeffStructure, curr_z_idx):
+        zlist = CoeffStructure.zintegral
+        Rs = CoeffStructure.Rtabsmoo
+        
+        Dg0 = CosmoParams.growthint(zlist[0])
+        Dg = CosmoParams.growthint(zlist[curr_z_idx])
+        if not self.SPHERIZE:
+            ion_field = np.any(self.density_smoothed_allr > (Dg0/Dg)*self.barrier[curr_z_idx, self._r_idx][:, None, None, None], axis=0)
+        else:
+            ion_field_Rs = np.zeros((len(self._r_idx),self.ncells,self.ncells,self.ncells))
+            for j in range(len(self._r_idx)):
+                curr_R = self._r_idx[j]
+                ion_field_oneR = self.density_smoothed_allr[j] > (Dg0/Dg)*self.barrier[curr_z_idx, curr_R]
+                ion_field_oneR_fft = np.fft.fftn(ion_field_oneR)
+                cutoff = 1/(4/3*np.pi*Rs[curr_R]**3)/2*(1+self.barrier[curr_z_idx, curr_R]) #comment
+                ion_spheres = z21_utilities.tophat_smooth(Rs[curr_R]/(1+self.barrier[curr_z_idx, curr_R])**(1/3), self._k, ion_field_oneR_fft) > cutoff
+                ion_field_Rs[j] = ion_spheres
+                    
+            ion_field = np.any(ion_field_Rs, axis=0)
+        return ion_field
     
-    ion_fields = np.array(ion_fields)
-
-    print('Done!')
-    print('\nTotal time:')
-    z21_utilities.print_timer(start_time)
-
-    if FLAG_return_densities == 0:
-        return ion_fields, ion_frac
-
-    elif FLAG_return_densities == 1:
-        return ion_fields, ion_frac, density_field
-
-    elif FLAG_return_densities == 2:
-        return ion_fields, ion_frac, density_field, smooth_density_fields
-
-    else:
-        print('WARNING: FLAG_return_densities is not set to (0, 1, or 2). Defaulting to 0.')
-        return ion_fields, ion_frac
-
-#look over this again
-def ionize(CosmoParams, zlist, Rs, curr_z_idx, smooth_density_fields, barrier, r_idx, klist3Dfft, spherize):
-    Dg0 = CosmoParams.growthint(zlist[0])
-    Dg = CosmoParams.growthint(zlist[curr_z_idx])
-    if not spherize:
-        ion_field = np.any(smooth_density_fields > (Dg0/Dg)*barrier[curr_z_idx, r_idx][:, None, None, None], axis=0)
-    else:
-        ion_field_Rs = []
-        for j in range(len(r_idx)):
-            curr_R = r_idx[j]
-            ion_field_oneR = smooth_density_fields[j] > (Dg0/Dg)*barrier[curr_z_idx, curr_R]
-            ionk = np.fft.fftn(ion_field_oneR)
-            cutoff = 1/(4/3*np.pi*Rs[curr_R]**3)/2*(1+barrier[curr_z_idx, curr_R]) #comment
-            ion_spheres = z21_utilities.tophat_smooth(Rs[curr_R]/(1+barrier[curr_z_idx, curr_R])**(1/3), klist3Dfft, ionk) > cutoff
-            ion_field_Rs.append(ion_spheres)
-                
-        ion_field = np.any(ion_field_Rs, axis=0)
-    return ion_field
\ No newline at end of file
+    def compute_massweighted_xHII(self, CosmoParams, lowres_massweighting=1):
+        if np.sum(self.density_allz[0, 0, 0]) == 0.:
+            self.generate_density_allz(CosmoParams)
+        self.lowres_massweighting = lowres_massweighting
+        if self.lowres_massweighting < 1:
+            raise Exeption('lowres_massweighting should be >=1.')
+        if not isinstance(self.lowres_massweighting, (int, np.int32, np.int64)):
+            raise Exeption('lowres_massweighting should be an integer.')
+        d_allz = self.density_allz[:, ::self.lowres_massweighting, ::self.lowres_massweighting, ::self.lowres_massweighting]
+        ion_allz = self.ion_field_allz[:, ::self.lowres_massweighting, ::self.lowres_massweighting, ::self.lowres_massweighting]
+        if self.PRINT_TIMER:
+            start_time = time.time()
+            print("Computing mass weighted ionized fraction...")
+        self.ion_frac_massweighted = np.average((1+d_allz) * ion_allz, axis=(1, 2, 3))
+        if self.PRINT_TIMER:
+            print_timer(start_time, text_before="    done in ")
+        return self.ion_frac_massweighted
+
diff --git a/zeus21/reionization.py b/zeus21/reionization.py
index 90917c3..6be8ff7 100644
--- a/zeus21/reionization.py
+++ b/zeus21/reionization.py
@@ -24,7 +24,7 @@ class BMF:
     
     """
     
-    def __init__(self, UserParams, CoeffStructure, HMFintclass, CosmoParams, AstroParams, R_linear_sigma_fit_input=10, FLAG_converge=True, max_iter=10, ZMAX_REION = 30):
+    def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, R_linear_sigma_fit_input=10, FLAG_converge=True, max_iter=10, ZMAX_REION = 30):
 
         self.ZMAX_REION = ZMAX_REION #max redshift up to which we calculate reionization observables
         self.zlist = CoeffStructure.zintegral
@@ -36,7 +36,6 @@ def __init__(self, UserParams, CoeffStructure, HMFintclass, CosmoParams, AstroPa
         self.sigma = CoeffStructure.sigmaofRtab
 
         self.Hz = cosmology.Hubinvyr(CosmoParams, self.zlist)
-        #self.Hzconv = (u.km/u.s/u.Mpc).to(1/u.yr) #multiply this times Hubble parameter to convert from km/s/Mpc to 1/yr
         self.trec0 = 1/(constants.alphaB * cosmology.n_H(CosmoParams,0) * AstroParams._clumping) #seconds
         self.trec = self.trec0/(1+self.zlist)**3/constants.yrTos #years
         self.niondot_avg = CoeffStructure.niondot_avg_II
@@ -47,7 +46,7 @@ def __init__(self, UserParams, CoeffStructure, HMFintclass, CosmoParams, AstroPa
 
         zr_mesh = np.meshgrid(np.arange(len(self.Rs)), np.arange(len(self.zlist)))
         self.nion_norm = self.nion_normalization(zr_mesh[1], zr_mesh[0])
-        self.barrier = self.compute_barrier(self.ion_frac)
+        self.barrier = self.compute_barrier(CosmoParams, self.ion_frac)
         self.barrier_initial = np.copy(self.barrier)
 
         self.R_linear_sigma_fit_idx = z21_utilities.find_nearest_idx(self.Rs, R_linear_sigma_fit_input)
@@ -56,14 +55,15 @@ def __init__(self, UserParams, CoeffStructure, HMFintclass, CosmoParams, AstroPa
         self.BMF = np.array([self.VRdn_dR(HMFintclass, np.arange(len(self.Rs)), i) for i in range(len(self.zlist))]) #initial bubble mass function
         self.BMF_initial = np.copy(self.BMF)
         self.ion_frac = np.nan_to_num([np.trapz(self.BMF[i], np.log(self.Rs)) for i in range(len(self.zlist))]) #ion_frac by integrating the BMF
-
+        self.ion_frac[self.barrier[:, -1]<=0] = 1
+        
         if FLAG_converge:
-            self.converge_BMF(HMFintclass, self.ion_frac, max_iter=max_iter)
+            self.converge_BMF(CosmoParams, HMFintclass, self.ion_frac, max_iter=max_iter)
         #two functions: compute BMF and iterate
         
 
 
-    def compute_barrier(self, ion_frac):
+    def compute_barrier(self, CosmoParams, ion_frac):
         """
         Computes the density barrier threshold for ionization.
         
@@ -71,6 +71,8 @@ def compute_barrier(self, ion_frac):
 
         Parameters
         ----------
+        CosmoParams: zeus21.Cosmo_Parameters class
+            Stores cosmology.
         ion_frac: 1D np.array
             The ionized fractions to be used to compute the number of recombinations. 
 
@@ -80,21 +82,20 @@ def compute_barrier(self, ion_frac):
             The resultant density threshold array. First dimension is each redshift, second dimension is each radius scale.
         """
         barrier = np.zeros((len(self.zlist), len(self.Rs)))
-        ds_array = np.linspace(-1, 2, 21)
-       
+        ds_array = np.linspace(-1, 5, 101)
         
         for ir in range(len(self.Rs)):
             #Compute nion_values and nrec_values for this 'ir'
-            nion_values = self.nion_delta_r_int(ds_array, ir)  #Shape (nd, nz)
-            nrec_values = self.nrec(ds_array, ion_frac)             #Shape (nd, nz)
+            nion_values = self.nion_delta_r_int(CosmoParams, ds_array, ir)  #Shape (nd, nz)
+            nrec_values = self.nrec(CosmoParams, ds_array, ion_frac)             #Shape (nd, nz)
         
-            total_values = np.log10(nion_values / (1 + nrec_values))   #taking difference in logspace to find zero-crossing 
+            total_values = np.log10(nion_values / (1 + nrec_values) + 1e-10)   #taking difference in logspace to find zero-crossing 
         
             #Loop over redshift indices
             for i in range(len(self.zlist)):
                 y_values = total_values[:, i]  #Shape (nd,)
         
-                # Find zero crossings
+                #Find zero crossings
                 sign_change = np.diff(np.sign(y_values))
                 idx = np.where(sign_change)[0]
                 if idx.size > 0:
@@ -107,22 +108,22 @@ def compute_barrier(self, ion_frac):
                     barrier[i, ir] = x_intersect[0]  #Assuming we take the first crossing
                 else:
                     barrier[i, ir] = np.nan #Never crosses
-
+        barrier = barrier * (CosmoParams.growthint(self.zlist)/CosmoParams.growthint(self.zlist[0]))[:, np.newaxis] #scale barrier with growth factor
         barrier[self.zlist > self.ZMAX_REION] = 100 #sets density to an unreachable barrier, as if reionization isn't happening
         return barrier
 
-
-
     #normalizing the nion/sfrd model
     def nion_normalization(self, z, R):
         return 1/np.sqrt(1-2*self.gamma2[z, R]*self.sigma[z, R]**2)*np.exp(self.gamma[z, R]**2 * self.sigma[z, R]**2 / (2-4*self.gamma2[z, R]*self.sigma[z, R]**2))
 
-    def nrec(self, d_array, ion_frac):
+    def nrec(self, CosmoParams, d_array, ion_frac):
         """
         Vectorized computation of nrec over an array of overdensities d_array.
 
         Parameters
         ----------
+        CosmoParams: zeus21.Cosmo_Parameters class
+            Stores cosmology.
         d_array: 1D np.array
             A list of sample overdensity values to evaluate nrec over.
         ion_frac: 1D np.array
@@ -142,21 +143,23 @@ def nrec(self, d_array, ion_frac):
     
         denom = -1 / (1 + z_rev) / Hz_rev / trec_rev
         integrand_base = denom * ion_frac_rev 
-    
-        # Broadcast over d_array
-        integrand = integrand_base[np.newaxis, :] * (1 + d_array[:, np.newaxis])
-        cumulative_integral = np.concatenate(
-            (np.zeros((len(d_array), 1)), cumulative_trapezoid(integrand, x=z_rev, axis=1)), axis=1
-        )
-        nrecs = cumulative_integral[:, ::-1]  # Reverse to match self.zlist order
+        Dg = CosmoParams.growthint(z_rev) #growth factor
+
+        nrecs = cumulative_trapezoid(integrand_base*(1+d_array[:, np.newaxis]*Dg/Dg[-1]), x=z_rev, initial=0) #(1+delta) rather than (1+delta)^2 because nrec and nion are per hydrogen atom 
+        
+        #TODO: nonlinear recombinations/higher order
+
+        nrecs = nrecs[:, ::-1]  # Reverse to match self.zlist order
         return nrecs
     
-    def niondot_delta_r(self, d_array, ir):
+    def niondot_delta_r(self, CosmoParams, d_array, ir):
         """
         Compute niondot over an array of overdensities d_array for a given ir.
 
         Parameters
         ----------
+        CosmoParams: zeus21.Cosmo_Parameters class
+            Stores cosmology.
         d_array: 1D np.array
             A list of sample overdensity values to evaluate niondot over.
         ir: int
@@ -168,7 +171,7 @@ def niondot_delta_r(self, d_array, ir):
             The rates of ionizing photon production. The first dimension is densities, the second dimension is redshifts.
         """
         
-        d_array = d_array[:, np.newaxis]
+        d_array = d_array[:, np.newaxis] * CosmoParams.growthint(self.zlist)[np.newaxis, :] / CosmoParams.growthint(self.zlist[0])
     
         gamma_ir = self.gamma[:, ir]   
         gamma2_ir = self.gamma2[:, ir]  
@@ -179,12 +182,14 @@ def niondot_delta_r(self, d_array, ir):
         
         return niondot
     
-    def nion_delta_r_int(self, d_array, ir):
+    def nion_delta_r_int(self, CosmoParams, d_array, ir):
         """
         Vectorized computation of nion over an array of overdensities d_array for a given ir.
 
         Parameters
         ----------
+        CosmoParams: zeus21.Cosmo_Parameters class
+            Stores cosmology.
         d_array: 1D np.array
             A list of sample overdensity values to evaluate niondot over.
         ir: int
@@ -200,13 +205,14 @@ def nion_delta_r_int(self, d_array, ir):
         z_rev = self.zlist[::-1]
         Hz_rev = self.Hz[::-1]
     
-        niondot_values = self.niondot_delta_r(d_array, ir)
+        niondot_values = self.niondot_delta_r(CosmoParams, d_array, ir)
     
         integrand = -1 / (1 + z_rev) / Hz_rev * niondot_values[:, ::-1]
-        cumulative_integral = np.concatenate(
-            (np.zeros((len(d_array), 1)), cumulative_trapezoid(integrand, x=z_rev, axis=1)), axis=1
-        )
-        nion = cumulative_integral[:, ::-1]  # Reverse back to match self.zlist order
+        #cumulative_integral = np.concatenate(
+        #    (np.zeros((len(d_array), 1)), cumulative_trapezoid(integrand, x=z_rev, axis=1)), axis=1
+        #)
+        nion = cumulative_trapezoid(integrand, x=z_rev, initial=0)[:, ::-1] #reverse back to match self.zlist order
+        #nion = cumulative_integral[:, ::-1]  # Reverse back to match self.zlist order
         return nion
 
     #calculating ionized fraction, put outside the class
@@ -261,19 +267,18 @@ def Rdn_dR(self, ir, iz):
         z = self.zlist[iz]
         return self.VRdn_dR(HMFintclass, ir, iz)*3/(4*np.pi*R**3)
 
-    def converge_BMF(self, HMFintclass, ion_frac_input, max_iter):
+    def converge_BMF(self, CosmoParams, HMFintclass, ion_frac_input, max_iter):
         self.ion_frac = ion_frac_input
         for j in trange(max_iter):
             ion_frac_prev = np.copy(self.ion_frac)
             
-            self.barrier = self.compute_barrier(self.ion_frac)
+            self.barrier = self.compute_barrier(CosmoParams, self.ion_frac)
             self.BMF = np.array([self.VRdn_dR(HMFintclass, np.arange(len(self.Rs)), i) for i in range(len(self.zlist))])
             self.ion_frac = np.nan_to_num([np.trapz(self.BMF[i], np.log(self.Rs)) for i in range(len(self.zlist))])
+            self.ion_frac[self.barrier[:, -1]<=0] = 1
 
             if np.allclose(ion_frac_prev, self.ion_frac):
                 print(f'SUCCESS: BMF converged in {j} iterations.')
                 return 
             
-        print(f"WARNING: BMF didn't converge within {max_iter} iterations.")
-            
-        
\ No newline at end of file
+        print(f"WARNING: BMF didn't converge within {max_iter} iterations.")
\ No newline at end of file
diff --git a/zeus21/sfrd.py b/zeus21/sfrd.py
index 101f538..c7a2975 100644
--- a/zeus21/sfrd.py
+++ b/zeus21/sfrd.py
@@ -114,7 +114,7 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
 
         #and EPS factors
         Nsigmad = 1.0 #how many sigmas we explore
-        Nds = 3 #how many deltas
+        Nds = 3 #how many deltas; must be an odd integer
         deltatab_norm = np.linspace(-Nsigmad,Nsigmad,Nds)
 
         #initialize Xrays
diff --git a/zeus21/z21_utilities.py b/zeus21/z21_utilities.py
index e05f231..9290700 100644
--- a/zeus21/z21_utilities.py
+++ b/zeus21/z21_utilities.py
@@ -37,12 +37,11 @@ def find_nearest_idx(array, values):
         idx.append((np.abs(array - values[i])).argmin())
     return np.unique(idx)
 
-def print_timer(start_time, text=""):
+def print_timer(start_time, text_before="", text_after=""):
     elapsed_time = time.time() - start_time
     mins = int(elapsed_time//60)
     secs = int(elapsed_time - mins*60)
-    print(f"\n{mins}min {secs}s\n")
-    print(text)
+    print(f"{text_before}{mins}min {secs}s{text_after}")
 
 def v2r(v):
     return (3/4/np.pi * v)**(1/3)

From 7b498f43730c5c134bd319ea4f637076f3bb5aef Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Thu, 24 Apr 2025 12:27:17 -0500
Subject: [PATCH 05/37] Fixed minor issues with reionization_maps.

---
 zeus21/maps.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index e94ccf2..57f7ce7 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -239,7 +239,7 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
             self.compute_massweighted_xHII(CosmoParams, self.lowres_massweighting)
         
         if self.PRINT_TIMER:
-            print_timer(self._start_time, text_before="Total computation time: ")
+            z21_utilities.print_timer(self._start_time, text_before="Total computation time: ")
         
 
     def generate_density(self, ClassyCosmo, CorrFClass):
@@ -260,7 +260,7 @@ def generate_density(self, ClassyCosmo, CorrFClass):
             pb = pbox.PowerBox(N=self.ncells, dim=3, pk=(lambda k: np.exp(pk_spl(np.log(k)))), boxlength=self.boxlength, seed=self.seed)
         density_field = pb.delta_x()
         if self.PRINT_TIMER:
-            print_timer(start_time, text_before="    done in ")
+            z21_utilities.print_timer(start_time, text_before="    done in ")
         return density_field
 
     def generate_density_allz(self, CosmoParams):
@@ -272,7 +272,7 @@ def generate_density_allz(self, CosmoParams):
         density_lastz = np.copy(self.density)
         self.density_allz = density_lastz[np.newaxis]*growthfactor_ratio
         if self.PRINT_TIMER:
-            print_timer(start_time, text_before="    done in ")
+            z21_utilities.print_timer(start_time, text_before="    done in ")
         return self.density_allz
 
     def compute_k(self):
@@ -287,7 +287,7 @@ def smooth_density(self):
         density_fft = np.fft.fftn(self.density)
         density_smoothed_allr = np.array([z21_utilities.tophat_smooth(rr, self._k, density_fft) for rr in self.r])
         if self.PRINT_TIMER:
-            print_timer(start_time, text_before="    done in ")
+            z21_utilities.print_timer(start_time, text_before="    done in ")
         return density_smoothed_allr
 
     def generate_xHII(self, CosmoParams, CoeffStructure, BMF):
@@ -302,7 +302,7 @@ def generate_xHII(self, CosmoParams, CoeffStructure, BMF):
             ion_field_allz[i] = ion_field
             ion_frac[i] = np.sum(ion_field)/(self.ncells**3)
         if self.PRINT_TIMER:
-            print_timer(start_time, text_before="    done in ")
+            z21_utilities.print_timer(start_time, text_before="    done in ")
         return ion_field_allz, ion_frac
 
     def ionize(self,CosmoParams, CoeffStructure, curr_z_idx):
@@ -341,6 +341,6 @@ def compute_massweighted_xHII(self, CosmoParams, lowres_massweighting=1):
             print("Computing mass weighted ionized fraction...")
         self.ion_frac_massweighted = np.average((1+d_allz) * ion_allz, axis=(1, 2, 3))
         if self.PRINT_TIMER:
-            print_timer(start_time, text_before="    done in ")
+            z21_utilities.print_timer(start_time, text_before="    done in ")
         return self.ion_frac_massweighted
 

From cb721cafbebea3e5855c6d872026a8cc9ffc83b7 Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Thu, 15 May 2025 13:08:31 -0500
Subject: [PATCH 06/37] Added zreion/treion maps computation.

---
 zeus21/cosmology.py | 32 ++++++++++++++++++++++++++++++++
 zeus21/maps.py      | 42 ++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/zeus21/cosmology.py b/zeus21/cosmology.py
index c2500b7..5bec74a 100644
--- a/zeus21/cosmology.py
+++ b/zeus21/cosmology.py
@@ -115,6 +115,38 @@ def runclass(CosmologyIn):
     
     return ClassCosmo
 
+def time_at_redshift(ClassyCosmo,z):
+    """
+    Returns the age of the Universe (in Gyrs) corresponding to a given redshift.
+
+    Parameters
+    ----------
+    ClassyCosmo: zeus21.runclass class
+        Sets up Class cosmology.
+    z: float
+        Redshift.
+    """
+    background = ClassyCosmo.get_background()
+    classy_t, classy_z = background['proper time [Gyr]'], background['z']
+    classy_tinterp = interpolate.interp1d(classy_z, classy_t)
+    return classy_tinterp(z)
+
+def redshift_at_time(ClassyCosmo,t):
+    """
+    Returns the redshift corresponding to a given age of the Universe (in Gyrs).
+
+    Parameters
+    ----------
+    ClassyCosmo: zeus21.runclass class
+        Sets up Class cosmology.
+    t: float
+        Age in Gyrs.
+    """
+    background = ClassyCosmo.get_background()
+    classy_t, classy_z = background['proper time [Gyr]'], background['z']
+    classy_tinterp = interpolate.interp1d(classy_t, classy_z)
+    return classy_tinterp(t)
+
 def Hub(Cosmo_Parameters, z):
 #Hubble(z) in km/s/Mpc
     return Cosmo_Parameters.h_fid * 100 * np.sqrt(Cosmo_Parameters.OmegaM * pow(1+z,3.)+Cosmo_Parameters.OmegaR * pow(1+z,4.)+Cosmo_Parameters.OmegaL)
diff --git a/zeus21/maps.py b/zeus21/maps.py
index 57f7ce7..807a66c 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -154,6 +154,8 @@ class reionization_maps:
         Whether to compute the mass weighted ionized fraction. If True, COMPUTE_DENSITY_AT_ALLZ will be forced to True, thus increasing dramatically computation time. Default is False.
     lowres_massweighting: int
         Compute the mass-weighted ionized fraction more efficiently by using lower resolution density and ionized fields. Has to be >=1 and an integer. Default is 1.
+    COMPUTE_ZREION: bool
+        Whether to compute the reionization redshift and time fields. Default is False.
 
     Attributes
     ----------
@@ -175,13 +177,18 @@ class reionization_maps:
         Volume weighted ionized fraction at all the redshifts asked by the user.
     ion_frac_massweighted: 1D np.array
         Mass weighted ionized fraction at all the redshifts asked by the user. Only computed if COMPUTE_MASSWEIGHTED_IONFRAC is True.
+    zreion: 3D np.array
+        Reionization redshift field.
+    treion: 3D np.array
+        Reionization time field.    
     """
     
     def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, input_z, 
                  input_boxlength=300., ncells=300, seed=1234, r_precision=1., barrier=None, 
                  PRINT_TIMER=True, ENFORCE_BMF_SCALE=True, 
                  LOGNORMAL_DENSITY=False, COMPUTE_DENSITY_AT_ALLZ=False, SPHERIZE=False, 
-                 COMPUTE_MASSWEIGHTED_IONFRAC=False, lowres_massweighting=1):
+                 COMPUTE_MASSWEIGHTED_IONFRAC=False, lowres_massweighting=1,
+                 COMPUTE_ZREION=False):
         #Measure time elapsed from start
         self._start_time = time.time()
         
@@ -237,6 +244,10 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
         self.ion_frac_massweighted = np.empty(len(self.z))
         if self.COMPUTE_MASSWEIGHTED_IONFRAC:
             self.compute_massweighted_xHII(CosmoParams, self.lowres_massweighting)
+
+        ### computing zreion and treion
+        self.zreion = self.compute_zreion_frombinaryxHII()
+        self.treion = self.compute_treion(ClassyCosmo)
         
         if self.PRINT_TIMER:
             z21_utilities.print_timer(self._start_time, text_before="Total computation time: ")
@@ -343,4 +354,31 @@ def compute_massweighted_xHII(self, CosmoParams, lowres_massweighting=1):
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
         return self.ion_frac_massweighted
-
+    
+    def compute_zreion_frombinaryxHII(self):
+        vectorized_zlist = np.vectorize(lambda iz: self.z[iz])
+        zreion = vectorized_zlist(np.argmin(self.ion_field_allz,axis=0)-1).reshape((self.ncells,self.ncells,self.ncells))
+        return zreion
+    
+    def compute_treion(ClassyCosmo):
+        return cosmology.time_at_redshift(ClassyCosmo,self.zreion)
+    
+    def _compute_ionfrac_from_zreion(self):
+        """
+        Way to compute the volume ionized fraction from zreion. Currently not used but there if needed.
+        """
+        zvalues = np.unique(self.zreion)
+        neutfrac = np.zeros(len(zvalues))
+        for i in range(len(zvalues)):
+            neutfrac[i] = np.sum(self.zreion<zvalues[i]) / self.ncells**3
+        return 1-neutfrac, zvalues
+
+    def _compute_ionfrac_from_treion(self):
+        """
+        Way to compute the volume ionized fraction from treion. Currently not used but there if needed.
+        """
+        tvalues = np.unique(self.treion)
+        neutfrac = np.zeros(len(tvalues))
+        for i in range(len(tvalues)):
+            neutfrac[i] = np.sum(self.treion>tvalues[i]) / self.ncells**3
+        return 1-neutfrac, tvalues

From 7068282111ef515ff2ccf3139b6c3c43c051d258 Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Thu, 15 May 2025 13:12:01 -0500
Subject: [PATCH 07/37] Fixes for zreion/treion maps.

---
 zeus21/maps.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index 807a66c..652755f 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -360,7 +360,7 @@ def compute_zreion_frombinaryxHII(self):
         zreion = vectorized_zlist(np.argmin(self.ion_field_allz,axis=0)-1).reshape((self.ncells,self.ncells,self.ncells))
         return zreion
     
-    def compute_treion(ClassyCosmo):
+    def compute_treion(self,ClassyCosmo):
         return cosmology.time_at_redshift(ClassyCosmo,self.zreion)
     
     def _compute_ionfrac_from_zreion(self):

From 39fe9654b755008df6d3853965ef5d096b8ce6f6 Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Thu, 15 May 2025 13:15:13 -0500
Subject: [PATCH 08/37] Fixes for zreion/treion maps.

---
 zeus21/cosmology.py |  4 ++--
 zeus21/maps.py      | 13 ++++++++++++-
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/zeus21/cosmology.py b/zeus21/cosmology.py
index 5bec74a..028533c 100644
--- a/zeus21/cosmology.py
+++ b/zeus21/cosmology.py
@@ -128,7 +128,7 @@ def time_at_redshift(ClassyCosmo,z):
     """
     background = ClassyCosmo.get_background()
     classy_t, classy_z = background['proper time [Gyr]'], background['z']
-    classy_tinterp = interpolate.interp1d(classy_z, classy_t)
+    classy_tinterp = interp1d(classy_z, classy_t)
     return classy_tinterp(z)
 
 def redshift_at_time(ClassyCosmo,t):
@@ -144,7 +144,7 @@ def redshift_at_time(ClassyCosmo,t):
     """
     background = ClassyCosmo.get_background()
     classy_t, classy_z = background['proper time [Gyr]'], background['z']
-    classy_tinterp = interpolate.interp1d(classy_t, classy_z)
+    classy_tinterp = interp1d(classy_t, classy_z)
     return classy_tinterp(t)
 
 def Hub(Cosmo_Parameters, z):
diff --git a/zeus21/maps.py b/zeus21/maps.py
index 652755f..07abea0 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -356,12 +356,23 @@ def compute_massweighted_xHII(self, CosmoParams, lowres_massweighting=1):
         return self.ion_frac_massweighted
     
     def compute_zreion_frombinaryxHII(self):
+        if self.PRINT_TIMER:
+            start_time = time.time()
+            print("Computing zreion map...")
         vectorized_zlist = np.vectorize(lambda iz: self.z[iz])
         zreion = vectorized_zlist(np.argmin(self.ion_field_allz,axis=0)-1).reshape((self.ncells,self.ncells,self.ncells))
+        if self.PRINT_TIMER:
+            z21_utilities.print_timer(start_time, text_before="    done in ")
         return zreion
     
     def compute_treion(self,ClassyCosmo):
-        return cosmology.time_at_redshift(ClassyCosmo,self.zreion)
+        if self.PRINT_TIMER:
+            start_time = time.time()
+            print("Computing zreion map...")
+        treion = cosmology.time_at_redshift(ClassyCosmo,self.zreion)
+        if self.PRINT_TIMER:
+            z21_utilities.print_timer(start_time, text_before="    done in ")
+        return treion
     
     def _compute_ionfrac_from_zreion(self):
         """

From a6ea097fba33caafb6515379920bf22f82f15db6 Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Thu, 15 May 2025 13:25:30 -0500
Subject: [PATCH 09/37] Fixes for zreion/treion maps.

---
 zeus21/maps.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index 07abea0..1daf534 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -246,8 +246,9 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
             self.compute_massweighted_xHII(CosmoParams, self.lowres_massweighting)
 
         ### computing zreion and treion
-        self.zreion = self.compute_zreion_frombinaryxHII()
-        self.treion = self.compute_treion(ClassyCosmo)
+        if COMPUTE_ZREION:
+            self.zreion = self.compute_zreion_frombinaryxHII()
+            self.treion = self.compute_treion(ClassyCosmo)
         
         if self.PRINT_TIMER:
             z21_utilities.print_timer(self._start_time, text_before="Total computation time: ")
@@ -368,7 +369,7 @@ def compute_zreion_frombinaryxHII(self):
     def compute_treion(self,ClassyCosmo):
         if self.PRINT_TIMER:
             start_time = time.time()
-            print("Computing zreion map...")
+            print("Computing treion map...")
         treion = cosmology.time_at_redshift(ClassyCosmo,self.zreion)
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")

From 94bbe4f0c2ce071cfc95bf2654bdad962d45a866 Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Thu, 22 May 2025 17:11:32 -0500
Subject: [PATCH 10/37] Include partial ionization in reionization_maps.

---
 zeus21/maps.py | 87 +++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 75 insertions(+), 12 deletions(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index 1daf534..f911467 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -147,13 +147,16 @@ class reionization_maps:
     LOGNORMAL_DENSITY: bool
         Whether to use lognormal (True) or Gaussian (False) density fields. Default is False.of radii. Default is False.
     COMPUTE_DENSITY_AT_ALLZ: bool
-        Whether to output the density field at all redshifts. If False, only the density at the lower input redshift is computed. If True, the computation time and memory usage dramatically increases. Default is False.
+        Whether to output the density field at all redshifts. If False, only the density at the lower input redshift is computed. If True, the computation time and memory usage dramatically increase. Default is False.
     SPHERIZE: bool
         Whether to flag spheres around ionized cells (True) instead of only central pixel flagging (False). Default is False. Central pixel flagging is generally more consistent with the bubble mass function than spherizing. Default is False.
     COMPUTE_MASSWEIGHTED_IONFRAC: bool
         Whether to compute the mass weighted ionized fraction. If True, COMPUTE_DENSITY_AT_ALLZ will be forced to True, thus increasing dramatically computation time. Default is False.
     lowres_massweighting: int
         Compute the mass-weighted ionized fraction more efficiently by using lower resolution density and ionized fields. Has to be >=1 and an integer. Default is 1.
+    INCLUDE_PARTIALION: bool
+        Whether to compute the ionized field while accounting for the ionization on the subpixels scale, so called partial ionization. Default is False.
+        This computation requires to know the density at all redshift. So if True, COMPUTE_DENSITY_AT_ALLZ will be forced to True, thus increasing dramatically computation time.
     COMPUTE_ZREION: bool
         Whether to compute the reionization redshift and time fields. Default is False.
 
@@ -173,10 +176,19 @@ class reionization_maps:
         Overdensity field at all the redshifts asked by the user. First dimension correponds to redshifts. Only computed if COMPUTE_DENSITY_AT_ALLZ is True.
     ion_field_allz: 4D np.array
         Ionized fraction field at all the redshifts asked by the user. First dimension correponds to redshifts.
+    ion_field_partial_allz: 4D np.array
+        Ionized fraction field (including partical ionization of pixels) at all the redshifts asked by the user. First dimension correponds to redshifts.
+        Only computed if INCLUDE_PARTIALION is True.
     ion_frac: 1D np.array
         Volume weighted ionized fraction at all the redshifts asked by the user.
+    ion_frac_subpixel: 1D np.array
+        Volume weighted ionized fraction (including subpixel ionization) at all the redshifts asked by the user.
+        Only computed if COMPUTE_MASSWEIGHTED_IONFRAC and INCLUDE_PARTIALION are True.
     ion_frac_massweighted: 1D np.array
         Mass weighted ionized fraction at all the redshifts asked by the user. Only computed if COMPUTE_MASSWEIGHTED_IONFRAC is True.
+    ion_frac_massweighted_subpixel: 1D np.array
+        Mass weighted ionized fraction (including subpixel ionization) at all the redshifts asked by the user. 
+        Only computed if COMPUTE_MASSWEIGHTED_IONFRAC and INCLUDE_PARTIALION are True.
     zreion: 3D np.array
         Reionization redshift field.
     treion: 3D np.array
@@ -188,6 +200,7 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
                  PRINT_TIMER=True, ENFORCE_BMF_SCALE=True, 
                  LOGNORMAL_DENSITY=False, COMPUTE_DENSITY_AT_ALLZ=False, SPHERIZE=False, 
                  COMPUTE_MASSWEIGHTED_IONFRAC=False, lowres_massweighting=1,
+                 INCLUDE_PARTIALION = False,
                  COMPUTE_ZREION=False):
         #Measure time elapsed from start
         self._start_time = time.time()
@@ -210,8 +223,10 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
         self.COMPUTE_DENSITY_AT_ALLZ = COMPUTE_DENSITY_AT_ALLZ
         self.SPHERIZE = SPHERIZE
         self.COMPUTE_MASSWEIGHTED_IONFRAC = COMPUTE_MASSWEIGHTED_IONFRAC
-        if self.COMPUTE_MASSWEIGHTED_IONFRAC:
+        self.INCLUDE_PARTIALION = INCLUDE_PARTIALION
+        if self.COMPUTE_MASSWEIGHTED_IONFRAC or self.INCLUDE_PARTIALION:
             self.COMPUTE_DENSITY_AT_ALLZ = True
+        self.COMPUTE_ZREION = COMPUTE_ZREION
 
         ### selecting redshifts and radii from available redshifts
         # redshifts
@@ -237,16 +252,23 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
         self.barrier = barrier
         if self.barrier is None:
             self.barrier = BMF.barrier
-        self.ion_field_allz, self.ion_frac = self.generate_xHII(CosmoParams, CoeffStructure, BMF)
+        self.ion_field_allz = self.generate_xHII(CosmoParams, CoeffStructure, BMF)
+        self.ion_frac = self.compute_volumeweigthed_xHII(self.ion_field_allz)
+        if self.INCLUDE_PARTIALION:
+            self.ion_field_partial_allz = self.generate_xHII_withPartialIonization(BMF)
+            self.ion_frac_subpixel = self.compute_volumeweigthed_xHII(self.ion_field_partial_allz)
 
         ### computing the mass weighted ionized fraction
         self.lowres_massweighting = lowres_massweighting
         self.ion_frac_massweighted = np.empty(len(self.z))
         if self.COMPUTE_MASSWEIGHTED_IONFRAC:
-            self.compute_massweighted_xHII(CosmoParams, self.lowres_massweighting)
+            self.ion_frac_massweighted = self.compute_massweighted_xHII(CosmoParams, self.ion_field_allz, self.lowres_massweighting)
+
+            if self.INCLUDE_PARTIALION:
+                self.ion_frac_massweighted_subpixel = self.compute_massweighted_xHII(CosmoParams, self.ion_field_partial_allz, self.lowres_massweighting)
 
         ### computing zreion and treion
-        if COMPUTE_ZREION:
+        if self.COMPUTE_ZREION:
             self.zreion = self.compute_zreion_frombinaryxHII()
             self.treion = self.compute_treion(ClassyCosmo)
         
@@ -271,6 +293,7 @@ def generate_density(self, ClassyCosmo, CorrFClass):
         else:
             pb = pbox.PowerBox(N=self.ncells, dim=3, pk=(lambda k: np.exp(pk_spl(np.log(k)))), boxlength=self.boxlength, seed=self.seed)
         density_field = pb.delta_x()
+
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
         return density_field
@@ -279,10 +302,12 @@ def generate_density_allz(self, CosmoParams):
         if self.PRINT_TIMER:
             start_time = time.time()
             print('Evolving density field...')
+
         Dg = CosmoParams.growthint(self.z)
         growthfactor_ratio = (Dg/Dg[0])[:, np.newaxis, np.newaxis, np.newaxis]
         density_lastz = np.copy(self.density)
         self.density_allz = density_lastz[np.newaxis]*growthfactor_ratio
+
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
         return self.density_allz
@@ -296,8 +321,10 @@ def smooth_density(self):
         if self.PRINT_TIMER:
             start_time = time.time()
             print("Smoothing density field...")
+
         density_fft = np.fft.fftn(self.density)
         density_smoothed_allr = np.array([z21_utilities.tophat_smooth(rr, self._k, density_fft) for rr in self.r])
+
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
         return density_smoothed_allr
@@ -306,16 +333,16 @@ def generate_xHII(self, CosmoParams, CoeffStructure, BMF):
         if self.PRINT_TIMER:
             start_time = time.time()
             print("Generating ionized field...")
+
         ion_field_allz = np.zeros((len(self.z),self.ncells,self.ncells,self.ncells))
-        ion_frac = np.zeros(len(self.z))
         for i in trange(len(self.z)):
             curr_z_idx = self._z_idx[i]
             ion_field = self.ionize(CosmoParams, CoeffStructure, curr_z_idx)
             ion_field_allz[i] = ion_field
-            ion_frac[i] = np.sum(ion_field)/(self.ncells**3)
+
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
-        return ion_field_allz, ion_frac
+        return ion_field_allz
 
     def ionize(self,CosmoParams, CoeffStructure, curr_z_idx):
         zlist = CoeffStructure.zintegral
@@ -338,7 +365,36 @@ def ionize(self,CosmoParams, CoeffStructure, curr_z_idx):
             ion_field = np.any(ion_field_Rs, axis=0)
         return ion_field
     
-    def compute_massweighted_xHII(self, CosmoParams, lowres_massweighting=1):
+    def generate_xHII_withPartialIonization(self, BMF, ir=0):
+        if self.PRINT_TIMER:
+            start_time = time.time()
+            print("Generating ionized field with partial ionization...")
+
+        partialion_field = np.empty(self.ion_field_allz.shape)
+        sample_d = np.linspace(-5, 5, 201)
+        for i in range(self.ncells):
+            nion_spl = spline(sample_d, BMF.nion_delta_r_int(CosmoParams, sample_d, ir).T[i])
+            nrec_spl = spline(sample_d, BMF.nrec(CosmoParams, sample_d, BMF.ion_frac).T[i])
+            partial_ion_spl = spline(sample_d, nion_spl(sample_d)/(1+nrec_spl(sample_d)))
+            partialion_field[i] = partial_ion_spl(self.density_allz)
+        ion_field_partial_allz = np.clip(self.ion_field_allz + partialion_field, 0, 1)
+
+        if self.PRINT_TIMER:
+            z21_utilities.print_timer(start_time, text_before="    done in ")
+        return ion_field_partial_allz
+    
+    def compute_volumeweigthed_xHII(self,xHII_field_allz):
+        if self.PRINT_TIMER:
+            start_time = time.time()
+            print("Computing volume weighted ionized fraction...")
+
+        ion_frac = np.sum(xHII_field_allz,axis=(1,2,3))/(self.ncells**3)
+
+        if self.PRINT_TIMER:
+            z21_utilities.print_timer(start_time, text_before="    done in ")
+        return ion_frac
+    
+    def compute_massweighted_xHII(self, CosmoParams, xHII_field_allz, lowres_massweighting=1):
         if np.sum(self.density_allz[0, 0, 0]) == 0.:
             self.generate_density_allz(CosmoParams)
         self.lowres_massweighting = lowres_massweighting
@@ -347,21 +403,26 @@ def compute_massweighted_xHII(self, CosmoParams, lowres_massweighting=1):
         if not isinstance(self.lowres_massweighting, (int, np.int32, np.int64)):
             raise Exeption('lowres_massweighting should be an integer.')
         d_allz = self.density_allz[:, ::self.lowres_massweighting, ::self.lowres_massweighting, ::self.lowres_massweighting]
-        ion_allz = self.ion_field_allz[:, ::self.lowres_massweighting, ::self.lowres_massweighting, ::self.lowres_massweighting]
+        ion_allz = xHII_field_allz[:, ::self.lowres_massweighting, ::self.lowres_massweighting, ::self.lowres_massweighting]
+
         if self.PRINT_TIMER:
             start_time = time.time()
             print("Computing mass weighted ionized fraction...")
-        self.ion_frac_massweighted = np.average((1+d_allz) * ion_allz, axis=(1, 2, 3))
+
+        ion_frac_massweighted = np.average((1+d_allz) * ion_allz, axis=(1, 2, 3))
+
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
-        return self.ion_frac_massweighted
+        return ion_frac_massweighted
     
     def compute_zreion_frombinaryxHII(self):
         if self.PRINT_TIMER:
             start_time = time.time()
             print("Computing zreion map...")
+
         vectorized_zlist = np.vectorize(lambda iz: self.z[iz])
         zreion = vectorized_zlist(np.argmin(self.ion_field_allz,axis=0)-1).reshape((self.ncells,self.ncells,self.ncells))
+
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
         return zreion
@@ -370,7 +431,9 @@ def compute_treion(self,ClassyCosmo):
         if self.PRINT_TIMER:
             start_time = time.time()
             print("Computing treion map...")
+
         treion = cosmology.time_at_redshift(ClassyCosmo,self.zreion)
+
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
         return treion

From 622055aeef8564512df415dcd2167abbe44fc19b Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Thu, 22 May 2025 17:14:31 -0500
Subject: [PATCH 11/37] Fixes for partial ionization.

---
 zeus21/maps.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index f911467..a3ac0b6 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -255,7 +255,7 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
         self.ion_field_allz = self.generate_xHII(CosmoParams, CoeffStructure, BMF)
         self.ion_frac = self.compute_volumeweigthed_xHII(self.ion_field_allz)
         if self.INCLUDE_PARTIALION:
-            self.ion_field_partial_allz = self.generate_xHII_withPartialIonization(BMF)
+            self.ion_field_partial_allz = self.generate_xHII_withPartialIonization(CosmoParams, BMF)
             self.ion_frac_subpixel = self.compute_volumeweigthed_xHII(self.ion_field_partial_allz)
 
         ### computing the mass weighted ionized fraction
@@ -365,7 +365,7 @@ def ionize(self,CosmoParams, CoeffStructure, curr_z_idx):
             ion_field = np.any(ion_field_Rs, axis=0)
         return ion_field
     
-    def generate_xHII_withPartialIonization(self, BMF, ir=0):
+    def generate_xHII_withPartialIonization(self, CosmoParams, BMF, ir=0):
         if self.PRINT_TIMER:
             start_time = time.time()
             print("Generating ionized field with partial ionization...")

From 347ea86773573c1c98e158647a485fd5ac5067b0 Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Thu, 22 May 2025 17:17:27 -0500
Subject: [PATCH 12/37] Fixes for partial ionization.

---
 zeus21/maps.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index a3ac0b6..df1acce 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -376,7 +376,7 @@ def generate_xHII_withPartialIonization(self, CosmoParams, BMF, ir=0):
             nion_spl = spline(sample_d, BMF.nion_delta_r_int(CosmoParams, sample_d, ir).T[i])
             nrec_spl = spline(sample_d, BMF.nrec(CosmoParams, sample_d, BMF.ion_frac).T[i])
             partial_ion_spl = spline(sample_d, nion_spl(sample_d)/(1+nrec_spl(sample_d)))
-            partialion_field[i] = partial_ion_spl(self.density_allz)
+            partialion_field[i] = partial_ion_spl(self.density_allz[i])
         ion_field_partial_allz = np.clip(self.ion_field_allz + partialion_field, 0, 1)
 
         if self.PRINT_TIMER:

From 323f38d7bb77ef45b3611e15172cfd387aa4b20d Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Thu, 22 May 2025 17:24:52 -0500
Subject: [PATCH 13/37] Fixes for partial ionization.

---
 zeus21/maps.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index df1acce..67559db 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -372,7 +372,7 @@ def generate_xHII_withPartialIonization(self, CosmoParams, BMF, ir=0):
 
         partialion_field = np.empty(self.ion_field_allz.shape)
         sample_d = np.linspace(-5, 5, 201)
-        for i in range(self.ncells):
+        for i in range(len(self.z)):
             nion_spl = spline(sample_d, BMF.nion_delta_r_int(CosmoParams, sample_d, ir).T[i])
             nrec_spl = spline(sample_d, BMF.nrec(CosmoParams, sample_d, BMF.ion_frac).T[i])
             partial_ion_spl = spline(sample_d, nion_spl(sample_d)/(1+nrec_spl(sample_d)))

From 2acf59524decc68c76f265948a0aedfd10b9887a Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Fri, 23 May 2025 10:58:05 -0500
Subject: [PATCH 14/37] Fixes for partial ionization.

---
 zeus21/maps.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index 67559db..5e956bd 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -373,8 +373,9 @@ def generate_xHII_withPartialIonization(self, CosmoParams, BMF, ir=0):
         partialion_field = np.empty(self.ion_field_allz.shape)
         sample_d = np.linspace(-5, 5, 201)
         for i in range(len(self.z)):
-            nion_spl = spline(sample_d, BMF.nion_delta_r_int(CosmoParams, sample_d, ir).T[i])
-            nrec_spl = spline(sample_d, BMF.nrec(CosmoParams, sample_d, BMF.ion_frac).T[i])
+            curr_z_idx = self._z_idx[i]
+            nion_spl = spline(sample_d, BMF.nion_delta_r_int(CosmoParams, sample_d, ir).T[curr_z_idx])
+            nrec_spl = spline(sample_d, BMF.nrec(CosmoParams, sample_d, BMF.ion_frac).T[curr_z_idx])
             partial_ion_spl = spline(sample_d, nion_spl(sample_d)/(1+nrec_spl(sample_d)))
             partialion_field[i] = partial_ion_spl(self.density_allz[i])
         ion_field_partial_allz = np.clip(self.ion_field_allz + partialion_field, 0, 1)

From 4285137332e749fcef27e166c55efee933952c4c Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Wed, 25 Jun 2025 11:04:55 -0500
Subject: [PATCH 15/37] Fix for partial ionization.

---
 zeus21/maps.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index 5e956bd..25b40bb 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -224,6 +224,8 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
         self.SPHERIZE = SPHERIZE
         self.COMPUTE_MASSWEIGHTED_IONFRAC = COMPUTE_MASSWEIGHTED_IONFRAC
         self.INCLUDE_PARTIALION = INCLUDE_PARTIALION
+        if self.INCLUDE_PARTIALION: # COMPUTE_MASSWEIGHTED_IONFRAC needs to be True so that partial ion computes on the actual density field at its appropriate redshift
+            self.COMPUTE_MASSWEIGHTED_IONFRAC = True 
         if self.COMPUTE_MASSWEIGHTED_IONFRAC or self.INCLUDE_PARTIALION:
             self.COMPUTE_DENSITY_AT_ALLZ = True
         self.COMPUTE_ZREION = COMPUTE_ZREION

From 6daacaaf4050b5b9113224bec6d305af00688556 Mon Sep 17 00:00:00 2001
From: Yonatan Sklansky <yonatan.sklansky@gmail.com>
Date: Tue, 21 Oct 2025 17:24:05 -0500
Subject: [PATCH 16/37] Changes to reionization.py:

-Functions take values instead of indexes
-Added interpolators
-Added Rmin and ZMAX
---
 zeus21/reionization.py | 237 ++++++++++++++++++++++++++---------------
 1 file changed, 152 insertions(+), 85 deletions(-)

diff --git a/zeus21/reionization.py b/zeus21/reionization.py
index 6be8ff7..4f225c7 100644
--- a/zeus21/reionization.py
+++ b/zeus21/reionization.py
@@ -12,8 +12,9 @@
 from . import cosmology
 from . import constants
 import numpy as np
-import astropy.units as u
 from scipy.integrate import cumulative_trapezoid
+from scipy.interpolate import interp1d
+from scipy.interpolate import RegularGridInterpolator
 from tqdm import trange
 
 
@@ -24,46 +25,60 @@ class BMF:
     
     """
     
-    def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, R_linear_sigma_fit_input=10, FLAG_converge=True, max_iter=10, ZMAX_REION = 30):
+    def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, ClassyCosmo, R_linear_sigma_fit_input=10, FLAG_converge=True, max_iter=10, ZMAX_REION = 30, Rmin=0.05):
 
         self.ZMAX_REION = ZMAX_REION #max redshift up to which we calculate reionization observables
         self.zlist = CoeffStructure.zintegral
         self.Rs = CoeffStructure.Rtabsmoo
+        self.Rs_BMF = np.logspace(np.log10(Rmin), np.log10(self.Rs[-1]), 100)
 
         
         self.gamma = CoeffStructure.gamma_niondot_II_index2D
         self.gamma2 = CoeffStructure.gamma2_niondot_II_index2D
         self.sigma = CoeffStructure.sigmaofRtab
+        
+        self.zr = [self.zlist, np.log(self.Rs)]
+        self.gamma_int = RegularGridInterpolator(self.zr, self.gamma, bounds_error = False, fill_value = np.nan)
+        self.gamma2_int = RegularGridInterpolator(self.zr, self.gamma2, bounds_error = False, fill_value = np.nan)
 
+        self.sigma_BMF = np.array([[ClassyCosmo.sigma(r, z) for z in self.zlist] for r in self.Rs_BMF]).T
+        self.zr_BMF = [self.zlist, np.log(self.Rs_BMF)]
+        self.sigma_int = RegularGridInterpolator(self.zr_BMF, self.sigma_BMF, bounds_error = False, fill_value = np.nan)
+        
         self.Hz = cosmology.Hubinvyr(CosmoParams, self.zlist)
         self.trec0 = 1/(constants.alphaB * cosmology.n_H(CosmoParams,0) * AstroParams._clumping) #seconds
         self.trec = self.trec0/(1+self.zlist)**3/constants.yrTos #years
+        self.trec_int = interp1d(self.zlist, self.trec, bounds_error = False, fill_value = np.nan)
+        
         self.niondot_avg = CoeffStructure.niondot_avg_II
+        self.niondot_avg_int = interp1d(self.zlist, self.niondot_avg, bounds_error = False, fill_value = np.nan)
 
-        #self.xHI = CoeffStructure.xHI_avg ----> need to figure out what's going on with this
-        self.ion_frac = np.fmin(1, [self.calc_Q(CosmoParams, i) for i in np.arange(len(self.zlist))])
+        self.ion_frac = np.fmin(1, [self.calc_Q(CosmoParams, z) for z in self.zlist])
         self.ion_frac_initial = np.copy(self.ion_frac)
 
         zr_mesh = np.meshgrid(np.arange(len(self.Rs)), np.arange(len(self.zlist)))
         self.nion_norm = self.nion_normalization(zr_mesh[1], zr_mesh[0])
-        self.barrier = self.compute_barrier(CosmoParams, self.ion_frac)
+        self.nion_norm_int = RegularGridInterpolator(self.zr, self.nion_norm, bounds_error = False, fill_value = np.nan)
+        
+        self.barrier = self.compute_barrier(CosmoParams, self.ion_frac, self.zlist, self.Rs)
         self.barrier_initial = np.copy(self.barrier)
+        self.barrier_int = RegularGridInterpolator(self.zr, self.barrier, bounds_error = False, fill_value = np.nan)
 
         self.R_linear_sigma_fit_idx = z21_utilities.find_nearest_idx(self.Rs, R_linear_sigma_fit_input)
         self.R_linear_sigma_fit = self.Rs[self.R_linear_sigma_fit_idx]
 
-        self.BMF = np.array([self.VRdn_dR(HMFintclass, np.arange(len(self.Rs)), i) for i in range(len(self.zlist))]) #initial bubble mass function
+        self.BMF = np.array([self.VRdn_dR(z, self.Rs_BMF) for z in self.zlist]) #initial bubble mass function
         self.BMF_initial = np.copy(self.BMF)
-        self.ion_frac = np.nan_to_num([np.trapz(self.BMF[i], np.log(self.Rs)) for i in range(len(self.zlist))]) #ion_frac by integrating the BMF
+        self.ion_frac = np.nan_to_num([np.trapz(self.BMF[i], np.log(self.Rs_BMF)) for i in range(len(self.zlist))]) #ion_frac by integrating the BMF
         self.ion_frac[self.barrier[:, -1]<=0] = 1
         
         if FLAG_converge:
-            self.converge_BMF(CosmoParams, HMFintclass, self.ion_frac, max_iter=max_iter)
+            self.converge_BMF(CosmoParams, self.ion_frac, max_iter=max_iter)
         #two functions: compute BMF and iterate
         
 
 
-    def compute_barrier(self, CosmoParams, ion_frac):
+    def compute_barrier(self, CosmoParams, ion_frac, z, R):
         """
         Computes the density barrier threshold for ionization.
         
@@ -81,13 +96,17 @@ def compute_barrier(self, CosmoParams, ion_frac):
         barrier: 2D np.array
             The resultant density threshold array. First dimension is each redshift, second dimension is each radius scale.
         """
-        barrier = np.zeros((len(self.zlist), len(self.Rs)))
+        barrier = np.zeros((len(z), len(R)))
         ds_array = np.linspace(-1, 5, 101)
+
+        zarg = np.argsort(z) #sort just in case
+        z = z[zarg]
+        ion_frac = ion_frac[zarg]
         
-        for ir in range(len(self.Rs)):
+        for ir in range(len(R)):
             #Compute nion_values and nrec_values for this 'ir'
-            nion_values = self.nion_delta_r_int(CosmoParams, ds_array, ir)  #Shape (nd, nz)
-            nrec_values = self.nrec(CosmoParams, ds_array, ion_frac)             #Shape (nd, nz)
+            nion_values = self.nion_delta_r_int(CosmoParams, ds_array, z, R[ir])  #Shape (nd, nz)
+            nrec_values = self.nrec(CosmoParams, ds_array, ion_frac, z)             #Shape (nd, nz)
         
             total_values = np.log10(nion_values / (1 + nrec_values) + 1e-10)   #taking difference in logspace to find zero-crossing 
         
@@ -116,7 +135,7 @@ def compute_barrier(self, CosmoParams, ion_frac):
     def nion_normalization(self, z, R):
         return 1/np.sqrt(1-2*self.gamma2[z, R]*self.sigma[z, R]**2)*np.exp(self.gamma[z, R]**2 * self.sigma[z, R]**2 / (2-4*self.gamma2[z, R]*self.sigma[z, R]**2))
 
-    def nrec(self, CosmoParams, d_array, ion_frac):
+    def nrec(self, CosmoParams, d_array, ion_frac, z):
         """
         Vectorized computation of nrec over an array of overdensities d_array.
 
@@ -134,11 +153,14 @@ def nrec(self, CosmoParams, d_array, ion_frac):
         nrecs: 2D np.array
             The total number of recombinations at each overdensity for a certain ionized fraction history at each redshift. The first dimension is densities, the second dimension is redshifts.
         """
+        zarg = np.argsort(z) #sort just in case
+        z = z[zarg]
+        ion_frac = ion_frac[zarg]
         
         #reverse the inputs to make the integral easier to compute
-        z_rev = self.zlist[::-1]
-        Hz_rev = self.Hz[::-1]
-        trec_rev = self.trec[::-1]
+        z_rev = z[::-1]
+        Hz_rev = cosmology.Hubinvyr(CosmoParams, z_rev)
+        trec_rev = self.trec_int(z_rev)
         ion_frac_rev = ion_frac[::-1]
     
         denom = -1 / (1 + z_rev) / Hz_rev / trec_rev
@@ -149,12 +171,12 @@ def nrec(self, CosmoParams, d_array, ion_frac):
         
         #TODO: nonlinear recombinations/higher order
 
-        nrecs = nrecs[:, ::-1]  # Reverse to match self.zlist order
+        nrecs = nrecs[:, ::-1]  #reverse back to increasing z order
         return nrecs
     
-    def niondot_delta_r(self, CosmoParams, d_array, ir):
+    def niondot_delta_r(self, CosmoParams, d_array, z, R):
         """
-        Compute niondot over an array of overdensities d_array for a given ir.
+        Compute niondot over an array of overdensities d_array for a given R.
 
         Parameters
         ----------
@@ -162,29 +184,32 @@ def niondot_delta_r(self, CosmoParams, d_array, ir):
             Stores cosmology.
         d_array: 1D np.array
             A list of sample overdensity values to evaluate niondot over.
-        ir: int
-            Index corresponding to a certain radius value from self.Rs.
+        R: float
+            Radius value (cMpc)
 
         Output
         ----------
         niondot: 2D np.array
             The rates of ionizing photon production. The first dimension is densities, the second dimension is redshifts.
         """
+
+        zarg = np.argsort(z) #sort just in case
+        z = z[zarg]
         
-        d_array = d_array[:, np.newaxis] * CosmoParams.growthint(self.zlist)[np.newaxis, :] / CosmoParams.growthint(self.zlist[0])
+        d_array = d_array[:, np.newaxis] * CosmoParams.growthint(z)[np.newaxis, :] / CosmoParams.growthint(z[0])
     
-        gamma_ir = self.gamma[:, ir]   
-        gamma2_ir = self.gamma2[:, ir]  
-        nion_norm_ir = self.nion_norm[:, ir] 
+        gamma_R = self.gammaz_int(z, R)   
+        gamma2_R = self.gamma2z_int(z, R)  
+        nion_norm_R = self.nion_normz_int(z, R)
     
-        exp_term = np.exp(gamma_ir[np.newaxis, :] * d_array + gamma2_ir[np.newaxis, :] * d_array**2)
-        niondot = (self.niondot_avg[np.newaxis, :] / nion_norm_ir[np.newaxis, :]) * exp_term
+        exp_term = np.exp(gamma_R[np.newaxis, :] * d_array + gamma2_R[np.newaxis, :] * d_array**2)
+        niondot = (self.niondot_avg_int(z)[np.newaxis, :] / nion_norm_R[np.newaxis, :]) * exp_term
         
         return niondot
     
-    def nion_delta_r_int(self, CosmoParams, d_array, ir):
+    def nion_delta_r_int(self, CosmoParams, d_array, z, R):
         """
-        Vectorized computation of nion over an array of overdensities d_array for a given ir.
+        Vectorized computation of nion over an array of overdensities d_array for a given R.
 
         Parameters
         ----------
@@ -192,93 +217,135 @@ def nion_delta_r_int(self, CosmoParams, d_array, ir):
             Stores cosmology.
         d_array: 1D np.array
             A list of sample overdensity values to evaluate niondot over.
-        ir: int
-            Index corresponding to a certain radius value from self.Rs.
+        R: float
+            Radius value (cMpc)
 
         Output
         ----------
         nion: 2D np.array
-            The total number of ionizing photons produced since z=50. The first dimension is densities, the second dimension is redshifts.
+            The total number of ionizing photons produced since z=zmax. The first dimension is densities, the second dimension is redshifts.
         """
+
+        z.sort() #sort if not sorted
         
         #reverse the inputs to make the integral easier to compute
-        z_rev = self.zlist[::-1]
-        Hz_rev = self.Hz[::-1]
+        z_rev = z[::-1]
+        Hz_rev = cosmology.Hubinvyr(CosmoParams, z_rev)
     
-        niondot_values = self.niondot_delta_r(CosmoParams, d_array, ir)
+        niondot_values = self.niondot_delta_r(CosmoParams, d_array, z, R)
     
         integrand = -1 / (1 + z_rev) / Hz_rev * niondot_values[:, ::-1]
-        #cumulative_integral = np.concatenate(
-        #    (np.zeros((len(d_array), 1)), cumulative_trapezoid(integrand, x=z_rev, axis=1)), axis=1
-        #)
-        nion = cumulative_trapezoid(integrand, x=z_rev, initial=0)[:, ::-1] #reverse back to match self.zlist order
-        #nion = cumulative_integral[:, ::-1]  # Reverse back to match self.zlist order
+        nion = cumulative_trapezoid(integrand, x=z_rev, initial=0)[:, ::-1] #reverse back to increasing z order
+        
         return nion
 
-    #calculating ionized fraction, put outside the class
-    def calc_Q(self, CosmoParams, iz):
-        z = self.zlist[iz]
-        dtdz = 1/cosmology.Hubinvyr(CosmoParams, self.zlist[iz:])/(1 + self.zlist[iz:])
-        tau0 = self.trec0 * np.sqrt(CosmoParams.OmegaM) * cosmology.Hubinvyr(CosmoParams,0) / constants.yrTos
-        exp = np.exp(2/3/tau0 * (np.power(1 + z, 3/2) - np.power(1 + self.zlist[iz:], 3/2))) #switched order around to be correct
-        
-        integrand = dtdz * self.niondot_avg[iz:] * exp
+    #calculating ionized fraction
+    def calc_Q(self, CosmoParams, z):
+        z_arr = np.logspace(np.log10(z), np.log10(self.zlist[-1]), 50)
+        dtdz = 1/cosmology.Hubinvyr(CosmoParams, z_arr)/(1 + z_arr)
+        tau0 = self.trec0 * np.sqrt(CosmoParams.OmegaM) * cosmology.Hubinvyr(CosmoParams, 0) / constants.yrTos
+        exp = np.exp(2/3/tau0 * (np.power(1 + z, 3/2) - np.power(1 + z_arr, 3/2))) #switched order around to be correct (typo in paper)
+
+        niondot_avgs = self.niondot_avg_int(z_arr)
+        integrand = dtdz * niondot_avgs * exp
     
-        return np.trapz(integrand, x = self.zlist[iz:])
+        return np.trapz(integrand, x = z_arr)
 
     #computing linear barrier
-    def B_1(self, HMFintclass, ir, iz):
-        sigmax = HMFintclass.sigmaR_int(self.Rs[self.R_linear_sigma_fit_idx+1], self.zlist[iz])
-        sigmin = HMFintclass.sigmaR_int(self.Rs[self.R_linear_sigma_fit_idx-1], self.zlist[iz])
-        barriermax = self.barrier[iz, self.R_linear_sigma_fit_idx+1]
-        barriermin = self.barrier[iz, self.R_linear_sigma_fit_idx-1]
+    def B_1(self, z):
+        sigmax = self.sigmaR_int(z, self.Rs[self.R_linear_sigma_fit_idx+1])
+        sigmin = self.sigmaR_int(z, self.Rs[self.R_linear_sigma_fit_idx-1])
+        barriermax = self.barrierR_int(z, self.Rs[self.R_linear_sigma_fit_idx+1])
+        barriermin = self.barrierR_int(z, self.Rs[self.R_linear_sigma_fit_idx-1])
         return (barriermax - barriermin)/(sigmax**2 - sigmin**2)
         
-    def B_0(self, HMFintclass, ir, iz):
-        sigmin = HMFintclass.sigmaR_int(self.Rs[self.R_linear_sigma_fit_idx-1], self.zlist[iz])
-        barriermin = self.barrier[iz, self.R_linear_sigma_fit_idx-1]
-        return barriermin - sigmin**2 * self.B_1(HMFintclass, ir, iz)
+    def B_0(self, z):
+        sigmin = self.sigmaR_int(z, self.Rs[self.R_linear_sigma_fit_idx-1])
+        barriermin = self.barrierR_int(z, self.Rs[self.R_linear_sigma_fit_idx-1])
+        return barriermin - sigmin**2 * self.B_1(z)
     
-    def B(self, HMFintclass, ir, iz): 
-        sig = HMFintclass.sigmaR_int(self.Rs[ir], self.zlist[iz])
-        return self.B_0(HMFintclass, ir, iz) + self.B_1(HMFintclass, ir, iz)*sig**2
-
+    def B(self, z, R): 
+        sig = self.sigmaR_int(z, R)
+        return self.B_0(z) + self.B_1(z)*sig**2
     
     #computing other terms in the BMF
-    def dsigma_dR(self, HMFintclass, ir, iz):
-        R = self.Rs[ir]
-        z = self.zlist[iz]
-        sigma = HMFintclass.sigmaR_int(R, z)
+    def dsigma_dR(self, z, R):
+        sigma = self.sigmaR_int(z, R)
         return sigma/R*np.gradient(np.log(sigma), np.log(R))
     
-    def dlogsigma_dlogR(self, HMFintclass, ir, iz):
-        R = self.Rs[ir]
-        z = self.zlist[iz]
-        return self.dsigma_dR(HMFintclass, ir, iz) * R/HMFintclass.sigmaR_int(R,z)
+    def dlogsigma_dlogR(self, z, R):
+        sigma = self.sigmaR_int(z, R)
+        return self.dsigma_dR(z, R) * R/sigma
     
-    def VRdn_dR(self, HMFintclass, ir, iz):
-        R = self.Rs[ir]
-        z = self.zlist[iz]
-        sig = HMFintclass.sigmaR_int(self.Rs[ir], self.zlist[iz])
-        return np.sqrt(2/np.pi) * np.abs(self.dlogsigma_dlogR(HMFintclass, ir, iz)) * np.abs(self.B_0(HMFintclass, ir, iz))/sig * np.exp(-self.B(HMFintclass, ir, iz)**2/2/sig**2)
+    def VRdn_dR(self, z, R):
+        sig = self.sigmaR_int(z, R)
+        return np.sqrt(2/np.pi) * np.abs(self.dlogsigma_dlogR(z, R)) * np.abs(self.B_0(z))/sig * np.exp(-self.B(z, R)**2/2/sig**2)
     
-    def Rdn_dR(self, ir, iz):
-        R = self.Rs[ir]
-        z = self.zlist[iz]
-        return self.VRdn_dR(HMFintclass, ir, iz)*3/(4*np.pi*R**3)
+    def Rdn_dR(self, z, R):
+        return self.VRdn_dR(z, R)*3/(4*np.pi*R**3)
 
-    def converge_BMF(self, CosmoParams, HMFintclass, ion_frac_input, max_iter):
+    def converge_BMF(self, CosmoParams, ion_frac_input, max_iter):
         self.ion_frac = ion_frac_input
         for j in trange(max_iter):
             ion_frac_prev = np.copy(self.ion_frac)
             
-            self.barrier = self.compute_barrier(CosmoParams, self.ion_frac)
-            self.BMF = np.array([self.VRdn_dR(HMFintclass, np.arange(len(self.Rs)), i) for i in range(len(self.zlist))])
-            self.ion_frac = np.nan_to_num([np.trapz(self.BMF[i], np.log(self.Rs)) for i in range(len(self.zlist))])
+            self.barrier = self.compute_barrier(CosmoParams, self.ion_frac, self.zlist, self.Rs)
+            self.barrier_int = RegularGridInterpolator(self.zr, self.barrier, bounds_error = False, fill_value = np.nan)
+
+            def barrierR_int(self, z, R):
+                return self.interpR(z, R, self.barrier_int)
+            def barrierz_int(self, z, R):
+                return self.interpz(z, R, self.barrier_int)
+            
+            self.BMF = np.array([self.VRdn_dR(z, self.Rs_BMF) for z in self.zlist])
+            self.ion_frac = np.nan_to_num([np.trapz(self.BMF[i], np.log(self.Rs_BMF)) for i in range(len(self.zlist))])
             self.ion_frac[self.barrier[:, -1]<=0] = 1
 
             if np.allclose(ion_frac_prev, self.ion_frac):
                 print(f'SUCCESS: BMF converged in {j} iterations.')
                 return 
             
-        print(f"WARNING: BMF didn't converge within {max_iter} iterations.")
\ No newline at end of file
+        print(f"WARNING: BMF didn't converge within {max_iter} iterations.")
+
+
+    #interpolators in z and R used in reionization.py
+    
+    def interpR(self, z, R, func):
+        "Interpolator to find func(z, R), designed to take a single z but an array of R in cMpc"
+        _logR = np.log(R)
+        logRvec = np.asarray([_logR]) if np.isscalar(_logR) else np.asarray(_logR)
+        inarray = np.array([[z, LR] for LR in logRvec])
+        return func(inarray)
+
+    def interpz(self, z, R, func):
+        "Interpolator to find func(z, R), designed to take a single R in cMpc but an array of z"
+        zvec = np.asarray([z]) if np.isscalar(z) else np.asarray(z)
+        inarray = np.array([[zz, np.log(R)] for zz in zvec])
+        return func(inarray)
+
+    #all instances of different (z, R) interpolators, named explicitly for clarity in the code
+
+    def sigmaR_int(self, z, R):
+        return self.interpR(z, R, self.sigma_int)
+    def sigmaz_int(self, z, R):
+        return self.interpz(z, R, self.sigma_int)
+
+    def barrierR_int(self, z, R): #unused
+        return self.interpR(z, R, self.barrier_int)
+    def barrierz_int(self, z, R): #unused
+        return self.interpz(z, R, self.barrier_int)
+
+    def gammaR_int(self, z, R):
+        return self.interpR(z, R, self.gamma_int)
+    def gammaz_int(self, z, R):
+        return self.interpz(z, R, self.gamma_int)
+
+    def gamma2R_int(self, z, R):
+        return self.interpR(z, R, self.gamma2_int)
+    def gamma2z_int(self, z, R):
+        return self.interpz(z, R, self.gamma2_int)
+
+    def nion_normR_int(self, z, R):
+        return self.interpR(z, R, self.nion_norm_int)
+    def nion_normz_int(self, z, R):
+        return self.interpz(z, R, self.nion_norm_int)
\ No newline at end of file

From bf415638673259f35d8712fec5ba766bb9574833 Mon Sep 17 00:00:00 2001
From: Yonatan Sklansky <yonatan.sklansky@gmail.com>
Date: Tue, 21 Oct 2025 17:30:24 -0500
Subject: [PATCH 17/37] Changes to maps.py

- Removed enforce_BMF_scale, allowing user to pick any box size and resolution, but forcing smoothing Rs.
- Allow input Rs but comes with a warning.
- Using linear barrier by default (uses new functions from reionization.py).
- Updated partial ionizations, and added their massweightings.
- Small fixes.
---
 zeus21/maps.py         | 258 ++++++++++++++++++++++++-----------------
 zeus21/reionization.py |   2 +-
 2 files changed, 151 insertions(+), 109 deletions(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index 25b40bb..02ee8cb 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -3,7 +3,7 @@
 Make maps! For fun and science
 
 Authors: Julian B. Muñoz, Yonatan Sklansky, Emilie Thelie
-UT Austin - February 2025
+UT Austin - October 2025
 
 """
 
@@ -15,7 +15,7 @@
 import powerbox as pbox
 from scipy.interpolate import interp1d
 from scipy.interpolate import InterpolatedUnivariateSpline as spline
-from pyfftw import empty_aligned as empty
+#from pyfftw import empty_aligned as empty
 from tqdm import trange
 import time
 
@@ -141,24 +141,21 @@ class reionization_maps:
     r_precision: float
         Allows to change the steps of the radii for faster computation. Default (and max) is 1, lower values make the computation faster at the cost of accuracy.
     barrier: function
-        Input density barrier to be used as the threshold for map generation. Takes z index //\\Change//\\ (relative to CoeffStructure.zintegral) as input and returns np.array of shape. Default is None.
+        Input density barrier to be used as the threshold for map generation. Takes z value as input and returns np.array of shape. Default is None.
     PRINT_TIMER: bool
         Whether to print the time elapsed along the process. Default is True.
     LOGNORMAL_DENSITY: bool
-        Whether to use lognormal (True) or Gaussian (False) density fields. Default is False.of radii. Default is False.
+        Whether to use lognormal (True) or Gaussian (False) density fields. Default is False.
     COMPUTE_DENSITY_AT_ALLZ: bool
-        Whether to output the density field at all redshifts. If False, only the density at the lower input redshift is computed. If True, the computation time and memory usage dramatically increase. Default is False.
+        Whether to output the density field at all redshifts. If False, only the density at the lower input redshift is computed. If True, the computation time and memory usage dramatically increases. Default is False.
     SPHERIZE: bool
         Whether to flag spheres around ionized cells (True) instead of only central pixel flagging (False). Default is False. Central pixel flagging is generally more consistent with the bubble mass function than spherizing. Default is False.
-    COMPUTE_MASSWEIGHTED_IONFRAC: bool
-        Whether to compute the mass weighted ionized fraction. If True, COMPUTE_DENSITY_AT_ALLZ will be forced to True, thus increasing dramatically computation time. Default is False.
+    COMPUTE_MASSWEIGHTED: bool
+        Whether to compute the mass weighted ionized field and fraction. If True, COMPUTE_DENSITY_AT_ALLZ will be forced to True, thus increasing computation time dramatically. Default is False.
     lowres_massweighting: int
-        Compute the mass-weighted ionized fraction more efficiently by using lower resolution density and ionized fields. Has to be >=1 and an integer. Default is 1.
-    INCLUDE_PARTIALION: bool
-        Whether to compute the ionized field while accounting for the ionization on the subpixels scale, so called partial ionization. Default is False.
-        This computation requires to know the density at all redshift. So if True, COMPUTE_DENSITY_AT_ALLZ will be forced to True, thus increasing dramatically computation time.
-    COMPUTE_ZREION: bool
-        Whether to compute the reionization redshift and time fields. Default is False.
+        Compute the mass-weighted ionized field and fraction more efficiently by using lower resolution density and ionized fields. Has to be >=1 and an integer. Default is 1.
+    COMPUTE_PARTIAL_IONIZATIONS: bool
+        Whether to compute the subpixel ionizations in the field and the ionized fractions.
 
     Attributes
     ----------
@@ -176,73 +173,67 @@ class reionization_maps:
         Overdensity field at all the redshifts asked by the user. First dimension correponds to redshifts. Only computed if COMPUTE_DENSITY_AT_ALLZ is True.
     ion_field_allz: 4D np.array
         Ionized fraction field at all the redshifts asked by the user. First dimension correponds to redshifts.
-    ion_field_partial_allz: 4D np.array
-        Ionized fraction field (including partical ionization of pixels) at all the redshifts asked by the user. First dimension correponds to redshifts.
-        Only computed if INCLUDE_PARTIALION is True.
     ion_frac: 1D np.array
         Volume weighted ionized fraction at all the redshifts asked by the user.
-    ion_frac_subpixel: 1D np.array
-        Volume weighted ionized fraction (including subpixel ionization) at all the redshifts asked by the user.
-        Only computed if COMPUTE_MASSWEIGHTED_IONFRAC and INCLUDE_PARTIALION are True.
     ion_frac_massweighted: 1D np.array
-        Mass weighted ionized fraction at all the redshifts asked by the user. Only computed if COMPUTE_MASSWEIGHTED_IONFRAC is True.
-    ion_frac_massweighted_subpixel: 1D np.array
-        Mass weighted ionized fraction (including subpixel ionization) at all the redshifts asked by the user. 
-        Only computed if COMPUTE_MASSWEIGHTED_IONFRAC and INCLUDE_PARTIALION are True.
-    zreion: 3D np.array
-        Reionization redshift field.
-    treion: 3D np.array
-        Reionization time field.    
+        Mass weighted ionized fraction at all the redshifts asked by the user. Only computed if COMPUTE_MASSWEIGHTED is True.
     """
     
     def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, input_z, 
-                 input_boxlength=300., ncells=300, seed=1234, r_precision=1., barrier=None, 
-                 PRINT_TIMER=True, ENFORCE_BMF_SCALE=True, 
+                 input_boxlength=300., ncells=300, seed=1234, r_precision=1., Rs=None, barrier=None, 
+                 PRINT_TIMER=True, 
                  LOGNORMAL_DENSITY=False, COMPUTE_DENSITY_AT_ALLZ=False, SPHERIZE=False, 
-                 COMPUTE_MASSWEIGHTED_IONFRAC=False, lowres_massweighting=1,
-                 INCLUDE_PARTIALION = False,
-                 COMPUTE_ZREION=False):
+                 COMPUTE_MASSWEIGHTED=False, lowres_massweighting=1, COMPUTE_PARTIAL_IONIZATIONS=False,
+                 COMPUTE_PARTIAL_AND_MASSWEIGHTED=False, COMPUTE_ZREION=False
+                ):
         #Measure time elapsed from start
         self._start_time = time.time()
         
         ### boxes parameters
         self.input_z = input_z
-        self.ENFORCE_BMF_SCALE = ENFORCE_BMF_SCALE
         self.ncells = ncells
         self.boxlength = input_boxlength
         self.dx = self.boxlength/self.ncells
-        if self.ENFORCE_BMF_SCALE:
-            self.dx = CoeffStructure.Rtabsmoo[0] * (4*np.pi/3)**(1/3) #change to constant
-            self.boxlength = self.dx * self.ncells
-            print(f"WARNING: with ENFORCE_BMF_SCALE=True, boxlength is changed to match BMF scale. Current value: {self.boxlength:.2f} cMpc.")
+
+        # radii
+        if Rs is None:
+            default_len = len(CoeffStructure.Rtabsmoo)
+            self.r_precision = r_precision
+            self.r = np.logspace(np.log10(self.dx * (3/4/np.pi)**(1/3)), np.log10(self.boxlength), int(default_len*self.r_precision))
+            self._r_idx = np.arange(int(default_len*self.r_precision))
+        else:
+            self.r_precision = r_precision
+            self.r = Rs
+            if self.r_precision > 1:
+                raise ValueError('r_precision cannot be greater than 1 if you input your own radii.')
+            self._r_idx = np.floor(np.arange(len(self.r), step=self.r_precision)).astype(int)
+            smallest_r = self.dx * (3/4/np.pi)**(1/3)
+            if self.r[0] < smallest_r:
+                print(f'WARNING: Your input radii are too small for the pixel size. The code will still run now.\nIn the future, for best performance and physical accuracy on this boxlength and ncells, the smallest smoothing radius should be no less than R=L/N * (4pi/3)^(-1/3), or approximately {smallest_r:.2f} cMpc.')
+        
         self.seed = seed
 
         ### FLAGS
         self.PRINT_TIMER = PRINT_TIMER
         self.LOGNORMAL_DENSITY = LOGNORMAL_DENSITY
         self.COMPUTE_DENSITY_AT_ALLZ = COMPUTE_DENSITY_AT_ALLZ
+        self._has_density = COMPUTE_DENSITY_AT_ALLZ
         self.SPHERIZE = SPHERIZE
-        self.COMPUTE_MASSWEIGHTED_IONFRAC = COMPUTE_MASSWEIGHTED_IONFRAC
-        self.INCLUDE_PARTIALION = INCLUDE_PARTIALION
-        if self.INCLUDE_PARTIALION: # COMPUTE_MASSWEIGHTED_IONFRAC needs to be True so that partial ion computes on the actual density field at its appropriate redshift
-            self.COMPUTE_MASSWEIGHTED_IONFRAC = True 
-        if self.COMPUTE_MASSWEIGHTED_IONFRAC or self.INCLUDE_PARTIALION:
+        self.COMPUTE_MASSWEIGHTED = COMPUTE_MASSWEIGHTED
+        self.COMPUTE_PARTIAL_IONIZATIONS = COMPUTE_PARTIAL_IONIZATIONS
+        self.COMPUTE_PARTIAL_AND_MASSWEIGHTED = COMPUTE_PARTIAL_AND_MASSWEIGHTED
+        if self.COMPUTE_MASSWEIGHTED or self.COMPUTE_PARTIAL_IONIZATIONS or self.COMPUTE_PARTIAL_AND_MASSWEIGHTED:
             self.COMPUTE_DENSITY_AT_ALLZ = True
-        self.COMPUTE_ZREION = COMPUTE_ZREION
 
         ### selecting redshifts and radii from available redshifts
         # redshifts
-        self._z_idx = z21_utilities.find_nearest_idx(CoeffStructure.zintegral, self.input_z)
-        self.z = CoeffStructure.zintegral[self._z_idx]
-        # radii
-        self.r_precision = r_precision
-        self._r_idx = np.linspace(0, len(CoeffStructure.Rtabsmoo)-1, int(len(CoeffStructure.Rtabsmoo)*self.r_precision), dtype=int)
-        self.r = CoeffStructure.Rtabsmoo[self._r_idx]
+        self._z_idx = np.arange(len(np.atleast_1d(input_z))) #z21_utilities.find_nearest_idx(CoeffStructure.zintegral, self.input_z)
+        self.z = np.atleast_1d(input_z) #CoeffStructure.zintegral[self._z_idx]
 
         ### generating the density field at the closest redshift to the lower one inputed
         self.z_of_density = self.z[0]
         self.density = self.generate_density(ClassyCosmo, CorrFClass)
-        self.density_allz = np.empty((len(self.z), self.ncells, self.ncells, self.ncells))
+        self.density_allz = np.empty((len(self.z), self.ncells, self.ncells, self.ncells), dtype=np.float32)
         if self.COMPUTE_DENSITY_AT_ALLZ:
             self.generate_density_allz(CosmoParams)
 
@@ -253,26 +244,28 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
         ### generating the ionized field, and computing the ionized fraction
         self.barrier = barrier
         if self.barrier is None:
-            self.barrier = BMF.barrier
-        self.ion_field_allz = self.generate_xHII(CosmoParams, CoeffStructure, BMF)
-        self.ion_frac = self.compute_volumeweigthed_xHII(self.ion_field_allz)
-        if self.INCLUDE_PARTIALION:
-            self.ion_field_partial_allz = self.generate_xHII_withPartialIonization(CosmoParams, BMF)
-            self.ion_frac_subpixel = self.compute_volumeweigthed_xHII(self.ion_field_partial_allz)
+            self.barrier = np.array([BMF.B(z, self.r) for z in self.z]) #BMF.barrier
+        self.ion_field_allz, self.ion_frac = self.generate_xHII(CosmoParams, CoeffStructure, BMF)
 
         ### computing the mass weighted ionized fraction
+
+        self._has_mw = False
         self.lowres_massweighting = lowres_massweighting
-        self.ion_frac_massweighted = np.empty(len(self.z))
-        if self.COMPUTE_MASSWEIGHTED_IONFRAC:
-            self.ion_frac_massweighted = self.compute_massweighted_xHII(CosmoParams, self.ion_field_allz, self.lowres_massweighting)
+        if self.COMPUTE_MASSWEIGHTED:
+            self.compute_massweighted(CosmoParams, self.lowres_massweighting)
 
-            if self.INCLUDE_PARTIALION:
-                self.ion_frac_massweighted_subpixel = self.compute_massweighted_xHII(CosmoParams, self.ion_field_partial_allz, self.lowres_massweighting)
+        self._has_p = False
+        if self.COMPUTE_PARTIAL_IONIZATIONS:
+            self.compute_partial(CosmoParams, BMF)
+
+        self._has_mwp = False
+        if self.COMPUTE_PARTIAL_AND_MASSWEIGHTED:
+            self.compute_partial_massweighted(CosmoParams, BMF)
 
-        ### computing zreion and treion
         if self.COMPUTE_ZREION:
             self.zreion = self.compute_zreion_frombinaryxHII()
             self.treion = self.compute_treion(ClassyCosmo)
+            
         
         if self.PRINT_TIMER:
             z21_utilities.print_timer(self._start_time, text_before="Total computation time: ")
@@ -294,8 +287,7 @@ def generate_density(self, ClassyCosmo, CorrFClass):
             pb = pbox.LogNormalPowerBox(N=self.ncells, dim=3, pk=(lambda k: np.exp(pk_spl(np.log(k)))), boxlength=self.boxlength, seed=self.seed)
         else:
             pb = pbox.PowerBox(N=self.ncells, dim=3, pk=(lambda k: np.exp(pk_spl(np.log(k)))), boxlength=self.boxlength, seed=self.seed)
-        density_field = pb.delta_x()
-
+        density_field = pb.delta_x().astype(np.float32, copy=False)
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
         return density_field
@@ -304,14 +296,15 @@ def generate_density_allz(self, CosmoParams):
         if self.PRINT_TIMER:
             start_time = time.time()
             print('Evolving density field...')
-
         Dg = CosmoParams.growthint(self.z)
         growthfactor_ratio = (Dg/Dg[0])[:, np.newaxis, np.newaxis, np.newaxis]
         density_lastz = np.copy(self.density)
         self.density_allz = density_lastz[np.newaxis]*growthfactor_ratio
-
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
+
+        self._has_density = True
+            
         return self.density_allz
 
     def compute_k(self):
@@ -323,10 +316,8 @@ def smooth_density(self):
         if self.PRINT_TIMER:
             start_time = time.time()
             print("Smoothing density field...")
-
         density_fft = np.fft.fftn(self.density)
         density_smoothed_allr = np.array([z21_utilities.tophat_smooth(rr, self._k, density_fft) for rr in self.r])
-
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
         return density_smoothed_allr
@@ -335,20 +326,23 @@ def generate_xHII(self, CosmoParams, CoeffStructure, BMF):
         if self.PRINT_TIMER:
             start_time = time.time()
             print("Generating ionized field...")
-
         ion_field_allz = np.zeros((len(self.z),self.ncells,self.ncells,self.ncells))
-        for i in trange(len(self.z)):
+        ion_frac = np.zeros(len(self.z))
+
+        iterator = trange(len(self.z)) if self.PRINT_TIMER else range(len(self.z))
+        
+        for i in iterator:
             curr_z_idx = self._z_idx[i]
             ion_field = self.ionize(CosmoParams, CoeffStructure, curr_z_idx)
             ion_field_allz[i] = ion_field
-
+            ion_frac[i] = np.sum(ion_field)/(self.ncells**3)
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
-        return ion_field_allz
+        return ion_field_allz, ion_frac
 
     def ionize(self,CosmoParams, CoeffStructure, curr_z_idx):
-        zlist = CoeffStructure.zintegral
-        Rs = CoeffStructure.Rtabsmoo
+        zlist = self.z #CoeffStructure.zintegral
+        Rs = self.r
         
         Dg0 = CosmoParams.growthint(zlist[0])
         Dg = CosmoParams.growthint(zlist[curr_z_idx])
@@ -367,57 +361,103 @@ def ionize(self,CosmoParams, CoeffStructure, curr_z_idx):
             ion_field = np.any(ion_field_Rs, axis=0)
         return ion_field
     
-    def generate_xHII_withPartialIonization(self, CosmoParams, BMF, ir=0):
+    def compute_massweighted(self, CosmoParams, lowres_massweighting=1):
+        if not self._has_mw:
+            self.ion_frac_massweighted = np.empty(len(self.z))
+            self.ion_field_massweighted_allz = np.empty_like(self.ion_field_allz)
+        if not self._has_density:
+            self.generate_density_allz(CosmoParams)
+        self.lowres_massweighting = lowres_massweighting
+        if self.lowres_massweighting < 1:
+            raise Exception('lowres_massweighting should be >=1.')
+        if not isinstance(self.lowres_massweighting, (int, np.int32, np.int64)):
+            raise Exception('lowres_massweighting should be an integer.')
+        d_allz = self.density_allz[:, ::self.lowres_massweighting, ::self.lowres_massweighting, ::self.lowres_massweighting]
+        ion_allz = self.ion_field_allz[:, ::self.lowres_massweighting, ::self.lowres_massweighting, ::self.lowres_massweighting]
         if self.PRINT_TIMER:
             start_time = time.time()
-            print("Generating ionized field with partial ionization...")
-
-        partialion_field = np.empty(self.ion_field_allz.shape)
-        sample_d = np.linspace(-5, 5, 201)
-        for i in range(len(self.z)):
-            curr_z_idx = self._z_idx[i]
-            nion_spl = spline(sample_d, BMF.nion_delta_r_int(CosmoParams, sample_d, ir).T[curr_z_idx])
-            nrec_spl = spline(sample_d, BMF.nrec(CosmoParams, sample_d, BMF.ion_frac).T[curr_z_idx])
-            partial_ion_spl = spline(sample_d, nion_spl(sample_d)/(1+nrec_spl(sample_d)))
-            partialion_field[i] = partial_ion_spl(self.density_allz[i])
-        ion_field_partial_allz = np.clip(self.ion_field_allz + partialion_field, 0, 1)
-
+            print("Computing mass-weighted field...")
+        self.ion_field_massweighted_allz = (1+d_allz) * ion_allz
+        if self.PRINT_TIMER:
+            print("Computing mass-weighted ionized fraction...")
+        self.ion_frac_massweighted = np.average(self.ion_field_massweighted_allz, axis=(1, 2, 3))
+        
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
-        return ion_field_partial_allz
-    
-    def compute_volumeweigthed_xHII(self,xHII_field_allz):
+
+        self._has_mw = True
+        
+        return self.ion_frac_massweighted, self.ion_field_massweighted_allz
+
+    def compute_partial(self, CosmoParams, BMF, r=None):
+        if r is None:
+            r = self.r[0]
+        if not self._has_p:
+            self.ion_frac_partial = np.empty(len(self.z))
+            self.ion_field_partial_allz = np.empty_like(self.ion_field_allz)
+        if not self._has_density:
+            self.generate_density_allz(CosmoParams)
+        sample_d = np.linspace(-5, 5, 51)
+
         if self.PRINT_TIMER:
             start_time = time.time()
-            print("Computing volume weighted ionized fraction...")
+            print("Computing partially ionized field...")
+
+        iterator = trange(len(self.z)) if self.PRINT_TIMER else range(len(self.z))
+        for i in iterator:
+            nion_spl = spline(sample_d, BMF.nion_delta_r_int(CosmoParams, sample_d, self.z, r)[:, i])
+            nrec_spl = spline(sample_d, BMF.nrec(CosmoParams, sample_d, BMF.ion_frac, self.z)[:, i])
+            partial_ion_spl = spline(sample_d, nion_spl(sample_d)/(1+nrec_spl(sample_d)))
 
-        ion_frac = np.sum(xHII_field_allz,axis=(1,2,3))/(self.ncells**3)
+            #if need to do by slices:
+            #np.array([partial_ion_noclip(maps.density_allz[:, i], maps.ion_field_allz[:, i], ir=0) for i in trange(300)]).transpose(1, 0, 2, 3)
+            
+            partialfield = np.abs(partial_ion_spl(self.density_allz[i]))
+            sumfield = self.ion_field_allz[i] + partialfield
+            self.ion_field_partial_allz[i] = np.clip(sumfield, 0, 1)
+
+        if self.PRINT_TIMER:
+            print("Computing partial ionized fraction...")
+
+        self.ion_frac_partial = np.average(self.ion_field_partial_allz, axis=(1, 2, 3))
 
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
-        return ion_frac
-    
-    def compute_massweighted_xHII(self, CosmoParams, xHII_field_allz, lowres_massweighting=1):
-        if np.sum(self.density_allz[0, 0, 0]) == 0.:
-            self.generate_density_allz(CosmoParams)
-        self.lowres_massweighting = lowres_massweighting
-        if self.lowres_massweighting < 1:
-            raise Exeption('lowres_massweighting should be >=1.')
-        if not isinstance(self.lowres_massweighting, (int, np.int32, np.int64)):
-            raise Exeption('lowres_massweighting should be an integer.')
-        d_allz = self.density_allz[:, ::self.lowres_massweighting, ::self.lowres_massweighting, ::self.lowres_massweighting]
-        ion_allz = xHII_field_allz[:, ::self.lowres_massweighting, ::self.lowres_massweighting, ::self.lowres_massweighting]
+
+        self._has_p = True
+            
+        return self.ion_frac_partial, self.ion_field_partial_allz
+
+    def compute_partial_massweighted(self, CosmoParams, BMF, r=None):
+        if not self._has_p:
+            self.compute_partial(CosmoParams, BMF, r)
+
+        if not self._has_mwp:
+            self.ion_frac_partial_massweighted = np.empty(len(self.z))
+            self.ion_field_partial_massweighted_allz = np.empty_like(self.ion_field_allz)
 
         if self.PRINT_TIMER:
             start_time = time.time()
-            print("Computing mass weighted ionized fraction...")
+            print("Computing mass-weighted partially ionized field...")
 
-        ion_frac_massweighted = np.average((1+d_allz) * ion_allz, axis=(1, 2, 3))
+        iterator = trange(len(self.z)) if self.PRINT_TIMER else range(len(self.z))
+        for i in iterator:
+            self.ion_field_partial_massweighted_allz[i] = (1+self.density_allz[i]) * self.ion_field_partial_allz[i]
+        
+        if self.PRINT_TIMER:
+            print("Computing mass-weighted partial ionized fraction...")
 
+        iterator = trange(len(self.z)) if self.PRINT_TIMER else range(len(self.z))
+        for i in iterator:
+            self.ion_frac_partial_massweighted[i] = np.average(self.ion_field_partial_massweighted_allz[i])
+        
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
-        return ion_frac_massweighted
-    
+
+        self._has_mwp = True
+        
+        return self.ion_frac_partial_massweighted, self.ion_field_partial_massweighted_allz
+
     def compute_zreion_frombinaryxHII(self):
         if self.PRINT_TIMER:
             start_time = time.time()
@@ -460,3 +500,5 @@ def _compute_ionfrac_from_treion(self):
         for i in range(len(tvalues)):
             neutfrac[i] = np.sum(self.treion>tvalues[i]) / self.ncells**3
         return 1-neutfrac, tvalues
+        
+
diff --git a/zeus21/reionization.py b/zeus21/reionization.py
index 4f225c7..7644cc8 100644
--- a/zeus21/reionization.py
+++ b/zeus21/reionization.py
@@ -4,7 +4,7 @@
 See Sklansky et al. (in prep)
 
 Authors: Yonatan Sklansky, Emilie Thelie
-UT Austin - January 2025
+UT Austin - October 2025
 
 """
 

From 1efe55cceaec5caebaa0b1fb872e68d96f50b5ee Mon Sep 17 00:00:00 2001
From: Yonatan Sklansky <yonatan.sklansky@gmail.com>
Date: Tue, 21 Oct 2025 17:39:04 -0500
Subject: [PATCH 18/37] Add ionized maps to __init__.py

---
 zeus21/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zeus21/__init__.py b/zeus21/__init__.py
index 7a9caee..97bfdd4 100644
--- a/zeus21/__init__.py
+++ b/zeus21/__init__.py
@@ -5,7 +5,7 @@
 from .sfrd import get_T21_coefficients
 from .xrays import Xray_class
 from .UVLFs import UVLF_binned
-from .maps import CoevalMaps
+from .maps import CoevalMaps, reionization_maps
 from .reionization import BMF
 
 import warnings

From ff54cc3121d4a855828bc4fcd0d6c6abe1377826 Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Fri, 24 Oct 2025 12:24:10 -0500
Subject: [PATCH 19/37] Update maps.py

---
 zeus21/maps.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index 02ee8cb..89fbef2 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -222,6 +222,7 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
         self.COMPUTE_MASSWEIGHTED = COMPUTE_MASSWEIGHTED
         self.COMPUTE_PARTIAL_IONIZATIONS = COMPUTE_PARTIAL_IONIZATIONS
         self.COMPUTE_PARTIAL_AND_MASSWEIGHTED = COMPUTE_PARTIAL_AND_MASSWEIGHTED
+        self.COMPUTE_ZREION = COMPUTE_ZREION
         if self.COMPUTE_MASSWEIGHTED or self.COMPUTE_PARTIAL_IONIZATIONS or self.COMPUTE_PARTIAL_AND_MASSWEIGHTED:
             self.COMPUTE_DENSITY_AT_ALLZ = True
 

From 100a0b07065277247c2a7642ae2f660575140f2b Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Fri, 24 Oct 2025 12:39:04 -0500
Subject: [PATCH 20/37] change np.trapz to np.trapezoid

---
 zeus21/reionization.py |  6 +++---
 zeus21/sfrd.py         | 34 +++++++++++++++++-----------------
 2 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/zeus21/reionization.py b/zeus21/reionization.py
index 7644cc8..5db95d0 100644
--- a/zeus21/reionization.py
+++ b/zeus21/reionization.py
@@ -69,7 +69,7 @@ def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, Classy
 
         self.BMF = np.array([self.VRdn_dR(z, self.Rs_BMF) for z in self.zlist]) #initial bubble mass function
         self.BMF_initial = np.copy(self.BMF)
-        self.ion_frac = np.nan_to_num([np.trapz(self.BMF[i], np.log(self.Rs_BMF)) for i in range(len(self.zlist))]) #ion_frac by integrating the BMF
+        self.ion_frac = np.nan_to_num([np.trapezoid(self.BMF[i], np.log(self.Rs_BMF)) for i in range(len(self.zlist))]) #ion_frac by integrating the BMF
         self.ion_frac[self.barrier[:, -1]<=0] = 1
         
         if FLAG_converge:
@@ -249,7 +249,7 @@ def calc_Q(self, CosmoParams, z):
         niondot_avgs = self.niondot_avg_int(z_arr)
         integrand = dtdz * niondot_avgs * exp
     
-        return np.trapz(integrand, x = z_arr)
+        return np.trapezoid(integrand, x = z_arr)
 
     #computing linear barrier
     def B_1(self, z):
@@ -298,7 +298,7 @@ def barrierz_int(self, z, R):
                 return self.interpz(z, R, self.barrier_int)
             
             self.BMF = np.array([self.VRdn_dR(z, self.Rs_BMF) for z in self.zlist])
-            self.ion_frac = np.nan_to_num([np.trapz(self.BMF[i], np.log(self.Rs_BMF)) for i in range(len(self.zlist))])
+            self.ion_frac = np.nan_to_num([np.trapezoid(self.BMF[i], np.log(self.Rs_BMF)) for i in range(len(self.zlist))])
             self.ion_frac[self.barrier[:, -1]<=0] = 1
 
             if np.allclose(ion_frac_prev, self.ion_frac):
diff --git a/zeus21/sfrd.py b/zeus21/sfrd.py
index ebf3c93..b618e65 100644
--- a/zeus21/sfrd.py
+++ b/zeus21/sfrd.py
@@ -129,14 +129,14 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
         zSFRD, mArray = np.meshgrid(zSFRDflat, HMF_interpolator.Mhtab, indexing = 'ij', sparse = True)
 
         J21LW_interp = interpolate.interp1d(zSFRDflat, np.zeros_like(zSFRDflat), kind = 'linear', bounds_error = False, fill_value = 0,) #no LW background. Controls only Mmol() function, NOT the individual Pop II and III LW background
-        SFRD_II_avg = np.trapz(SFRD_II_integrand(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, mArray, zSFRD, zSFRD), HMF_interpolator.logtabMh, axis = 1) #never changes with J_LW
+        SFRD_II_avg = np.trapezoid(SFRD_II_integrand(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, mArray, zSFRD, zSFRD), HMF_interpolator.logtabMh, axis = 1) #never changes with J_LW
         SFRD_II_interp = interpolate.interp1d(zSFRDflat, SFRD_II_avg, kind = 'cubic', bounds_error = False, fill_value = 0,)
 
         J21LW_II = 1e21 * J_LW(Astro_Parameters, Cosmo_Parameters, SFRD_II_avg, zSFRDflat, 2) #this never changes; only Pop III Quanties change
         self.J_21_LW_II = interpolate.interp1d(zSFRDflat, J21LW_II, kind = 'cubic')(self.zintegral) #different from J21LW_interp
 
         if Astro_Parameters.USE_POPIII == True:
-            SFRD_III_Iter_Matrix = [np.trapz(SFRD_III_integrand(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, mArray, J21LW_interp, zSFRD, zSFRD, ClassCosmo.pars['v_avg']), HMF_interpolator.logtabMh, axis = 1)] #changes with each iteration
+            SFRD_III_Iter_Matrix = [np.trapezoid(SFRD_III_integrand(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, mArray, J21LW_interp, zSFRD, zSFRD, ClassCosmo.pars['v_avg']), HMF_interpolator.logtabMh, axis = 1)] #changes with each iteration
 
             errorTolerance = 0.001 # 0.1 percent accuracy
             recur_iterate_Flag = True
@@ -144,7 +144,7 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
                 J21LW_III_iter = 1e21 * J_LW(Astro_Parameters, Cosmo_Parameters, SFRD_III_Iter_Matrix[-1], zSFRDflat, 3)
                 J21LW_interp = interpolate.interp1d(zSFRDflat, J21LW_II + J21LW_III_iter, kind = 'linear', fill_value = 0, bounds_error = False)
 
-                SFRD_III_avg_n = np.trapz(SFRD_III_integrand(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, mArray, J21LW_interp, zSFRD, zSFRD, ClassCosmo.pars['v_avg']), HMF_interpolator.logtabMh, axis = 1)
+                SFRD_III_avg_n = np.trapezoid(SFRD_III_integrand(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, mArray, J21LW_interp, zSFRD, zSFRD, ClassCosmo.pars['v_avg']), HMF_interpolator.logtabMh, axis = 1)
                 SFRD_III_Iter_Matrix.append(SFRD_III_avg_n)
 
                 if max(SFRD_III_Iter_Matrix[-1]/SFRD_III_Iter_Matrix[-2]) < 1.0 + errorTolerance and min(SFRD_III_Iter_Matrix[-1]/SFRD_III_Iter_Matrix[-2]) > 1.0 - errorTolerance:
@@ -170,8 +170,8 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
             zpTable, tempTable, mTable = np.meshgrid(self.zintegral, self.Rtabsmoo, HMF_interpolator.Mhtab, indexing = 'ij', sparse = True)
             zppTable = self.zGreaterMatrix.reshape((len(self.zintegral), len(self.Rtabsmoo), 1))
 
-            self.SFRDbar2D_II = np.trapz(SFRD_II_integrand(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, mTable, zppTable, zpTable), HMF_interpolator.logtabMh, axis = 2)
-            self.SFRDbar2D_III = np.trapz(SFRD_III_integrand(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, mTable, J21LW_interp, zppTable, zpTable, ClassCosmo.pars['v_avg']), HMF_interpolator.logtabMh, axis = 2)
+            self.SFRDbar2D_II = np.trapezoid(SFRD_II_integrand(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, mTable, zppTable, zpTable), HMF_interpolator.logtabMh, axis = 2)
+            self.SFRDbar2D_III = np.trapezoid(SFRD_III_integrand(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, mTable, J21LW_interp, zppTable, zpTable, ClassCosmo.pars['v_avg']), HMF_interpolator.logtabMh, axis = 2)
 
             self.SFRDbar2D_II[np.isnan(self.SFRDbar2D_II)] = 0.0
             self.SFRDbar2D_III[np.isnan(self.SFRDbar2D_III)] = 0.0
@@ -228,8 +228,8 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
 
         ########
         # Compute SFRD and niondot quantities
-        SFRD_II_dR = np.trapz(integrand_II, HMF_interpolator.logtabMh, axis = 2)
-        niondot_II_dR = np.trapz(integrand_II*fesctab_II[None, None, :, None], HMF_interpolator.logtabMh, axis = 2)
+        SFRD_II_dR = np.trapezoid(integrand_II, HMF_interpolator.logtabMh, axis = 2)
+        niondot_II_dR = np.trapezoid(integrand_II*fesctab_II[None, None, :, None], HMF_interpolator.logtabMh, axis = 2)
 
         ###
         if Astro_Parameters.USE_POPIII == True:
@@ -238,8 +238,8 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
             elif(Cosmo_Parameters.Flag_emulate_21cmfast==True):
                 integrand_III = PS_HMF_corr * SFR_III(Astro_Parameters, Cosmo_Parameters, ClassCosmo, HMF_interpolator, mArray, J21LW_interp, zGreaterArray, zGreaterArray, ClassCosmo.pars['v_avg']) * mArray
 
-            SFRD_III_dR = np.trapz(integrand_III, HMF_interpolator.logtabMh, axis = 2)
-            niondot_III_dR = np.trapz(integrand_III*fesctab_III[None, None, :, None], HMF_interpolator.logtabMh, axis = 2)
+            SFRD_III_dR = np.trapezoid(integrand_III, HMF_interpolator.logtabMh, axis = 2)
+            niondot_III_dR = np.trapezoid(integrand_III*fesctab_III[None, None, :, None], HMF_interpolator.logtabMh, axis = 2)
         else:
             SFRD_III_dR = np.zeros_like(SFRD_II_dR)
             
@@ -347,7 +347,7 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
                 else:
                     print("ERROR: Need to set FLAG_EMULATE_21CMFAST at True or False in the self.gamma_index2D calculation.")
 
-                SFRD_III_dR_V = np.trapz(integrand_III, HMF_interpolator.logtabMh, axis = 2)
+                SFRD_III_dR_V = np.trapezoid(integrand_III, HMF_interpolator.logtabMh, axis = 2)
 
                 SFRDIII_Ratio = SFRD_III_dR_V / SFRD_III_dR_V[:,:,len(vAvg_array)//2].reshape((len(self.zintegral), len(self.Rtabsmoo), 1))
                 SFRDIII_Ratio[np.isnan(SFRDIII_Ratio)] = 0.0
@@ -471,7 +471,7 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
 
         opticalDepthIntegrand = 1 / cosmology.HubinvMpc(Cosmo_Parameters, zPPCube) / (1+zPPCube) * sigmatot * cosmology.n_H(Cosmo_Parameters, zPPCube) * constants.Mpctocm #this uses atom fractions of 1 for HI and x_He for HeI
 #        opticalDepthIntegrand = 1 / cosmology.HubinvMpc(Cosmo_Parameters, zPPCube) / (1+zPPCube) * sigmatot * cosmology.n_baryon(Cosmo_Parameters, zPPCube) * constants.Mpctocm
-        tauCube = np.trapz(opticalDepthIntegrand, zPPCube, axis = 3)
+        tauCube = np.trapezoid(opticalDepthIntegrand, zPPCube, axis = 3)
 
         indextautoolarge = np.array(tauCube>=Xrays.TAUMAX)
         tauCube[indextautoolarge] = Xrays.TAUMAX
@@ -629,8 +629,8 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete
         integrand_II_table = SFRD_II_integrand(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, mArray, zArray, zArray)
         integrand_III_table = SFRD_III_integrand(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, mArray, J21LW_interp, zArray, zArray, ClassCosmo.pars['v_avg'])
         
-        self.niondot_avg_II = Astro_Parameters.N_ion_perbaryon_II/cosmology.rho_baryon(Cosmo_Parameters,0.) * np.trapz(integrand_II_table * fesctab_II, HMF_interpolator.logtabMh, axis = 1)
-        self.niondot_avg_III = Astro_Parameters.N_ion_perbaryon_II/cosmology.rho_baryon(Cosmo_Parameters,0.) * np.trapz(integrand_III_table * fesctab_III, HMF_interpolator.logtabMh, axis = 1)
+        self.niondot_avg_II = Astro_Parameters.N_ion_perbaryon_II/cosmology.rho_baryon(Cosmo_Parameters,0.) * np.trapezoid(integrand_II_table * fesctab_II, HMF_interpolator.logtabMh, axis = 1)
+        self.niondot_avg_III = Astro_Parameters.N_ion_perbaryon_II/cosmology.rho_baryon(Cosmo_Parameters,0.) * np.trapezoid(integrand_III_table * fesctab_III, HMF_interpolator.logtabMh, axis = 1)
         self.niondot_avg = self.niondot_avg_II + self.niondot_avg_III
 
         if(Cosmo_Parameters.Flag_emulate_21cmfast==False): #regular calculation, integrating over time and accounting for recombinations in the exponent
@@ -678,7 +678,7 @@ def tau_reio(Cosmo_Parameters, T21_coefficients):
     _nelistlowz = cosmology.n_H(Cosmo_Parameters,_zlistlowz)*(1.0 + Cosmo_Parameters.x_He + Cosmo_Parameters.x_He * np.heaviside(constants.zHeIIreio - _zlistlowz,0.5))
 #    _nelistlowz = cosmology.n_baryon(Cosmo_Parameters,_zlistlowz)*(Cosmo_Parameters.f_H + Cosmo_Parameters.f_He + Cosmo_Parameters.f_He * np.heaviside(constants.zHeIIreio - _zlistlowz,0.5))
     _distlistlowz = 1.0/cosmology.HubinvMpc(Cosmo_Parameters,_zlistlowz)/(1+_zlistlowz)
-    _lowzint = constants.sigmaT * np.trapz(_nelistlowz*_distlistlowz,_zlistlowz) * constants.Mpctocm
+    _lowzint = constants.sigmaT * np.trapezoid(_nelistlowz*_distlistlowz,_zlistlowz) * constants.Mpctocm
 
     _zlisthiz = T21_coefficients.zintegral
     
@@ -686,7 +686,7 @@ def tau_reio(Cosmo_Parameters, T21_coefficients):
 #    _nelistlhiz = cosmology.n_baryon(Cosmo_Parameters,_zlisthiz) * (1.0 - T21_coefficients.xHI_avg)
     _distlisthiz = 1.0/cosmology.HubinvMpc(Cosmo_Parameters,_zlisthiz)/(1+_zlisthiz)
 
-    _hizint = constants.sigmaT * np.trapz(_nelistlhiz*_distlisthiz,_zlisthiz) * constants.Mpctocm
+    _hizint = constants.sigmaT * np.trapezoid(_nelistlhiz*_distlisthiz,_zlisthiz) * constants.Mpctocm
 
     return(_lowzint + _hizint)
 
@@ -768,7 +768,7 @@ def J_LW(Astro_Parameters, Cosmo_Parameters, sfrdIter, z, pop):
 #    integrandLW *= (1+z)**3 / cosmology.Hubinvyr(Cosmo_Parameters,zIntMatrix) / (1+zIntMatrix) #HAC: delete this and comment above back in!!!
     integrandLW *= Nlw * Elw / massProton / deltaNulw
     integrandLW = integrandLW * sfrdIterMatrix_LW * (1 /u.Mpc**2).to(1/u.cm**2).value #broadcasting doesn't like augmented assignment operations (like *=) for some reason
-    return np.trapz(integrandLW, x = zIntMatrix, axis = 0)
+    return np.trapezoid(integrandLW, x = zIntMatrix, axis = 0)
 
 
 def SFRD_II_integrand(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, massVector, z, z2):
@@ -902,7 +902,7 @@ def dSFRDIII_dJ(Astro_Parameters, Cosmo_Parameters, HMF_interpolator, J21LW_inte
     integrand_III = HMF_curr * SFRtab_currIII * HMF_interpolator.Mhtab
     integrand_III *= Astro_Parameters.A_LW * Astro_Parameters.beta_LW * J21LW_interp(z)**(Astro_Parameters.beta_LW - 1)
     integrand_III *= -1 * Mmol_vcb(Astro_Parameters, Cosmo_Parameters, z, Cosmo_Parameters.vcb_avg)/ HMF_interpolator.Mhtab
-    return np.trapz(integrand_III, HMF_interpolator.logtabMh)
+    return np.trapezoid(integrand_III, HMF_interpolator.logtabMh)
 
 
 def fesc_II(Astro_Parameters, Mh):

From c4798f7856fd460b350cae8453bd2192d30345d6 Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Fri, 24 Oct 2025 13:08:46 -0500
Subject: [PATCH 21/37] Fixed some test calls

---
 tests/test_astrophysics.py | 2 +-
 tests/test_maps.py         | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/test_astrophysics.py b/tests/test_astrophysics.py
index 9b23d20..4fbc238 100644
--- a/tests/test_astrophysics.py
+++ b/tests/test_astrophysics.py
@@ -43,7 +43,7 @@
 AstroParams_21cmfast = zeus21.Astro_Parameters(UserParams,CosmoParams_21cmfast, astromodel = 1)
 
 #and for bubbles:
-BMF_class = zeus21.BMF(UserParams, Coeffs, HMFintclass, CosmoParams, AstroParams)
+BMF_class = zeus21.BMF(Coeffs, HMFintclass, CosmoParams, AstroParams, ClassyCosmo)
 
 
 ztest = 20.
diff --git a/tests/test_maps.py b/tests/test_maps.py
index 4b37381..977abd6 100644
--- a/tests/test_maps.py
+++ b/tests/test_maps.py
@@ -11,7 +11,8 @@
 import zeus21
 import numpy as np
 
-from zeus21.maps import CoevalMaps, powerboxCtoR
+from zeus21.maps import CoevalMaps
+from zeus21.z21_utilities import powerboxCtoR
 
 def test_coevalmaps_initialization():
     """Test that CoevalMaps initializes correctly"""

From 85a0634d7a1eb181a8efa03e91cec41f1e5759de Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Fri, 24 Oct 2025 13:21:23 -0500
Subject: [PATCH 22/37] More trapz to trapezoid, and also uncommented import
 empty

---
 tests/test_inputs.py   | 8 ++++----
 zeus21/UVLFs.py        | 4 ++--
 zeus21/correlations.py | 4 ++--
 zeus21/cosmology.py    | 8 ++++----
 zeus21/maps.py         | 2 +-
 zeus21/xrays.py        | 4 ++--
 6 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/tests/test_inputs.py b/tests/test_inputs.py
index 4a3102d..606f143 100644
--- a/tests/test_inputs.py
+++ b/tests/test_inputs.py
@@ -96,23 +96,23 @@ def test_inputs():
     #test Pop II Xray SED
     Energylisttest = np.logspace(2,np.log10(AstroParams.Emax_xray_norm),100)
     SEDXtab_test = AstroParams.SED_XRAY(Energylisttest, 2) #same in both models
-    normalization_XraySED = np.trapz(Energylisttest * SEDXtab_test,Energylisttest)
+    normalization_XraySED = np.trapezoid(Energylisttest * SEDXtab_test,Energylisttest)
     assert( normalization_XraySED == pytest.approx(1.0, 0.05) ) #5% is enough here
     
     #test Pop III Xray SED
     SEDXtab_test = AstroParams.SED_XRAY(Energylisttest, 3) #same in both models
-    normalization_XraySED = np.trapz(Energylisttest * SEDXtab_test,Energylisttest)
+    normalization_XraySED = np.trapezoid(Energylisttest * SEDXtab_test,Energylisttest)
     assert( normalization_XraySED == pytest.approx(1.0, 0.05) ) #5% is enough here
 
 
     #test Pop II LyA SED
     nulisttest = np.linspace(zeus21.constants.freqLyA, zeus21.constants.freqLyCont, 100)
     SEDLtab_test = AstroParams.SED_LyA(nulisttest, 2) #same in both models
-    normalization_LyASED = np.trapz(SEDLtab_test,nulisttest)
+    normalization_LyASED = np.trapezoid(SEDLtab_test,nulisttest)
     assert( normalization_LyASED == pytest.approx(1.0, 0.05) ) #5% is enough here
     
     #test Pop III LyA SED
     nulisttest = np.linspace(zeus21.constants.freqLyA, zeus21.constants.freqLyCont, 100)
     SEDLtab_test = AstroParams.SED_LyA(nulisttest, 3) #same in both models
-    normalization_LyASED = np.trapz(SEDLtab_test,nulisttest)
+    normalization_LyASED = np.trapezoid(SEDLtab_test,nulisttest)
     assert( normalization_LyASED == pytest.approx(1.0, 0.05) ) #5% is enough here
diff --git a/zeus21/UVLFs.py b/zeus21/UVLFs.py
index b654c53..ec08a94 100644
--- a/zeus21/UVLFs.py
+++ b/zeus21/UVLFs.py
@@ -82,7 +82,7 @@ def UVLF_binned(Astro_Parameters,Cosmo_Parameters,HMF_interpolator, zcenter, zwi
     xlo = np.subtract.outer(MUVcutlo, currMUV )/(np.sqrt(2) * sigmaUV)
     weights = (erf(xhi) - erf(xlo)).T/(2.0 * MUVwidths)
     
-    UVLF_filtered = np.trapz(weights.T * HMFcurr, HMF_interpolator.Mhtab, axis=-1)
+    UVLF_filtered = np.trapezoid(weights.T * HMFcurr, HMF_interpolator.Mhtab, axis=-1)
 
     if(Astro_Parameters.USE_POPIII==False):
         return UVLF_filtered
@@ -98,7 +98,7 @@ def UVLF_binned(Astro_Parameters,Cosmo_Parameters,HMF_interpolator, zcenter, zwi
         xlo = np.subtract.outer(MUVcutlo, MUVbarlist_III)/(np.sqrt(2) * sigmaUV)
         weights = (erf(xhi) - erf(xlo)).T/(2.0 * MUVwidths)
 
-        UVLF_filtered_III = np.trapz(weights.T * HMFcurr, HMF_interpolator.Mhtab, axis=-1)
+        UVLF_filtered_III = np.trapezoid(weights.T * HMFcurr, HMF_interpolator.Mhtab, axis=-1)
     
         return UVLF_filtered, UVLF_filtered_III
 
diff --git a/zeus21/correlations.py b/zeus21/correlations.py
index 8e3287a..7653ecf 100644
--- a/zeus21/correlations.py
+++ b/zeus21/correlations.py
@@ -374,7 +374,7 @@ def __init__(self, User_Parameters, Cosmo_Parameters, Astro_Parameters, ClassCos
             #the z>zmax part of the integral we do aside. Assume Tk=Tadiabatic from CLASS.
             _zlisthighz_ = np.linspace(T21_coefficients.zintegral[-1], 99., 100) #beyond z=100 need to explictly tell CLASS to save growth
             _dgrowthhighz_ = cosmology.dgrowth_dz(Cosmo_Parameters, _zlisthighz_)
-            _hizintegral = np.trapz(cosmology.Tadiabatic(Cosmo_Parameters,_zlisthighz_)
+            _hizintegral = np.trapezoid(cosmology.Tadiabatic(Cosmo_Parameters,_zlisthighz_)
             /(1+_zlisthighz_)**2 * _dgrowthhighz_, _zlisthighz_)
 
         self._betaTad_ = -2./3. * _factor_adi_/self._lingrowthd * (np.cumsum(_integrand_adi[::-1])[::-1] + _hizintegral) #units of Tk_avg. Internal sum goes from high to low z (backwards), minus sign accounts for it properly so it's positive.
@@ -1124,7 +1124,7 @@ def get_Pk_from_xi(self, rsinput, xiinput):
 #
 #                 Probdtab = np.exp(-dtab**2/sigmaRRsq/2.0)
 #
-#                 norm = np.trapz(NionEPS * Probdtab, dtab)
+#                 norm = np.trapezoid(NionEPS * Probdtab, dtab)
 #                 NionEPS/=norm
 #
 #                 bindex = min(range(len(NionEPS)), key=lambda i: abs(NionEPS[i]-_invQbar))
diff --git a/zeus21/cosmology.py b/zeus21/cosmology.py
index 028533c..564e34f 100644
--- a/zeus21/cosmology.py
+++ b/zeus21/cosmology.py
@@ -75,13 +75,13 @@ def runclass(CosmologyIn):
         theta_b = velTransFunc['t_b']
         theta_c = velTransFunc['t_cdm']
 
-        sigma_vcb = np.sqrt(np.trapz(CosmologyIn.As * (kVel/0.05)**(CosmologyIn.ns-1) /kVel * (theta_b - theta_c)**2/kVel**2, kVel)) * constants.c_kms
+        sigma_vcb = np.sqrt(np.trapezoid(CosmologyIn.As * (kVel/0.05)**(CosmologyIn.ns-1) /kVel * (theta_b - theta_c)**2/kVel**2, kVel)) * constants.c_kms
         ClassCosmo.pars['sigma_vcb'] = sigma_vcb
         
         ###HAC: now computing average velocity assuming a Maxwell-Boltzmann distribution of velocities
         velArr = np.geomspace(0.01, constants.c_kms, 1000) #in km/s
         vavgIntegrand = (3 / (2 * np.pi * sigma_vcb**2))**(3/2) * 4 * np.pi * velArr**2 * np.exp(-3 * velArr**2 / (2 * sigma_vcb**2))
-        ClassCosmo.pars['v_avg'] = np.trapz(vavgIntegrand * velArr, velArr)
+        ClassCosmo.pars['v_avg'] = np.trapezoid(vavgIntegrand * velArr, velArr)
         
         ###HAC: Computing Vcb Power Spectrum
         ClassCosmo.pars['k_vcb'] = kVel
@@ -99,8 +99,8 @@ def runclass(CosmologyIn):
         j0bessel = lambda x: np.sin(x)/x
         j2bessel = lambda x: (3 / x**2 - 1) * np.sin(x)/x - 3*np.cos(x)/x**2
         
-        psi0 = 1 / 3 / (sigma_vcb/constants.c_kms)**2 * np.trapz(kVelIntp**2 / 2 / np.pi**2 * p_vcb_intp(np.log(kVelIntp)) * j0bessel(kVelIntp * np.transpose([rVelIntp])), kVelIntp, axis = 1)
-        psi2 = -2 / 3 / (sigma_vcb/constants.c_kms)**2 * np.trapz(kVelIntp**2 / 2 / np.pi**2 * p_vcb_intp(np.log(kVelIntp)) * j2bessel(kVelIntp * np.transpose([rVelIntp])), kVelIntp, axis = 1)
+        psi0 = 1 / 3 / (sigma_vcb/constants.c_kms)**2 * np.trapezoid(kVelIntp**2 / 2 / np.pi**2 * p_vcb_intp(np.log(kVelIntp)) * j0bessel(kVelIntp * np.transpose([rVelIntp])), kVelIntp, axis = 1)
+        psi2 = -2 / 3 / (sigma_vcb/constants.c_kms)**2 * np.trapezoid(kVelIntp**2 / 2 / np.pi**2 * p_vcb_intp(np.log(kVelIntp)) * j2bessel(kVelIntp * np.transpose([rVelIntp])), kVelIntp, axis = 1)
         
         k_eta, P_eta = mcfit.xi2P(rVelIntp, l=0, lowring = True)((6 * psi0**2 + 3 * psi2**2), extrap = False)
         
diff --git a/zeus21/maps.py b/zeus21/maps.py
index 89fbef2..d04109b 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -15,7 +15,7 @@
 import powerbox as pbox
 from scipy.interpolate import interp1d
 from scipy.interpolate import InterpolatedUnivariateSpline as spline
-#from pyfftw import empty_aligned as empty
+from pyfftw import empty_aligned as empty
 from tqdm import trange
 import time
 
diff --git a/zeus21/xrays.py b/zeus21/xrays.py
index 6666014..78cdaa1 100644
--- a/zeus21/xrays.py
+++ b/zeus21/xrays.py
@@ -43,7 +43,7 @@ def optical_depth(self, User_Parameters, Cosmo_Parameters, En,z,zp):
         # divided by factor of H(z')(1+z') because of variable of integration change from proper distance to redshift
         integrand = 1.0/HubinvMpc(Cosmo_Parameters, zinttau)/(1+zinttau) * sigmatot * n_H(Cosmo_Parameters, zinttau) * constants.Mpctocm
 #        integrand = 1.0/HubinvMpc(Cosmo_Parameters, zinttau)/(1+zinttau) * sigmatot * n_baryon(Cosmo_Parameters, zinttau) * constants.Mpctocm
-        taulist = np.trapz(integrand, zinttau, axis=1)
+        taulist = np.trapezoid(integrand, zinttau, axis=1)
 
         #OLD: kept for reference only.
         # taulist = 1.0*np.zeros_like(Envec)
@@ -55,7 +55,7 @@ def optical_depth(self, User_Parameters, Cosmo_Parameters, En,z,zp):
         #
         #     integrand = 1.0/HubinvMpc(Cosmo_Parameters, zinttau)/(1+zinttau) * sigmatot * n_baryon(Cosmo_Parameters, zinttau) * constants.Mpctocm
         #
-        #     taulist[iE] = np.trapz(integrand, zinttau)
+        #     taulist[iE] = np.trapezoid(integrand, zinttau)
 
         indextautoolarge = np.array(taulist>=self.TAUMAX)
         taulist [indextautoolarge] = self.TAUMAX

From c4c06ff931159e9fd635f23238b73c35a6395d4b Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Fri, 24 Oct 2025 13:30:32 -0500
Subject: [PATCH 23/37] Added required import

---
 zeus21/inputs.py        | 2 +-
 zeus21/z21_utilities.py | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/zeus21/inputs.py b/zeus21/inputs.py
index 4ceadc7..b55a859 100644
--- a/zeus21/inputs.py
+++ b/zeus21/inputs.py
@@ -36,7 +36,7 @@ class User_Parameters:
         0 to do standard calculation, 1 to force linearization of correlation function.
     MIN_R_NONLINEAR: float
         Minimum radius R/cMpc in which we start doing the nonlinear calculation. 
-        Below ~1 it will blow up because sigma > 1 eventually, and our exp(\delta) approximation breaks. 
+        Below ~1 it will blow up because sigma > 1 eventually, and our exp(delta) approximation breaks. 
         Check if you play with it and if you change Window().
     MAX_R_NONLINEAR: float
         Maximum radius R/cMpc in which we start doing the nonlinear calculation (above this it is very linear)
diff --git a/zeus21/z21_utilities.py b/zeus21/z21_utilities.py
index 9290700..9e9474a 100644
--- a/zeus21/z21_utilities.py
+++ b/zeus21/z21_utilities.py
@@ -7,6 +7,7 @@
 """
 
 import numpy as np
+from pyfftw import empty_aligned as empty
 import time
 
 def powerboxCtoR(pbobject,mapkin = None):

From 3a4e2fe411e8b8001922a0ba83cdcb4535648123 Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Fri, 24 Oct 2025 13:37:34 -0500
Subject: [PATCH 24/37] pbox import

---
 zeus21/z21_utilities.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/zeus21/z21_utilities.py b/zeus21/z21_utilities.py
index 9e9474a..4bdd2ab 100644
--- a/zeus21/z21_utilities.py
+++ b/zeus21/z21_utilities.py
@@ -7,6 +7,7 @@
 """
 
 import numpy as np
+import powerbox as pbox
 from pyfftw import empty_aligned as empty
 import time
 

From 7858c4821f9c038dd1b312036b34ecf997d68195 Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Fri, 24 Oct 2025 14:57:36 -0500
Subject: [PATCH 25/37] Added reionization map tests

---
 tests/test_maps.py | 47 +++++++++++++++++++++++++++++++++++++++++++++-
 zeus21/maps.py     |  2 +-
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/tests/test_maps.py b/tests/test_maps.py
index 977abd6..9ee1af0 100644
--- a/tests/test_maps.py
+++ b/tests/test_maps.py
@@ -157,4 +157,49 @@ def test_powerboxCtoR():
     # Test with default parameter (None)
     real_field2 = powerboxCtoR(pb)
     assert np.isreal(real_field2).all()
-    assert real_field2.shape == (20, 20, 20)
\ No newline at end of file
+    assert real_field2.shape == (20, 20, 20)
+
+def test_reionization_maps():
+    """Test reionization maps generation"""
+    # Set up the necessary objects
+    UserParams = zeus21.User_Parameters()
+    CosmoParams_input = zeus21.Cosmo_Parameters_Input(kmax_CLASS=100.) # Use higher kmax_CLASS as in test_astrophysics.py
+    ClassyCosmo = zeus21.runclass(CosmoParams_input)
+    CosmoParams = zeus21.Cosmo_Parameters(UserParams, CosmoParams_input, ClassyCosmo)
+    
+    AstroParams = zeus21.Astro_Parameters(UserParams, CosmoParams)
+    HMFintclass = zeus21.HMF_interpolator(UserParams, CosmoParams, ClassyCosmo)
+    CorrFClass = zeus21.Correlations(UserParams, CosmoParams, ClassyCosmo)
+    
+    # Generate T21 coefficients
+    ZMIN = 20.0  # Use same ZMIN as in test_astrophysics.py
+    Coeffs = zeus21.get_T21_coefficients(UserParams, CosmoParams, ClassyCosmo, AstroParams, HMFintclass, zmin=ZMIN)
+    
+    # Generate BMF
+    BMF = zeus21.BMF(Coeffs, HMFintclass, CosmoParams, AstroParams, ClassyCosmo)
+
+    # Test redshift
+    ztest = 10.0  # A redshift during reionization
+    
+    # Initialize ionized field map
+    map_obj = zeus21.reionization_maps(CosmoParams, ClassyCosmo, CorrFClass, Coeffs, BMF, ztest, input_boxlength=300, ncells=50, seed=12345, PRINT_TIMER=False, COMPUTE_PARTIAL_AND_MASSWEIGHTED=True)
+
+    # Check that ionization map exists
+    assert map_obj.ion_field_allz is not None
+    assert map_obj.ion_field_partial_massweighted_allz is not None
+    assert map_obj.ion_frac is not None
+    assert map_obj.ion_frac_partial_massweighted is not None
+    
+    # Check dimensions
+    assert map_obj.ion_field_allz.shape == (1, 50, 50, 50)
+    assert map_obj.ion_field_partial_massweighted_allz.shape == (1, 50, 50, 50)
+    assert map_obj.ion_frac.shape == (1,)
+    assert map_obj.ion_frac_partial_massweighted.shape == (1,)
+    assert map_obj.barrier.shape == (1, len(map_obj.r))
+    
+    # Check that ionization fractions are between 0 and 1
+    assert (map_obj.ion_field_allz >= 0).all() and (map_obj.ion_field_allz <= 1).all()
+    assert (map_obj.ion_field_partial_massweighted_allz >= 0).all() and (map_obj.ion_field_partial_massweighted_allz <= np.max(map_obj.density_allz+1)).all()
+    assert 0.0 <= map_obj.ion_frac <= 1.0
+    assert 0.0 <= map_obj.ion_frac_partial_massweighted <= 1.0
+    
\ No newline at end of file
diff --git a/zeus21/maps.py b/zeus21/maps.py
index d04109b..4956b59 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -245,7 +245,7 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
         ### generating the ionized field, and computing the ionized fraction
         self.barrier = barrier
         if self.barrier is None:
-            self.barrier = np.array([BMF.B(z, self.r) for z in self.z]) #BMF.barrier
+            self.barrier = np.array([BMF.B(z, self.r) for z in self.z]) #BMF linear barrier
         self.ion_field_allz, self.ion_frac = self.generate_xHII(CosmoParams, CoeffStructure, BMF)
 
         ### computing the mass weighted ionized fraction

From 2f32de52b18f375e7db06f8eb9c7926f9d73c5b0 Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Tue, 18 Nov 2025 15:29:05 -0600
Subject: [PATCH 26/37] Fixd single-z map computation and small R partials

Fixed single-z map computation by precomputing partial ionizations as a function of density in reionization.py

Allowed extrapolation of small R values in partial ionizations. This lets the user make more highly resolved maps than before.
---
 zeus21/maps.py         |  7 +---
 zeus21/reionization.py | 87 +++++++++++++++++++++++++++++++++---------
 2 files changed, 71 insertions(+), 23 deletions(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index 4956b59..e0ee561 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -406,12 +406,7 @@ def compute_partial(self, CosmoParams, BMF, r=None):
 
         iterator = trange(len(self.z)) if self.PRINT_TIMER else range(len(self.z))
         for i in iterator:
-            nion_spl = spline(sample_d, BMF.nion_delta_r_int(CosmoParams, sample_d, self.z, r)[:, i])
-            nrec_spl = spline(sample_d, BMF.nrec(CosmoParams, sample_d, BMF.ion_frac, self.z)[:, i])
-            partial_ion_spl = spline(sample_d, nion_spl(sample_d)/(1+nrec_spl(sample_d)))
-
-            #if need to do by slices:
-            #np.array([partial_ion_noclip(maps.density_allz[:, i], maps.ion_field_allz[:, i], ir=0) for i in trange(300)]).transpose(1, 0, 2, 3)
+            partial_ion_spl = spline(sample_d, BMF.prebarrier_xHII_int_grid(sample_d, self.z[i], r)) #spline is faster than RGI, so build a spline on sample densities
             
             partialfield = np.abs(partial_ion_spl(self.density_allz[i]))
             sumfield = self.ion_field_allz[i] + partialfield
diff --git a/zeus21/reionization.py b/zeus21/reionization.py
index 5db95d0..6c5ce80 100644
--- a/zeus21/reionization.py
+++ b/zeus21/reionization.py
@@ -31,6 +31,7 @@ def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, Classy
         self.zlist = CoeffStructure.zintegral
         self.Rs = CoeffStructure.Rtabsmoo
         self.Rs_BMF = np.logspace(np.log10(Rmin), np.log10(self.Rs[-1]), 100)
+        self.ds_array = np.linspace(-1, 5, 101)
 
         
         self.gamma = CoeffStructure.gamma_niondot_II_index2D
@@ -59,11 +60,15 @@ def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, Classy
         zr_mesh = np.meshgrid(np.arange(len(self.Rs)), np.arange(len(self.zlist)))
         self.nion_norm = self.nion_normalization(zr_mesh[1], zr_mesh[0])
         self.nion_norm_int = RegularGridInterpolator(self.zr, self.nion_norm, bounds_error = False, fill_value = np.nan)
-        
+
+        self.prebarrier_xHII = np.empty((len(self.ds_array), len(self.zlist), len(self.Rs)))
         self.barrier = self.compute_barrier(CosmoParams, self.ion_frac, self.zlist, self.Rs)
         self.barrier_initial = np.copy(self.barrier)
         self.barrier_int = RegularGridInterpolator(self.zr, self.barrier, bounds_error = False, fill_value = np.nan)
 
+        self.dzr = [self.ds_array, self.zlist, np.log(self.Rs)]
+        self.prebarrier_xHII_int = RegularGridInterpolator(self.dzr, self.prebarrier_xHII, bounds_error = False, fill_value = None) #allow extrapolation
+
         self.R_linear_sigma_fit_idx = z21_utilities.find_nearest_idx(self.Rs, R_linear_sigma_fit_input)
         self.R_linear_sigma_fit = self.Rs[self.R_linear_sigma_fit_idx]
 
@@ -77,6 +82,16 @@ def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, Classy
         #two functions: compute BMF and iterate
         
 
+    def compute_prebarrier_xHII(self, CosmoParams, ion_frac, z, R):
+        """
+        
+        """
+        nion_values = self.nion_delta_r_int(CosmoParams, z, R)  #Shape (nd, nz)
+        nrec_values = self.nrec(CosmoParams, ion_frac, z)       #Shape (nd, nz)
+        
+        prebarrier_xHII = nion_values / (1 + nrec_values)
+
+        return prebarrier_xHII
 
     def compute_barrier(self, CosmoParams, ion_frac, z, R):
         """
@@ -97,7 +112,6 @@ def compute_barrier(self, CosmoParams, ion_frac, z, R):
             The resultant density threshold array. First dimension is each redshift, second dimension is each radius scale.
         """
         barrier = np.zeros((len(z), len(R)))
-        ds_array = np.linspace(-1, 5, 101)
 
         zarg = np.argsort(z) #sort just in case
         z = z[zarg]
@@ -105,28 +119,26 @@ def compute_barrier(self, CosmoParams, ion_frac, z, R):
         
         for ir in range(len(R)):
             #Compute nion_values and nrec_values for this 'ir'
-            nion_values = self.nion_delta_r_int(CosmoParams, ds_array, z, R[ir])  #Shape (nd, nz)
-            nrec_values = self.nrec(CosmoParams, ds_array, ion_frac, z)             #Shape (nd, nz)
-        
-            total_values = np.log10(nion_values / (1 + nrec_values) + 1e-10)   #taking difference in logspace to find zero-crossing 
+            self.prebarrier_xHII[:, :, ir] =  self.compute_prebarrier_xHII(CosmoParams, ion_frac, z, R[ir])
+            total_values = np.log10(self.prebarrier_xHII[:, :, ir] + 1e-10)
         
             #Loop over redshift indices
-            for i in range(len(self.zlist)):
-                y_values = total_values[:, i]  #Shape (nd,)
+            for iz in range(len(self.zlist)):
+                y_values = total_values[:, iz]  #Shape (nd,)
         
                 #Find zero crossings
                 sign_change = np.diff(np.sign(y_values))
                 idx = np.where(sign_change)[0]
                 if idx.size > 0:
                     #Linear interpolation to find zero crossings
-                    x0 = ds_array[idx]
-                    x1 = ds_array[idx + 1]
+                    x0 = self.ds_array[idx]
+                    x1 = self.ds_array[idx + 1]
                     y0 = y_values[idx]
                     y1 = y_values[idx + 1]
                     x_intersect = x0 - y0 * (x1 - x0) / (y1 - y0)
-                    barrier[i, ir] = x_intersect[0]  #Assuming we take the first crossing
+                    barrier[iz, ir] = x_intersect[0]  #Assuming we take the first crossing
                 else:
-                    barrier[i, ir] = np.nan #Never crosses
+                    barrier[iz, ir] = np.nan #Never crosses
         barrier = barrier * (CosmoParams.growthint(self.zlist)/CosmoParams.growthint(self.zlist[0]))[:, np.newaxis] #scale barrier with growth factor
         barrier[self.zlist > self.ZMAX_REION] = 100 #sets density to an unreachable barrier, as if reionization isn't happening
         return barrier
@@ -135,7 +147,7 @@ def compute_barrier(self, CosmoParams, ion_frac, z, R):
     def nion_normalization(self, z, R):
         return 1/np.sqrt(1-2*self.gamma2[z, R]*self.sigma[z, R]**2)*np.exp(self.gamma[z, R]**2 * self.sigma[z, R]**2 / (2-4*self.gamma2[z, R]*self.sigma[z, R]**2))
 
-    def nrec(self, CosmoParams, d_array, ion_frac, z):
+    def nrec(self, CosmoParams, ion_frac, z, d_array=None):
         """
         Vectorized computation of nrec over an array of overdensities d_array.
 
@@ -156,6 +168,9 @@ def nrec(self, CosmoParams, d_array, ion_frac, z):
         zarg = np.argsort(z) #sort just in case
         z = z[zarg]
         ion_frac = ion_frac[zarg]
+
+        if d_array is None:
+            d_array = self.ds_array
         
         #reverse the inputs to make the integral easier to compute
         z_rev = z[::-1]
@@ -174,7 +189,7 @@ def nrec(self, CosmoParams, d_array, ion_frac, z):
         nrecs = nrecs[:, ::-1]  #reverse back to increasing z order
         return nrecs
     
-    def niondot_delta_r(self, CosmoParams, d_array, z, R):
+    def niondot_delta_r(self, CosmoParams, z, R, d_array=None):
         """
         Compute niondot over an array of overdensities d_array for a given R.
 
@@ -195,6 +210,9 @@ def niondot_delta_r(self, CosmoParams, d_array, z, R):
 
         zarg = np.argsort(z) #sort just in case
         z = z[zarg]
+
+        if d_array is None:
+            d_array = self.ds_array
         
         d_array = d_array[:, np.newaxis] * CosmoParams.growthint(z)[np.newaxis, :] / CosmoParams.growthint(z[0])
     
@@ -207,7 +225,7 @@ def niondot_delta_r(self, CosmoParams, d_array, z, R):
         
         return niondot
     
-    def nion_delta_r_int(self, CosmoParams, d_array, z, R):
+    def nion_delta_r_int(self, CosmoParams, z, R, d_array=None):
         """
         Vectorized computation of nion over an array of overdensities d_array for a given R.
 
@@ -227,12 +245,15 @@ def nion_delta_r_int(self, CosmoParams, d_array, z, R):
         """
 
         z.sort() #sort if not sorted
+
+        if d_array is None:
+            d_array = self.ds_array
         
         #reverse the inputs to make the integral easier to compute
         z_rev = z[::-1]
         Hz_rev = cosmology.Hubinvyr(CosmoParams, z_rev)
     
-        niondot_values = self.niondot_delta_r(CosmoParams, d_array, z, R)
+        niondot_values = self.niondot_delta_r(CosmoParams, z, R, d_array)
     
         integrand = -1 / (1 + z_rev) / Hz_rev * niondot_values[:, ::-1]
         nion = cumulative_trapezoid(integrand, x=z_rev, initial=0)[:, ::-1] #reverse back to increasing z order
@@ -348,4 +369,36 @@ def gamma2z_int(self, z, R):
     def nion_normR_int(self, z, R):
         return self.interpR(z, R, self.nion_norm_int)
     def nion_normz_int(self, z, R):
-        return self.interpz(z, R, self.nion_norm_int)
\ No newline at end of file
+        return self.interpz(z, R, self.nion_norm_int)
+
+    def prebarrier_xHII_int_grid(self, d, z, R):
+        """
+        Evaluate prebarrier xHII on a density field d(x),
+        at fixed redshift z and smoothing radius R.
+
+        Parameters
+        ----------
+        d: np.ndarray
+            Density/overdensity field. Can be any shape (...).
+        z: float
+            Redshift.
+        R: float
+            Smoothing radius (cMpc).
+
+        Output
+        ----------
+        values: np.ndarray
+            xHII field with the same shape as d.
+        """
+        
+        d = np.asarray(d, dtype=float)
+
+        z_arr   = np.full_like(d, float(z), dtype=float)
+        logr_arr = np.full_like(d, np.log(float(R)), dtype=float)
+
+        #stack into points (..., 3) where last axis is (delta, z, logR)
+        points = np.stack([d, z_arr, logr_arr], axis=-1)
+
+        values = self.prebarrier_xHII_int(points)
+
+        return values
\ No newline at end of file

From 547115ddd16c46cac78e1e8f15f2d0880067b783 Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Tue, 18 Nov 2025 15:44:26 -0600
Subject: [PATCH 27/37] Added print option for successful run

---
 zeus21/reionization.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/zeus21/reionization.py b/zeus21/reionization.py
index 6c5ce80..90979af 100644
--- a/zeus21/reionization.py
+++ b/zeus21/reionization.py
@@ -25,14 +25,14 @@ class BMF:
     
     """
     
-    def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, ClassyCosmo, R_linear_sigma_fit_input=10, FLAG_converge=True, max_iter=10, ZMAX_REION = 30, Rmin=0.05):
+    def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, ClassyCosmo, R_linear_sigma_fit_input=10, FLAG_converge=True, max_iter=10, ZMAX_REION = 30, Rmin=0.05, PRINT_SUCCESS=True):
 
+        self.PRINT_SUCCESS = PRINT_SUCCESS
         self.ZMAX_REION = ZMAX_REION #max redshift up to which we calculate reionization observables
         self.zlist = CoeffStructure.zintegral
         self.Rs = CoeffStructure.Rtabsmoo
         self.Rs_BMF = np.logspace(np.log10(Rmin), np.log10(self.Rs[-1]), 100)
         self.ds_array = np.linspace(-1, 5, 101)
-
         
         self.gamma = CoeffStructure.gamma_niondot_II_index2D
         self.gamma2 = CoeffStructure.gamma2_niondot_II_index2D
@@ -307,7 +307,8 @@ def Rdn_dR(self, z, R):
 
     def converge_BMF(self, CosmoParams, ion_frac_input, max_iter):
         self.ion_frac = ion_frac_input
-        for j in trange(max_iter):
+        iterator = trange(max_iter) if self.PRINT_SUCCESS else range(max_iter)
+        for j in iterator:
             ion_frac_prev = np.copy(self.ion_frac)
             
             self.barrier = self.compute_barrier(CosmoParams, self.ion_frac, self.zlist, self.Rs)
@@ -323,7 +324,8 @@ def barrierz_int(self, z, R):
             self.ion_frac[self.barrier[:, -1]<=0] = 1
 
             if np.allclose(ion_frac_prev, self.ion_frac):
-                print(f'SUCCESS: BMF converged in {j} iterations.')
+                if self.PRINT_SUCCESS:
+                    print(f'SUCCESS: BMF converged in {j} iterations.')
                 return 
             
         print(f"WARNING: BMF didn't converge within {max_iter} iterations.")

From f5f1b78dce4b1073bc47dc0f10b677be69ad1b6d Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Thu, 20 Nov 2025 11:12:38 -0600
Subject: [PATCH 28/37] Extend density power in CoevalMaps

---
 zeus21/maps.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index e0ee561..3605592 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -45,13 +45,14 @@ def __init__(self, T21_coefficients, Power_Spectrum, z, Lbox=600, Nbox=200, KIND
         if (KIND == 0): #just T21, ~gaussian
                 
             P21 = Power_Spectrum.Deltasq_T21_lin[_iz]/k3over2pi2
-            P21norminterp = interp1d(klist,P21/self.T21global**2,fill_value=0.0,bounds_error=False)
+            #P21norminterp = interp1d(klist,P21/self.T21global**2,fill_value=0.0,bounds_error=False) OLD
+            P21_spl = spline(np.log(klist), np.log(P21/self.T21global**2)) #spline over log values
 
 
             pb = pbox.PowerBox(
                 N=self.Nbox,                     
                 dim=3,                     
-                pk = lambda k: P21norminterp(k), 
+                pk = lambda k: np.exp(P21_spl(np.log(k))), 
                 boxlength = self.Lbox,           
                 seed = self.seed                
             )
@@ -63,12 +64,13 @@ def __init__(self, T21_coefficients, Power_Spectrum, z, Lbox=600, Nbox=200, KIND
             
         elif (KIND == 1):
             Pd = Power_Spectrum.Deltasq_d_lin[_iz,:]/k3over2pi2
-            Pdinterp = interp1d(klist,Pd,fill_value=0.0,bounds_error=False)
+            #Pdinterp = interp1d(klist,Pd,fill_value=0.0,bounds_error=False) OLD
+            Pd_spl = spline(np.log(klist), np.log(Pd))
 
             pb = pbox.PowerBox(
                 N=self.Nbox,                     
                 dim=3,                     
-                pk = lambda k: Pdinterp(k), 
+                pk = lambda k: np.exp(Pd_spl(np.log(k))), 
                 boxlength = self.Lbox,           
                 seed = self.seed               
             )
@@ -78,12 +80,13 @@ def __init__(self, T21_coefficients, Power_Spectrum, z, Lbox=600, Nbox=200, KIND
             #then we make a map of the linear T21 fluctuation, better to use the cross to keep sign, at linear level same 
             PdT21 = Power_Spectrum.Deltasq_dT21[_iz]/k3over2pi2
 
-            powerratioint = interp1d(klist,PdT21/Pd,fill_value=0.0,bounds_error=False)
+            #powerratioint = interp1d(klist,PdT21/Pd,fill_value=0.0,bounds_error=False) OLD
+            powerratio_spl = spline(klist, PdT21/Pd) #cross can be negative, so can't interpolate over log values
 
 
             deltak = pb.delta_k()
 
-            powerratio = powerratioint(pb.k())
+            powerratio = powerratio_spl(pb.k())
             T21lin_k = powerratio * deltak
             self.T21maplin= self.T21global + z21_utilities.powerboxCtoR(pb,mapkin = T21lin_k)
 

From 7788c8743d428b529c46f1f96b5f125cd4b178ad Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Mon, 9 Feb 2026 15:44:06 -0600
Subject: [PATCH 29/37] Altered partial ionizations

maps.py: density array used for partial ionizations and massweighting now the density field smoothed by the pixel scale.

reionization.py: minor tweaks for self-consistency
---
 zeus21/maps.py         | 12 +++++++-----
 zeus21/reionization.py |  4 ++--
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index 3605592..559ef56 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -237,14 +237,16 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
         ### generating the density field at the closest redshift to the lower one inputed
         self.z_of_density = self.z[0]
         self.density = self.generate_density(ClassyCosmo, CorrFClass)
-        self.density_allz = np.empty((len(self.z), self.ncells, self.ncells, self.ncells), dtype=np.float32)
-        if self.COMPUTE_DENSITY_AT_ALLZ:
-            self.generate_density_allz(CosmoParams)
 
         ### smoothing the density field
         self._k = self.compute_k()
         self.density_smoothed_allr = self.smooth_density()
 
+        ### evolving density
+        self.density_allz = np.empty((len(self.z), self.ncells, self.ncells, self.ncells), dtype=np.float32)
+        if self.COMPUTE_DENSITY_AT_ALLZ:
+            self.generate_density_allz(CosmoParams)
+
         ### generating the ionized field, and computing the ionized fraction
         self.barrier = barrier
         if self.barrier is None:
@@ -302,7 +304,7 @@ def generate_density_allz(self, CosmoParams):
             print('Evolving density field...')
         Dg = CosmoParams.growthint(self.z)
         growthfactor_ratio = (Dg/Dg[0])[:, np.newaxis, np.newaxis, np.newaxis]
-        density_lastz = np.copy(self.density)
+        density_lastz = np.copy(self.density_smoothed_allr[0])
         self.density_allz = density_lastz[np.newaxis]*growthfactor_ratio
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
@@ -411,7 +413,7 @@ def compute_partial(self, CosmoParams, BMF, r=None):
         for i in iterator:
             partial_ion_spl = spline(sample_d, BMF.prebarrier_xHII_int_grid(sample_d, self.z[i], r)) #spline is faster than RGI, so build a spline on sample densities
             
-            partialfield = np.abs(partial_ion_spl(self.density_allz[i]))
+            partialfield = np.abs(partial_ion_spl(self.density_smoothed_allr[0])) #abs just in case, but it never actually triggers afaik
             sumfield = self.ion_field_allz[i] + partialfield
             self.ion_field_partial_allz[i] = np.clip(sumfield, 0, 1)
 
diff --git a/zeus21/reionization.py b/zeus21/reionization.py
index 90979af..48699f4 100644
--- a/zeus21/reionization.py
+++ b/zeus21/reionization.py
@@ -36,7 +36,7 @@ def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, Classy
         
         self.gamma = CoeffStructure.gamma_niondot_II_index2D
         self.gamma2 = CoeffStructure.gamma2_niondot_II_index2D
-        self.sigma = CoeffStructure.sigmaofRtab
+        self.sigma = np.array([[ClassyCosmo.sigma(r, z) for z in self.zlist] for r in self.Rs]).T#CoeffStructure.sigmaofRtab
         
         self.zr = [self.zlist, np.log(self.Rs)]
         self.gamma_int = RegularGridInterpolator(self.zr, self.gamma, bounds_error = False, fill_value = np.nan)
@@ -44,7 +44,7 @@ def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, Classy
 
         self.sigma_BMF = np.array([[ClassyCosmo.sigma(r, z) for z in self.zlist] for r in self.Rs_BMF]).T
         self.zr_BMF = [self.zlist, np.log(self.Rs_BMF)]
-        self.sigma_int = RegularGridInterpolator(self.zr_BMF, self.sigma_BMF, bounds_error = False, fill_value = np.nan)
+        self.sigma_int = RegularGridInterpolator(self.zr_BMF, self.sigma_BMF, bounds_error = False, fill_value = None)
         
         self.Hz = cosmology.Hubinvyr(CosmoParams, self.zlist)
         self.trec0 = 1/(constants.alphaB * cosmology.n_H(CosmoParams,0) * AstroParams._clumping) #seconds

From 8a5d1f6d4e2c5f1bdc363054774bc73fd2cdf14b Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Tue, 31 Mar 2026 14:58:22 -0500
Subject: [PATCH 30/37] Updated Modeling of Maps and Bubbles

Maps:
- Added mass variance correction to account for lack of ergodicity in smaller simulations
- Fixed partial ionization computation accounting for this sigma correction.

Bubbles:
- Changed all interpolators to allow for linear extrapolation
- Added new moving pivot for linear barrier to account for the typical bubble size changing at different redshifts. Now the typical bubbles are the ones best approximated in the linear barrier.
- Now computing the ionized fraction from the analytic solution to the integral of the BMF. This accounts for contribution from bubbles outside the set radius range rather than integrating numerically.
- Increased tolerance for wiggle between iterations of the BMF.
---
 zeus21/maps.py         | 11 +++++--
 zeus21/reionization.py | 70 +++++++++++++++++++++++++++---------------
 2 files changed, 53 insertions(+), 28 deletions(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index 559ef56..08410a3 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -3,7 +3,7 @@
 Make maps! For fun and science
 
 Authors: Julian B. Muñoz, Yonatan Sklansky, Emilie Thelie
-UT Austin - October 2025
+UT Austin - March 2026
 
 """
 
@@ -237,6 +237,7 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
         ### generating the density field at the closest redshift to the lower one inputed
         self.z_of_density = self.z[0]
         self.density = self.generate_density(ClassyCosmo, CorrFClass)
+        self.density /= self.sigma_correction(ClassyCosmo) #ergodicity correction
 
         ### smoothing the density field
         self._k = self.compute_k()
@@ -304,7 +305,7 @@ def generate_density_allz(self, CosmoParams):
             print('Evolving density field...')
         Dg = CosmoParams.growthint(self.z)
         growthfactor_ratio = (Dg/Dg[0])[:, np.newaxis, np.newaxis, np.newaxis]
-        density_lastz = np.copy(self.density_smoothed_allr[0])
+        density_lastz = np.copy(self.density)
         self.density_allz = density_lastz[np.newaxis]*growthfactor_ratio
         if self.PRINT_TIMER:
             z21_utilities.print_timer(start_time, text_before="    done in ")
@@ -328,6 +329,10 @@ def smooth_density(self):
             z21_utilities.print_timer(start_time, text_before="    done in ")
         return density_smoothed_allr
 
+    def sigma_correction(self, ClassyCosmo):
+        sigma_ratio = np.std(self.density)/ClassyCosmo.sigma(self.r[0], self.z_of_density)
+        return sigma_ratio
+
     def generate_xHII(self, CosmoParams, CoeffStructure, BMF):
         if self.PRINT_TIMER:
             start_time = time.time()
@@ -413,7 +418,7 @@ def compute_partial(self, CosmoParams, BMF, r=None):
         for i in iterator:
             partial_ion_spl = spline(sample_d, BMF.prebarrier_xHII_int_grid(sample_d, self.z[i], r)) #spline is faster than RGI, so build a spline on sample densities
             
-            partialfield = np.abs(partial_ion_spl(self.density_smoothed_allr[0])) #abs just in case, but it never actually triggers afaik
+            partialfield = np.abs(partial_ion_spl(self.density)) #abs just in case, but it never actually triggers afaik
             sumfield = self.ion_field_allz[i] + partialfield
             self.ion_field_partial_allz[i] = np.clip(sumfield, 0, 1)
 
diff --git a/zeus21/reionization.py b/zeus21/reionization.py
index 48699f4..415ac5b 100644
--- a/zeus21/reionization.py
+++ b/zeus21/reionization.py
@@ -15,6 +15,7 @@
 from scipy.integrate import cumulative_trapezoid
 from scipy.interpolate import interp1d
 from scipy.interpolate import RegularGridInterpolator
+from scipy.special import erfc
 from tqdm import trange
 
 
@@ -39,8 +40,8 @@ def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, Classy
         self.sigma = np.array([[ClassyCosmo.sigma(r, z) for z in self.zlist] for r in self.Rs]).T#CoeffStructure.sigmaofRtab
         
         self.zr = [self.zlist, np.log(self.Rs)]
-        self.gamma_int = RegularGridInterpolator(self.zr, self.gamma, bounds_error = False, fill_value = np.nan)
-        self.gamma2_int = RegularGridInterpolator(self.zr, self.gamma2, bounds_error = False, fill_value = np.nan)
+        self.gamma_int = RegularGridInterpolator(self.zr, self.gamma, bounds_error = False, fill_value = None)
+        self.gamma2_int = RegularGridInterpolator(self.zr, self.gamma2, bounds_error = False, fill_value = None)
 
         self.sigma_BMF = np.array([[ClassyCosmo.sigma(r, z) for z in self.zlist] for r in self.Rs_BMF]).T
         self.zr_BMF = [self.zlist, np.log(self.Rs_BMF)]
@@ -49,22 +50,22 @@ def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, Classy
         self.Hz = cosmology.Hubinvyr(CosmoParams, self.zlist)
         self.trec0 = 1/(constants.alphaB * cosmology.n_H(CosmoParams,0) * AstroParams._clumping) #seconds
         self.trec = self.trec0/(1+self.zlist)**3/constants.yrTos #years
-        self.trec_int = interp1d(self.zlist, self.trec, bounds_error = False, fill_value = np.nan)
+        self.trec_int = interp1d(self.zlist, self.trec, bounds_error = False, fill_value = None)
         
         self.niondot_avg = CoeffStructure.niondot_avg_II
-        self.niondot_avg_int = interp1d(self.zlist, self.niondot_avg, bounds_error = False, fill_value = np.nan)
+        self.niondot_avg_int = interp1d(self.zlist, self.niondot_avg, bounds_error = False, fill_value = None)
 
         self.ion_frac = np.fmin(1, [self.calc_Q(CosmoParams, z) for z in self.zlist])
         self.ion_frac_initial = np.copy(self.ion_frac)
 
         zr_mesh = np.meshgrid(np.arange(len(self.Rs)), np.arange(len(self.zlist)))
         self.nion_norm = self.nion_normalization(zr_mesh[1], zr_mesh[0])
-        self.nion_norm_int = RegularGridInterpolator(self.zr, self.nion_norm, bounds_error = False, fill_value = np.nan)
+        self.nion_norm_int = RegularGridInterpolator(self.zr, self.nion_norm, bounds_error = False, fill_value = None)
 
         self.prebarrier_xHII = np.empty((len(self.ds_array), len(self.zlist), len(self.Rs)))
         self.barrier = self.compute_barrier(CosmoParams, self.ion_frac, self.zlist, self.Rs)
         self.barrier_initial = np.copy(self.barrier)
-        self.barrier_int = RegularGridInterpolator(self.zr, self.barrier, bounds_error = False, fill_value = np.nan)
+        self.barrier_int = RegularGridInterpolator(self.zr, self.barrier, bounds_error = False, fill_value = None)
 
         self.dzr = [self.ds_array, self.zlist, np.log(self.Rs)]
         self.prebarrier_xHII_int = RegularGridInterpolator(self.dzr, self.prebarrier_xHII, bounds_error = False, fill_value = None) #allow extrapolation
@@ -72,9 +73,17 @@ def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, Classy
         self.R_linear_sigma_fit_idx = z21_utilities.find_nearest_idx(self.Rs, R_linear_sigma_fit_input)
         self.R_linear_sigma_fit = self.Rs[self.R_linear_sigma_fit_idx]
 
-        self.BMF = np.array([self.VRdn_dR(z, self.Rs_BMF) for z in self.zlist]) #initial bubble mass function
+        #fake bubble mass function to impose peak around R_linear_sigma_fit for the initial linear barriers
+        #looks something like [0, 0, ..., 1, ..., 0, 0]*(number of redshifts)
+        self.BMF = np.repeat([np.eye(len(self.Rs_BMF))[self.R_linear_sigma_fit_idx]], len(self.zlist), axis=0)
+
+        #second computation of BMF using the initial guess peaks
+        self.BMF = np.array([self.VRdn_dR(z, self.Rs_BMF) for z in self.zlist])
         self.BMF_initial = np.copy(self.BMF)
-        self.ion_frac = np.nan_to_num([np.trapezoid(self.BMF[i], np.log(self.Rs_BMF)) for i in range(len(self.zlist))]) #ion_frac by integrating the BMF
+        
+        #self.ion_frac = np.nan_to_num([np.trapezoid(self.BMF[i], np.log(self.Rs_BMF)) for i in range(len(self.zlist))]) #ion_frac by numerically integrating the BMF
+        self.ion_frac = np.nan_to_num([self.analytic_Q(ClassyCosmo, z)[0] for z in self.zlist]) #ion_frac by analytic integral of BMF
+        
         self.ion_frac[self.barrier[:, -1]<=0] = 1
         
         if FLAG_converge:
@@ -274,15 +283,17 @@ def calc_Q(self, CosmoParams, z):
 
     #computing linear barrier
     def B_1(self, z):
-        sigmax = self.sigmaR_int(z, self.Rs[self.R_linear_sigma_fit_idx+1])
-        sigmin = self.sigmaR_int(z, self.Rs[self.R_linear_sigma_fit_idx-1])
-        barriermax = self.barrierR_int(z, self.Rs[self.R_linear_sigma_fit_idx+1])
-        barriermin = self.barrierR_int(z, self.Rs[self.R_linear_sigma_fit_idx-1])
+        R_pivot = np.max([0.1, self.BMF_peak_R(z)])
+        sigmax = self.sigmaR_int(z, R_pivot*1.1)
+        sigmin = self.sigmaR_int(z, R_pivot*0.9)
+        barriermax = self.barrierR_int(z, R_pivot*1.1)
+        barriermin = self.barrierR_int(z, R_pivot*0.9)
         return (barriermax - barriermin)/(sigmax**2 - sigmin**2)
         
     def B_0(self, z):
-        sigmin = self.sigmaR_int(z, self.Rs[self.R_linear_sigma_fit_idx-1])
-        barriermin = self.barrierR_int(z, self.Rs[self.R_linear_sigma_fit_idx-1])
+        R_pivot = np.max([0.1, self.BMF_peak_R(z)])
+        sigmin = self.sigmaR_int(z, R_pivot*0.9)
+        barriermin = self.barrierR_int(z, R_pivot*0.9)
         return barriermin - sigmin**2 * self.B_1(z)
     
     def B(self, z, R): 
@@ -305,6 +316,19 @@ def VRdn_dR(self, z, R):
     def Rdn_dR(self, z, R):
         return self.VRdn_dR(z, R)*3/(4*np.pi*R**3)
 
+    def BMF_peak_R(self, z):
+        iz = z21_utilities.find_nearest_idx(self.zlist, z)
+        ir = np.argmax(self.BMF[iz])
+        return self.Rs_BMF[ir]
+
+    def analytic_Q(self, ClassyCosmo, z): #analytically integrating the BMF to get Q
+        Rmin = 1e-10 #arbitrarily small
+        B0 = self.B_0(z)
+        B1 = self.B_1(z)
+        sigmin = ClassyCosmo.sigma(Rmin, z)
+        s2 = sigmin**2
+        return 0.5*np.exp(-2*B0*B1)*erfc((B0-B1*s2)/np.sqrt(2*s2)) + 0.5*erfc((B0+B1*s2)/np.sqrt(2*s2))
+
     def converge_BMF(self, CosmoParams, ion_frac_input, max_iter):
         self.ion_frac = ion_frac_input
         iterator = trange(max_iter) if self.PRINT_SUCCESS else range(max_iter)
@@ -312,20 +336,15 @@ def converge_BMF(self, CosmoParams, ion_frac_input, max_iter):
             ion_frac_prev = np.copy(self.ion_frac)
             
             self.barrier = self.compute_barrier(CosmoParams, self.ion_frac, self.zlist, self.Rs)
-            self.barrier_int = RegularGridInterpolator(self.zr, self.barrier, bounds_error = False, fill_value = np.nan)
-
-            def barrierR_int(self, z, R):
-                return self.interpR(z, R, self.barrier_int)
-            def barrierz_int(self, z, R):
-                return self.interpz(z, R, self.barrier_int)
+            self.barrier_int = RegularGridInterpolator(self.zr, self.barrier, bounds_error = False, fill_value = None)
             
             self.BMF = np.array([self.VRdn_dR(z, self.Rs_BMF) for z in self.zlist])
-            self.ion_frac = np.nan_to_num([np.trapezoid(self.BMF[i], np.log(self.Rs_BMF)) for i in range(len(self.zlist))])
+            self.ion_frac = np.nan_to_num([self.analytic_Q(ClassyCosmo, z)[0] for z in self.zlist])
             self.ion_frac[self.barrier[:, -1]<=0] = 1
 
-            if np.allclose(ion_frac_prev, self.ion_frac):
+            if np.allclose(ion_frac_prev, self.ion_frac, rtol=1e-1, atol=1e-2):
                 if self.PRINT_SUCCESS:
-                    print(f'SUCCESS: BMF converged in {j} iterations.')
+                    print(f"SUCCESS: BMF converged after {j+1} iteration{'s' if j > 0 else ''}.")
                 return 
             
         print(f"WARNING: BMF didn't converge within {max_iter} iterations.")
@@ -353,9 +372,9 @@ def sigmaR_int(self, z, R):
     def sigmaz_int(self, z, R):
         return self.interpz(z, R, self.sigma_int)
 
-    def barrierR_int(self, z, R): #unused
+    def barrierR_int(self, z, R):
         return self.interpR(z, R, self.barrier_int)
-    def barrierz_int(self, z, R): #unused
+    def barrierz_int(self, z, R):
         return self.interpz(z, R, self.barrier_int)
 
     def gammaR_int(self, z, R):
@@ -373,6 +392,7 @@ def nion_normR_int(self, z, R):
     def nion_normz_int(self, z, R):
         return self.interpz(z, R, self.nion_norm_int)
 
+
     def prebarrier_xHII_int_grid(self, d, z, R):
         """
         Evaluate prebarrier xHII on a density field d(x),

From f6b6c1e9cb4f1af2b3003f6ba55e95fdd7533560 Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Tue, 31 Mar 2026 15:08:03 -0500
Subject: [PATCH 31/37] Small fix

Forgot to add input to function
---
 zeus21/reionization.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/zeus21/reionization.py b/zeus21/reionization.py
index 415ac5b..ea5b45f 100644
--- a/zeus21/reionization.py
+++ b/zeus21/reionization.py
@@ -87,7 +87,7 @@ def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, Classy
         self.ion_frac[self.barrier[:, -1]<=0] = 1
         
         if FLAG_converge:
-            self.converge_BMF(CosmoParams, self.ion_frac, max_iter=max_iter)
+            self.converge_BMF(CosmoParams, ClassyCosmo, self.ion_frac, max_iter=max_iter)
         #two functions: compute BMF and iterate
         
 
@@ -329,7 +329,7 @@ def analytic_Q(self, ClassyCosmo, z): #analytically integrating the BMF to get Q
         s2 = sigmin**2
         return 0.5*np.exp(-2*B0*B1)*erfc((B0-B1*s2)/np.sqrt(2*s2)) + 0.5*erfc((B0+B1*s2)/np.sqrt(2*s2))
 
-    def converge_BMF(self, CosmoParams, ion_frac_input, max_iter):
+    def converge_BMF(self, CosmoParams, ClassyCosmo, ion_frac_input, max_iter):
         self.ion_frac = ion_frac_input
         iterator = trange(max_iter) if self.PRINT_SUCCESS else range(max_iter)
         for j in iterator:

From 24aa05b905a5dde826ecf66d263eb9f115d749ae Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Tue, 7 Apr 2026 16:08:16 -0500
Subject: [PATCH 32/37] Added clumping as arg to AstroParams, inputs for
 densities in reionization_maps, and delete_class_attributes to z21_utilities.
 Removed density_allz from partials in reionization_maps. Created
 astro_variations class in reionization_maps. Changed CosmoParams.Rsmmin to
 0.05.

---
 zeus21/inputs.py        |   5 +-
 zeus21/maps.py          | 291 ++++++++++++++++++++++++++++++++++++++--
 zeus21/z21_utilities.py |   8 +-
 3 files changed, 291 insertions(+), 13 deletions(-)

diff --git a/zeus21/inputs.py b/zeus21/inputs.py
index b55a859..f0fe7a3 100644
--- a/zeus21/inputs.py
+++ b/zeus21/inputs.py
@@ -179,7 +179,7 @@ def __init__(self, UserParams, CosmoParams_input, ClassCosmo):
 
 
         #and define the shells that we integrate over at each z.
-        self.Rsmmin = 0.5
+        self.Rsmmin = 0.05
         self.Rsmmax = 2000.
 
         if(self.Flag_emulate_21cmfast==True):
@@ -234,6 +234,7 @@ def __init__(self, UserParams, Cosmo_Parameters,
                     E0_xray = 500.,
                     alpha_xray = -1.0,
                     Emax_xray_norm=2000,
+                    _clumping = 3,
 
                     Nalpha_lyA_II = 9690,
                     Nalpha_lyA_III = 17900,
@@ -334,7 +335,7 @@ def __init__(self, UserParams, Cosmo_Parameters,
         #fesc(M) parameter. Power law normalized (fesc10) at M=1e10 Msun with index alphaesc
         self.fesc10 = fesc10
         self.alphaesc = alphaesc
-        self._clumping = 3.0 #clumping factor, z-independent and fixed for now
+        self._clumping = _clumping #clumping factor, z-independent and fixed for now
         if(Cosmo_Parameters.Flag_emulate_21cmfast==True):
             self._clumping = 2.0 #this is the 21cmFAST value
 
diff --git a/zeus21/maps.py b/zeus21/maps.py
index 08410a3..bc17118 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -10,6 +10,9 @@
 from . import cosmology
 from . import constants
 from . import z21_utilities
+from . import inputs
+from . import sfrd
+from . import reionization
 
 import numpy as np
 import powerbox as pbox
@@ -18,6 +21,8 @@
 from pyfftw import empty_aligned as empty
 from tqdm import trange
 import time
+from dataclasses import dataclass
+from collections.abc import Sequence
 
 
 class CoevalMaps:
@@ -187,7 +192,8 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
                  PRINT_TIMER=True, 
                  LOGNORMAL_DENSITY=False, COMPUTE_DENSITY_AT_ALLZ=False, SPHERIZE=False, 
                  COMPUTE_MASSWEIGHTED=False, lowres_massweighting=1, COMPUTE_PARTIAL_IONIZATIONS=False,
-                 COMPUTE_PARTIAL_AND_MASSWEIGHTED=False, COMPUTE_ZREION=False
+                 COMPUTE_PARTIAL_AND_MASSWEIGHTED=False, COMPUTE_ZREION=False,
+                 input_density=None, input_density_smoothed_allr=None, input_density_allz=None
                 ):
         #Measure time elapsed from start
         self._start_time = time.time()
@@ -226,7 +232,8 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
         self.COMPUTE_PARTIAL_IONIZATIONS = COMPUTE_PARTIAL_IONIZATIONS
         self.COMPUTE_PARTIAL_AND_MASSWEIGHTED = COMPUTE_PARTIAL_AND_MASSWEIGHTED
         self.COMPUTE_ZREION = COMPUTE_ZREION
-        if self.COMPUTE_MASSWEIGHTED or self.COMPUTE_PARTIAL_IONIZATIONS or self.COMPUTE_PARTIAL_AND_MASSWEIGHTED:
+        #if self.COMPUTE_MASSWEIGHTED or self.COMPUTE_PARTIAL_IONIZATIONS or self.COMPUTE_PARTIAL_AND_MASSWEIGHTED:
+        if self.COMPUTE_MASSWEIGHTED or self.COMPUTE_PARTIAL_AND_MASSWEIGHTED:
             self.COMPUTE_DENSITY_AT_ALLZ = True
 
         ### selecting redshifts and radii from available redshifts
@@ -236,17 +243,42 @@ def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, in
 
         ### generating the density field at the closest redshift to the lower one inputed
         self.z_of_density = self.z[0]
-        self.density = self.generate_density(ClassyCosmo, CorrFClass)
-        self.density /= self.sigma_correction(ClassyCosmo) #ergodicity correction
+        if input_density is None:
+            self.density = self.generate_density(ClassyCosmo, CorrFClass)
+            self.density /= self.sigma_correction(ClassyCosmo) #ergodicity correction
+        else:
+            if input_density.shape==(self.ncells,self.ncells,self.ncells):
+                self.density = input_density # no np.copy here to reduce live memory
+            else:
+                raise ValueError(f"The input density should have a shape {(self.ncells,self.ncells,self.ncells)}.")
+            if PRINT_TIMER:
+                z21_utilities.print_timer(self._start_time, text_before="input density in ")
 
         ### smoothing the density field
         self._k = self.compute_k()
-        self.density_smoothed_allr = self.smooth_density()
+        if input_density_smoothed_allr is None:
+            self.density_smoothed_allr = self.smooth_density()
+        else:
+            if input_density_smoothed_allr.shape==(len(self.r),self.ncells,self.ncells,self.ncells):
+                self.density_smoothed_allr = input_density_smoothed_allr # no np.copy here to reduce live memory
+            else:
+                raise ValueError(f"The input smooth densities should have a shape {(len(self.r),self.ncells,self.ncells,self.ncells)}.")
+            if PRINT_TIMER:
+                z21_utilities.print_timer(self._start_time, text_before="input smoothed density in ")
 
         ### evolving density
-        self.density_allz = np.empty((len(self.z), self.ncells, self.ncells, self.ncells), dtype=np.float32)
-        if self.COMPUTE_DENSITY_AT_ALLZ:
-            self.generate_density_allz(CosmoParams)
+        if input_density_allz is None:
+            self.density_allz = np.empty((len(self.z), self.ncells, self.ncells, self.ncells), dtype=np.float32)
+            if self.COMPUTE_DENSITY_AT_ALLZ:
+                self.generate_density_allz(CosmoParams)
+        else:
+            if input_density_allz.shape==(len(self.z),self.ncells,self.ncells,self.ncells):
+                self.density_allz = input_density_allz # no np.copy here to reduce live memory
+                self._has_density = True
+            else:
+                raise ValueError(f"The input all-z-densities should have a shape {(len(self.z),self.ncells,self.ncells,self.ncells)}.")
+            if PRINT_TIMER:
+                z21_utilities.print_timer(self._start_time, text_before="input all z densities in ")
 
         ### generating the ionized field, and computing the ionized fraction
         self.barrier = barrier
@@ -406,8 +438,8 @@ def compute_partial(self, CosmoParams, BMF, r=None):
         if not self._has_p:
             self.ion_frac_partial = np.empty(len(self.z))
             self.ion_field_partial_allz = np.empty_like(self.ion_field_allz)
-        if not self._has_density:
-            self.generate_density_allz(CosmoParams)
+        #if not self._has_density:
+        #    self.generate_density_allz(CosmoParams)
         sample_d = np.linspace(-5, 5, 51)
 
         if self.PRINT_TIMER:
@@ -508,3 +540,242 @@ def _compute_ionfrac_from_treion(self):
         return 1-neutfrac, tvalues
         
 
+
+
+##### Running reionization_maps multiple times with variations in astrophysics
+# dataclasses for inputing arguments to astro_variations
+@dataclass(frozen=True)
+class AstroParamsConfig:
+    """
+    All arguments of Astro_Parameters that have default values
+    """
+    astromodel: int = 0
+    accretion_model: int = 0
+    alphastar: float = 0.5
+    betastar: float = -0.5
+    epsstar: float = 0.1
+    Mc: float = 3e11
+    dlog10epsstardz: float = 0.0
+    
+    fesc10: float = 0.1
+    alphaesc: float = 0.0
+    L40_xray: float = 3.0
+    E0_xray: float = 500.
+    alpha_xray: float = -1.0
+    Emax_xray_norm: float = 2000
+    _clumping: float = 3
+    
+    Nalpha_lyA_II: float = 9690
+    Nalpha_lyA_III: float = 17900
+    
+    Mturn_fixed: float | None = None
+    FLAG_MTURN_SHARP: bool = False
+    
+    C0dust: float = 4.43
+    C1dust: float = 1.99
+    
+    sigmaUV: float = 0.5
+    
+    USE_POPIII: bool = False
+    
+    alphastar_III: float = 0
+    betastar_III: float = 0
+    fstar_III: float = 10**(-2.5)
+    Mc_III: float = 1e7
+    dlog10epsstardz_III: float = 0.0
+    
+    fesc7_III: float = 10**(-1.35)
+    alphaesc_III: float = -0.3
+    L40_xray_III: float = 3.0
+    alpha_xray_III: float = -1.0
+    
+    USE_LW_FEEDBACK: bool = True
+    A_LW: float = 2.0
+    beta_LW: float = 0.6
+    
+    A_vcb: float = 1.0
+    beta_vcb: float = 1.8,
+
+@dataclass
+class BMFConfig:
+    """
+    All arguments of BMF that have default values.
+    """
+    R_linear_sigma_fit_input: float = 10
+    FLAG_converge: bool = True
+    max_iter: int = 10
+    ZMAX_REION: float = 30
+    Rmin: float = 0.05
+    PRINT_SUCCESS: bool = True
+
+@dataclass
+class ReioMapsConfig:
+    """
+    All arguments of reionization_maps that have default values, 
+    except the input_densities since they are changed dynamically in astro_variations
+    """
+    r_precision: float = 1.
+    Rs: list | np.ndarray | None = None
+    barrier: function = None
+    PRINT_TIMER: bool = True
+    LOGNORMAL_DENSITY: bool = False
+    COMPUTE_DENSITY_AT_ALLZ: bool = False
+    SPHERIZE: bool = False
+    COMPUTE_MASSWEIGHTED: bool = False
+    lowres_massweighting: int = 1
+    COMPUTE_PARTIAL_IONIZATIONS: bool = False
+    COMPUTE_PARTIAL_AND_MASSWEIGHTED: bool = False
+    COMPUTE_ZREION: bool = False
+
+# running class that will do the astro variations
+class astro_variations:
+    """
+    Efficiently rerun the reionization_maps class while varying astrophysical parameters.
+    The efficiency comes from:
+        - making use of the different input density boxes in reionization_maps and not reruning them if multiple models are run.
+        - deleting the instances of the reionization_map class generated at each run after saving the attributes the user wants to keep.
+
+    Parameters
+    ----------
+    UserParams: zeus21.User_Parameters class
+        Stores the user parameters.
+    CosmoParams: zeus21.Cosmo_Parameters class
+        Stores the cosmology.
+    ClassyCosmo: zeus21.runclass class
+        Sets up Class cosmology.
+    HMFintclass: zeus21.HMF_interpolator class
+        Stores the HMF interpolator.
+    CorrFClass: zeus21.Correlations class
+        Stores the correlations.
+    input_boxlength: float
+        Comoving physical side length of the box.
+    ncells: int
+        Number of cells on a side.
+    seed: int
+        Sets the predetermined generation of maps. Default is 1234.
+    AstroParams_configs: sequence of AstroParamsConfig instances
+        Sets the different astro models to run.
+    BMF_quantities_to_save: list of str of the names of BMF attributes
+        Specifiy which quantities to save as attributes in astro_variations. 
+        They will be saved in lists named self.BMF_attr (e.g.: self.BMF_ion_frac).
+        Each element of these lists correspond to the corresponding AstroParams_configs element. 
+    ReioMaps_quantities_to_save: list of str of the names of reionization_maps attributes
+        Specifiy which quantities to save as attributes in astro_variations.
+        They will be saved in lists named self.ReioMaps_attr (e.g.: self.ReioMaps_ion_frac).
+        Each element of these lists correspond to the corresponding AstroParams_configs element. 
+    BMF_config: BMFConfig class
+        Sets up the config for the BMF class. Default is None and corresponds to default BMF arguments.
+    ReioMaps_config: ReioMapsConfig class
+        Sets up the config for the reionization_maps class. Default is None and corresponds to default reionization_maps arguments.
+    ZMIN: float
+        Minimum redshift to which zeus21.get_T21_coefficients will do its integrations. Default is 5.0.
+    input_z: list of float
+        Redshift list over which the reionization maps are made. Default is None, corresponding to zeus21.get_T21_coefficients.zintegral.
+    seed: int
+        Sets the seed used to generate the boxes. Default is 1234.
+
+    Attributes
+    ----------
+    density: 3D np.array
+        Saves the density maps generated in the first run.
+    density_smoothed_allr: 4D np.array
+        Saves the density_smoothed_allr maps generated in the first run.
+    density_allz: 4D np.array
+        Saves the density_allz maps generated in the first run.
+    BMF_... and ReioMaps_...:
+        All attributes from the BMF and reionization_map classes the user wants to save.
+    """
+
+    def __init__(self, UserParams, CosmoParams, ClassyCosmo, HMFintclass, CorrFClass, 
+                 input_boxlength, ncells, 
+                 AstroParams_configs: Sequence[AstroParamsConfig],
+                 BMF_quantities_to_save, ReioMaps_quantities_to_save, 
+                 BMF_config: BMFConfig | None = None,
+                 ReioMaps_config: ReioMapsConfig | None = None,
+                 ZMIN = 5., input_z=None, seed=1234):
+        
+        # zeus21 classes
+        self.UserParams = UserParams
+        self.CosmoParams, self.ClassyCosmo = CosmoParams, ClassyCosmo
+        self.HMFintclass, self.CorrFClass = HMFintclass, CorrFClass
+        self.ZMIN = ZMIN
+        
+        # maps input
+        self.input_z = input_z
+        self.input_boxlength, self.ncells = input_boxlength, ncells
+        self.seed = seed
+        
+        # BMF and reionization_maps arguments
+        self.AstroParams_configs = tuple(AstroParams_configs)
+        self.BMF_config = BMF_config or BMFConfig()
+        self.ReioMaps_config = ReioMaps_config or ReioMapsConfig()
+        
+        # saved density boxes for time and memory efficiency
+        self.density = None
+        self.density_smoothed_allr = None
+        self.density_allz = None
+
+        # initialize physical quantities to save
+        self.BMF_quantities_to_save = BMF_quantities_to_save
+        self.ReioMaps_quantities_to_save = ReioMaps_quantities_to_save
+        for q in BMF_quantities_to_save:
+            if q in vars(reionization.BMF)["__static_attributes__"]:
+                setattr(self, "BMF_"+q, [])
+            else:
+                raise ValueError(f"{q} is not an attribute of the BMF class.")
+        for q in ReioMaps_quantities_to_save:
+            if q in vars(reionization_maps)["__static_attributes__"]:
+                setattr(self, "ReioMaps_"+q, [])
+            else:
+                raise ValueError(f"{q} is not an attribute of the reionization_maps class.")
+
+
+    def run_1model(self, AstroParams_config):
+        AstroParams = inputs.Astro_Parameters(self.UserParams, self.CosmoParams, **vars(AstroParams_config))
+        CoeffStructure = sfrd.get_T21_coefficients(self.UserParams, self.CosmoParams, self.ClassyCosmo, AstroParams, self.HMFintclass, zmin=self.ZMIN)
+    
+        BMF = reionization.BMF(CoeffStructure, self.HMFintclass, self.CosmoParams, AstroParams, self.ClassyCosmo,
+                            **vars(self.BMF_config))
+        
+        if self.input_z is None:
+            self.input_z = CoeffStructure.zintegral
+        
+        reio_maps = reionization_maps(self.CosmoParams, self.ClassyCosmo, self.CorrFClass, CoeffStructure, BMF, self.input_z,
+                                      input_boxlength=self.input_boxlength, ncells=self.ncells, seed=self.seed,
+                                      **vars(self.ReioMaps_config),
+                                      input_density=self.density, 
+                                      input_density_smoothed_allr=self.density_smoothed_allr, 
+                                      input_density_allz=self.density_allz)
+        
+        return BMF, reio_maps
+    
+    
+    def save_quantities(self,BMF,reio_maps):
+        for q in self.BMF_quantities_to_save:
+            getattr(self, "BMF_" + q).append(getattr(BMF, q))
+        for q in self.ReioMaps_quantities_to_save:
+            getattr(self, "ReioMaps_" + q).append(getattr(reio_maps, q))
+
+
+    def run(self):
+        ### run the first model and save the densities for the next ones
+        # run model
+        BMF, reio_maps = self.run_1model(self.AstroParams_configs[0])
+        self.density = reio_maps.density # no need of np.copy here (reducing a bit the memory) since we won't modify the reio_maps.density
+        self.density_smoothed_allr = reio_maps.density_smoothed_allr
+        if reio_maps.COMPUTE_DENSITY_AT_ALLZ:
+            self.density_allz = reio_maps.density_allz
+        # save quantities
+        self.save_quantities(BMF,reio_maps)
+        # deallocate the reio maps class
+        z21_utilities.delete_class_attributes(reio_maps) # deleting reio_maps doesn't mean that self.density will get deleted since self.density is in fact a reference to the array
+
+        ### run all remaining models
+        for AstroParams_config in self.AstroParams_configs[1:]:
+            # run model
+            BMF, reio_maps = self.run_1model(AstroParams_config)
+            # save quantities
+            self.save_quantities(BMF,reio_maps)
+            # deallocate the reoi maps class
+            z21_utilities.delete_class_attributes(reio_maps)
+
diff --git a/zeus21/z21_utilities.py b/zeus21/z21_utilities.py
index 4bdd2ab..ccbef91 100644
--- a/zeus21/z21_utilities.py
+++ b/zeus21/z21_utilities.py
@@ -10,6 +10,7 @@
 import powerbox as pbox
 from pyfftw import empty_aligned as empty
 import time
+import gc
 
 def powerboxCtoR(pbobject,mapkin = None):
     'Function to convert a complex field to real 3D (eg density, T21...) on the powerbox notation'
@@ -49,4 +50,9 @@ def v2r(v):
     return (3/4/np.pi * v)**(1/3)
     
 def r2v(r):
-    return 4/3 * np.pi * r**3
\ No newline at end of file
+    return 4/3 * np.pi * r**3
+
+def delete_class_attributes(class_instance): # delete all attributes of the class instance
+    for attr in list(class_instance.__dict__):    
+        delattr(class_instance, attr)
+    gc.collect()
\ No newline at end of file

From 6b87804017efae48c983517902250550e945e7fc Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Tue, 7 Apr 2026 16:39:15 -0500
Subject: [PATCH 33/37] Fixes for the type for the barrier input in
 reionization_maps.

---
 zeus21/inputs.py | 6 ++++--
 zeus21/maps.py   | 5 +++--
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/zeus21/inputs.py b/zeus21/inputs.py
index f0fe7a3..788cbbf 100644
--- a/zeus21/inputs.py
+++ b/zeus21/inputs.py
@@ -234,7 +234,7 @@ def __init__(self, UserParams, Cosmo_Parameters,
                     E0_xray = 500.,
                     alpha_xray = -1.0,
                     Emax_xray_norm=2000,
-                    _clumping = 3,
+                    clumping = 3.0,
 
                     Nalpha_lyA_II = 9690,
                     Nalpha_lyA_III = 17900,
@@ -335,9 +335,11 @@ def __init__(self, UserParams, Cosmo_Parameters,
         #fesc(M) parameter. Power law normalized (fesc10) at M=1e10 Msun with index alphaesc
         self.fesc10 = fesc10
         self.alphaesc = alphaesc
-        self._clumping = _clumping #clumping factor, z-independent and fixed for now
+        self._clumping = clumping #clumping factor, z-independent and fixed for now
         if(Cosmo_Parameters.Flag_emulate_21cmfast==True):
             self._clumping = 2.0 #this is the 21cmFAST value
+            if clumping != 3.0:
+                self._clumping = clumping
 
 
 
diff --git a/zeus21/maps.py b/zeus21/maps.py
index bc17118..889b64a 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -23,6 +23,7 @@
 import time
 from dataclasses import dataclass
 from collections.abc import Sequence
+from typing import Callable
 
 
 class CoevalMaps:
@@ -148,7 +149,7 @@ class reionization_maps:
         Sets the predetermined generation of maps. Default is 1234.
     r_precision: float
         Allows to change the steps of the radii for faster computation. Default (and max) is 1, lower values make the computation faster at the cost of accuracy.
-    barrier: function
+    barrier: 2D np.array
         Input density barrier to be used as the threshold for map generation. Takes z value as input and returns np.array of shape. Default is None.
     PRINT_TIMER: bool
         Whether to print the time elapsed along the process. Default is True.
@@ -616,7 +617,7 @@ class ReioMapsConfig:
     """
     r_precision: float = 1.
     Rs: list | np.ndarray | None = None
-    barrier: function = None
+    barrier: np.ndarray = None
     PRINT_TIMER: bool = True
     LOGNORMAL_DENSITY: bool = False
     COMPUTE_DENSITY_AT_ALLZ: bool = False

From b2f1c0ac08a008baaa9e8a659261aab254fd824d Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Thu, 9 Apr 2026 16:05:32 -0500
Subject: [PATCH 34/37] Improve astro_variations + increase kmax in
 test_correlations + remove classy.class from cosmo_wrapper

---
 tests/test_correlations.py |   2 +-
 zeus21/cosmology.py        |   4 +-
 zeus21/maps.py             | 246 +++++++++++++++++++++----------------
 3 files changed, 143 insertions(+), 109 deletions(-)

diff --git a/tests/test_correlations.py b/tests/test_correlations.py
index f35ce13..ce7f226 100644
--- a/tests/test_correlations.py
+++ b/tests/test_correlations.py
@@ -19,7 +19,7 @@
 
 UserParams = zeus21.User_Parameters()
 
-CosmoParams_input = zeus21.Cosmo_Parameters_Input(kmax_CLASS = 10., zmax_CLASS = 10.) #to speed up
+CosmoParams_input = zeus21.Cosmo_Parameters_Input(kmax_CLASS = 500, zmax_CLASS = 10.) #to speed up
 ClassyCosmo = zeus21.runclass(CosmoParams_input)
 CosmoParams = zeus21.Cosmo_Parameters(UserParams, CosmoParams_input, ClassyCosmo)
 
diff --git a/zeus21/cosmology.py b/zeus21/cosmology.py
index 564e34f..653c89d 100644
--- a/zeus21/cosmology.py
+++ b/zeus21/cosmology.py
@@ -27,8 +27,8 @@ def cosmo_wrapper(User_Parameters, Cosmo_Parameters_Input):
     Cosmo_Parameters, Class_Cosmo, Correlations, HMF_interpolator
     """
 
-    ClassCosmo = Class()
-    ClassCosmo.compute()
+    #ClassCosmo = Class()
+    #ClassCosmo.compute()
 
     ClassyCosmo = runclass(Cosmo_Parameters_Input)
     CosmoParams = Cosmo_Parameters(User_Parameters, Cosmo_Parameters_Input, ClassyCosmo) 
diff --git a/zeus21/maps.py b/zeus21/maps.py
index 889b64a..9ae620b 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -8,22 +8,20 @@
 """
 
 from . import cosmology
-from . import constants
 from . import z21_utilities
 from . import inputs
+from . import correlations
 from . import sfrd
 from . import reionization
+import classy
 
 import numpy as np
 import powerbox as pbox
 from scipy.interpolate import interp1d
 from scipy.interpolate import InterpolatedUnivariateSpline as spline
-from pyfftw import empty_aligned as empty
 from tqdm import trange
 import time
-from dataclasses import dataclass
-from collections.abc import Sequence
-from typing import Callable
+from dataclasses import dataclass, field as _field
 
 
 class CoevalMaps:
@@ -543,9 +541,35 @@ def _compute_ionfrac_from_treion(self):
 
 
 
-##### Running reionization_maps multiple times with variations in astrophysics
+##### Wrapper for maps generation
 # dataclasses for inputing arguments to astro_variations
-@dataclass(frozen=True)
+# caveat: if the default are changed in the definition of the classes in zeus21, they will not change here...
+# --> this should be improved later
+@dataclass(kw_only=True)
+class UserParamsConfig:
+    precisionboost: float = 1.0
+    FLAG_FORCE_LINEAR_CF: int = 0
+    MIN_R_NONLINEAR: float = 2.0
+    MAX_R_NONLINEAR: float = 100.0
+    FLAG_DO_DENS_NL: bool = False
+    FLAG_WF_ITERATIVE: bool = True
+
+@dataclass(kw_only=True)
+class CosmoParamsInputConfig:
+    omegab: float = 0.0223828
+    omegac: float = 0.1201075
+    h_fid: float = 0.67810
+    As: float = 2.100549e-09
+    ns: float = 0.9660499
+    tau_fid: float = 0.05430842
+    kmax_CLASS: float = 500.
+    zmax_CLASS: float = 50.
+    zmin_CLASS: float = 5.
+    Flag_emulate_21cmfast: bool = False
+    USE_RELATIVE_VELOCITIES: bool = False
+    HMF_CHOICE: str = "ST"
+
+@dataclass(kw_only=True)
 class AstroParamsConfig:
     """
     All arguments of Astro_Parameters that have default values
@@ -564,7 +588,7 @@ class AstroParamsConfig:
     E0_xray: float = 500.
     alpha_xray: float = -1.0
     Emax_xray_norm: float = 2000
-    _clumping: float = 3
+    clumping: float = 3
     
     Nalpha_lyA_II: float = 9690
     Nalpha_lyA_III: float = 17900
@@ -595,9 +619,16 @@ class AstroParamsConfig:
     beta_LW: float = 0.6
     
     A_vcb: float = 1.0
-    beta_vcb: float = 1.8,
+    beta_vcb: float = 1.8
 
-@dataclass
+@dataclass(kw_only=True)
+class getT21coeffConfig:
+    """
+    All arguments of reionization_maps that have default values
+    """
+    zmin: float = 10.0
+
+@dataclass(kw_only=True)
 class BMFConfig:
     """
     All arguments of BMF that have default values.
@@ -609,12 +640,14 @@ class BMFConfig:
     Rmin: float = 0.05
     PRINT_SUCCESS: bool = True
 
-@dataclass
+@dataclass(kw_only=True)
 class ReioMapsConfig:
     """
-    All arguments of reionization_maps that have default values, 
-    except the input_densities since they are changed dynamically in astro_variations
+    All arguments of reionization_maps that have default values
     """
+    input_boxlength: float = 300.
+    ncells: int = 300
+    seed: int = 1234
     r_precision: float = 1.
     Rs: list | np.ndarray | None = None
     barrier: np.ndarray = None
@@ -627,8 +660,12 @@ class ReioMapsConfig:
     COMPUTE_PARTIAL_IONIZATIONS: bool = False
     COMPUTE_PARTIAL_AND_MASSWEIGHTED: bool = False
     COMPUTE_ZREION: bool = False
+    input_density: np.ndarray | None = None
+    input_density_smoothed_allr: np.ndarray | None = None
+    input_density_allz: np.ndarray | None = None
 
-# running class that will do the astro variations
+# class to run reionization_maps multiple times with variations in astrophysics
+@dataclass(kw_only=True)
 class astro_variations:
     """
     Efficiently rerun the reionization_maps class while varying astrophysical parameters.
@@ -638,145 +675,142 @@ class astro_variations:
 
     Parameters
     ----------
-    UserParams: zeus21.User_Parameters class
-        Stores the user parameters.
-    CosmoParams: zeus21.Cosmo_Parameters class
-        Stores the cosmology.
-    ClassyCosmo: zeus21.runclass class
-        Sets up Class cosmology.
-    HMFintclass: zeus21.HMF_interpolator class
-        Stores the HMF interpolator.
-    CorrFClass: zeus21.Correlations class
-        Stores the correlations.
-    input_boxlength: float
-        Comoving physical side length of the box.
-    ncells: int
-        Number of cells on a side.
-    seed: int
-        Sets the predetermined generation of maps. Default is 1234.
     AstroParams_configs: sequence of AstroParamsConfig instances
         Sets the different astro models to run.
     BMF_quantities_to_save: list of str of the names of BMF attributes
         Specifiy which quantities to save as attributes in astro_variations. 
-        They will be saved in lists named self.BMF_attr (e.g.: self.BMF_ion_frac).
+        They will be saved in lists in a dict named self.BMF_quantities.
         Each element of these lists correspond to the corresponding AstroParams_configs element. 
     ReioMaps_quantities_to_save: list of str of the names of reionization_maps attributes
         Specifiy which quantities to save as attributes in astro_variations.
-        They will be saved in lists named self.ReioMaps_attr (e.g.: self.ReioMaps_ion_frac).
+        They will be saved in lists named self.ReioMaps_quantities.
         Each element of these lists correspond to the corresponding AstroParams_configs element. 
+    input_z: list of float
+        Redshift list over which the reionization maps are made. Default is None, corresponding to zeus21.get_T21_coefficients.zintegral.
+    UserParams_config: UserParamsConfig class
+        Sets up the config for the User_Parameters class. Default is None and corresponds to default User_Parameters arguments.
+    CosmoParamsInput_config: CosmoParamsInputConfig class
+        Sets up the config for the Cosmo_Parameters_Input class. Default is None and corresponds to default Cosmo_Parameters_Input arguments.
+    getT21coeff_config: getT21coeffConfig class
+        Sets up the config for the get_T21_coefficients class. Default is None and corresponds to default get_T21_coefficients arguments.
     BMF_config: BMFConfig class
         Sets up the config for the BMF class. Default is None and corresponds to default BMF arguments.
     ReioMaps_config: ReioMapsConfig class
         Sets up the config for the reionization_maps class. Default is None and corresponds to default reionization_maps arguments.
-    ZMIN: float
-        Minimum redshift to which zeus21.get_T21_coefficients will do its integrations. Default is 5.0.
-    input_z: list of float
-        Redshift list over which the reionization maps are made. Default is None, corresponding to zeus21.get_T21_coefficients.zintegral.
-    seed: int
-        Sets the seed used to generate the boxes. Default is 1234.
 
     Attributes
     ----------
-    density: 3D np.array
-        Saves the density maps generated in the first run.
-    density_smoothed_allr: 4D np.array
-        Saves the density_smoothed_allr maps generated in the first run.
-    density_allz: 4D np.array
-        Saves the density_allz maps generated in the first run.
-    BMF_... and ReioMaps_...:
-        All attributes from the BMF and reionization_map classes the user wants to save.
+    BMF_quantities:
+        dict containing all the quantities the user wants to save from the BMF class.
+        Each quantity has the same name in the dict than in the class and is a list, of which
+        each element is the output of zeus21 for the given astro model.
+    ReioMaps_quantities:
+        dict containing all the quantities the user wants to save from the reionization_map class.
+        Each quantity has the same name in the dict than in the class and is a list, of which
+        each element is the output of zeus21 for the given astro model.
+    UserParams: User_Parameters 
+        Stores the user parameters.
+    CosmoParams_input: Cosmo_Parameters_Input
+        Stores the input for cosmology.
+    CosmoParams: Cosmo_Parameters
+        Stores the cosmology.
+    ClassyCosmo: classy.Class
+        Stores the Classy cosmology.
+    CorrFClass: Correlations
+        Stores the correlations.
+    HMFintclass: HMF_interpolator
+        Stores the HMF interpolator.
     """
-
-    def __init__(self, UserParams, CosmoParams, ClassyCosmo, HMFintclass, CorrFClass, 
-                 input_boxlength, ncells, 
-                 AstroParams_configs: Sequence[AstroParamsConfig],
-                 BMF_quantities_to_save, ReioMaps_quantities_to_save, 
-                 BMF_config: BMFConfig | None = None,
-                 ReioMaps_config: ReioMapsConfig | None = None,
-                 ZMIN = 5., input_z=None, seed=1234):
-        
-        # zeus21 classes
-        self.UserParams = UserParams
-        self.CosmoParams, self.ClassyCosmo = CosmoParams, ClassyCosmo
-        self.HMFintclass, self.CorrFClass = HMFintclass, CorrFClass
-        self.ZMIN = ZMIN
-        
-        # maps input
-        self.input_z = input_z
-        self.input_boxlength, self.ncells = input_boxlength, ncells
-        self.seed = seed
-        
-        # BMF and reionization_maps arguments
-        self.AstroParams_configs = tuple(AstroParams_configs)
-        self.BMF_config = BMF_config or BMFConfig()
-        self.ReioMaps_config = ReioMaps_config or ReioMapsConfig()
-        
-        # saved density boxes for time and memory efficiency
-        self.density = None
-        self.density_smoothed_allr = None
-        self.density_allz = None
-
-        # initialize physical quantities to save
-        self.BMF_quantities_to_save = BMF_quantities_to_save
-        self.ReioMaps_quantities_to_save = ReioMaps_quantities_to_save
-        for q in BMF_quantities_to_save:
+    # arguments with no default: what astrophysics and which quantities to save from BMF and reionization_maps
+    AstroParams_configs: list[AstroParamsConfig]
+    BMF_quantities_to_save: list[str] 
+    ReioMaps_quantities_to_save: list[str]
+
+    # maps input
+    input_z: list | None = _field(default=None)
+
+    # zeus21 configuration
+    UserParams_config: UserParamsConfig = _field(default_factory=UserParamsConfig)
+    CosmoParamsInput_config: CosmoParamsInputConfig = _field(default_factory=CosmoParamsInputConfig)
+    getT21coeff_config: getT21coeffConfig = _field(default_factory=getT21coeffConfig)
+    BMF_config: BMFConfig = _field(default_factory=BMFConfig)
+    ReioMaps_config: ReioMapsConfig = _field(default_factory=ReioMapsConfig)
+
+    # initialize quantities dict and zeus21
+    BMF_quantities: dict[str, list] = _field(init=False)
+    ReioMaps_quantities: dict[str, list] = _field(init=False)
+    UserParams: inputs.User_Parameters = _field(init=False)
+    CosmoParams_input: inputs.Cosmo_Parameters_Input = _field(init=False)
+    CosmoParams: inputs.Cosmo_Parameters = _field(init=False)
+    ClassyCosmo: classy.Class = _field(init=False)
+    CorrFClass: correlations.Correlations = _field(init=False)
+    HMFintclass: cosmology.HMF_interpolator = _field(init=False)
+
+    def __post_init__(self):
+        # initialize BMF quantities to save
+        self.BMF_quantities = {}
+        for q in self.BMF_quantities_to_save:
             if q in vars(reionization.BMF)["__static_attributes__"]:
-                setattr(self, "BMF_"+q, [])
+                self.BMF_quantities[q] = []
             else:
                 raise ValueError(f"{q} is not an attribute of the BMF class.")
-        for q in ReioMaps_quantities_to_save:
+        
+        # initialize reionization_maps quantities to save
+        self.ReioMaps_quantities = {}
+        for q in self.ReioMaps_quantities_to_save:
             if q in vars(reionization_maps)["__static_attributes__"]:
-                setattr(self, "ReioMaps_"+q, [])
+                self.ReioMaps_quantities[q] = []
             else:
-                raise ValueError(f"{q} is not an attribute of the reionization_maps class.")
+                raise ValueError(f"{q} is not an attribute of the reionization_maps class.")     
+
+        # initialize zeus21
+        self.UserParams = inputs.User_Parameters(**vars(self.UserParams_config))
+        self.CosmoParams_input = inputs.Cosmo_Parameters_Input(**vars(self.CosmoParamsInput_config)) 
+        self.CosmoParams, self.ClassyCosmo, self.CorrFClass , self.HMFintclass =  cosmology.cosmo_wrapper(self.UserParams, self.CosmoParams_input)
 
 
     def run_1model(self, AstroParams_config):
         AstroParams = inputs.Astro_Parameters(self.UserParams, self.CosmoParams, **vars(AstroParams_config))
-        CoeffStructure = sfrd.get_T21_coefficients(self.UserParams, self.CosmoParams, self.ClassyCosmo, AstroParams, self.HMFintclass, zmin=self.ZMIN)
-    
-        BMF = reionization.BMF(CoeffStructure, self.HMFintclass, self.CosmoParams, AstroParams, self.ClassyCosmo,
-                            **vars(self.BMF_config))
-        
+        CoeffStructure = sfrd.get_T21_coefficients(self.UserParams, self.CosmoParams, self.ClassyCosmo, AstroParams, self.HMFintclass, 
+                                                   **vars(self.getT21coeff_config))
         if self.input_z is None:
             self.input_z = CoeffStructure.zintegral
+    
+        bmf = reionization.BMF(CoeffStructure, self.HMFintclass, self.CosmoParams, AstroParams, self.ClassyCosmo, 
+                               **vars(self.BMF_config))
         
-        reio_maps = reionization_maps(self.CosmoParams, self.ClassyCosmo, self.CorrFClass, CoeffStructure, BMF, self.input_z,
-                                      input_boxlength=self.input_boxlength, ncells=self.ncells, seed=self.seed,
-                                      **vars(self.ReioMaps_config),
-                                      input_density=self.density, 
-                                      input_density_smoothed_allr=self.density_smoothed_allr, 
-                                      input_density_allz=self.density_allz)
+        reio_maps = reionization_maps(self.CosmoParams, self.ClassyCosmo, self.CorrFClass, CoeffStructure, bmf, self.input_z, 
+                                      **vars(self.ReioMaps_config))
         
-        return BMF, reio_maps
+        return bmf, reio_maps
     
     
-    def save_quantities(self,BMF,reio_maps):
+    def save_quantities(self,bmf,reio_maps):
         for q in self.BMF_quantities_to_save:
-            getattr(self, "BMF_" + q).append(getattr(BMF, q))
+            self.BMF_quantities[q].append(getattr(bmf, q))
         for q in self.ReioMaps_quantities_to_save:
-            getattr(self, "ReioMaps_" + q).append(getattr(reio_maps, q))
+            self.ReioMaps_quantities[q].append(getattr(reio_maps, q))
 
 
     def run(self):
         ### run the first model and save the densities for the next ones
         # run model
-        BMF, reio_maps = self.run_1model(self.AstroParams_configs[0])
-        self.density = reio_maps.density # no need of np.copy here (reducing a bit the memory) since we won't modify the reio_maps.density
-        self.density_smoothed_allr = reio_maps.density_smoothed_allr
+        bmf, reio_maps = self.run_1model(self.AstroParams_configs[0])
+        self.ReioMaps_config.input_density = reio_maps.density # no need of np.copy here (reducing a bit the memory) since we won't modify the reio_maps.density
+        self.ReioMaps_config.input_density_smoothed_allr = reio_maps.density_smoothed_allr
         if reio_maps.COMPUTE_DENSITY_AT_ALLZ:
-            self.density_allz = reio_maps.density_allz
+            self.ReioMaps_config.input_density_allz = reio_maps.density_allz
         # save quantities
-        self.save_quantities(BMF,reio_maps)
+        self.save_quantities(bmf,reio_maps)
         # deallocate the reio maps class
         z21_utilities.delete_class_attributes(reio_maps) # deleting reio_maps doesn't mean that self.density will get deleted since self.density is in fact a reference to the array
 
         ### run all remaining models
         for AstroParams_config in self.AstroParams_configs[1:]:
             # run model
-            BMF, reio_maps = self.run_1model(AstroParams_config)
+            bmf, reio_maps = self.run_1model(AstroParams_config)
             # save quantities
-            self.save_quantities(BMF,reio_maps)
+            self.save_quantities(bmf,reio_maps)
             # deallocate the reoi maps class
             z21_utilities.delete_class_attributes(reio_maps)
 

From a9564e47ca004b796affca9f3d829b9d6d97b3ad Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Thu, 9 Apr 2026 19:31:14 -0500
Subject: [PATCH 35/37] Small fix in astro_variations.

---
 zeus21/maps.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index 9ae620b..edcb0dd 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -750,7 +750,7 @@ def __post_init__(self):
         # initialize BMF quantities to save
         self.BMF_quantities = {}
         for q in self.BMF_quantities_to_save:
-            if q in vars(reionization.BMF)["__static_attributes__"]:
+            if q in getattr(reionization.BMF, "__static_attributes__"):
                 self.BMF_quantities[q] = []
             else:
                 raise ValueError(f"{q} is not an attribute of the BMF class.")
@@ -758,7 +758,7 @@ def __post_init__(self):
         # initialize reionization_maps quantities to save
         self.ReioMaps_quantities = {}
         for q in self.ReioMaps_quantities_to_save:
-            if q in vars(reionization_maps)["__static_attributes__"]:
+            if q in getattr(reionization_maps, "__static_attributes__"):
                 self.ReioMaps_quantities[q] = []
             else:
                 raise ValueError(f"{q} is not an attribute of the reionization_maps class.")     

From 20c4e13340ed25354125bdd3dd371053fcdaa238 Mon Sep 17 00:00:00 2001
From: Emilie Thelie <emilie77220@gmail.com>
Date: Fri, 10 Apr 2026 15:53:07 -0500
Subject: [PATCH 36/37] Small fix in astro_variations.

---
 zeus21/maps.py | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/zeus21/maps.py b/zeus21/maps.py
index edcb0dd..38ccf56 100644
--- a/zeus21/maps.py
+++ b/zeus21/maps.py
@@ -749,19 +749,27 @@ class astro_variations:
     def __post_init__(self):
         # initialize BMF quantities to save
         self.BMF_quantities = {}
-        for q in self.BMF_quantities_to_save:
-            if q in getattr(reionization.BMF, "__static_attributes__"):
-                self.BMF_quantities[q] = []
-            else:
-                raise ValueError(f"{q} is not an attribute of the BMF class.")
+        __static_attributes__ = getattr(reionization.BMF, "__static_attributes__", None)
+        if __static_attributes__ is None:
+            print(f"Warning: cannot check whether BMF_quantities are attributes of the BMF class: errors could show up after generating the boxes.")
+        else:
+            for q in self.BMF_quantities_to_save:
+                if q in __static_attributes__:
+                    self.BMF_quantities[q] = []
+                else:
+                    raise ValueError(f"{q} is not an attribute of the BMF class.")
         
         # initialize reionization_maps quantities to save
         self.ReioMaps_quantities = {}
-        for q in self.ReioMaps_quantities_to_save:
-            if q in getattr(reionization_maps, "__static_attributes__"):
-                self.ReioMaps_quantities[q] = []
-            else:
-                raise ValueError(f"{q} is not an attribute of the reionization_maps class.")     
+        __static_attributes__ = getattr(reionization_maps, "__static_attributes__", None)
+        if __static_attributes__ is None:
+            print(f"Warning: cannot check whether ReioMaps_quantities are attributes of the reionization_maps class: errors could show up after generating the boxes.")
+        else:
+            for q in self.ReioMaps_quantities_to_save:
+                if q in __static_attributes__:
+                    self.ReioMaps_quantities[q] = []
+                else:
+                    raise ValueError(f"{q} is not an attribute of the reionization_maps class.")     
 
         # initialize zeus21
         self.UserParams = inputs.User_Parameters(**vars(self.UserParams_config))

From d6976a9b0284374d5a9b6f642da543c8cb9d0c18 Mon Sep 17 00:00:00 2001
From: yonboyage <59982772+yonboyage@users.noreply.github.com>
Date: Wed, 29 Apr 2026 17:14:17 -0500
Subject: [PATCH 37/37] Added broadcasting to most code

- Broadcasting used to speed up computation by 3x, almost everywhere
- Several BMF functions now accept arrays as well as scalars
- Added moving pivot for barrier fits, so that now the BMF best represents the bubbles around the peak, rather than a static fit at one bubble size that has increasing errors farther from the pivot; i.e. the most common bubbles are always well-modeled by the new moving pivot.
- Added a few interpolators, including for peak R of BMF
---
 zeus21/reionization.py | 234 ++++++++++++++++++++++++++++++-----------
 1 file changed, 172 insertions(+), 62 deletions(-)

diff --git a/zeus21/reionization.py b/zeus21/reionization.py
index ea5b45f..2accc9e 100644
--- a/zeus21/reionization.py
+++ b/zeus21/reionization.py
@@ -15,9 +15,12 @@
 from scipy.integrate import cumulative_trapezoid
 from scipy.interpolate import interp1d
 from scipy.interpolate import RegularGridInterpolator
+from scipy.interpolate import UnivariateSpline
 from scipy.special import erfc
 from tqdm import trange
 
+import time
+
 
 class BMF:
     """
@@ -25,7 +28,6 @@ class BMF:
 
     
     """
-    
     def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, ClassyCosmo, R_linear_sigma_fit_input=10, FLAG_converge=True, max_iter=10, ZMAX_REION = 30, Rmin=0.05, PRINT_SUCCESS=True):
 
         self.PRINT_SUCCESS = PRINT_SUCCESS
@@ -34,16 +36,20 @@ def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, Classy
         self.Rs = CoeffStructure.Rtabsmoo
         self.Rs_BMF = np.logspace(np.log10(Rmin), np.log10(self.Rs[-1]), 100)
         self.ds_array = np.linspace(-1, 5, 101)
+
+        self._r_array, self._z_array = np.meshgrid(self.Rs, self.zlist, sparse=True, indexing='ij')
+        self._rb_array, self._z_array = np.meshgrid(self.Rs_BMF, self.zlist, sparse=True, indexing='ij')
         
         self.gamma = CoeffStructure.gamma_niondot_II_index2D
         self.gamma2 = CoeffStructure.gamma2_niondot_II_index2D
-        self.sigma = np.array([[ClassyCosmo.sigma(r, z) for z in self.zlist] for r in self.Rs]).T#CoeffStructure.sigmaofRtab
+        #self.sigma = np.array([[ClassyCosmo.sigma(r, z) for z in self.zlist] for r in self.Rs]).T#CoeffStructure.sigmaofRtab
+        self.sigma =  HMFintclass.sigmaRintlog((np.log(self._r_array), self._z_array)).T
         
         self.zr = [self.zlist, np.log(self.Rs)]
         self.gamma_int = RegularGridInterpolator(self.zr, self.gamma, bounds_error = False, fill_value = None)
         self.gamma2_int = RegularGridInterpolator(self.zr, self.gamma2, bounds_error = False, fill_value = None)
 
-        self.sigma_BMF = np.array([[ClassyCosmo.sigma(r, z) for z in self.zlist] for r in self.Rs_BMF]).T
+        self.sigma_BMF = HMFintclass.sigmaRintlog((np.log(self._rb_array), self._z_array)).T #might need to make a new interpolator for different R range
         self.zr_BMF = [self.zlist, np.log(self.Rs_BMF)]
         self.sigma_int = RegularGridInterpolator(self.zr_BMF, self.sigma_BMF, bounds_error = False, fill_value = None)
         
@@ -55,7 +61,7 @@ def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, Classy
         self.niondot_avg = CoeffStructure.niondot_avg_II
         self.niondot_avg_int = interp1d(self.zlist, self.niondot_avg, bounds_error = False, fill_value = None)
 
-        self.ion_frac = np.fmin(1, [self.calc_Q(CosmoParams, z) for z in self.zlist])
+        self.ion_frac = np.fmin(1, self.Madau_Q(CosmoParams, self.zlist))
         self.ion_frac_initial = np.copy(self.ion_frac)
 
         zr_mesh = np.meshgrid(np.arange(len(self.Rs)), np.arange(len(self.zlist)))
@@ -70,19 +76,22 @@ def __init__(self, CoeffStructure, HMFintclass, CosmoParams, AstroParams, Classy
         self.dzr = [self.ds_array, self.zlist, np.log(self.Rs)]
         self.prebarrier_xHII_int = RegularGridInterpolator(self.dzr, self.prebarrier_xHII, bounds_error = False, fill_value = None) #allow extrapolation
 
-        self.R_linear_sigma_fit_idx = z21_utilities.find_nearest_idx(self.Rs, R_linear_sigma_fit_input)
+        self.R_linear_sigma_fit_idx = z21_utilities.find_nearest_idx(self.Rs, R_linear_sigma_fit_input)[0]
         self.R_linear_sigma_fit = self.Rs[self.R_linear_sigma_fit_idx]
 
         #fake bubble mass function to impose peak around R_linear_sigma_fit for the initial linear barriers
         #looks something like [0, 0, ..., 1, ..., 0, 0]*(number of redshifts)
         self.BMF = np.repeat([np.eye(len(self.Rs_BMF))[self.R_linear_sigma_fit_idx]], len(self.zlist), axis=0)
 
+        self.peakRofz = np.array([self.BMF_peak_R(z) for z in self.zlist])
+        self.peakRofz_int = interp1d(self.zlist, self.peakRofz, bounds_error = False, fill_value = None)
+
         #second computation of BMF using the initial guess peaks
-        self.BMF = np.array([self.VRdn_dR(z, self.Rs_BMF) for z in self.zlist])
+        self.BMF = self.VRdn_dR(self.zlist, self.Rs_BMF)
         self.BMF_initial = np.copy(self.BMF)
         
         #self.ion_frac = np.nan_to_num([np.trapezoid(self.BMF[i], np.log(self.Rs_BMF)) for i in range(len(self.zlist))]) #ion_frac by numerically integrating the BMF
-        self.ion_frac = np.nan_to_num([self.analytic_Q(ClassyCosmo, z)[0] for z in self.zlist]) #ion_frac by analytic integral of BMF
+        self.ion_frac = np.nan_to_num(self.analytic_Q(CosmoParams, ClassyCosmo, self.zlist)) #ion_frac by analytic integral of BMF
         
         self.ion_frac[self.barrier[:, -1]<=0] = 1
         
@@ -95,8 +104,8 @@ def compute_prebarrier_xHII(self, CosmoParams, ion_frac, z, R):
         """
         
         """
-        nion_values = self.nion_delta_r_int(CosmoParams, z, R)  #Shape (nd, nz)
-        nrec_values = self.nrec(CosmoParams, ion_frac, z)       #Shape (nd, nz)
+        nion_values = self.nion_delta_r_int(CosmoParams, z, R)  #Shape (nd, nz, nR)
+        nrec_values = self.nrec(CosmoParams, ion_frac, z)[:, :, None]       #Shape (nd, nz) * (1, 1, nR)
         
         prebarrier_xHII = nion_values / (1 + nrec_values)
 
@@ -125,15 +134,15 @@ def compute_barrier(self, CosmoParams, ion_frac, z, R):
         zarg = np.argsort(z) #sort just in case
         z = z[zarg]
         ion_frac = ion_frac[zarg]
+
+        #Compute nion_values and nrec_values based on (re)computed ion_frac
+        self.prebarrier_xHII =  self.compute_prebarrier_xHII(CosmoParams, ion_frac, z, R)
+        total_values = np.log10(self.prebarrier_xHII + 1e-10)
         
         for ir in range(len(R)):
-            #Compute nion_values and nrec_values for this 'ir'
-            self.prebarrier_xHII[:, :, ir] =  self.compute_prebarrier_xHII(CosmoParams, ion_frac, z, R[ir])
-            total_values = np.log10(self.prebarrier_xHII[:, :, ir] + 1e-10)
-        
             #Loop over redshift indices
             for iz in range(len(self.zlist)):
-                y_values = total_values[:, iz]  #Shape (nd,)
+                y_values = total_values[:, iz, ir]  #Shape (nd,)
         
                 #Find zero crossings
                 sign_change = np.diff(np.sign(y_values))
@@ -148,7 +157,7 @@ def compute_barrier(self, CosmoParams, ion_frac, z, R):
                     barrier[iz, ir] = x_intersect[0]  #Assuming we take the first crossing
                 else:
                     barrier[iz, ir] = np.nan #Never crosses
-        barrier = barrier * (CosmoParams.growthint(self.zlist)/CosmoParams.growthint(self.zlist[0]))[:, np.newaxis] #scale barrier with growth factor
+        barrier = barrier * (CosmoParams.growthint(self.zlist)/CosmoParams.growthint(self.zlist[0]))[:, None] #scale barrier with growth factor
         barrier[self.zlist > self.ZMAX_REION] = 100 #sets density to an unreachable barrier, as if reionization isn't happening
         return barrier
 
@@ -217,20 +226,23 @@ def niondot_delta_r(self, CosmoParams, z, R, d_array=None):
             The rates of ionizing photon production. The first dimension is densities, the second dimension is redshifts.
         """
 
-        zarg = np.argsort(z) #sort just in case
-        z = z[zarg]
+        z1d = np.copy(z)
+        R1d = np.copy(R)
+        
+        z = z[None, :, None]
+        R = R[None, None, :]
 
         if d_array is None:
-            d_array = self.ds_array
-        
-        d_array = d_array[:, np.newaxis] * CosmoParams.growthint(z)[np.newaxis, :] / CosmoParams.growthint(z[0])
+            d_array = self.ds_array[:, None, None]
+
+        d_array = d_array * CosmoParams.growthint(z) / CosmoParams.growthint(z1d[0])
     
-        gamma_R = self.gammaz_int(z, R)   
-        gamma2_R = self.gamma2z_int(z, R)  
-        nion_norm_R = self.nion_normz_int(z, R)
+        gamma = self.gamma_zR_int(z1d[:, None], R1d[None, :])[None, :, :]
+        gamma2 = self.gamma2_zR_int(z1d[:, None], R1d[None, :])[None, :, :]
+        nion_norm = self.nion_norm_zR_int(z1d[:, None], R1d[None, :])[None, :, :]
     
-        exp_term = np.exp(gamma_R[np.newaxis, :] * d_array + gamma2_R[np.newaxis, :] * d_array**2)
-        niondot = (self.niondot_avg_int(z)[np.newaxis, :] / nion_norm_R[np.newaxis, :]) * exp_term
+        exp_term = np.exp(gamma * d_array + gamma2 * d_array**2)
+        niondot = (self.niondot_avg_int(z) / nion_norm) * exp_term
         
         return niondot
     
@@ -256,7 +268,7 @@ def nion_delta_r_int(self, CosmoParams, z, R, d_array=None):
         z.sort() #sort if not sorted
 
         if d_array is None:
-            d_array = self.ds_array
+            d_array = self.ds_array[:, None, None]
         
         #reverse the inputs to make the integral easier to compute
         z_rev = z[::-1]
@@ -264,14 +276,15 @@ def nion_delta_r_int(self, CosmoParams, z, R, d_array=None):
     
         niondot_values = self.niondot_delta_r(CosmoParams, z, R, d_array)
     
-        integrand = -1 / (1 + z_rev) / Hz_rev * niondot_values[:, ::-1]
-        nion = cumulative_trapezoid(integrand, x=z_rev, initial=0)[:, ::-1] #reverse back to increasing z order
+        integrand = -1 / (1 + z_rev[None, :, None]) / Hz_rev[None, :, None] * niondot_values[:, ::-1]
+        nion = cumulative_trapezoid(integrand, x=z_rev, initial=0, axis=1)[:, ::-1] #reverse back to increasing z order
         
         return nion
 
-    #calculating ionized fraction
-    def calc_Q(self, CosmoParams, z):
-        z_arr = np.logspace(np.log10(z), np.log10(self.zlist[-1]), 50)
+    #calculating naive ionized fraction
+    def Madau_Q(self, CosmoParams, z):
+        z = np.atleast_1d(z) #accepts scalar or array
+        z_arr = np.geomspace(z, self.zlist[-1], len(self.zlist))
         dtdz = 1/cosmology.Hubinvyr(CosmoParams, z_arr)/(1 + z_arr)
         tau0 = self.trec0 * np.sqrt(CosmoParams.OmegaM) * cosmology.Hubinvyr(CosmoParams, 0) / constants.yrTos
         exp = np.exp(2/3/tau0 * (np.power(1 + z, 3/2) - np.power(1 + z_arr, 3/2))) #switched order around to be correct (typo in paper)
@@ -279,53 +292,102 @@ def calc_Q(self, CosmoParams, z):
         niondot_avgs = self.niondot_avg_int(z_arr)
         integrand = dtdz * niondot_avgs * exp
     
-        return np.trapezoid(integrand, x = z_arr)
+        return np.trapezoid(integrand, x = z_arr, axis = 0)
 
     #computing linear barrier
     def B_1(self, z):
-        R_pivot = np.max([0.1, self.BMF_peak_R(z)])
-        sigmax = self.sigmaR_int(z, R_pivot*1.1)
-        sigmin = self.sigmaR_int(z, R_pivot*0.9)
-        barriermax = self.barrierR_int(z, R_pivot*1.1)
-        barriermin = self.barrierR_int(z, R_pivot*0.9)
+        R_pivot = self.peakRofz_int(z)
+        sigmax = np.diagonal(self.sigma_zR_int(z[:, None], (R_pivot*1.1)[None, :]))
+        sigmin = np.diagonal(self.sigma_zR_int(z[:, None], (R_pivot*0.9)[None, :]))
+        barriermax = np.diagonal(self.barrier_zR_int(z[:, None], (R_pivot*1.1)[None, :]))
+        barriermin = np.diagonal(self.barrier_zR_int(z[:, None], (R_pivot*0.9)[None, :]))
         return (barriermax - barriermin)/(sigmax**2 - sigmin**2)
         
     def B_0(self, z):
-        R_pivot = np.max([0.1, self.BMF_peak_R(z)])
-        sigmin = self.sigmaR_int(z, R_pivot*0.9)
-        barriermin = self.barrierR_int(z, R_pivot*0.9)
+        R_pivot = self.peakRofz_int(z)
+        sigmin = np.diagonal(self.sigma_zR_int(z[:, None], (R_pivot*0.9)[None, :]))
+        barriermin = np.diagonal(self.barrier_zR_int(z[:, None], (R_pivot*0.9)[None, :]))
         return barriermin - sigmin**2 * self.B_1(z)
     
-    def B(self, z, R): 
-        sig = self.sigmaR_int(z, R)
-        return self.B_0(z) + self.B_1(z)*sig**2
+    def B(self, z, R, sig):
+        B0 = self.B_0(z)
+        B1 = self.B_1(z)
+        return B0[:, None] + B1[:, None]*sig**2
     
     #computing other terms in the BMF
-    def dsigma_dR(self, z, R):
-        sigma = self.sigmaR_int(z, R)
-        return sigma/R*np.gradient(np.log(sigma), np.log(R))
-    
-    def dlogsigma_dlogR(self, z, R):
-        sigma = self.sigmaR_int(z, R)
-        return self.dsigma_dR(z, R) * R/sigma
+    def dlogsigma_dlogR(self, z, R, sig):
+        return np.gradient(np.log(sig), np.log(R), axis=1)
     
     def VRdn_dR(self, z, R):
-        sig = self.sigmaR_int(z, R)
-        return np.sqrt(2/np.pi) * np.abs(self.dlogsigma_dlogR(z, R)) * np.abs(self.B_0(z))/sig * np.exp(-self.B(z, R)**2/2/sig**2)
+        z = np.atleast_1d(z)
+        sig = self.sigma_zR_int(z[:, None], R[None, :])
+        B0 = self.B_0(z)
+        B1 = self.B_1(z)
+        return np.sqrt(2/np.pi) * np.abs(self.dlogsigma_dlogR(z, R, sig)) * np.abs(B0[:, None])/sig * np.exp(-(B0[:, None]+B1[:, None]*sig**2)**2/2/sig**2)
     
     def Rdn_dR(self, z, R):
-        return self.VRdn_dR(z, R)*3/(4*np.pi*R**3)
+        return self.VRdn_dR(z, R)*3/(4*np.pi*R[None, :]**3)
+#
+#    def BMF_peak_R(self, z):
+#        iz = z21_utilities.find_nearest_idx(self.zlist, z)
+#        ir = np.argmax(self.BMF[iz])
+#        return self.Rs_BMF[ir]
+
+    def BMF_peak_R(self, z, fit_window=5, max_bubble=100, min_bubble = 0.2): #written and commented by Claude
+        iz = z21_utilities.find_nearest_idx(self.zlist, z)[0]
+
+        # Find the coarse peak index
+        ir_peak = np.argmax(self.BMF[iz])
+
+        # Slice a window around the peak
+        i_lo = max(0, ir_peak - fit_window)
+        i_hi = min(len(self.Rs_BMF), ir_peak + fit_window + 1)
+        
+        R_window = self.Rs_BMF[i_lo:i_hi]
+        BMF_row = self.BMF[iz, :]
+        BMF_window = BMF_row[i_lo:i_hi]
+        
+        # If the peak is within fit_window of either edge, the true peak may
+        # be at the boundary — skip the spline and return the coarse peak
+        peak_at_left_edge  = (ir_peak - fit_window <= 0)
+        peak_at_right_edge = (ir_peak + fit_window >= len(self.Rs_BMF) - 1)
+        
+        if peak_at_left_edge or peak_at_right_edge:
+            return np.clip(self.Rs_BMF[ir_peak], min_bubble, max_bubble)
+        
+        # Also guard against a window that's too small to fit a degree-4 spline
+        # (need at least k+1 = 5 points)
+        if len(R_window) < 5:
+            return np.clip(self.Rs_BMF[ir_peak], min_bubble, max_bubble)
+        
+        # Fit a spline and find its maximum
+        spline = UnivariateSpline(R_window, BMF_window, k=4, s=0)
+        roots = spline.derivative().roots()
+        
+        # Keep only roots that are local maxima (second derivative < 0)
+        # and lie within the window bounds
+        d2 = spline.derivative(n=2)
+        valid_roots = [
+            r for r in roots
+            if d2(r) < 0 and R_window[0] <= r <= R_window[-1]
+        ]
+        
+        # Return the valid root closest to the coarse peak, or fall back
+        if len(valid_roots) == 0:
+            return np.clip(self.Rs_BMF[ir_peak], min_bubble, max_bubble)
+        
+        ir_peak_R = self.Rs_BMF[ir_peak]
+
+        peak_R = valid_roots[np.argmin(np.abs(np.array(valid_roots) - ir_peak_R))]
 
-    def BMF_peak_R(self, z):
-        iz = z21_utilities.find_nearest_idx(self.zlist, z)
-        ir = np.argmax(self.BMF[iz])
-        return self.Rs_BMF[ir]
+        return np.clip(peak_R, min_bubble, max_bubble) #peak can't be outside the allowed bounds
 
-    def analytic_Q(self, ClassyCosmo, z): #analytically integrating the BMF to get Q
+    def analytic_Q(self, CosmoParams, ClassyCosmo, z): #analytically integrating the BMF to get Q
+        z = np.atleast_1d(z)
         Rmin = 1e-10 #arbitrarily small
         B0 = self.B_0(z)
         B1 = self.B_1(z)
-        sigmin = ClassyCosmo.sigma(Rmin, z)
+        sigmin = ClassyCosmo.sigma(Rmin, z[0])*CosmoParams.growthint(z)/CosmoParams.growthint(z[0])
         s2 = sigmin**2
         return 0.5*np.exp(-2*B0*B1)*erfc((B0-B1*s2)/np.sqrt(2*s2)) + 0.5*erfc((B0+B1*s2)/np.sqrt(2*s2))
 
@@ -334,12 +396,15 @@ def converge_BMF(self, CosmoParams, ClassyCosmo, ion_frac_input, max_iter):
         iterator = trange(max_iter) if self.PRINT_SUCCESS else range(max_iter)
         for j in iterator:
             ion_frac_prev = np.copy(self.ion_frac)
-            
+
             self.barrier = self.compute_barrier(CosmoParams, self.ion_frac, self.zlist, self.Rs)
             self.barrier_int = RegularGridInterpolator(self.zr, self.barrier, bounds_error = False, fill_value = None)
             
-            self.BMF = np.array([self.VRdn_dR(z, self.Rs_BMF) for z in self.zlist])
-            self.ion_frac = np.nan_to_num([self.analytic_Q(ClassyCosmo, z)[0] for z in self.zlist])
+            self.BMF = self.VRdn_dR(self.zlist, self.Rs_BMF)
+            self.peakRofz = np.array([self.BMF_peak_R(z) for z in self.zlist])
+            self.peakRofz_int = interp1d(self.zlist, self.peakRofz, bounds_error = False, fill_value = None)
+
+            self.ion_frac = np.nan_to_num(self.analytic_Q(CosmoParams, ClassyCosmo, self.zlist))
             self.ion_frac[self.barrier[:, -1]<=0] = 1
 
             if np.allclose(ion_frac_prev, self.ion_frac, rtol=1e-1, atol=1e-2):
@@ -392,6 +457,51 @@ def nion_normR_int(self, z, R):
     def nion_normz_int(self, z, R):
         return self.interpz(z, R, self.nion_norm_int)
 
+    def interp_zR(self, z, R, func):
+        """
+        Evaluate a RegularGridInterpolator defined on (z, logR).
+    
+        Accepts scalar, 1D, 2D, or ND z and R.
+        z and R are broadcast against each other.
+    
+        Examples
+        --------
+        scalar z, vector R:
+            out.shape == R.shape
+    
+        vector z, scalar R:
+            out.shape == z.shape
+    
+        z[:, None], R[None, :]:
+            out.shape == (nz, nR)
+        """
+        z = np.asarray(z, dtype=float)
+        R = np.asarray(R, dtype=float)
+    
+        z_b, R_b = np.broadcast_arrays(z, R)
+    
+        points = np.column_stack([
+            z_b.ravel(),
+            np.log(R_b).ravel()
+        ])
+    
+        out = func(points)
+        return out.reshape(z_b.shape)
+
+    def sigma_zR_int(self, z, R):
+        return self.interp_zR(z, R, self.sigma_int)
+    
+    def barrier_zR_int(self, z, R):
+        return self.interp_zR(z, R, self.barrier_int)
+    
+    def gamma_zR_int(self, z, R):
+        return self.interp_zR(z, R, self.gamma_int)
+    
+    def gamma2_zR_int(self, z, R):
+        return self.interp_zR(z, R, self.gamma2_int)
+    
+    def nion_norm_zR_int(self, z, R):
+        return self.interp_zR(z, R, self.nion_norm_int)
 
     def prebarrier_xHII_int_grid(self, d, z, R):
         """