diff --git a/.gitignore b/.gitignore index 5c2a373..460612c 100644 --- a/.gitignore +++ b/.gitignore @@ -132,3 +132,5 @@ dmypy.json # Pyre type checker .pyre/ +.worktrees/ +.worktrees/ diff --git a/README.md b/README.md index 369c0c6..c2821a1 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,27 @@ python setup.py install --user (modifying the Makefile to your `gcc` as needed) +### Multiprocessing + +Zeus21 now includes built-in process-level caching for CLASS cosmology objects. When running many simulations in parallel, each worker process will only initialize CLASS once and reuse the cached object for subsequent calls. This provides a significant speedup for large-scale inference or Monte-Carlo sampling. + +Because CLASS uses C-global state, **always use `spawn` (not `fork`)** when creating a multiprocessing pool: + +```python +import multiprocessing as mp +import zeus21 + +ctx = mp.get_context("spawn") +with ctx.Pool(processes=4) as pool: + pool.map(my_zeus21_worker, params_list) +``` + +See `examples/benchmark_multiprocess.py` for a runnable before/after benchmark. + +### NumPy 2.0 Compatibility + +Zeus21 is now compatible with NumPy 2.0+. All deprecated `np.trapz` calls have been replaced with `np.trapezoid`. + ## Citation If you find this code useful please cite: diff --git a/examples/benchmark_multiprocess.py b/examples/benchmark_multiprocess.py new file mode 100644 index 0000000..682e408 --- /dev/null +++ b/examples/benchmark_multiprocess.py @@ -0,0 +1,80 @@ +""" +Benchmark script demonstrating Zeus21 multiprocess performance with process-level caching. + +This script compares the time to run multiple simulations with and without +the built-in CLASS caching in `zeus21.cosmology.runclass()`. + +Usage: + python examples/benchmark_multiprocess.py + +Requirements: + - zeus21 installed (including CLASS) + - numpy +""" + +import time +import multiprocessing as mp +import numpy as np +import zeus21 + + +def _sim_worker_cached(_): + """Worker that reuses the per-process CLASS cache.""" + user_params = zeus21.User_Parameters() + cosmo_input = zeus21.Cosmo_Parameters_Input() + classy_cosmo = zeus21.runclass(cosmo_input) + cosmo_params = zeus21.Cosmo_Parameters(user_params, cosmo_input, classy_cosmo) + hmf_interp = zeus21.HMF_interpolator(user_params, cosmo_params, classy_cosmo) + astro_params = zeus21.Astro_Parameters(user_params, cosmo_params) + coeffs = zeus21.get_T21_coefficients( + user_params, cosmo_params, classy_cosmo, astro_params, hmf_interp, zmin=10.0 + ) + return float(coeffs.T21avg.sum()) + + +def main(): + N_SIMS = 8 + N_WORKERS = 4 + + print("=" * 60) + print("Zeus21 Multiprocess Caching Benchmark") + print("=" * 60) + + # Warmup: ensure CLASS is compiled/loaded in the main process + print("\n[Warmup] Running first simulation...") + _sim_worker_cached(None) + print("Done.") + + # Single-process benchmark + print("\n[Single-process benchmark]") + t0 = time.time() + for _ in range(N_SIMS): + _sim_worker_cached(None) + t_single = time.time() - t0 + print(f" {N_SIMS} simulations in {t_single:.1f}s ({t_single / N_SIMS:.2f}s/sim)") + + # Multiprocess benchmark with spawn (safe for CLASS) + print("\n[Multiprocess benchmark]") + print(f" Using spawn context with {N_WORKERS} workers for {N_SIMS} simulations") + ctx = mp.get_context("spawn") + t0 = time.time() + with ctx.Pool(processes=N_WORKERS) as pool: + _ = pool.map(_sim_worker_cached, range(N_SIMS)) + t_multi = time.time() - t0 + print(f" {N_SIMS} simulations in {t_multi:.1f}s ({t_multi / N_SIMS:.2f}s/sim)") + + print("\n" + "=" * 60) + print("Summary") + print("=" * 60) + print(f"Single-process: {t_single:.1f}s total") + print(f"Multiprocess : {t_multi:.1f}s total") + if t_multi > 0: + print(f"Speedup : {t_single / t_multi:.1f}x") + print("\nNote: The first simulation in each new process incurs the") + print("CLASS initialization cost (~few seconds). Subsequent calls in") + print("the same process reuse the cached CLASS object automatically.") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/tests/test_inputs.py b/tests/test_inputs.py index 4a3102d..606f143 100644 --- a/tests/test_inputs.py +++ b/tests/test_inputs.py @@ -96,23 +96,23 @@ def test_inputs(): #test Pop II Xray SED Energylisttest = np.logspace(2,np.log10(AstroParams.Emax_xray_norm),100) SEDXtab_test = AstroParams.SED_XRAY(Energylisttest, 2) #same in both models - normalization_XraySED = np.trapz(Energylisttest * SEDXtab_test,Energylisttest) + normalization_XraySED = np.trapezoid(Energylisttest * SEDXtab_test,Energylisttest) assert( normalization_XraySED == pytest.approx(1.0, 0.05) ) #5% is enough here #test Pop III Xray SED SEDXtab_test = AstroParams.SED_XRAY(Energylisttest, 3) #same in both models - normalization_XraySED = np.trapz(Energylisttest * SEDXtab_test,Energylisttest) + normalization_XraySED = np.trapezoid(Energylisttest * SEDXtab_test,Energylisttest) assert( normalization_XraySED == pytest.approx(1.0, 0.05) ) #5% is enough here #test Pop II LyA SED nulisttest = np.linspace(zeus21.constants.freqLyA, zeus21.constants.freqLyCont, 100) SEDLtab_test = AstroParams.SED_LyA(nulisttest, 2) #same in both models - normalization_LyASED = np.trapz(SEDLtab_test,nulisttest) + normalization_LyASED = np.trapezoid(SEDLtab_test,nulisttest) assert( normalization_LyASED == pytest.approx(1.0, 0.05) ) #5% is enough here #test Pop III LyA SED nulisttest = np.linspace(zeus21.constants.freqLyA, zeus21.constants.freqLyCont, 100) SEDLtab_test = AstroParams.SED_LyA(nulisttest, 3) #same in both models - normalization_LyASED = np.trapz(SEDLtab_test,nulisttest) + normalization_LyASED = np.trapezoid(SEDLtab_test,nulisttest) assert( normalization_LyASED == pytest.approx(1.0, 0.05) ) #5% is enough here diff --git a/zeus21/UVLFs.py b/zeus21/UVLFs.py index b654c53..ec08a94 100644 --- a/zeus21/UVLFs.py +++ b/zeus21/UVLFs.py @@ -82,7 +82,7 @@ def UVLF_binned(Astro_Parameters,Cosmo_Parameters,HMF_interpolator, zcenter, zwi xlo = np.subtract.outer(MUVcutlo, currMUV )/(np.sqrt(2) * sigmaUV) weights = (erf(xhi) - erf(xlo)).T/(2.0 * MUVwidths) - UVLF_filtered = np.trapz(weights.T * HMFcurr, HMF_interpolator.Mhtab, axis=-1) + UVLF_filtered = np.trapezoid(weights.T * HMFcurr, HMF_interpolator.Mhtab, axis=-1) if(Astro_Parameters.USE_POPIII==False): return UVLF_filtered @@ -98,7 +98,7 @@ def UVLF_binned(Astro_Parameters,Cosmo_Parameters,HMF_interpolator, zcenter, zwi xlo = np.subtract.outer(MUVcutlo, MUVbarlist_III)/(np.sqrt(2) * sigmaUV) weights = (erf(xhi) - erf(xlo)).T/(2.0 * MUVwidths) - UVLF_filtered_III = np.trapz(weights.T * HMFcurr, HMF_interpolator.Mhtab, axis=-1) + UVLF_filtered_III = np.trapezoid(weights.T * HMFcurr, HMF_interpolator.Mhtab, axis=-1) return UVLF_filtered, UVLF_filtered_III diff --git a/zeus21/correlations.py b/zeus21/correlations.py index 8e3287a..7653ecf 100644 --- a/zeus21/correlations.py +++ b/zeus21/correlations.py @@ -374,7 +374,7 @@ def __init__(self, User_Parameters, Cosmo_Parameters, Astro_Parameters, ClassCos #the z>zmax part of the integral we do aside. Assume Tk=Tadiabatic from CLASS. _zlisthighz_ = np.linspace(T21_coefficients.zintegral[-1], 99., 100) #beyond z=100 need to explictly tell CLASS to save growth _dgrowthhighz_ = cosmology.dgrowth_dz(Cosmo_Parameters, _zlisthighz_) - _hizintegral = np.trapz(cosmology.Tadiabatic(Cosmo_Parameters,_zlisthighz_) + _hizintegral = np.trapezoid(cosmology.Tadiabatic(Cosmo_Parameters,_zlisthighz_) /(1+_zlisthighz_)**2 * _dgrowthhighz_, _zlisthighz_) self._betaTad_ = -2./3. * _factor_adi_/self._lingrowthd * (np.cumsum(_integrand_adi[::-1])[::-1] + _hizintegral) #units of Tk_avg. Internal sum goes from high to low z (backwards), minus sign accounts for it properly so it's positive. @@ -1124,7 +1124,7 @@ def get_Pk_from_xi(self, rsinput, xiinput): # # Probdtab = np.exp(-dtab**2/sigmaRRsq/2.0) # -# norm = np.trapz(NionEPS * Probdtab, dtab) +# norm = np.trapezoid(NionEPS * Probdtab, dtab) # NionEPS/=norm # # bindex = min(range(len(NionEPS)), key=lambda i: abs(NionEPS[i]-_invQbar)) diff --git a/zeus21/cosmology.py b/zeus21/cosmology.py index c2500b7..bee6c2e 100644 --- a/zeus21/cosmology.py +++ b/zeus21/cosmology.py @@ -11,6 +11,7 @@ """ import numpy as np +import os from classy import Class from scipy.interpolate import RegularGridInterpolator from scipy.interpolate import interp1d @@ -21,6 +22,10 @@ from .inputs import Cosmo_Parameters, Cosmo_Parameters_Input from .correlations import Correlations +# Process-level cache for CLASS cosmology objects. Keyed by cosmological parameters +# and os.getpid() so that each spawned worker only initializes CLASS once. +_COSMO_CACHE = {} + def cosmo_wrapper(User_Parameters, Cosmo_Parameters_Input): """ Wrapper function for all the cosmology. It takes Cosmo_Parameters_Input and returns: @@ -41,6 +46,17 @@ def cosmo_wrapper(User_Parameters, Cosmo_Parameters_Input): def runclass(CosmologyIn): "Set up CLASS cosmology. Takes CosmologyIn class input and returns CLASS Cosmology object" + cache_key = ( + CosmologyIn.omegab, CosmologyIn.omegac, CosmologyIn.h_fid, + CosmologyIn.As, CosmologyIn.ns, CosmologyIn.tau_fid, + CosmologyIn.kmax_CLASS, CosmologyIn.zmax_CLASS, + CosmologyIn.zmin_CLASS, CosmologyIn.Flag_emulate_21cmfast, + CosmologyIn.USE_RELATIVE_VELOCITIES, CosmologyIn.HMF_CHOICE, + os.getpid(), + ) + if cache_key in _COSMO_CACHE: + return _COSMO_CACHE[cache_key] + ClassCosmo = Class() ClassCosmo.set({'omega_b': CosmologyIn.omegab,'omega_cdm': CosmologyIn.omegac, 'h': CosmologyIn.h_fid,'A_s': CosmologyIn.As,'n_s': CosmologyIn.ns,'tau_reio': CosmologyIn.tau_fid}) @@ -75,13 +91,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 +115,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) @@ -112,7 +128,8 @@ def runclass(CosmologyIn): else: ClassCosmo.pars['v_avg'] = 0.0 ClassCosmo.pars['sigma_vcb'] = 1.0 #Avoids excess computation, but doesn't matter what value we set it to because the flag in inputs.py sets all feedback parameters to zero - + + _COSMO_CACHE[cache_key] = ClassCosmo return ClassCosmo def Hub(Cosmo_Parameters, z): diff --git a/zeus21/sfrd.py b/zeus21/sfrd.py index b36c4f9..5601a20 100644 --- a/zeus21/sfrd.py +++ b/zeus21/sfrd.py @@ -130,14 +130,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 @@ -145,7 +145,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: @@ -171,8 +171,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 @@ -229,7 +229,7 @@ def __init__(self, User_Parameters, Cosmo_Parameters, ClassCosmo, Astro_Paramete ######## # Compute SFRD quantities - SFRD_II_dR = np.trapz(integrand_II, HMF_interpolator.logtabMh, axis = 2) + SFRD_II_dR = np.trapezoid(integrand_II, HMF_interpolator.logtabMh, axis = 2) ### if Astro_Parameters.USE_POPIII == True: @@ -238,7 +238,7 @@ 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) else: SFRD_III_dR = np.zeros_like(SFRD_II_dR) @@ -314,7 +314,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 @@ -438,7 +438,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 @@ -588,8 +588,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 @@ -638,7 +638,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 @@ -646,7 +646,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) @@ -728,7 +728,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): @@ -862,7 +862,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