Add working bubble model for reionization

requirements.txt

@@ -4,7 +4,9 @@ scipy
 mcfit
 numexpr
 astropy
-powerbox
+tqdm
+matplotlib
 pyfftw
+powerbox
 sphinx
-myst_parser
+myst_parser

setup.py

@@ -20,5 +20,8 @@
            "classy",
            "numexpr",
            "astropy",
+           "powerbox",
+           "pyfftw",
+           "tqdm"
        ],
 )

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(Coeffs, HMFintclass, CosmoParams, AstroParams, ClassyCosmo)
+
 
 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())
+
 
 
 

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)
 

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

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"""
@@ -156,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)
+    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
+

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
 

zeus21/__init__.py

@@ -5,7 +5,8 @@
 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
 warnings.filterwarnings("ignore", category=UserWarning) #to silence unnecessary warning in mcfit

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))

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) 
@@ -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)
 
@@ -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 = 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 = 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)

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)
@@ -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.0,
 
                     Nalpha_lyA_II = 9690,
                     Nalpha_lyA_III = 17900,
@@ -334,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 = 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
+            if clumping != 3.0:
+                self._clumping = clumping