diff --git a/requirements.txt b/requirements.txt index 54f6d49..73ff78f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,9 @@ scipy mcfit numexpr astropy -powerbox +tqdm +matplotlib pyfftw +powerbox sphinx -myst_parser \ No newline at end of file +myst_parser 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/tests/test_astrophysics.py b/tests/test_astrophysics.py index 99e1065..4fbc238 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(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()) + 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/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/tests/test_maps.py b/tests/test_maps.py index 4b37381..9ee1af0 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""" @@ -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) \ 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/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/__init__.py b/zeus21/__init__.py index 656a475..97bfdd4 100644 --- a/zeus21/__init__.py +++ b/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 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 c2500b7..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) @@ -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) diff --git a/zeus21/inputs.py b/zeus21/inputs.py index 4ceadc7..788cbbf 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) @@ -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 diff --git a/zeus21/maps.py b/zeus21/maps.py index e8bcd4c..38ccf56 100644 --- a/zeus21/maps.py +++ b/zeus21/maps.py @@ -2,18 +2,26 @@ 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 - March 2026 """ 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 pyfftw import empty_aligned as empty +from scipy.interpolate import InterpolatedUnivariateSpline as spline +from tqdm import trange +import time +from dataclasses import dataclass, field as _field class CoevalMaps: @@ -41,13 +49,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 ) @@ -59,12 +68,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 ) @@ -74,14 +84,15 @@ 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 + 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 +119,706 @@ def __init__(self, T21_coefficients, Power_Spectrum, z, Lbox=600, Nbox=200, KIND print('ERROR, KIND not implemented yet!') +class reionization_maps: + """ + 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. + + 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. + 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. + 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. + 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. + 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. + LOGNORMAL_DENSITY: bool + 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 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: 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 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 + ---------- + 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 + 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 is True. + """ + + def __init__(self, CosmoParams, ClassyCosmo, CorrFClass, CoeffStructure, BMF, input_z, + 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=False, lowres_massweighting=1, COMPUTE_PARTIAL_IONIZATIONS=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() + + ### boxes parameters + self.input_z = input_z + self.ncells = ncells + self.boxlength = input_boxlength + self.dx = self.boxlength/self.ncells + + # 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 = 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: + if self.COMPUTE_MASSWEIGHTED or self.COMPUTE_PARTIAL_AND_MASSWEIGHTED: + self.COMPUTE_DENSITY_AT_ALLZ = True + + ### selecting redshifts and radii from available redshifts + # redshifts + 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] + 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() + 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 + 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 + if self.barrier is None: + 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 + + self._has_mw = False + self.lowres_massweighting = lowres_massweighting + if self.COMPUTE_MASSWEIGHTED: + self.compute_massweighted(CosmoParams, 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) + + 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: ") + + + 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)) + + #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().astype(np.float32, copy=False) + if self.PRINT_TIMER: + z21_utilities.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: + z21_utilities.print_timer(start_time, text_before=" done in ") + + self._has_density = True + + 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: + 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() + 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)) + + 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, ion_frac + + def ionize(self,CosmoParams, CoeffStructure, curr_z_idx): + zlist = self.z #CoeffStructure.zintegral + Rs = self.r + + 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 + + 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("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 ") + + 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 partially ionized field...") + + iterator = trange(len(self.z)) if self.PRINT_TIMER else range(len(self.z)) + 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)) #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) + + if self.PRINT_TIMER: + print("Computing partial ionized fraction...") + + self.ion_frac_partial = np.average(self.ion_field_partial_allz, axis=(1, 2, 3)) -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...' + if self.PRINT_TIMER: + z21_utilities.print_timer(start_time, text_before=" done in ") - 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) + 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 partially ionized field...") + + 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 ") + + 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() + 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): + 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 + + 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.zreiontvalues[i]) / self.ncells**3 + return 1-neutfrac, tvalues + + + + +##### Wrapper for maps generation +# dataclasses for inputing arguments to astro_variations +# 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 + """ + 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(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. + """ + 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(kw_only=True) +class ReioMapsConfig: + """ + 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 + 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 + input_density: np.ndarray | None = None + input_density_smoothed_allr: np.ndarray | None = None + input_density_allz: np.ndarray | None = None + +# 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. + 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 + ---------- + 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 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_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. + + Attributes + ---------- + 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. + """ + # 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 = {} + __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 = {} + __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)) + 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, + **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, + **vars(self.ReioMaps_config)) + + return bmf, reio_maps + + + def save_quantities(self,bmf,reio_maps): + for q in self.BMF_quantities_to_save: + self.BMF_quantities[q].append(getattr(bmf, q)) + for q in self.ReioMaps_quantities_to_save: + 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.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.ReioMaps_config.input_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) - return realmap \ No newline at end of file diff --git a/zeus21/reionization.py b/zeus21/reionization.py new file mode 100644 index 0000000..2accc9e --- /dev/null +++ b/zeus21/reionization.py @@ -0,0 +1,536 @@ +""" + +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 - October 2025 + +""" + +from . import z21_utilities +from . import cosmology +from . import constants +import numpy as np +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: + """ + Computes the bubble mass function (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 + 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._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 = 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 = 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) + + 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 = None) + + 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.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))) + 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 = 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 = 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 + + 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 = 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(CosmoParams, ClassyCosmo, self.zlist)) #ion_frac by analytic integral of BMF + + self.ion_frac[self.barrier[:, -1]<=0] = 1 + + if FLAG_converge: + self.converge_BMF(CosmoParams, ClassyCosmo, self.ion_frac, max_iter=max_iter) + #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, nR) + nrec_values = self.nrec(CosmoParams, ion_frac, z)[:, :, None] #Shape (nd, nz) * (1, 1, nR) + + prebarrier_xHII = nion_values / (1 + nrec_values) + + return prebarrier_xHII + + def compute_barrier(self, CosmoParams, ion_frac, z, R): + """ + 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 + ---------- + CosmoParams: zeus21.Cosmo_Parameters class + Stores cosmology. + 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(z), len(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)): + #Loop over redshift indices + for iz in range(len(self.zlist)): + y_values = total_values[:, iz, ir] #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 = 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[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]))[:, 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 + + #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, CosmoParams, ion_frac, z, d_array=None): + """ + 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 + 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. + """ + 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] + 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 + integrand_base = denom * ion_frac_rev + 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 back to increasing z order + return nrecs + + def niondot_delta_r(self, CosmoParams, z, R, d_array=None): + """ + Compute niondot over an array of overdensities d_array for a given R. + + Parameters + ---------- + CosmoParams: zeus21.Cosmo_Parameters class + Stores cosmology. + d_array: 1D np.array + A list of sample overdensity values to evaluate niondot over. + 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. + """ + + 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[:, None, None] + + d_array = d_array * CosmoParams.growthint(z) / CosmoParams.growthint(z1d[0]) + + 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 * d_array + gamma2 * d_array**2) + niondot = (self.niondot_avg_int(z) / nion_norm) * exp_term + + return niondot + + 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. + + Parameters + ---------- + CosmoParams: zeus21.Cosmo_Parameters class + Stores cosmology. + d_array: 1D np.array + A list of sample overdensity values to evaluate niondot over. + R: float + Radius value (cMpc) + + Output + ---------- + nion: 2D np.array + 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 + + if d_array is None: + d_array = self.ds_array[:, None, None] + + #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, z, R, d_array) + + 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 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) + + niondot_avgs = self.niondot_avg_int(z_arr) + integrand = dtdz * niondot_avgs * exp + + return np.trapezoid(integrand, x = z_arr, axis = 0) + + #computing linear barrier + def B_1(self, z): + 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 = 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): + B0 = self.B_0(z) + B1 = self.B_1(z) + return B0[:, None] + B1[:, None]*sig**2 + + #computing other terms in the BMF + def dlogsigma_dlogR(self, z, R, sig): + return np.gradient(np.log(sig), np.log(R), axis=1) + + def VRdn_dR(self, z, R): + 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[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))] + + return np.clip(peak_R, min_bubble, max_bubble) #peak can't be outside the allowed bounds + + 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[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)) + + 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: + 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 = 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): + if self.PRINT_SUCCESS: + 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.") + + + #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): + return self.interpR(z, R, self.barrier_int) + def barrierz_int(self, z, R): + 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) + + 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): + """ + 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 diff --git a/zeus21/sfrd.py b/zeus21/sfrd.py index 62a635b..b618e65 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; must be an odd integer deltatab_norm = np.linspace(-Nsigmad,Nsigmad,Nds) #initialize Xrays @@ -127,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 @@ -142,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: @@ -168,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 @@ -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 - SFRD_II_dR = np.trapz(integrand_II, HMF_interpolator.logtabMh, axis = 2) + # Compute SFRD and niondot quantities + 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: @@ -235,20 +238,53 @@ 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) + 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) - #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 #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: - 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(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) ##################################################################################################### ### STEP 3: Computing lambdas in velocity anisotropies @@ -311,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 @@ -435,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 @@ -548,10 +584,18 @@ 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 + #ADD AFTER CHECKING W JULIAN + #collision coefficient fh(1-x_e) + #self.xc_HH = (1.0 - self.xe_avg) * cosmology.n_H(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_H(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 + self.coeff_Ja_xa = self._coeff_Ja_xa_0 * Salpha_exp(self.zintegral, self.Tk_avg, self.xe_avg) self.xa_avg = self.coeff_Ja_xa * self.Jalpha_avg self.invTcol_avg = 1.0 / self.Tk_avg - self._invTs_avg = (1.0/self.T_CMB+self.xa_avg*self.invTcol_avg)/(1+self.xa_avg) + self._invTs_avg = (1.0/self.T_CMB+(self.xa_avg)*self.invTcol_avg)/(1+self.xa_avg) + #AFTER CHECKING ALSO UNCOMMENT THIS: + #self._invTs_avg = (1.0/self.T_CMB+(self.xa_avg+self.xc_avg)*self.invTcol_avg)/(1+self.xa_avg+self.xc_avg) if(User_Parameters.FLAG_WF_ITERATIVE==True): #iteratively find Tcolor and Ts. Could initialize one to zero, but this should converge faster ### iteration routine to find Tcolor and Ts _invTs_tryfirst = 1.0/self.T_CMB @@ -585,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 @@ -607,7 +651,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 @@ -635,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 @@ -643,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) @@ -725,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): @@ -859,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): 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 diff --git a/zeus21/z21_utilities.py b/zeus21/z21_utilities.py new file mode 100644 index 0000000..ccbef91 --- /dev/null +++ b/zeus21/z21_utilities.py @@ -0,0 +1,58 @@ +""" +Helper functions to be used across zeus21 + +Authors: Yonatan Sklansky, Emilie Thelie +UT Austin - February 2025 + +""" + +import numpy as np +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' + '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_before="", text_after=""): + elapsed_time = time.time() - start_time + mins = int(elapsed_time//60) + secs = int(elapsed_time - mins*60) + print(f"{text_before}{mins}min {secs}s{text_after}") + +def v2r(v): + return (3/4/np.pi * v)**(1/3) + +def r2v(r): + 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