From c29577155f1f069c2dc65e151af24ef032610fef Mon Sep 17 00:00:00 2001 From: Pierre Burger Date: Fri, 20 Sep 2024 08:40:04 -0400 Subject: [PATCH 01/13] Add CSMF to abacusutilities Add line of code that are able to compute CSMF HOD --- abacusnbody/hod/GRAND_HOD.py | 1755 +++++++++++++++------------------ abacusnbody/hod/abacus_hod.py | 46 +- 2 files changed, 854 insertions(+), 947 deletions(-) diff --git a/abacusnbody/hod/GRAND_HOD.py b/abacusnbody/hod/GRAND_HOD.py index 739790dd..4ab7222f 100644 --- a/abacusnbody/hod/GRAND_HOD.py +++ b/abacusnbody/hod/GRAND_HOD.py @@ -10,7 +10,7 @@ from astropy.table import Table from numba import njit, types from numba.typed import Dict - +import scipy.special as sc # import yaml # config = yaml.safe_load(open('config/abacus_hod.yaml')) # numba.set_num_threads(16) @@ -19,18 +19,211 @@ G = 4.302e-6 # in kpc/Msol (km.s)^2 +@njit(fastmath=True) +def n_cen_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): + """ + Standard Cacciato et al. (2008) centrals HOD parametrization for CSMF + """ + M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) + x_low = np.log10(Mstar_low/M_c_value)/(1.41421356*sigma_c) + x_up = np.log10(Mstar_up/M_c_value)/(1.41421356*sigma_c) + + return 0.5*(math.erf(x_up) - math.erf(x_low)) + +@njit(fastmath=True) +def CSMF_centrals(M_h, Mstar, M_1, M_0, gamma1, gamma2, sigma_c): + """ + Eq. (34) from Cacciato et al. (2008) + """ + + M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) + + return 1/(1.41421356*np.sqrt(np.pi)*np.log(10)*sigma_c*Mstar)*np.exp(-(np.log10(Mstar)-np.log10(M_c_value))**2/(2*sigma_c**2)) + +@njit(fastmath=True) +def M_c(M_h, M_1, M_0, gamma1, gamma2): + """ + Eq. (37) from Cacciato et al. (2008) + """ + return M_0 * (M_h/M_1)**gamma1/(1+M_h/M_1)**(gamma1-gamma2) + + + +@njit(fastmath=True) +def get_random_cen_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): + + nbins = 1000 + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + + CSMF_cen = CSMF_centrals(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c) + + cdf = [] + cdf_current_value = 0 + for k in range(len(delta_stellar_masses)): + cdf_current_value = cdf_current_value + CSMF_cen[k]*delta_stellar_masses[k] + cdf.append(cdf_current_value) + cdf=np.array(cdf) + + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin_clostest = (np.abs(cdf - random_rv)).argmin() + + return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) + + + + +@njit(fastmath=True) +def get_random_cen_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): + + nbins = 1000 + # stellarmass = np.linspace(Mstar_low,Mstar_up,nbins) + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + + CSMF_cen = CSMF_centrals(M_h, stellarmass, M_1, M_0, gamma1, gamma2, sigma_c) + + delta=stellarmass[1:]-stellarmass[:-1] + + cdf = [] + cdf_current_value = 0 + for i in range(len(delta)): + cdf_current_value = cdf_current_value + CSMF_cen[i]*delta[i] + cdf.append(cdf_current_value) + cdf=np.array(cdf) + + # cdf = np.cumsum(CSMF_sat*delta) + + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin = np.where(cdf>random_rv)[0][0] + + m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) + return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + + +@njit(fastmath=True) +def n_sat_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c, a1, a2, M2, b0, b1, b2, delta1, delta2): + """ + Standard Cacciato et al. (2008) satellite HOD parametrization for CSMF + """ + nbins = 1000 + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + + CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + # delta=stellarmass[1:]-stellarmass[:-1] + # nsat = np.sum(CSMF_sat[:-1]*delta) + + # ncen = n_cen_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c) + nsat = 0 + for i in range(nbins-1): + nsat += (CSMF_sat[i+1]-CSMF_sat[i])*(stellarmass[i+1]-stellarmass[i])/2 + (stellarmass[i+1]-stellarmass[i])*CSMF_sat[i] + + return nsat#*ncen + + +@njit(fastmath=True) +def CSMF_satelites(M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): + """ + Eq. (36) from Cacciato et al. (2008) + """ + M_s_value = M_s(M_h, M_1, M_0, gamma1, gamma2) + alpha_s_value = alpha_s(M_h, a1, a2, M2) + phi_s_value = phi_s(M_h, b0, b1, b2) + + delta = 10 ** (delta1 + delta2 * (np.log10(M_h) - 12)) + + return phi_s_value/M_s_value * (Mstar/M_s_value)**alpha_s_value * np.exp(-delta*(Mstar/M_s_value)**2) + +@njit(fastmath=True) +def M_s(M_h, M_1, M_0, gamma1, gamma2): + """ + Eq. (38) from Cacciato et al. (2008) + """ + return 0.562 * M_c(M_h, M_1, M_0, gamma1, gamma2) + +@njit(fastmath=True) +def alpha_s(M_h, a1, a2, M2): + """ + Eq. (39) from Cacciato et al. (2008) + """ + return -2.0 + a1 * (1-2/np.pi*np.arctan(a2*np.log10(M_h/M2))) + +@njit(fastmath=True) +def phi_s(M_h, b0, b1, b2): + """ + Eq. (40) from Cacciato et al. (2008) + """ + M12 = M_h/1e12 + log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12)**2 + return 10**log_phi_s + + +@njit(fastmath=True) +def get_random_sat_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): + + nbins = 1000 + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + + CSMF_sat = CSMF_satelites(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + + + cdf = [] + cdf_current_value = 0 + for k in range(len(delta_stellar_masses)): + cdf_current_value = cdf_current_value + CSMF_sat[k]*delta_stellar_masses[k] + cdf.append(cdf_current_value) + cdf=np.array(cdf) + + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin_clostest = (np.abs(cdf - random_rv)).argmin() + + return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) + + +@njit(fastmath=True) +def get_random_sat_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): + + nbins = 100 + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + + CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + + delta=stellarmass[1:]-stellarmass[:-1] + + cdf = [] + cdf_current_value = 0 + for i in range(len(delta)): + cdf_current_value = cdf_current_value + CSMF_sat[i]*delta[i] + cdf.append(cdf_current_value) + cdf=np.array(cdf) + + # cdf = np.cumsum(CSMF_sat*delta) + + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin = np.where(cdf>random_rv)[0][0] + + m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) + return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + + @njit(fastmath=True) def n_sat_LRG_modified(M_h, logM_cut, M_cut, M_1, sigma, alpha, kappa): """ Standard Zheng et al. (2005) satellite HOD parametrization for LRGs, modified with n_cent_LRG """ - if M_h - kappa * M_cut < 0: + if M_h - kappa*M_cut < 0: return 0 - return ( - ((M_h - kappa * M_cut) / M_1) ** alpha - * 0.5 - * math.erfc((logM_cut - np.log10(M_h)) / (1.41421356 * sigma)) - ) + return ((M_h - kappa*M_cut)/M_1)**alpha*0.5*math.erfc((logM_cut - np.log10(M_h))/(1.41421356*sigma)) @njit(fastmath=True) @@ -38,44 +231,36 @@ def n_cen_LRG(M_h, logM_cut, sigma): """ Standard Zheng et al. (2005) central HOD parametrization for LRGs. """ - return 0.5 * math.erfc((logM_cut - np.log10(M_h)) / (1.41421356 * sigma)) - + return 0.5*math.erfc((logM_cut - np.log10(M_h))/(1.41421356*sigma)) @njit(fastmath=True) -def N_sat_generic(M_h, M_cut, kappa, M_1, alpha, A_s=1.0): +def N_sat_generic(M_h, M_cut, kappa, M_1, alpha, A_s=1.): """ Standard Zheng et al. (2005) satellite HOD parametrization for all tracers with an optional amplitude parameter, A_s. """ - if M_h - kappa * M_cut < 0: + if M_h - kappa*M_cut < 0: return 0 - return A_s * ((M_h - kappa * M_cut) / M_1) ** alpha - + return A_s*((M_h-kappa*M_cut)/M_1)**alpha @njit(fastmath=True) -def N_sat_elg(M_h, M_cut, kappa, M_1, alpha, A_s=1.0, alpha1=0.0, beta=0.0): +def N_sat_elg(M_h, M_cut, kappa, M_1, alpha, A_s=1., alpha1 = 0., beta = 0.): """ Standard power law modulated by an exponential fall off at small M """ # return (M_h/M_1)**alpha/(1+np.exp(-A_s*(np.log10(M_h)-np.log10(kappa*M_cut)))) + beta*(M_h/M_1)**(-alpha1)/100 - if M_h - kappa * M_cut < 0: + if M_h - kappa*M_cut < 0: return 0 - return ( - A_s * ((M_h - kappa * M_cut) / M_1) ** alpha - ) # + beta*(M_h/M_1)**(-alpha1)/100 - + return A_s*((M_h-kappa*M_cut)/M_1)**alpha # + beta*(M_h/M_1)**(-alpha1)/100 @njit(fastmath=True) -def N_cen_ELG_v1(M_h, p_max, Q, logM_cut, sigma, gamma, Anorm=1): +def N_cen_ELG_v1(M_h, p_max, Q, logM_cut, sigma, gamma, Anorm = 1): """ HOD function for ELG centrals taken from arXiv:1910.05095. """ logM_h = np.log10(M_h) phi = phi_fun(logM_h, logM_cut, sigma) Phi = Phi_fun(logM_h, logM_cut, sigma, gamma) - return ( - 2.0 * (p_max - 1.0 / Q) * phi * Phi / Anorm - ) # + 0.5/Q*(1 + math.erf((logM_h-logM_cut-0.8)*3)) - + return 2.*(p_max-1./Q)*phi*Phi/Anorm # + 0.5/Q*(1 + math.erf((logM_h-logM_cut-0.8)*3)) @njit(fastmath=True) def N_cen_ELG_v2(M_h, p_max, logM_cut, sigma, gamma): @@ -84,17 +269,16 @@ def N_cen_ELG_v2(M_h, p_max, logM_cut, sigma, gamma): """ logM_h = np.log10(M_h) if logM_h <= logM_cut: - return p_max * Gaussian_fun(logM_h, logM_cut, sigma) + return p_max*Gaussian_fun(logM_h, logM_cut, sigma) else: - return p_max * (M_h / 10**logM_cut) ** gamma / (2.5066283 * sigma) - + return p_max*(M_h/10**logM_cut)**gamma/(2.5066283*sigma) @njit(fastmath=True) def N_cen_QSO(M_h, logM_cut, sigma): """ HOD function (Zheng et al. (2005) with p_max) for QSO centrals taken from arXiv:2007.09012. """ - return 0.5 * (1 + math.erf((np.log10(M_h) - logM_cut) / 1.41421356 / sigma)) + return 0.5*(1 + math.erf((np.log10(M_h)-logM_cut)/1.41421356/sigma)) @njit(fastmath=True) @@ -105,60 +289,37 @@ def phi_fun(logM_h, logM_cut, sigma): phi = Gaussian_fun(logM_h, logM_cut, sigma) return phi - @njit(fastmath=True) def Phi_fun(logM_h, logM_cut, sigma, gamma): """ Aiding function for N_cen_ELG_v1(). """ - x = gamma * (logM_h - logM_cut) / sigma - Phi = 0.5 * (1 + math.erf(x / np.sqrt(2))) + x = gamma*(logM_h-logM_cut)/sigma + Phi = 0.5*(1 + math.erf(x/np.sqrt(2))) return Phi - @njit(fastmath=True) def Gaussian_fun(x, mean, sigma): """ Gaussian function with centered at `mean' with standard deviation `sigma'. """ - return 0.3989422804014327 / sigma * np.exp(-((x - mean) ** 2) / 2 / sigma**2) + return 0.3989422804014327/sigma*np.exp(-(x - mean)**2/2/sigma**2) @njit(fastmath=True) def wrap(x, L): - """Fast scalar mod implementation""" - L2 = L / 2 + '''Fast scalar mod implementation''' + L2 = L/2 if x >= L2: return x - L elif x < -L2: return x + L return x - @njit(parallel=True, fastmath=True) -def gen_cent( - pos, - vel, - mass, - ids, - multis, - randoms, - vdev, - deltac, - fenv, - shear, - LRG_hod_dict, - ELG_hod_dict, - QSO_hod_dict, - rsd, - inv_velz2kms, - lbox, - want_LRG, - want_ELG, - want_QSO, - Nthread, - origin, -): +def gen_cent(pos, vel, mass, ids, multis, randoms, vdev, deltac, fenv, shear, + LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, CSMF_hod_dict, + rsd, inv_velz2kms, lbox, want_LRG, want_ELG, want_QSO, want_CSMF, Nthread, origin): """ Generate central galaxies in place in memory with a two pass numba parallel implementation. """ @@ -166,47 +327,30 @@ def gen_cent( if want_LRG: # parse out the hod parameters logM_cut_L, sigma_L = LRG_hod_dict['logM_cut'], LRG_hod_dict['sigma'] - ic_L, alpha_c_L, Ac_L, Bc_L = ( - LRG_hod_dict['ic'], - LRG_hod_dict['alpha_c'], - LRG_hod_dict['Acent'], - LRG_hod_dict['Bcent'], - ) + ic_L, alpha_c_L, Ac_L, Bc_L = LRG_hod_dict['ic'], LRG_hod_dict['alpha_c'], LRG_hod_dict['Acent'], LRG_hod_dict['Bcent'] if want_ELG: - pmax_E, Q_E, logM_cut_E, sigma_E, gamma_E = ( - ELG_hod_dict['p_max'], - ELG_hod_dict['Q'], - ELG_hod_dict['logM_cut'], - ELG_hod_dict['sigma'], - ELG_hod_dict['gamma'], - ) - alpha_c_E, Ac_E, Bc_E, Cc_E, ic_E = ( - ELG_hod_dict['alpha_c'], - ELG_hod_dict['Acent'], - ELG_hod_dict['Bcent'], - ELG_hod_dict['Ccent'], - ELG_hod_dict['ic'], - ) + pmax_E, Q_E, logM_cut_E, sigma_E, gamma_E = \ + ELG_hod_dict['p_max'], ELG_hod_dict['Q'], ELG_hod_dict['logM_cut'], ELG_hod_dict['sigma'], ELG_hod_dict['gamma'] + alpha_c_E, Ac_E, Bc_E, Cc_E, ic_E = ELG_hod_dict['alpha_c'], ELG_hod_dict['Acent'], ELG_hod_dict['Bcent'], ELG_hod_dict['Ccent'], ELG_hod_dict['ic'] if want_QSO: logM_cut_Q, sigma_Q = QSO_hod_dict['logM_cut'], QSO_hod_dict['sigma'] - alpha_c_Q, Ac_Q, Bc_Q, ic_Q = ( - QSO_hod_dict['alpha_c'], - QSO_hod_dict['Acent'], - QSO_hod_dict['Bcent'], - QSO_hod_dict['ic'], - ) + alpha_c_Q, Ac_Q, Bc_Q, ic_Q = QSO_hod_dict['alpha_c'], QSO_hod_dict['Acent'], QSO_hod_dict['Bcent'], QSO_hod_dict['ic'] + + if want_CSMF: + Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C = CSMF_hod_dict['Mstar_low'], CSMF_hod_dict['Mstar_up'], CSMF_hod_dict['M_1'], CSMF_hod_dict['M_0'], CSMF_hod_dict['gamma_1'], CSMF_hod_dict['gamma_2'], CSMF_hod_dict['sigma_c'] + ic_C, alpha_c_C, Ac_C, Bc_C = CSMF_hod_dict['ic'], CSMF_hod_dict['alpha_c'], CSMF_hod_dict['Acent'], CSMF_hod_dict['Bcent'] + + # print(Ac_C,Bc_C) H = len(mass) - + numba.set_num_threads(Nthread) - Nout = np.zeros((Nthread, 3, 8), dtype=np.int64) - hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype( - np.int64 - ) # starting index of each thread + Nout = np.zeros((Nthread, 4, 8), dtype = np.int64) + hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype(np.int64) # starting index of each thread - keep = np.empty(H, dtype=np.int8) # mask array tracking which halos to keep + keep = np.empty(H, dtype = np.int8) # mask array tracking which halos to keep # figuring out the number of halos kept for each thread for tid in numba.prange(Nthread): @@ -216,174 +360,210 @@ def gen_cent( if want_LRG: # do assembly bias and secondary bias logM_cut_L_temp = logM_cut_L + Ac_L * deltac[i] + Bc_L * fenv[i] - LRG_marker += ( - n_cen_LRG(mass[i], logM_cut_L_temp, sigma_L) * ic_L * multis[i] - ) + LRG_marker += n_cen_LRG(mass[i], logM_cut_L_temp, sigma_L) * ic_L * multis[i] ELG_marker = LRG_marker if want_ELG: - logM_cut_E_temp = ( - logM_cut_E + Ac_E * deltac[i] + Bc_E * fenv[i] + Cc_E * shear[i] - ) - ELG_marker += ( - N_cen_ELG_v1( - mass[i], pmax_E, Q_E, logM_cut_E_temp, sigma_E, gamma_E - ) - * ic_E - * multis[i] - ) + logM_cut_E_temp = logM_cut_E + Ac_E * deltac[i] + Bc_E * fenv[i] + Cc_E * shear[i] + ELG_marker += N_cen_ELG_v1(mass[i], pmax_E, Q_E, logM_cut_E_temp, sigma_E, gamma_E) * ic_E * multis[i] QSO_marker = ELG_marker if want_QSO: logM_cut_Q_temp = logM_cut_Q + Ac_Q * deltac[i] + Bc_Q * fenv[i] - QSO_marker += ( - N_cen_QSO(mass[i], logM_cut_Q_temp, sigma_Q) * ic_Q * multis[i] - ) - + QSO_marker += N_cen_QSO(mass[i], logM_cut_Q_temp, sigma_Q) * ic_Q * multis[i] + CSMF_marker = QSO_marker + if want_CSMF: + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + + ncen = n_cen_CSMF(mass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) + + CSMF_marker += ncen * ic_C * multis[i] + if randoms[i] <= LRG_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 elif randoms[i] <= ELG_marker: - Nout[tid, 1, 0] += 1 # counting + Nout[tid, 1, 0] += 1 # counting keep[i] = 2 elif randoms[i] <= QSO_marker: - Nout[tid, 2, 0] += 1 # counting + Nout[tid, 2, 0] += 1 # counting keep[i] = 3 + elif randoms[i] <= CSMF_marker: + Nout[tid, 3, 0] += 1 # counting + keep[i] = 4 else: keep[i] = 0 - + # compose galaxy array, first create array of galaxy starting indices for the threads - gstart = np.empty((Nthread + 1, 3), dtype=np.int64) + gstart = np.empty((Nthread + 1, 4), dtype = np.int64) gstart[0, :] = 0 gstart[1:, 0] = Nout[:, 0, 0].cumsum() gstart[1:, 1] = Nout[:, 1, 0].cumsum() gstart[1:, 2] = Nout[:, 2, 0].cumsum() + gstart[1:, 3] = Nout[:, 3, 0].cumsum() # galaxy arrays N_lrg = gstart[-1, 0] - lrg_x = np.empty(N_lrg, dtype=mass.dtype) - lrg_y = np.empty(N_lrg, dtype=mass.dtype) - lrg_z = np.empty(N_lrg, dtype=mass.dtype) - lrg_vx = np.empty(N_lrg, dtype=mass.dtype) - lrg_vy = np.empty(N_lrg, dtype=mass.dtype) - lrg_vz = np.empty(N_lrg, dtype=mass.dtype) - lrg_mass = np.empty(N_lrg, dtype=mass.dtype) - lrg_id = np.empty(N_lrg, dtype=ids.dtype) + lrg_x = np.empty(N_lrg, dtype = mass.dtype) + lrg_y = np.empty(N_lrg, dtype = mass.dtype) + lrg_z = np.empty(N_lrg, dtype = mass.dtype) + lrg_vx = np.empty(N_lrg, dtype = mass.dtype) + lrg_vy = np.empty(N_lrg, dtype = mass.dtype) + lrg_vz = np.empty(N_lrg, dtype = mass.dtype) + lrg_mass = np.empty(N_lrg, dtype = mass.dtype) + lrg_id = np.empty(N_lrg, dtype = ids.dtype) # galaxy arrays N_elg = gstart[-1, 1] - elg_x = np.empty(N_elg, dtype=mass.dtype) - elg_y = np.empty(N_elg, dtype=mass.dtype) - elg_z = np.empty(N_elg, dtype=mass.dtype) - elg_vx = np.empty(N_elg, dtype=mass.dtype) - elg_vy = np.empty(N_elg, dtype=mass.dtype) - elg_vz = np.empty(N_elg, dtype=mass.dtype) - elg_mass = np.empty(N_elg, dtype=mass.dtype) - elg_id = np.empty(N_elg, dtype=ids.dtype) + elg_x = np.empty(N_elg, dtype = mass.dtype) + elg_y = np.empty(N_elg, dtype = mass.dtype) + elg_z = np.empty(N_elg, dtype = mass.dtype) + elg_vx = np.empty(N_elg, dtype = mass.dtype) + elg_vy = np.empty(N_elg, dtype = mass.dtype) + elg_vz = np.empty(N_elg, dtype = mass.dtype) + elg_mass = np.empty(N_elg, dtype = mass.dtype) + elg_id = np.empty(N_elg, dtype = ids.dtype) # galaxy arrays N_qso = gstart[-1, 2] - qso_x = np.empty(N_qso, dtype=mass.dtype) - qso_y = np.empty(N_qso, dtype=mass.dtype) - qso_z = np.empty(N_qso, dtype=mass.dtype) - qso_vx = np.empty(N_qso, dtype=mass.dtype) - qso_vy = np.empty(N_qso, dtype=mass.dtype) - qso_vz = np.empty(N_qso, dtype=mass.dtype) - qso_mass = np.empty(N_qso, dtype=mass.dtype) - qso_id = np.empty(N_qso, dtype=ids.dtype) + qso_x = np.empty(N_qso, dtype = mass.dtype) + qso_y = np.empty(N_qso, dtype = mass.dtype) + qso_z = np.empty(N_qso, dtype = mass.dtype) + qso_vx = np.empty(N_qso, dtype = mass.dtype) + qso_vy = np.empty(N_qso, dtype = mass.dtype) + qso_vz = np.empty(N_qso, dtype = mass.dtype) + qso_mass = np.empty(N_qso, dtype = mass.dtype) + qso_id = np.empty(N_qso, dtype = ids.dtype) + + # galaxy arrays + N_CSMF = gstart[-1, 3] + CSMF_x = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_y = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_z = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_id = np.empty(N_CSMF, dtype = ids.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): - j1, j2, j3 = gstart[tid] + j1, j2, j3, j4 = gstart[tid] for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: # loop thru three directions to assign galaxy velocities and positions - lrg_x[j1] = pos[i, 0] - lrg_vx[j1] = vel[i, 0] + alpha_c_L * vdev[i, 0] # velocity bias - lrg_y[j1] = pos[i, 1] - lrg_vy[j1] = vel[i, 1] + alpha_c_L * vdev[i, 1] # velocity bias - lrg_z[j1] = pos[i, 2] - lrg_vz[j1] = vel[i, 2] + alpha_c_L * vdev[i, 2] # velocity bias + lrg_x[j1] = pos[i,0] + lrg_vx[j1] = vel[i,0] + alpha_c_L * vdev[i,0] # velocity bias + lrg_y[j1] = pos[i,1] + lrg_vy[j1] = vel[i,1] + alpha_c_L * vdev[i,1] # velocity bias + lrg_z[j1] = pos[i,2] + lrg_vz[j1] = vel[i,2] + alpha_c_L * vdev[i,2] # velocity bias # rsd only applies to the z direction if rsd and origin is not None: nx = lrg_x[j1] - origin[0] ny = lrg_y[j1] - origin[1] nz = lrg_z[j1] - origin[2] - inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms * ( - lrg_vx[j1] * nx + lrg_vy[j1] * ny + lrg_vz[j1] * nz - ) - lrg_x[j1] = lrg_x[j1] + proj * nx - lrg_y[j1] = lrg_y[j1] + proj * ny - lrg_z[j1] = lrg_z[j1] + proj * nz + proj = inv_velz2kms * (lrg_vx[j1]*nx + lrg_vy[j1]*ny + lrg_vz[j1]*nz) + lrg_x[j1] = lrg_x[j1]+proj*nx + lrg_y[j1] = lrg_y[j1]+proj*ny + lrg_z[j1] = lrg_z[j1]+proj*nz elif rsd: - lrg_z[j1] = wrap(pos[i, 2] + lrg_vz[j1] * inv_velz2kms, lbox) + lrg_z[j1] = wrap(pos[i,2] + lrg_vz[j1] * inv_velz2kms, lbox) lrg_mass[j1] = mass[i] lrg_id[j1] = ids[i] j1 += 1 elif keep[i] == 2: # loop thru three directions to assign galaxy velocities and positions - elg_x[j2] = pos[i, 0] - elg_vx[j2] = vel[i, 0] + alpha_c_E * vdev[i, 0] # velocity bias - elg_y[j2] = pos[i, 1] - elg_vy[j2] = vel[i, 1] + alpha_c_E * vdev[i, 1] # velocity bias - elg_z[j2] = pos[i, 2] - elg_vz[j2] = vel[i, 2] + alpha_c_E * vdev[i, 2] # velocity bias + elg_x[j2] = pos[i,0] + elg_vx[j2] = vel[i,0] + alpha_c_E * vdev[i,0] # velocity bias + elg_y[j2] = pos[i,1] + elg_vy[j2] = vel[i,1] + alpha_c_E * vdev[i,1] # velocity bias + elg_z[j2] = pos[i,2] + elg_vz[j2] = vel[i,2] + alpha_c_E * vdev[i,2] # velocity bias # rsd only applies to the z direction if rsd and origin is not None: nx = elg_x[j2] - origin[0] ny = elg_y[j2] - origin[1] nz = elg_z[j2] - origin[2] - inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms * ( - elg_vx[j2] * nx + elg_vy[j2] * ny + elg_vz[j2] * nz - ) - elg_x[j2] = elg_x[j2] + proj * nx - elg_y[j2] = elg_y[j2] + proj * ny - elg_z[j2] = elg_z[j2] + proj * nz + proj = inv_velz2kms*(elg_vx[j2]*nx+elg_vy[j2]*ny+elg_vz[j2]*nz) + elg_x[j2] = elg_x[j2]+proj*nx + elg_y[j2] = elg_y[j2]+proj*ny + elg_z[j2] = elg_z[j2]+proj*nz elif rsd: - elg_z[j2] = wrap(pos[i, 2] + elg_vz[j2] * inv_velz2kms, lbox) + elg_z[j2] = wrap(pos[i,2] + elg_vz[j2] * inv_velz2kms, lbox) elg_mass[j2] = mass[i] elg_id[j2] = ids[i] j2 += 1 elif keep[i] == 3: # loop thru three directions to assign galaxy velocities and positions - qso_x[j3] = pos[i, 0] - qso_vx[j3] = vel[i, 0] + alpha_c_Q * vdev[i, 0] # velocity bias - qso_y[j3] = pos[i, 1] - qso_vy[j3] = vel[i, 1] + alpha_c_Q * vdev[i, 1] # velocity bias - qso_z[j3] = pos[i, 2] - qso_vz[j3] = vel[i, 2] + alpha_c_Q * vdev[i, 2] # velocity bias + qso_x[j3] = pos[i,0] + qso_vx[j3] = vel[i,0] + alpha_c_Q * vdev[i,0] # velocity bias + qso_y[j3] = pos[i,1] + qso_vy[j3] = vel[i,1] + alpha_c_Q * vdev[i,1] # velocity bias + qso_z[j3] = pos[i,2] + qso_vz[j3] = vel[i,2] + alpha_c_Q * vdev[i,2] # velocity bias # rsd only applies to the z direction if rsd and origin is not None: nx = qso_x[j3] - origin[0] ny = qso_y[j3] - origin[1] nz = qso_z[j3] - origin[2] - inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms * ( - qso_vx[j3] * nx + qso_vy[j3] * ny + qso_vz[j3] * nz - ) - qso_x[j3] = qso_x[j3] + proj * nx - qso_y[j3] = qso_y[j3] + proj * ny - qso_z[j3] = qso_z[j3] + proj * nz + proj = inv_velz2kms*(qso_vx[j3]*nx+qso_vy[j3]*ny+qso_vz[j3]*nz) + qso_x[j3] = qso_x[j3]+proj*nx + qso_y[j3] = qso_y[j3]+proj*ny + qso_z[j3] = qso_z[j3]+proj*nz elif rsd: - qso_z[j3] = wrap(pos[i, 2] + qso_vz[j3] * inv_velz2kms, lbox) + qso_z[j3] = wrap(pos[i,2] + qso_vz[j3] * inv_velz2kms, lbox) qso_mass[j3] = mass[i] qso_id[j3] = ids[i] j3 += 1 + elif keep[i] == 4: + # loop thru three directions to assign galaxy velocities and positions + CSMF_x[j4] = pos[i,0] + CSMF_vx[j4] = vel[i,0] + alpha_c_C * vdev[i,0] # velocity bias + CSMF_y[j4] = pos[i,1] + CSMF_vy[j4] = vel[i,1] + alpha_c_C * vdev[i,1] # velocity bias + CSMF_z[j4] = pos[i,2] + CSMF_vz[j4] = vel[i,2] + alpha_c_C * vdev[i,2] # velocity bias + + # rsd only applies to the z direction + if rsd and origin is not None: + nx = CSMF_x[j4] - origin[0] + ny = CSMF_y[j4] - origin[1] + nz = CSMF_z[j4] - origin[2] + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + nx *= inv_norm + ny *= inv_norm + nz *= inv_norm + proj = inv_velz2kms*(CSMF_vx[j4]*nx+CSMF_vy[j4]*ny+CSMF_vz[j4]*nz) + CSMF_x[j4] = CSMF_x[j4]+proj*nx + CSMF_y[j4] = CSMF_y[j4]+proj*ny + CSMF_z[j4] = CSMF_z[j4]+proj*nz + elif rsd: + CSMF_z[j4] = wrap(pos[i,2] + CSMF_vz[j4] * inv_velz2kms, lbox) + + CSMF_mass[j4] = mass[i] + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + CSMF_stellarmass[j4] = get_random_cen_stellarmass_linearinterpolation(mass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) + CSMF_id[j4] = ids[i] + j4 += 1 # assert j == gstart[tid + 1] - LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) + LRG_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + ELG_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + QSO_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + ID_dict = Dict.empty(key_type = types.unicode_type, value_type = int_array) LRG_dict['x'] = lrg_x LRG_dict['y'] = lrg_y LRG_dict['z'] = lrg_z @@ -410,18 +590,27 @@ def gen_cent( QSO_dict['vz'] = qso_vz QSO_dict['mass'] = qso_mass ID_dict['QSO'] = qso_id - return LRG_dict, ELG_dict, QSO_dict, ID_dict, keep - + + CSMF_dict['x'] = CSMF_x + CSMF_dict['y'] = CSMF_y + CSMF_dict['z'] = CSMF_z + CSMF_dict['vx'] = CSMF_vx + CSMF_dict['vy'] = CSMF_vy + CSMF_dict['vz'] = CSMF_vz + CSMF_dict['mass'] = CSMF_mass + CSMF_dict['stellarmass'] = CSMF_stellarmass + ID_dict['CSMF'] = CSMF_id + return LRG_dict, ELG_dict, QSO_dict, CSMF_dict, ID_dict, keep @njit(parallel=True, fastmath=True) def getPointsOnSphere(nPoints, Nthread, seed=None): - """ + ''' --- Aiding function for NFW computation, generate random points in a sphere - """ + ''' numba.set_num_threads(Nthread) ind = min(Nthread, nPoints) # starting index of each thread - hstart = np.rint(np.linspace(0, nPoints, ind + 1)) + hstart = np.rint(np.linspace(0, nPoints, ind+1)) ur = np.zeros((nPoints, 3), dtype=np.float64) cmin = -1 cmax = +1 @@ -431,38 +620,19 @@ def getPointsOnSphere(nPoints, Nthread, seed=None): np.random.seed(seed[tid]) for i in range(hstart[tid], hstart[tid + 1]): u1, u2 = np.random.uniform(0, 1), np.random.uniform(0, 1) - ra = 0 + u1 * (2 * np.pi - 0) - dec = np.pi - (np.arccos(cmin + u2 * (cmax - cmin))) + ra = 0 + u1*(2*np.pi-0) + dec = np.pi - (np.arccos(cmin+u2*(cmax-cmin))) ur[i, 0] = np.sin(dec) * np.cos(ra) ur[i, 1] = np.sin(dec) * np.sin(ra) ur[i, 2] = np.cos(dec) return ur +@njit(fastmath=True, parallel=True) # parallel=True, +def compute_fast_NFW(NFW_draw, h_id, x_h, y_h, z_h, vx_h, vy_h, vz_h, vrms_h, c, M, Rvir, + rd_pos, num_sat, f_sigv, vel_sat = 'rd_normal', Nthread = 16, + exp_frac=0, exp_scale=1, nfw_rescale=1): -@njit(fastmath=True, parallel=True) # parallel=True, -def compute_fast_NFW( - NFW_draw, - h_id, - x_h, - y_h, - z_h, - vx_h, - vy_h, - vz_h, - vrms_h, - c, - M, - Rvir, - rd_pos, - num_sat, - f_sigv, - vel_sat='rd_normal', - Nthread=16, - exp_frac=0, - exp_scale=1, - nfw_rescale=1, -): """ --- Compute NFW positions and velocities for satelitte galaxies c: r98/r25 @@ -493,57 +663,36 @@ def compute_fast_NFW( for tid in numba.prange(Nthread): for i in range(int(hstart[tid]), int(hstart[tid + 1])): ind = i - # while (NFW_draw[ind] > c[i]): + #while (NFW_draw[ind] > c[i]): # ind = np.random.randint(0, len(NFW_draw)) - # etaVir = NFW_draw[ind]/c[i] # =r/rvir - if np.random.uniform(0, 1) < exp_frac: - tt = np.random.exponential(exp_scale, size=1)[0] - etaVir = tt / c[i] + #etaVir = NFW_draw[ind]/c[i] # =r/rvir + if np.random.uniform(0,1) c[i]: + # ind = np.random.choice(np.where(NFW_draw c[i]): ind = np.random.randint(0, len(NFW_draw)) - etaVir = NFW_draw[ind] / c[i] * nfw_rescale + etaVir = NFW_draw[ind]/c[i]*nfw_rescale p = etaVir * Rvir[i] x_sat[i] = x_h[i] + rd_pos[i, 0] * p y_sat[i] = y_h[i] + rd_pos[i, 1] * p z_sat[i] = z_h[i] + rd_pos[i, 2] * p - if vel_sat == 'rd_normal': - sig = vrms_h[i] * 0.577 * f_sigv - vx_sat[i] = np.random.normal(loc=vx_h[i], scale=sig) - vy_sat[i] = np.random.normal(loc=vy_h[i], scale=sig) - vz_sat[i] = np.random.normal(loc=vz_h[i], scale=sig) - else: - raise ValueError('Wrong vel_sat argument only "rd_normal"') + # if vel_sat == 'rd_normal': + sig = vrms_h[i]*0.577*f_sigv + vx_sat[i] = np.random.normal(loc=vx_h[i], scale=sig) + vy_sat[i] = np.random.normal(loc=vy_h[i], scale=sig) + vz_sat[i] = np.random.normal(loc=vz_h[i], scale=sig) + # else: + # raise ValueError( + # 'Wrong vel_sat argument only "rd_normal"') return h_id, x_sat, y_sat, z_sat, vx_sat, vy_sat, vz_sat, M - -@njit(fastmath=True, parallel=True) -def gen_sats_nfw( - NFW_draw, - hpos, - hvel, - hmass, - hid, - hdeltac, - hfenv, - hshear, - hvrms, - hc, - hrvir, - LRG_hod_dict, - ELG_hod_dict, - QSO_hod_dict, - want_LRG, - want_ELG, - want_QSO, - rsd, - inv_velz2kms, - lbox, - keep_cent, - vel_sat='rd_normal', - Nthread=16, -): +@njit(fastmath = True, parallel=True) +def gen_sats_nfw(NFW_draw, hpos, hvel, hmass, hid, hdeltac, hfenv, hshear, hvrms, hc, hrvir, + LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, CSMF_hod_dict, want_LRG, want_ELG, want_QSO, want_CSMF, + rsd, inv_velz2kms, lbox, keep_cent, vel_sat = 'rd_normal', Nthread = 16): """ Generate satellite galaxies on an NFW profile, with option for an extended profile. See Rocher et al. 2023. @@ -552,245 +701,130 @@ def gen_sats_nfw( """ if want_LRG: - logM_cut_L, logM1_L, sigma_L, alpha_L, kappa_L = ( - LRG_hod_dict['logM_cut'], - LRG_hod_dict['logM1'], - LRG_hod_dict['sigma'], - LRG_hod_dict['alpha'], - LRG_hod_dict['kappa'], - ) - Ac_L, As_L, Bc_L, Bs_L, ic_L = ( - LRG_hod_dict['Acent'], - LRG_hod_dict['Asat'], - LRG_hod_dict['Bcent'], - LRG_hod_dict['Bsat'], - LRG_hod_dict['ic'], - ) + logM_cut_L, logM1_L, sigma_L, alpha_L, kappa_L = \ + LRG_hod_dict['logM_cut'], LRG_hod_dict['logM1'], LRG_hod_dict['sigma'], LRG_hod_dict['alpha'], LRG_hod_dict['kappa'] + Ac_L, As_L, Bc_L, Bs_L, ic_L = \ + LRG_hod_dict['Acent'], LRG_hod_dict['Asat'], LRG_hod_dict['Bcent'], \ + LRG_hod_dict['Bsat'], LRG_hod_dict['ic'] f_sigv_L = LRG_hod_dict['f_sigv'] if want_ELG: - logM_cut_E, kappa_E, logM1_E, alpha_E, A_E = ( - ELG_hod_dict['logM_cut'], - ELG_hod_dict['kappa'], - ELG_hod_dict['logM1'], - ELG_hod_dict['alpha'], - ELG_hod_dict['A_s'], - ) - ( - Ac_E, - As_E, - Bc_E, - Bs_E, - Cc_E, - Cs_E, - ic_E, - ) = ( - ELG_hod_dict['Acent'], - ELG_hod_dict['Asat'], - ELG_hod_dict['Bcent'], - ELG_hod_dict['Bsat'], - ELG_hod_dict['Ccent'], - ELG_hod_dict['Csat'], - ELG_hod_dict['ic'], - ) - logM1_EE, alpha_EE, logM1_EL, alpha_EL = ( - ELG_hod_dict['logM1_EE'], - ELG_hod_dict['alpha_EE'], - ELG_hod_dict['logM1_EL'], - ELG_hod_dict['alpha_EL'], - ) + logM_cut_E, kappa_E, logM1_E, alpha_E, A_E = \ + ELG_hod_dict['logM_cut'], ELG_hod_dict['kappa'], ELG_hod_dict['logM1'], ELG_hod_dict['alpha'], ELG_hod_dict['A_s'] + Ac_E, As_E, Bc_E, Bs_E, Cc_E, Cs_E, ic_E, = \ + ELG_hod_dict['Acent'], ELG_hod_dict['Asat'], ELG_hod_dict['Bcent'], \ + ELG_hod_dict['Bsat'], ELG_hod_dict['Ccent'], ELG_hod_dict['Csat'], ELG_hod_dict['ic'] + logM1_EE, alpha_EE, logM1_EL, alpha_EL = \ + ELG_hod_dict['logM1_EE'], ELG_hod_dict['alpha_EE'], \ + ELG_hod_dict['logM1_EL'], ELG_hod_dict['alpha_EL'] f_sigv_E = ELG_hod_dict['f_sigv'] exp_frac = ELG_hod_dict['exp_frac'] exp_scale = ELG_hod_dict['exp_scale'] nfw_rescale = ELG_hod_dict['nfw_rescale'] if want_QSO: - logM_cut_Q, kappa_Q, logM1_Q, alpha_Q = ( - QSO_hod_dict['logM_cut'], - QSO_hod_dict['kappa'], - QSO_hod_dict['logM1'], - QSO_hod_dict['alpha'], - ) - Ac_Q, As_Q, Bc_Q, Bs_Q, ic_Q = ( - QSO_hod_dict['Acent'], - QSO_hod_dict['Asat'], - QSO_hod_dict['Bcent'], - QSO_hod_dict['Bsat'], - QSO_hod_dict['ic'], - ) + logM_cut_Q, kappa_Q, logM1_Q, alpha_Q = \ + QSO_hod_dict['logM_cut'], QSO_hod_dict['kappa'], QSO_hod_dict['logM1'], QSO_hod_dict['alpha'] + Ac_Q, As_Q, Bc_Q, Bs_Q, ic_Q = \ + QSO_hod_dict['Acent'], QSO_hod_dict['Asat'], QSO_hod_dict['Bcent'], \ + QSO_hod_dict['Bsat'], QSO_hod_dict['ic'] f_sigv_Q = QSO_hod_dict['f_sigv'] + + if want_CSMF: + Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = \ + CSMF_hod_dict['Mstar_low'], CSMF_hod_dict['Mstar_up'], CSMF_hod_dict['M_1'], CSMF_hod_dict['M_0'], CSMF_hod_dict['gamma_1'], \ + CSMF_hod_dict['gamma_2'], CSMF_hod_dict['sigma_c'], CSMF_hod_dict['a_1'], CSMF_hod_dict['a_2'], CSMF_hod_dict['M_2'], \ + CSMF_hod_dict['b_0'], CSMF_hod_dict['b_1'], CSMF_hod_dict['b_2'], CSMF_hod_dict['delta_1'], CSMF_hod_dict['delta_2'] + Ac_C, As_C, Bc_C, Bs_C, ic_C = \ + CSMF_hod_dict['Acent'], CSMF_hod_dict['Asat'], CSMF_hod_dict['Bcent'], \ + CSMF_hod_dict['Bsat'], CSMF_hod_dict['ic'] + f_sigv_C = CSMF_hod_dict['f_sigv'] numba.set_num_threads(Nthread) # compute nsate for each halo # figuring out the number of particles kept for each thread - num_sats_L = np.zeros(len(hid), dtype=np.int64) - num_sats_E = np.zeros(len(hid), dtype=np.int64) - num_sats_Q = np.zeros(len(hid), dtype=np.int64) - hstart = np.rint(np.linspace(0, len(hid), Nthread + 1)).astype( - np.int64 - ) # starting index of each thread + num_sats_L = np.zeros(len(hid), dtype = np.int64) + num_sats_E = np.zeros(len(hid), dtype = np.int64) + num_sats_Q = np.zeros(len(hid), dtype = np.int64) + num_sats_C = np.zeros(len(hid), dtype = np.int64) + stellarmass_C = np.zeros(len(hid), dtype = np.int64) + hstart = np.rint(np.linspace(0, len(hid), Nthread + 1)).astype(np.int64) # starting index of each thread for tid in range(Nthread): for i in range(hstart[tid], hstart[tid + 1]): if want_LRG: - M1_L_temp = 10 ** (logM1_L + As_L * hdeltac[i] + Bs_L * hfenv[i]) + M1_L_temp = 10**(logM1_L + As_L * hdeltac[i] + Bs_L * hfenv[i]) logM_cut_L_temp = logM_cut_L + Ac_L * hdeltac[i] + Bc_L * hfenv[i] - base_p_L = ( - n_sat_LRG_modified( - hmass[i], - logM_cut_L_temp, - 10**logM_cut_L_temp, - M1_L_temp, - sigma_L, - alpha_L, - kappa_L, - ) - * ic_L - ) + base_p_L = n_sat_LRG_modified(hmass[i], logM_cut_L_temp, + 10**logM_cut_L_temp, M1_L_temp, sigma_L, alpha_L, kappa_L) * ic_L num_sats_L[i] = np.random.poisson(base_p_L) if want_ELG: - M1_E_temp = 10 ** ( - logM1_E + As_E * hdeltac[i] + Bs_E * hfenv[i] + Cs_E * hshear[i] - ) - logM_cut_E_temp = ( - logM_cut_E + Ac_E * hdeltac[i] + Bc_E * hfenv[i] + Cc_E * hshear[i] - ) - base_p_E = ( - N_sat_elg( - hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_E, A_E - ) - * ic_E - ) + M1_E_temp = 10**(logM1_E + As_E * hdeltac[i] + Bs_E * hfenv[i] + Cs_E * hshear[i]) + logM_cut_E_temp = logM_cut_E + Ac_E * hdeltac[i] + Bc_E * hfenv[i] + Cc_E * hshear[i] + base_p_E = N_sat_elg( + hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_E, A_E) * ic_E # elg conformity if keep_cent[i] == 1: - M1_E_temp = 10 ** (logM1_EL + As_E * hdeltac[i] + Bs_E * hfenv[i]) - base_p_E = ( - N_sat_elg( - hmass[i], - 10**logM_cut_E_temp, - kappa_E, - M1_E_temp, - alpha_EL, - A_E, - ) - * ic_E - ) + M1_E_temp = 10**(logM1_EL + As_E * hdeltac[i] + Bs_E * hfenv[i]) + base_p_E = N_sat_elg( + hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_EL, A_E) * ic_E elif keep_cent[i] == 2: - M1_E_temp = 10 ** ( - logM1_EE + As_E * hdeltac[i] + Bs_E * hfenv[i] - ) # M1_E_temp*10**delta_M1 - base_p_E = ( - N_sat_elg( - hmass[i], - 10**logM_cut_E_temp, - kappa_E, - M1_E_temp, - alpha_EE, - A_E, - ) - * ic_E - ) + M1_E_temp = 10**(logM1_EE + As_E * hdeltac[i] + Bs_E * hfenv[i]) # M1_E_temp*10**delta_M1 + base_p_E = N_sat_elg( + hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_EE, A_E) * ic_E num_sats_E[i] = np.random.poisson(base_p_E) if want_QSO: - M1_Q_temp = 10 ** (logM1_Q + As_Q * hdeltac[i] + Bs_Q * hfenv[i]) + M1_Q_temp = 10**(logM1_Q + As_Q * hdeltac[i] + Bs_Q * hfenv[i]) logM_cut_Q_temp = logM_cut_Q + Ac_Q * hdeltac[i] + Bc_Q * hfenv[i] - base_p_Q = ( - N_sat_generic( - hmass[i], 10**logM_cut_Q_temp, kappa_Q, M1_Q_temp, alpha_Q - ) - * ic_Q - ) + base_p_Q = N_sat_generic( + hmass[i], 10**logM_cut_Q_temp, kappa_Q, M1_Q_temp, alpha_Q) * ic_Q num_sats_Q[i] = np.random.poisson(base_p_Q) - + + if want_CSMF: + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] + base_p_C = n_sat_CSMF(hmass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C_temp, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) * ic_C + + num_sats_C[i] = np.random.poisson(base_p_C) + # generate rdpos rd_pos_L = getPointsOnSphere(np.sum(num_sats_L), Nthread) rd_pos_E = getPointsOnSphere(np.sum(num_sats_E), Nthread) rd_pos_Q = getPointsOnSphere(np.sum(num_sats_Q), Nthread) + rd_pos_C = getPointsOnSphere(np.sum(num_sats_C), Nthread) # put satellites on NFW - h_id_L, x_sat_L, y_sat_L, z_sat_L, vx_sat_L, vy_sat_L, vz_sat_L, M_L = ( - compute_fast_NFW( - NFW_draw, - hid, - hpos[:, 0], - hpos[:, 1], - hpos[:, 2], - hvel[:, 0], - hvel[:, 1], - hvel[:, 2], - hvrms, - hc, - hmass, - hrvir, - rd_pos_L, - num_sats_L, - f_sigv_L, - vel_sat, - Nthread, - exp_frac, - exp_scale, - nfw_rescale, - ) - ) - h_id_E, x_sat_E, y_sat_E, z_sat_E, vx_sat_E, vy_sat_E, vz_sat_E, M_E = ( - compute_fast_NFW( - NFW_draw, - hid, - hpos[:, 0], - hpos[:, 1], - hpos[:, 2], - hvel[:, 0], - hvel[:, 1], - hvel[:, 2], - hvrms, - hc, - hmass, - hrvir, - rd_pos_E, - num_sats_E, - f_sigv_E, - vel_sat, - Nthread, - exp_frac, - exp_scale, - nfw_rescale, - ) - ) - h_id_Q, x_sat_Q, y_sat_Q, z_sat_Q, vx_sat_Q, vy_sat_Q, vz_sat_Q, M_Q = ( - compute_fast_NFW( - NFW_draw, - hid, - hpos[:, 0], - hpos[:, 1], - hpos[:, 2], - hvel[:, 0], - hvel[:, 1], - hvel[:, 2], - hvrms, - hc, - hmass, - hrvir, - rd_pos_Q, - num_sats_Q, - f_sigv_Q, - vel_sat, - Nthread, - exp_frac, - exp_scale, - nfw_rescale, - ) - ) + h_id_L, x_sat_L, y_sat_L, z_sat_L, vx_sat_L, vy_sat_L, vz_sat_L, M_L\ + = compute_fast_NFW(NFW_draw, hid, hpos[:, 0], hpos[:, 1], hpos[:, 2], hvel[:, 0], hvel[:, 1], hvel[:, 2], + hvrms, hc, hmass, hrvir, rd_pos_L, num_sats_L, f_sigv_L, vel_sat, Nthread, + exp_frac, exp_scale, nfw_rescale) + h_id_E, x_sat_E, y_sat_E, z_sat_E, vx_sat_E, vy_sat_E, vz_sat_E, M_E\ + = compute_fast_NFW(NFW_draw, hid, hpos[:, 0], hpos[:, 1], hpos[:, 2], hvel[:, 0], hvel[:, 1], hvel[:, 2], + hvrms, hc, hmass, hrvir, rd_pos_E, num_sats_E, f_sigv_E, vel_sat, Nthread, + exp_frac, exp_scale, nfw_rescale) + h_id_Q, x_sat_Q, y_sat_Q, z_sat_Q, vx_sat_Q, vy_sat_Q, vz_sat_Q, M_Q\ + = compute_fast_NFW(NFW_draw, hid, hpos[:, 0], hpos[:, 1], hpos[:, 2], hvel[:, 0], hvel[:, 1], hvel[:, 2], + hvrms, hc, hmass, hrvir, rd_pos_Q, num_sats_Q, f_sigv_Q, vel_sat, Nthread, + exp_frac, exp_scale, nfw_rescale) + h_id_C, x_sat_C, y_sat_C, z_sat_C, vx_sat_C, vy_sat_C, vz_sat_C, M_C\ + = compute_fast_NFW(NFW_draw, hid, hpos[:, 0], hpos[:, 1], hpos[:, 2], hvel[:, 0], hvel[:, 1], hvel[:, 2], + hvrms, hc, hmass, hrvir, rd_pos_C, num_sats_C, f_sigv_C, vel_sat, Nthread, + exp_frac, exp_scale, nfw_rescale) + + # do rsd if rsd: z_sat_L = (z_sat_L + vz_sat_L * inv_velz2kms) % lbox z_sat_E = (z_sat_E + vz_sat_E * inv_velz2kms) % lbox z_sat_Q = (z_sat_Q + vz_sat_Q * inv_velz2kms) % lbox + z_sat_C = (z_sat_C + vz_sat_C * inv_velz2kms) % lbox - LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) + LRG_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + ELG_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + QSO_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + ID_dict = Dict.empty(key_type = types.unicode_type, value_type = int_array) LRG_dict['x'] = x_sat_L LRG_dict['y'] = y_sat_L LRG_dict['z'] = z_sat_L @@ -817,170 +851,93 @@ def gen_sats_nfw( QSO_dict['vz'] = vz_sat_Q QSO_dict['mass'] = M_Q ID_dict['QSO'] = h_id_Q + + CSMF_dict['x'] = x_sat_C + CSMF_dict['y'] = y_sat_C + CSMF_dict['z'] = z_sat_C + CSMF_dict['vx'] = vx_sat_C + CSMF_dict['vy'] = vy_sat_C + CSMF_dict['vz'] = vz_sat_C + CSMF_dict['mass'] = M_C + stellarmass_C = np.empty_like(M_C) + ## compute stellarmass of all the satelites + if want_CSMF: + hstart = np.rint(np.linspace(0, num_sats_C.sum(), Nthread + 1)) + for tid in numba.prange(Nthread): + for i in range(int(hstart[tid]), int(hstart[tid + 1])): + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] + stellarmass_C[i] = get_random_sat_stellarmass_linearinterpolation(M_C[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, a1_C_temp, s, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) + CSMF_dict['stellarmass'] = stellarmass_C + ID_dict['CSMF'] = h_id_C + + return LRG_dict, ELG_dict, QSO_dict, CSMF_dict, ID_dict + +@njit(parallel = True, fastmath = True) +def gen_sats(ppos, pvel, hvel, hmass, hid, weights, randoms, hdeltac, hfenv, hshear, + enable_ranks, ranks, ranksv, ranksp, ranksr, ranksc, + LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, CSMF_hod_dict, + rsd, inv_velz2kms, lbox, Mpart, want_LRG, want_ELG, want_QSO, want_CSMF, Nthread, origin, keep_cent): - return LRG_dict, ELG_dict, QSO_dict, ID_dict - - -@njit(parallel=True, fastmath=True) -def gen_sats( - ppos, - pvel, - hvel, - hmass, - hid, - weights, - randoms, - hdeltac, - hfenv, - hshear, - enable_ranks, - ranks, - ranksv, - ranksp, - ranksr, - ranksc, - LRG_hod_dict, - ELG_hod_dict, - QSO_hod_dict, - rsd, - inv_velz2kms, - lbox, - Mpart, - want_LRG, - want_ELG, - want_QSO, - Nthread, - origin, - keep_cent, -): """ Generate satellite galaxies in place in memory with a two pass numba parallel implementation. """ if want_LRG: - logM_cut_L, logM1_L, sigma_L, alpha_L, kappa_L = ( - LRG_hod_dict['logM_cut'], - LRG_hod_dict['logM1'], - LRG_hod_dict['sigma'], - LRG_hod_dict['alpha'], - LRG_hod_dict['kappa'], - ) - alpha_s_L, s_L, s_v_L, s_p_L, s_r_L, Ac_L, As_L, Bc_L, Bs_L, ic_L = ( - LRG_hod_dict['alpha_s'], - LRG_hod_dict['s'], - LRG_hod_dict['s_v'], - LRG_hod_dict['s_p'], - LRG_hod_dict['s_r'], - LRG_hod_dict['Acent'], - LRG_hod_dict['Asat'], - LRG_hod_dict['Bcent'], - LRG_hod_dict['Bsat'], - LRG_hod_dict['ic'], - ) + logM_cut_L, logM1_L, sigma_L, alpha_L, kappa_L = \ + LRG_hod_dict['logM_cut'], LRG_hod_dict['logM1'], LRG_hod_dict['sigma'], LRG_hod_dict['alpha'], LRG_hod_dict['kappa'] + alpha_s_L, s_L, s_v_L, s_p_L, s_r_L, Ac_L, As_L, Bc_L, Bs_L, ic_L = \ + LRG_hod_dict['alpha_s'], LRG_hod_dict['s'], LRG_hod_dict['s_v'], LRG_hod_dict['s_p'], \ + LRG_hod_dict['s_r'], LRG_hod_dict['Acent'], LRG_hod_dict['Asat'], LRG_hod_dict['Bcent'], \ + LRG_hod_dict['Bsat'], LRG_hod_dict['ic'] if want_ELG: - logM_cut_E, kappa_E, logM1_E, alpha_E, A_E = ( - ELG_hod_dict['logM_cut'], - ELG_hod_dict['kappa'], - ELG_hod_dict['logM1'], - ELG_hod_dict['alpha'], - ELG_hod_dict['A_s'], - ) - ( - alpha_s_E, - s_E, - s_v_E, - s_p_E, - s_r_E, - Ac_E, - As_E, - Bc_E, - Bs_E, - Cc_E, - Cs_E, - ic_E, - logM1_EE, - alpha_EE, - logM1_EL, - alpha_EL, - ) = ( - ELG_hod_dict['alpha_s'], - ELG_hod_dict['s'], - ELG_hod_dict['s_v'], - ELG_hod_dict['s_p'], - ELG_hod_dict['s_r'], - ELG_hod_dict['Acent'], - ELG_hod_dict['Asat'], - ELG_hod_dict['Bcent'], - ELG_hod_dict['Bsat'], - ELG_hod_dict['Ccent'], - ELG_hod_dict['Csat'], - ELG_hod_dict['ic'], - ELG_hod_dict['logM1_EE'], - ELG_hod_dict['alpha_EE'], - ELG_hod_dict['logM1_EL'], - ELG_hod_dict['alpha_EL'], - ) + logM_cut_E, kappa_E, logM1_E, alpha_E, A_E = \ + ELG_hod_dict['logM_cut'], ELG_hod_dict['kappa'], ELG_hod_dict['logM1'], ELG_hod_dict['alpha'], ELG_hod_dict['A_s'] + alpha_s_E, s_E, s_v_E, s_p_E, s_r_E, Ac_E, As_E, Bc_E, Bs_E, Cc_E, Cs_E, ic_E, logM1_EE, alpha_EE, logM1_EL, alpha_EL = \ + ELG_hod_dict['alpha_s'], ELG_hod_dict['s'], ELG_hod_dict['s_v'], ELG_hod_dict['s_p'], \ + ELG_hod_dict['s_r'], ELG_hod_dict['Acent'], ELG_hod_dict['Asat'], ELG_hod_dict['Bcent'], \ + ELG_hod_dict['Bsat'], ELG_hod_dict['Ccent'], ELG_hod_dict['Csat'], ELG_hod_dict['ic'], \ + ELG_hod_dict['logM1_EE'], ELG_hod_dict['alpha_EE'], ELG_hod_dict['logM1_EL'], ELG_hod_dict['alpha_EL'] if want_QSO: - logM_cut_Q, kappa_Q, logM1_Q, alpha_Q = ( - QSO_hod_dict['logM_cut'], - QSO_hod_dict['kappa'], - QSO_hod_dict['logM1'], - QSO_hod_dict['alpha'], - ) - alpha_s_Q, s_Q, s_v_Q, s_p_Q, s_r_Q, Ac_Q, As_Q, Bc_Q, Bs_Q, ic_Q = ( - QSO_hod_dict['alpha_s'], - QSO_hod_dict['s'], - QSO_hod_dict['s_v'], - QSO_hod_dict['s_p'], - QSO_hod_dict['s_r'], - QSO_hod_dict['Acent'], - QSO_hod_dict['Asat'], - QSO_hod_dict['Bcent'], - QSO_hod_dict['Bsat'], - QSO_hod_dict['ic'], - ) - - H = len(hmass) # num of particles + logM_cut_Q, kappa_Q, logM1_Q, alpha_Q = \ + QSO_hod_dict['logM_cut'], QSO_hod_dict['kappa'], QSO_hod_dict['logM1'], QSO_hod_dict['alpha'] + alpha_s_Q, s_Q, s_v_Q, s_p_Q, s_r_Q, Ac_Q, As_Q, Bc_Q, Bs_Q, ic_Q = \ + QSO_hod_dict['alpha_s'], QSO_hod_dict['s'], QSO_hod_dict['s_v'], QSO_hod_dict['s_p'], \ + QSO_hod_dict['s_r'], QSO_hod_dict['Acent'], QSO_hod_dict['Asat'], QSO_hod_dict['Bcent'], \ + QSO_hod_dict['Bsat'], QSO_hod_dict['ic'] + + if want_CSMF: + Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = \ + CSMF_hod_dict['Mstar_low'], CSMF_hod_dict['Mstar_up'], CSMF_hod_dict['M_1'], CSMF_hod_dict['M_0'], CSMF_hod_dict['gamma_1'], \ + CSMF_hod_dict['gamma_2'], CSMF_hod_dict['sigma_c'], CSMF_hod_dict['a_1'], CSMF_hod_dict['a_2'], CSMF_hod_dict['M_2'], \ + CSMF_hod_dict['b_0'], CSMF_hod_dict['b_1'], CSMF_hod_dict['b_2'], CSMF_hod_dict['delta_1'], CSMF_hod_dict['delta_2'] + alpha_s_C, s_C, s_v_C, s_p_C, s_r_C, Ac_C, As_C, Bc_C, Bs_C, ic_C = CSMF_hod_dict['alpha_s'], CSMF_hod_dict['s'],\ + CSMF_hod_dict['s_v'], CSMF_hod_dict['s_p'], CSMF_hod_dict['s_r'], CSMF_hod_dict['Acent'],\ + CSMF_hod_dict['Asat'], CSMF_hod_dict['Bcent'], CSMF_hod_dict['Bsat'], CSMF_hod_dict['ic'] + + # print(Ac_C, As_C, Bc_C, Bs_C) + H = len(hmass) # num of particles numba.set_num_threads(Nthread) - Nout = np.zeros((Nthread, 3, 8), dtype=np.int64) - hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype( - np.int64 - ) # starting index of each thread + Nout = np.zeros((Nthread, 4, 8), dtype = np.int64) + hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype(np.int64) # starting index of each thread - keep = np.empty(H, dtype=np.int8) # mask array tracking which halos to keep + keep = np.empty(H, dtype = np.int8) # mask array tracking which halos to keep # figuring out the number of particles kept for each thread - for tid in numba.prange(Nthread): # numba.prange(Nthread): + for tid in numba.prange(Nthread): #numba.prange(Nthread): for i in range(hstart[tid], hstart[tid + 1]): # print(logM1, As, hdeltac[i], Bs, hfenv[i]) LRG_marker = 0 if want_LRG: - M1_L_temp = 10 ** (logM1_L + As_L * hdeltac[i] + Bs_L * hfenv[i]) + M1_L_temp = 10**(logM1_L + As_L * hdeltac[i] + Bs_L * hfenv[i]) logM_cut_L_temp = logM_cut_L + Ac_L * hdeltac[i] + Bc_L * hfenv[i] - base_p_L = ( - n_sat_LRG_modified( - hmass[i], - logM_cut_L_temp, - 10**logM_cut_L_temp, - M1_L_temp, - sigma_L, - alpha_L, - kappa_L, - ) - * weights[i] - * ic_L - ) + base_p_L = n_sat_LRG_modified(hmass[i], logM_cut_L_temp, + 10**logM_cut_L_temp, M1_L_temp, sigma_L, alpha_L, kappa_L) * weights[i] * ic_L if enable_ranks: - decorator_L = ( - 1 - + s_L * ranks[i] - + s_v_L * ranksv[i] - + s_p_L * ranksp[i] - + s_r_L * ranksr[i] - ) + decorator_L = 1 + s_L * ranks[i] + s_v_L * ranksv[i] + s_p_L * ranksp[i] + s_r_L * ranksr[i] exp_sat = base_p_L * decorator_L else: exp_sat = base_p_L @@ -988,50 +945,19 @@ def gen_sats( ELG_marker = LRG_marker if want_ELG: - M1_E_temp = 10 ** ( - logM1_E + As_E * hdeltac[i] + Bs_E * hfenv[i] + Cs_E * hshear[i] - ) - logM_cut_E_temp = ( - logM_cut_E + Ac_E * hdeltac[i] + Bc_E * hfenv[i] + Cc_E * hshear[i] - ) - base_p_E = ( - N_sat_elg( - hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_E, A_E - ) - * weights[i] - * ic_E - ) + M1_E_temp = 10**(logM1_E + As_E * hdeltac[i] + Bs_E * hfenv[i] + Cs_E * hshear[i]) + logM_cut_E_temp = logM_cut_E + Ac_E * hdeltac[i] + Bc_E * hfenv[i] + Cc_E * hshear[i] + base_p_E = N_sat_elg( + hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_E, A_E) * weights[i] * ic_E # elg conformity if keep_cent[i] == 1: - M1_E_temp = 10 ** (logM1_EL + As_E * hdeltac[i] + Bs_E * hfenv[i]) - base_p_E = ( - N_sat_elg( - hmass[i], - 10**logM_cut_E_temp, - kappa_E, - M1_E_temp, - alpha_EL, - A_E, - ) - * weights[i] - * ic_E - ) + M1_E_temp = 10**(logM1_EL + As_E * hdeltac[i] + Bs_E * hfenv[i]) + base_p_E = N_sat_elg( + hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_EL, A_E) * weights[i] * ic_E elif keep_cent[i] == 2: - M1_E_temp = 10 ** ( - logM1_EE + As_E * hdeltac[i] + Bs_E * hfenv[i] - ) # M1_E_temp*10**delta_M1 - base_p_E = ( - N_sat_elg( - hmass[i], - 10**logM_cut_E_temp, - kappa_E, - M1_E_temp, - alpha_EE, - A_E, - ) - * weights[i] - * ic_E - ) + M1_E_temp = 10**(logM1_EE + As_E * hdeltac[i] + Bs_E * hfenv[i]) # M1_E_temp*10**delta_M1 + base_p_E = N_sat_elg( + hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_EE, A_E) * weights[i] * ic_E # if base_p_E > 1: # print("ExE new p", base_p_E, np.log10(hmass[i]), N_sat_elg( @@ -1039,124 +965,134 @@ def gen_sats( # rank mods if enable_ranks: - decorator_E = ( - 1 - + s_E * ranks[i] - + s_v_E * ranksv[i] - + s_p_E * ranksp[i] - + s_r_E * ranksr[i] - ) + decorator_E = 1 + s_E * ranks[i] + s_v_E * ranksv[i] + s_p_E * ranksp[i] + s_r_E * ranksr[i] base_p_E = base_p_E * decorator_E ELG_marker += base_p_E QSO_marker = ELG_marker if want_QSO: - M1_Q_temp = 10 ** (logM1_Q + As_Q * hdeltac[i] + Bs_Q * hfenv[i]) + M1_Q_temp = 10**(logM1_Q + As_Q * hdeltac[i] + Bs_Q * hfenv[i]) logM_cut_Q_temp = logM_cut_Q + Ac_Q * hdeltac[i] + Bc_Q * hfenv[i] - base_p_Q = ( - N_sat_generic( - hmass[i], 10**logM_cut_Q_temp, kappa_Q, M1_Q_temp, alpha_Q - ) - * weights[i] - * ic_Q - ) + base_p_Q = N_sat_generic( + hmass[i], 10**logM_cut_Q_temp, kappa_Q, M1_Q_temp, alpha_Q) * weights[i] * ic_Q if enable_ranks: - decorator_Q = ( - 1 - + s_Q * ranks[i] - + s_v_Q * ranksv[i] - + s_p_Q * ranksp[i] - + s_r_Q * ranksr[i] - ) + decorator_Q = 1 + s_Q * ranks[i] + s_v_Q * ranksv[i] + s_p_Q * ranksp[i] + s_r_Q * ranksr[i] exp_sat = base_p_Q * decorator_Q else: exp_sat = base_p_Q QSO_marker += exp_sat + + CSMF_marker = QSO_marker + if want_CSMF: + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] + base_p_C = n_sat_CSMF(hmass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C_temp, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) * weights[i] * ic_C + if enable_ranks: + decorator_C = 1 + s_C * ranks[i] + s_v_C * ranksv[i] + s_p_C * ranksp[i] + s_r_C * ranksr[i] + exp_sat = base_p_C * decorator_C + else: + exp_sat = base_p_C + CSMF_marker += exp_sat if randoms[i] <= LRG_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 elif randoms[i] <= ELG_marker: - Nout[tid, 1, 0] += 1 # counting + Nout[tid, 1, 0] += 1 # counting keep[i] = 2 elif randoms[i] <= QSO_marker: - Nout[tid, 2, 0] += 1 # counting + Nout[tid, 2, 0] += 1 # counting keep[i] = 3 + elif randoms[i] <= CSMF_marker: + Nout[tid, 3, 0] += 1 # counting + keep[i] = 4 else: keep[i] = 0 +# if(0 not in keep): +# print("Waring: All particles are used for satelites") +# else: +# unused_haloes = sum(keep == 0) +# print("From "+str(H)+" particles we used:",(H-unused_haloes)/H) + # compose galaxy array, first create array of galaxy starting indices for the threads - gstart = np.empty((Nthread + 1, 3), dtype=np.int64) + gstart = np.empty((Nthread + 1, 4), dtype = np.int64) gstart[0, :] = 0 gstart[1:, 0] = Nout[:, 0, 0].cumsum() gstart[1:, 1] = Nout[:, 1, 0].cumsum() gstart[1:, 2] = Nout[:, 2, 0].cumsum() + gstart[1:, 3] = Nout[:, 3, 0].cumsum() # galaxy arrays N_lrg = gstart[-1, 0] - lrg_x = np.empty(N_lrg, dtype=hmass.dtype) - lrg_y = np.empty(N_lrg, dtype=hmass.dtype) - lrg_z = np.empty(N_lrg, dtype=hmass.dtype) - lrg_vx = np.empty(N_lrg, dtype=hmass.dtype) - lrg_vy = np.empty(N_lrg, dtype=hmass.dtype) - lrg_vz = np.empty(N_lrg, dtype=hmass.dtype) - lrg_mass = np.empty(N_lrg, dtype=hmass.dtype) - lrg_id = np.empty(N_lrg, dtype=hid.dtype) + lrg_x = np.empty(N_lrg, dtype = hmass.dtype) + lrg_y = np.empty(N_lrg, dtype = hmass.dtype) + lrg_z = np.empty(N_lrg, dtype = hmass.dtype) + lrg_vx = np.empty(N_lrg, dtype = hmass.dtype) + lrg_vy = np.empty(N_lrg, dtype = hmass.dtype) + lrg_vz = np.empty(N_lrg, dtype = hmass.dtype) + lrg_mass = np.empty(N_lrg, dtype = hmass.dtype) + lrg_id = np.empty(N_lrg, dtype = hid.dtype) # galaxy arrays N_elg = gstart[-1, 1] - elg_x = np.empty(N_elg, dtype=hmass.dtype) - elg_y = np.empty(N_elg, dtype=hmass.dtype) - elg_z = np.empty(N_elg, dtype=hmass.dtype) - elg_vx = np.empty(N_elg, dtype=hmass.dtype) - elg_vy = np.empty(N_elg, dtype=hmass.dtype) - elg_vz = np.empty(N_elg, dtype=hmass.dtype) - elg_mass = np.empty(N_elg, dtype=hmass.dtype) - elg_id = np.empty(N_elg, dtype=hid.dtype) + elg_x = np.empty(N_elg, dtype = hmass.dtype) + elg_y = np.empty(N_elg, dtype = hmass.dtype) + elg_z = np.empty(N_elg, dtype = hmass.dtype) + elg_vx = np.empty(N_elg, dtype = hmass.dtype) + elg_vy = np.empty(N_elg, dtype = hmass.dtype) + elg_vz = np.empty(N_elg, dtype = hmass.dtype) + elg_mass = np.empty(N_elg, dtype = hmass.dtype) + elg_id = np.empty(N_elg, dtype = hid.dtype) # galaxy arrays N_qso = gstart[-1, 2] - qso_x = np.empty(N_qso, dtype=hmass.dtype) - qso_y = np.empty(N_qso, dtype=hmass.dtype) - qso_z = np.empty(N_qso, dtype=hmass.dtype) - qso_vx = np.empty(N_qso, dtype=hmass.dtype) - qso_vy = np.empty(N_qso, dtype=hmass.dtype) - qso_vz = np.empty(N_qso, dtype=hmass.dtype) - qso_mass = np.empty(N_qso, dtype=hmass.dtype) - qso_id = np.empty(N_qso, dtype=hid.dtype) + qso_x = np.empty(N_qso, dtype = hmass.dtype) + qso_y = np.empty(N_qso, dtype = hmass.dtype) + qso_z = np.empty(N_qso, dtype = hmass.dtype) + qso_vx = np.empty(N_qso, dtype = hmass.dtype) + qso_vy = np.empty(N_qso, dtype = hmass.dtype) + qso_vz = np.empty(N_qso, dtype = hmass.dtype) + qso_mass = np.empty(N_qso, dtype = hmass.dtype) + qso_id = np.empty(N_qso, dtype = hid.dtype) + + # galaxy arrays + N_CSMF = gstart[-1, 3] + CSMF_x = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_y = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_z = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_id = np.empty(N_CSMF, dtype = hid.dtype) + # fill in the galaxy arrays for tid in numba.prange(Nthread): - j1, j2, j3 = gstart[tid] + j1, j2, j3, j4 = gstart[tid] for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: lrg_x[j1] = ppos[i, 0] - lrg_vx[j1] = hvel[i, 0] + alpha_s_L * ( - pvel[i, 0] - hvel[i, 0] - ) # velocity bias + lrg_vx[j1] = hvel[i, 0] + alpha_s_L * (pvel[i, 0] - hvel[i, 0]) # velocity bias lrg_y[j1] = ppos[i, 1] - lrg_vy[j1] = hvel[i, 1] + alpha_s_L * ( - pvel[i, 1] - hvel[i, 1] - ) # velocity bias + lrg_vy[j1] = hvel[i, 1] + alpha_s_L * (pvel[i, 1] - hvel[i, 1]) # velocity bias lrg_z[j1] = ppos[i, 2] - lrg_vz[j1] = hvel[i, 2] + alpha_s_L * ( - pvel[i, 2] - hvel[i, 2] - ) # velocity bias + lrg_vz[j1] = hvel[i, 2] + alpha_s_L * (pvel[i, 2] - hvel[i, 2]) # velocity bias if rsd and origin is not None: nx = lrg_x[j1] - origin[0] ny = lrg_y[j1] - origin[1] nz = lrg_z[j1] - origin[2] - inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms * ( - lrg_vx[j1] * nx + lrg_vy[j1] * ny + lrg_vz[j1] * nz - ) - lrg_x[j1] = lrg_x[j1] + proj * nx - lrg_y[j1] = lrg_y[j1] + proj * ny - lrg_z[j1] = lrg_z[j1] + proj * nz + proj = inv_velz2kms*(lrg_vx[j1]*nx+lrg_vy[j1]*ny+lrg_vz[j1]*nz) + lrg_x[j1] = lrg_x[j1]+proj*nx + lrg_y[j1] = lrg_y[j1]+proj*ny + lrg_z[j1] = lrg_z[j1]+proj*nz elif rsd: lrg_z[j1] = wrap(lrg_z[j1] + lrg_vz[j1] * inv_velz2kms, lbox) lrg_mass[j1] = hmass[i] @@ -1164,31 +1100,23 @@ def gen_sats( j1 += 1 elif keep[i] == 2: elg_x[j2] = ppos[i, 0] - elg_vx[j2] = hvel[i, 0] + alpha_s_E * ( - pvel[i, 0] - hvel[i, 0] - ) # velocity bias + elg_vx[j2] = hvel[i, 0] + alpha_s_E * (pvel[i, 0] - hvel[i, 0]) # velocity bias elg_y[j2] = ppos[i, 1] - elg_vy[j2] = hvel[i, 1] + alpha_s_E * ( - pvel[i, 1] - hvel[i, 1] - ) # velocity bias + elg_vy[j2] = hvel[i, 1] + alpha_s_E * (pvel[i, 1] - hvel[i, 1]) # velocity bias elg_z[j2] = ppos[i, 2] - elg_vz[j2] = hvel[i, 2] + alpha_s_E * ( - pvel[i, 2] - hvel[i, 2] - ) # velocity bias + elg_vz[j2] = hvel[i, 2] + alpha_s_E * (pvel[i, 2] - hvel[i, 2]) # velocity bias if rsd and origin is not None: nx = elg_x[j2] - origin[0] ny = elg_y[j2] - origin[1] nz = elg_z[j2] - origin[2] - inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms * ( - elg_vx[j2] * nx + elg_vy[j2] * ny + elg_vz[j2] * nz - ) - elg_x[j2] = elg_x[j2] + proj * nx - elg_y[j2] = elg_y[j2] + proj * ny - elg_z[j2] = elg_z[j2] + proj * nz + proj = inv_velz2kms*(elg_vx[j2]*nx+elg_vy[j2]*ny+elg_vz[j2]*nz) + elg_x[j2] = elg_x[j2]+proj*nx + elg_y[j2] = elg_y[j2]+proj*ny + elg_z[j2] = elg_z[j2]+proj*nz elif rsd: elg_z[j2] = wrap(elg_z[j2] + elg_vz[j2] * inv_velz2kms, lbox) elg_mass[j2] = hmass[i] @@ -1196,42 +1124,70 @@ def gen_sats( j2 += 1 elif keep[i] == 3: qso_x[j3] = ppos[i, 0] - qso_vx[j3] = hvel[i, 0] + alpha_s_Q * ( - pvel[i, 0] - hvel[i, 0] - ) # velocity bias + qso_vx[j3] = hvel[i, 0] + alpha_s_Q * (pvel[i, 0] - hvel[i, 0]) # velocity bias qso_y[j3] = ppos[i, 1] - qso_vy[j3] = hvel[i, 1] + alpha_s_Q * ( - pvel[i, 1] - hvel[i, 1] - ) # velocity bias + qso_vy[j3] = hvel[i, 1] + alpha_s_Q * (pvel[i, 1] - hvel[i, 1]) # velocity bias qso_z[j3] = ppos[i, 2] - qso_vz[j3] = hvel[i, 2] + alpha_s_Q * ( - pvel[i, 2] - hvel[i, 2] - ) # velocity bias + qso_vz[j3] = hvel[i, 2] + alpha_s_Q * (pvel[i, 2] - hvel[i, 2]) # velocity bias if rsd and origin is not None: nx = qso_x[j3] - origin[0] ny = qso_y[j3] - origin[1] nz = qso_z[j3] - origin[2] - inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms * ( - qso_vx[j3] * nx + qso_vy[j3] * ny + qso_vz[j3] * nz - ) - qso_x[j3] = qso_x[j3] + proj * nx - qso_y[j3] = qso_y[j3] + proj * ny - qso_z[j3] = qso_z[j3] + proj * nz + proj = inv_velz2kms*(qso_vx[j3]*nx+qso_vy[j3]*ny+qso_vz[j3]*nz) + qso_x[j3] = qso_x[j3]+proj*nx + qso_y[j3] = qso_y[j3]+proj*ny + qso_z[j3] = qso_z[j3]+proj*nz elif rsd: qso_z[j3] = wrap(qso_z[j3] + qso_vz[j3] * inv_velz2kms, lbox) qso_mass[j3] = hmass[i] qso_id[j3] = hid[i] j3 += 1 + elif keep[i] == 4: + CSMF_x[j4] = ppos[i, 0] + CSMF_vx[j4] = hvel[i, 0] + alpha_s_C * (pvel[i, 0] - hvel[i, 0]) # velocity bias + CSMF_y[j4] = ppos[i, 1] + CSMF_vy[j4] = hvel[i, 1] + alpha_s_C * (pvel[i, 1] - hvel[i, 1]) # velocity bias + CSMF_z[j4] = ppos[i, 2] + CSMF_vz[j4] = hvel[i, 2] + alpha_s_C * (pvel[i, 2] - hvel[i, 2]) # velocity bias + if rsd and origin is not None: + nx = CSMF_x[j4] - origin[0] + ny = CSMF_y[j4] - origin[1] + nz = CSMF_z[j4] - origin[2] + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + nx *= inv_norm + ny *= inv_norm + nz *= inv_norm + proj = inv_velz2kms*(CSMF_vx[j4]*nx+CSMF_vy[j4]*ny+CSMF_vz[j4]*nz) + CSMF_x[j4] = CSMF_x[j4]+proj*nx + CSMF_y[j4] = CSMF_y[j4]+proj*ny + CSMF_z[j4] = CSMF_z[j4]+proj*nz + elif rsd: + CSMF_z[j4] = wrap(CSMF_z[j4] + CSMF_vz[j4] * inv_velz2kms, lbox) + + + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] + # stellar_mass = get_random_sat_stellarmass_linearinterpolation(hmass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, a1_C_temp, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) + stellar_mass = get_random_sat_stellarmass(hmass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, a1_C_temp, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) + # if(j4>16000): + + CSMF_stellarmass[j4] = stellar_mass + CSMF_mass[j4] = hmass[i] + # CSMF_stellarmass[j4] = hmass[i] + CSMF_id[j4] = hid[i] + j4 += 1 # assert j == gstart[tid + 1] - LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) + LRG_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + ELG_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + QSO_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + ID_dict = Dict.empty(key_type = types.unicode_type, value_type = int_array) + LRG_dict['x'] = lrg_x LRG_dict['y'] = lrg_y LRG_dict['z'] = lrg_z @@ -1258,12 +1214,22 @@ def gen_sats( QSO_dict['vz'] = qso_vz QSO_dict['mass'] = qso_mass ID_dict['QSO'] = qso_id - return LRG_dict, ELG_dict, QSO_dict, ID_dict - - -@njit(parallel=True, fastmath=True) + + CSMF_dict['x'] = CSMF_x + CSMF_dict['y'] = CSMF_y + CSMF_dict['z'] = CSMF_z + CSMF_dict['vx'] = CSMF_vx + CSMF_dict['vy'] = CSMF_vy + CSMF_dict['vz'] = CSMF_vz + CSMF_dict['mass'] = CSMF_mass + CSMF_dict['stellarmass'] = CSMF_stellarmass + ID_dict['CSMF'] = CSMF_id + + return LRG_dict, ELG_dict, QSO_dict, CSMF_dict, ID_dict + +@njit(parallel = True, fastmath = True) def fast_concatenate(array1, array2, Nthread): - """Fast concatenate with numba parallel""" + ''' Fast concatenate with numba parallel''' N1 = len(array1) N2 = len(array2) @@ -1272,7 +1238,7 @@ def fast_concatenate(array1, array2, Nthread): elif N2 == 0: return array1 - final_array = np.empty(N1 + N2, dtype=array1.dtype) + final_array = np.empty(N1 + N2, dtype = array1.dtype) # if one thread, then no need to parallel if Nthread == 1: for i in range(N1): @@ -1287,7 +1253,7 @@ def fast_concatenate(array1, array2, Nthread): hstart1 = np.rint(np.linspace(0, N1, Nthread1 + 1)).astype(np.int64) hstart2 = np.rint(np.linspace(0, N2, Nthread2 + 1)).astype(np.int64) + N1 - for tid in numba.prange(Nthread): # numba.prange(Nthread): + for tid in numba.prange(Nthread): #numba.prange(Nthread): if tid < Nthread1: for i in range(hstart1[tid], hstart1[tid + 1]): final_array[i] = array1[i] @@ -1298,18 +1264,7 @@ def fast_concatenate(array1, array2, Nthread): return final_array -def gen_gals( - halos_array, - subsample, - tracers, - params, - Nthread, - enable_ranks, - rsd, - verbose, - nfw, - NFW_draw=None, -): +def gen_gals(halos_array, subsample, tracers, params, Nthread, enable_ranks, rsd, verbose, nfw, NFW_draw = None): """ parse hod parameters, pass them on to central and satellite generators and then format the results @@ -1345,26 +1300,26 @@ def gen_gals( ELG_HOD = tracers[tracer] if tracer == 'QSO': QSO_HOD = tracers[tracer] + if tracer == 'CSMF': + CSMF_HOD = tracers[tracer] + + # print(CSMF_HOD) if 'LRG' in tracers.keys(): want_LRG = True - LRG_hod_dict = nb.typed.Dict.empty( - key_type=nb.types.unicode_type, value_type=nb.types.float64 - ) + LRG_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) for key, value in LRG_HOD.items(): LRG_hod_dict[key] = value # LRG design and decorations logM_cut_L, logM1_L = map(LRG_HOD.get, ('logM_cut', 'logM1')) # z-evolving HOD - Delta_a = 1.0 / (1 + params['z']) - 1.0 / ( - 1 + LRG_HOD.get('z_pivot', params['z']) - ) + Delta_a = 1./(1+params['z']) - 1./(1+LRG_HOD.get('z_pivot', params['z'])) logM_cut_pr = LRG_HOD.get('logM_cut_pr', 0.0) logM1_pr = LRG_HOD.get('logM1_pr', 0.0) - logM_cut_L = logM_cut_L + logM_cut_pr * Delta_a - logM1_L = logM1_L + logM1_pr * Delta_a + logM_cut_L = logM_cut_L + logM_cut_pr*Delta_a + logM1_L = logM1_L + logM1_pr*Delta_a LRG_hod_dict['logM_cut'] = logM_cut_L LRG_hod_dict['logM1'] = logM1_L @@ -1379,30 +1334,25 @@ def gen_gals( else: want_LRG = False - LRG_hod_dict = nb.typed.Dict.empty( - key_type=nb.types.unicode_type, value_type=nb.types.float64 - ) + LRG_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + if 'ELG' in tracers.keys(): # ELG design want_ELG = True - ELG_hod_dict = nb.typed.Dict.empty( - key_type=nb.types.unicode_type, value_type=nb.types.float64 - ) + ELG_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) for key, value in ELG_HOD.items(): ELG_hod_dict[key] = value logM_cut_E, logM1_E = map(ELG_HOD.get, ('logM_cut', 'logM1')) # z-evolving HOD - Delta_a = 1.0 / (1 + params['z']) - 1.0 / ( - 1 + ELG_HOD.get('z_pivot', params['z']) - ) + Delta_a = 1./(1+params['z']) - 1./(1+ELG_HOD.get('z_pivot', params['z'])) logM_cut_pr = ELG_HOD.get('logM_cut_pr', 0.0) logM1_pr = ELG_HOD.get('logM1_pr', 0.0) - logM_cut_E = logM_cut_E + logM_cut_pr * Delta_a - logM1_E = logM1_E + logM1_pr * Delta_a + logM_cut_E = logM_cut_E + logM_cut_pr*Delta_a + logM1_E = logM1_E + logM1_pr*Delta_a ELG_hod_dict['logM_cut'] = logM_cut_E ELG_hod_dict['logM1'] = logM1_E @@ -1426,28 +1376,22 @@ def gen_gals( else: want_ELG = False - ELG_hod_dict = nb.typed.Dict.empty( - key_type=nb.types.unicode_type, value_type=nb.types.float64 - ) + ELG_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) if 'QSO' in tracers.keys(): # QSO design want_QSO = True - QSO_hod_dict = nb.typed.Dict.empty( - key_type=nb.types.unicode_type, value_type=nb.types.float64 - ) + QSO_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) for key, value in QSO_HOD.items(): QSO_hod_dict[key] = value logM_cut_Q, logM1_Q = map(QSO_HOD.get, ('logM_cut', 'logM1')) # z-evolving HOD - Delta_a = 1.0 / (1 + params['z']) - 1.0 / ( - 1 + QSO_HOD.get('z_pivot', params['z']) - ) + Delta_a = 1./(1+params['z']) - 1./(1+QSO_HOD.get('z_pivot', params['z'])) logM_cut_pr = QSO_HOD.get('logM_cut_pr', 0.0) logM1_pr = QSO_HOD.get('logM1_pr', 0.0) - logM_cut_Q = logM_cut_Q + logM_cut_pr * Delta_a - logM1_Q = logM1_Q + logM1_pr * Delta_a + logM_cut_Q = logM_cut_Q + logM_cut_pr*Delta_a + logM1_Q = logM1_Q + logM1_pr*Delta_a QSO_hod_dict['logM_cut'] = logM_cut_Q QSO_hod_dict['logM1'] = logM1_Q @@ -1462,150 +1406,96 @@ def gen_gals( else: want_QSO = False - QSO_hod_dict = nb.typed.Dict.empty( - key_type=nb.types.unicode_type, value_type=nb.types.float64 - ) + QSO_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + + + if 'CSMF' in tracers.keys(): + want_CSMF = True + + CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + for key, value in CSMF_HOD.items(): + CSMF_hod_dict[key] = value + + CSMF_hod_dict['Acent'] = CSMF_HOD.get('Acent', 0.0) + CSMF_hod_dict['Asat'] = CSMF_HOD.get('Asat', 0.0) + CSMF_hod_dict['Bcent'] = CSMF_HOD.get('Bcent', 0.0) + CSMF_hod_dict['Bsat'] = CSMF_HOD.get('Bsat', 0.0) + CSMF_hod_dict['ic'] = CSMF_HOD.get('ic', 1.0) + CSMF_hod_dict['f_sigv'] = CSMF_HOD.get('f_sigv', 0) + else: + want_CSMF = False + CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + + # print(CSMF_hod_dict) start = time.time() velz2kms = params['velz2kms'] - inv_velz2kms = 1 / velz2kms + inv_velz2kms = 1/velz2kms lbox = params['Lbox'] origin = params['origin'] - LRG_dict_cent, ELG_dict_cent, QSO_dict_cent, ID_dict_cent, keep_cent = gen_cent( - halos_array['hpos'], - halos_array['hvel'], - halos_array['hmass'], - halos_array['hid'], - halos_array['hmultis'], - halos_array['hrandoms'], - halos_array['hveldev'], - halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), - halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), - halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), - LRG_hod_dict, - ELG_hod_dict, - QSO_hod_dict, - rsd, - inv_velz2kms, - lbox, - want_LRG, - want_ELG, - want_QSO, - Nthread, - origin, - ) + LRG_dict_cent, ELG_dict_cent, QSO_dict_cent, CSMF_dict_cent, ID_dict_cent, keep_cent = \ + gen_cent(halos_array['hpos'], halos_array['hvel'], halos_array['hmass'], halos_array['hid'], halos_array['hmultis'], + halos_array['hrandoms'], halos_array['hveldev'], + halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), + halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), + halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), + LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, CSMF_hod_dict, rsd, inv_velz2kms, lbox, want_LRG, want_ELG, want_QSO, want_CSMF, Nthread, origin) + if verbose: - print('generating centrals took ', time.time() - start) + print("generating centrals took ", time.time() - start) start = time.time() if nfw: - warnings.warn( - 'NFW profile is unoptimized. It has different velocity bias. It does not support lightcone.' - ) - LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, ID_dict_sat = gen_sats_nfw( - NFW_draw, - halos_array['hpos'], - halos_array['hvel'], - halos_array['hmass'], - halos_array['hid'], - halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), - halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), - halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), - halos_array['hsigma3d'], - halos_array['hc'], - halos_array['hrvir'], - LRG_hod_dict, - ELG_hod_dict, - QSO_hod_dict, - want_LRG, - want_ELG, - want_QSO, - rsd, - inv_velz2kms, - lbox, - keep_cent, - Nthread=Nthread, - ) + warnings.warn("NFW profile is unoptimized. It has different velocity bias. It does not support lightcone.") + LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, CSMF_dict_sat, ID_dict_sat = \ + gen_sats_nfw(NFW_draw, halos_array['hpos'], halos_array['hvel'], halos_array['hmass'], halos_array['hid'], + halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), + halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), + halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), + halos_array['hsigma3d'], halos_array['hc'], halos_array['hrvir'], + LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, CSMF_hod_dict, want_LRG, want_ELG, want_QSO, want_CSMF, + rsd, inv_velz2kms, lbox, keep_cent, Nthread = Nthread) else: - LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, ID_dict_sat = gen_sats( - subsample['ppos'], - subsample['pvel'], - subsample['phvel'], - subsample['phmass'], - subsample['phid'], - subsample['pweights'], - subsample['prandoms'], - subsample.get('pdeltac', np.zeros(len(subsample['phid']))), - subsample.get('pfenv', np.zeros(len(subsample['phid']))), - subsample.get('pshear', np.zeros(len(subsample['phid']))), - enable_ranks, - subsample['pranks'], - subsample['pranksv'], - subsample['pranksp'], - subsample['pranksr'], - subsample['pranksc'], - LRG_hod_dict, - ELG_hod_dict, - QSO_hod_dict, - rsd, - inv_velz2kms, - lbox, - params['Mpart'], - want_LRG, - want_ELG, - want_QSO, - Nthread, - origin, - keep_cent[subsample['pinds']], - ) + LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, CSMF_dict_sat, ID_dict_sat = \ + gen_sats(subsample['ppos'], subsample['pvel'], subsample['phvel'], subsample['phmass'], subsample['phid'], + subsample['pweights'], subsample['prandoms'], + subsample.get('pdeltac', np.zeros(len(subsample['phid']))), + subsample.get('pfenv', np.zeros(len(subsample['phid']))), + subsample.get('pshear', np.zeros(len(subsample['phid']))), + enable_ranks, subsample['pranks'], subsample['pranksv'], subsample['pranksp'], subsample['pranksr'], subsample['pranksc'], + LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, CSMF_hod_dict, rsd, inv_velz2kms, lbox, params['Mpart'], + want_LRG, want_ELG, want_QSO, want_CSMF, Nthread, origin, keep_cent[subsample['pinds']]) if verbose: - print('generating satellites took ', time.time() - start) - + print("generating satellites took ", time.time() - start) + + # B.H. TODO: need a for loop above so we don't need to do this by hand - HOD_dict_sat = {'LRG': LRG_dict_sat, 'ELG': ELG_dict_sat, 'QSO': QSO_dict_sat} - HOD_dict_cent = {'LRG': LRG_dict_cent, 'ELG': ELG_dict_cent, 'QSO': QSO_dict_cent} + HOD_dict_sat = {'LRG': LRG_dict_sat, 'ELG': ELG_dict_sat, 'QSO': QSO_dict_sat, 'CSMF': CSMF_dict_sat} + HOD_dict_cent = {'LRG': LRG_dict_cent, 'ELG': ELG_dict_cent, 'QSO': QSO_dict_cent, 'CSMF': CSMF_dict_cent} + # do a concatenate in numba parallel start = time.time() HOD_dict = {} for tracer in tracers: - tracer_dict = {'Ncent': len(HOD_dict_cent[tracer]['x'])} + tracer_dict = {'Ncent':len(HOD_dict_cent[tracer]['x'])} for k in HOD_dict_cent[tracer]: - tracer_dict[k] = fast_concatenate( - HOD_dict_cent[tracer][k], HOD_dict_sat[tracer][k], Nthread - ) - tracer_dict['id'] = fast_concatenate( - ID_dict_cent[tracer], ID_dict_sat[tracer], Nthread - ) + tracer_dict[k] = fast_concatenate(HOD_dict_cent[tracer][k], HOD_dict_sat[tracer][k], Nthread) + tracer_dict['id'] = fast_concatenate(ID_dict_cent[tracer], ID_dict_sat[tracer], Nthread) if verbose: - print(tracer, 'number of galaxies ', len(tracer_dict['x'])) - print( - 'satellite fraction ', - len(HOD_dict_sat[tracer]['x']) / len(tracer_dict['x']), - ) + print(tracer, "number of galaxies ", len(tracer_dict['x'])) + print("satellite fraction ", len(HOD_dict_sat[tracer]['x'])/len(tracer_dict['x'])) HOD_dict[tracer] = tracer_dict if verbose: - print('organizing outputs took ', time.time() - start) + print("organizing outputs took ", time.time() - start) return HOD_dict -def gen_gal_cat( - halo_data, - particle_data, - tracers, - params, - Nthread=16, - enable_ranks=False, - rsd=True, - nfw=False, - NFW_draw=None, - write_to_disk=False, - savedir='./', - verbose=False, - fn_ext=None, -): +def gen_gal_cat(halo_data, particle_data, tracers, params, Nthread = 16, + enable_ranks = False, rsd = True, nfw = False, NFW_draw = None, + write_to_disk = False, savedir = "./", verbose = False, fn_ext = None): """ pass on inputs to the gen_gals function and takes care of I/O @@ -1655,67 +1545,42 @@ def gen_gal_cat( """ if not isinstance(rsd, bool): - raise ValueError('Error: rsd has to be a boolean') + raise ValueError("Error: rsd has to be a boolean") # find the halos, populate them with galaxies and write them to files - HOD_dict = gen_gals( - halo_data, - particle_data, - tracers, - params, - Nthread, - enable_ranks, - rsd, - verbose, - nfw, - NFW_draw, - ) + HOD_dict = gen_gals(halo_data, particle_data, tracers, params, Nthread, enable_ranks, rsd, verbose, nfw, NFW_draw) # how many galaxies were generated and write them to disk for tracer in tracers.keys(): Ncent = HOD_dict[tracer]['Ncent'] if verbose: - print( - 'generated %ss:' % tracer, - len(HOD_dict[tracer]['x']), - 'satellite fraction ', - 1 - Ncent / len(HOD_dict[tracer]['x']), - ) + print("generated %ss:"%tracer, len(HOD_dict[tracer]['x']), + "satellite fraction ", 1 - Ncent/len(HOD_dict[tracer]['x'])) if write_to_disk: if verbose: - print('outputting galaxies to disk') + print("outputting galaxies to disk") if rsd: - rsd_string = '_rsd' + rsd_string = "_rsd" else: - rsd_string = '' + rsd_string = "" if fn_ext is None: - outdir = (savedir) / ('galaxies' + rsd_string) + outdir = (savedir) / ("galaxies"+rsd_string) else: - outdir = (savedir) / ('galaxies' + rsd_string + fn_ext) + outdir = (savedir) / ("galaxies"+rsd_string+fn_ext) # create directories if not existing - os.makedirs(outdir, exist_ok=True) + os.makedirs(outdir, exist_ok = True) # save to file # outdict = HOD_dict[tracer].pop('Ncent', None) - table = Table( - HOD_dict[tracer], - meta={'Ncent': Ncent, 'Gal_type': tracer, **tracers[tracer]}, - ) + table = Table(HOD_dict[tracer], meta = {'Ncent': Ncent, 'Gal_type': tracer, **tracers[tracer]}) if params['chunk'] == -1: - ascii.write( - table, outdir / (f'{tracer}s.dat'), overwrite=True, format='ecsv' - ) + ascii.write(table, outdir / (f"{tracer}s.dat"), overwrite = True, format = 'ecsv') else: - ascii.write( - table, - outdir / (f"{tracer}s_chunk{params['chunk']:d}.dat"), - overwrite=True, - format='ecsv', - ) + ascii.write(table, outdir / (f"{tracer}s_chunk{params['chunk']:d}.dat"), overwrite = True, format = 'ecsv') return HOD_dict diff --git a/abacusnbody/hod/abacus_hod.py b/abacusnbody/hod/abacus_hod.py index de49a210..6d7c615b 100644 --- a/abacusnbody/hod/abacus_hod.py +++ b/abacusnbody/hod/abacus_hod.py @@ -303,14 +303,15 @@ def staging(self): params['numslabs'] = end - start self.lbox = header['BoxSize'] - # count ther number of halos and particles + # count the number of halos and particles Nhalos = np.zeros(params['numslabs']) Nparts = np.zeros(params['numslabs']) for eslab in range(start, end): if ( ('ELG' not in self.tracers.keys()) and ('QSO' not in self.tracers.keys()) - and (not self.force_mt) + and (not self.force_mt) + and ('CSMF' not in self.tracers.keys()) ): halofilename = subsample_dir / ( 'halos_xcom_%d_seed600_abacushod_oldfenv' % eslab @@ -892,6 +893,17 @@ def compute_ngal(self, tracers=None, Nthread=16): ) ngal_dict[etracer] = newngal[0] + newngal[1] fsat_dict[etracer] = newngal[1] / (newngal[0] + newngal[1]) + + + elif etracer == 'CSMF': + newngal = AbacusHOD._compute_ngal_CSMF( + self.logMbins, self.deltacbins, self.fenvbins, self.halo_mass_func, + tracer_hod['Mstar_low'], tracer_hod['Mstar_up'], tracer_hod['M_1'], tracer_hod['M_0'], tracer_hod['gamma_1'], tracer_hod['gamma_2'], + tracer_hod['sigma_c'], tracer_hod['a_1'], tracer_hod['a_2'], tracer_hod['M_2'], tracer_hod['b_0'], tracer_hod['b_1'], + tracer_hod['b_2'], tracer_hod.get('delta_1', 0), tracer_hod.get('delta_2', 0), tracer_hod.get('Acent', 0), + tracer_hod.get('Asat', 0), tracer_hod.get('Bcent', 0), tracer_hod.get('Bsat', 0), tracer_hod.get('ic', 1), Nthread) + ngal_dict[etracer] = newngal[0] + newngal[1] + fsat_dict[etracer] = newngal[1] / (newngal[0] + newngal[1]) return ngal_dict, fsat_dict @staticmethod @@ -1094,6 +1106,36 @@ def _compute_ngal_qso( ngal_cent += halo_mass_func[i, j, k] * ncent_temp * ic ngal_sat += halo_mass_func[i, j, k] * nsat_temp * ic return ngal_cent, ngal_sat + + + @staticmethod + @njit(fastmath = True, parallel = True) + def _compute_ngal_CSMF(logMbins, deltacbins, fenvbins, halo_mass_func, + Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c, a1, a2, M2, b0, b1, b2, delta1, delta2, Acent, Asat, Bcent, Bsat, ic, Nthread): + """ + internal helper to compute number of CSMFs + """ + numba.set_num_threads(Nthread) + + logMs = 0.5*(logMbins[1:] + logMbins[:-1]) + deltacs = 0.5*(deltacbins[1:] + deltacbins[:-1]) + fenvs = 0.5*(fenvbins[1:] + fenvbins[:-1]) + ngal_cent = 0 + ngal_sat = 0 + + for i in numba.prange(len(logMbins) - 1): + for j in range(len(deltacbins) - 1): + for k in range(len(fenvbins) - 1): + Mh_temp = 10**logMs[i] + M_1_temp = 10**(np.log10(M_1)) #+ Acent * deltacs[j] + Bcent * fenvs[k]) + M2_temp = 10**(np.log10(M2)) #+ Asat * deltacs[j] + Bsat * fenvs[k]) + + ncent_temp = n_cen_CSMF(Mh_temp, Mstar_low, Mstar_up, M_1_temp, M_0, gamma1, gamma2, sigma_c) + nsat_temp = n_sat_CSMF(Mh_temp, Mstar_low, Mstar_up, M_1_temp, M_0, gamma1, gamma2, sigma_c, a1, a2, M2_temp, b0, b1, b2, delta1, delta2) + ngal_cent += halo_mass_func[i, j, k] * ncent_temp * ic + ngal_sat += halo_mass_func[i, j, k] * nsat_temp * ic + + return ngal_cent, ngal_sat def compute_clustering(self, mock_dict, *args, **kwargs): """ From 2abddc83e9b4b60f4e9131ddfc5c86cc24a71dd0 Mon Sep 17 00:00:00 2001 From: Pierre Burger Date: Fri, 20 Sep 2024 08:51:56 -0400 Subject: [PATCH 02/13] back to origin back to origin --- abacusnbody/hod/GRAND_HOD.py | 1755 ++++++++++++++++++--------------- abacusnbody/hod/abacus_hod.py | 3 + 2 files changed, 948 insertions(+), 810 deletions(-) diff --git a/abacusnbody/hod/GRAND_HOD.py b/abacusnbody/hod/GRAND_HOD.py index 4ab7222f..739790dd 100644 --- a/abacusnbody/hod/GRAND_HOD.py +++ b/abacusnbody/hod/GRAND_HOD.py @@ -10,7 +10,7 @@ from astropy.table import Table from numba import njit, types from numba.typed import Dict -import scipy.special as sc + # import yaml # config = yaml.safe_load(open('config/abacus_hod.yaml')) # numba.set_num_threads(16) @@ -19,211 +19,18 @@ G = 4.302e-6 # in kpc/Msol (km.s)^2 -@njit(fastmath=True) -def n_cen_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): - """ - Standard Cacciato et al. (2008) centrals HOD parametrization for CSMF - """ - M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - x_low = np.log10(Mstar_low/M_c_value)/(1.41421356*sigma_c) - x_up = np.log10(Mstar_up/M_c_value)/(1.41421356*sigma_c) - - return 0.5*(math.erf(x_up) - math.erf(x_low)) - -@njit(fastmath=True) -def CSMF_centrals(M_h, Mstar, M_1, M_0, gamma1, gamma2, sigma_c): - """ - Eq. (34) from Cacciato et al. (2008) - """ - - M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - - return 1/(1.41421356*np.sqrt(np.pi)*np.log(10)*sigma_c*Mstar)*np.exp(-(np.log10(Mstar)-np.log10(M_c_value))**2/(2*sigma_c**2)) - -@njit(fastmath=True) -def M_c(M_h, M_1, M_0, gamma1, gamma2): - """ - Eq. (37) from Cacciato et al. (2008) - """ - return M_0 * (M_h/M_1)**gamma1/(1+M_h/M_1)**(gamma1-gamma2) - - - -@njit(fastmath=True) -def get_random_cen_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): - - nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - - CSMF_cen = CSMF_centrals(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c) - - cdf = [] - cdf_current_value = 0 - for k in range(len(delta_stellar_masses)): - cdf_current_value = cdf_current_value + CSMF_cen[k]*delta_stellar_masses[k] - cdf.append(cdf_current_value) - cdf=np.array(cdf) - - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) - bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) - - - - -@njit(fastmath=True) -def get_random_cen_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): - - nbins = 1000 - # stellarmass = np.linspace(Mstar_low,Mstar_up,nbins) - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - - CSMF_cen = CSMF_centrals(M_h, stellarmass, M_1, M_0, gamma1, gamma2, sigma_c) - - delta=stellarmass[1:]-stellarmass[:-1] - - cdf = [] - cdf_current_value = 0 - for i in range(len(delta)): - cdf_current_value = cdf_current_value + CSMF_cen[i]*delta[i] - cdf.append(cdf_current_value) - cdf=np.array(cdf) - - # cdf = np.cumsum(CSMF_sat*delta) - - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) - bin = np.where(cdf>random_rv)[0][0] - - m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) - return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] - - -@njit(fastmath=True) -def n_sat_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c, a1, a2, M2, b0, b1, b2, delta1, delta2): - """ - Standard Cacciato et al. (2008) satellite HOD parametrization for CSMF - """ - nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - - CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) - # delta=stellarmass[1:]-stellarmass[:-1] - # nsat = np.sum(CSMF_sat[:-1]*delta) - - # ncen = n_cen_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c) - nsat = 0 - for i in range(nbins-1): - nsat += (CSMF_sat[i+1]-CSMF_sat[i])*(stellarmass[i+1]-stellarmass[i])/2 + (stellarmass[i+1]-stellarmass[i])*CSMF_sat[i] - - return nsat#*ncen - - -@njit(fastmath=True) -def CSMF_satelites(M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): - """ - Eq. (36) from Cacciato et al. (2008) - """ - M_s_value = M_s(M_h, M_1, M_0, gamma1, gamma2) - alpha_s_value = alpha_s(M_h, a1, a2, M2) - phi_s_value = phi_s(M_h, b0, b1, b2) - - delta = 10 ** (delta1 + delta2 * (np.log10(M_h) - 12)) - - return phi_s_value/M_s_value * (Mstar/M_s_value)**alpha_s_value * np.exp(-delta*(Mstar/M_s_value)**2) - -@njit(fastmath=True) -def M_s(M_h, M_1, M_0, gamma1, gamma2): - """ - Eq. (38) from Cacciato et al. (2008) - """ - return 0.562 * M_c(M_h, M_1, M_0, gamma1, gamma2) - -@njit(fastmath=True) -def alpha_s(M_h, a1, a2, M2): - """ - Eq. (39) from Cacciato et al. (2008) - """ - return -2.0 + a1 * (1-2/np.pi*np.arctan(a2*np.log10(M_h/M2))) - -@njit(fastmath=True) -def phi_s(M_h, b0, b1, b2): - """ - Eq. (40) from Cacciato et al. (2008) - """ - M12 = M_h/1e12 - log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12)**2 - return 10**log_phi_s - - -@njit(fastmath=True) -def get_random_sat_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): - - nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - - CSMF_sat = CSMF_satelites(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) - - - cdf = [] - cdf_current_value = 0 - for k in range(len(delta_stellar_masses)): - cdf_current_value = cdf_current_value + CSMF_sat[k]*delta_stellar_masses[k] - cdf.append(cdf_current_value) - cdf=np.array(cdf) - - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) - bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) - - -@njit(fastmath=True) -def get_random_sat_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): - - nbins = 100 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - - CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) - - delta=stellarmass[1:]-stellarmass[:-1] - - cdf = [] - cdf_current_value = 0 - for i in range(len(delta)): - cdf_current_value = cdf_current_value + CSMF_sat[i]*delta[i] - cdf.append(cdf_current_value) - cdf=np.array(cdf) - - # cdf = np.cumsum(CSMF_sat*delta) - - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) - bin = np.where(cdf>random_rv)[0][0] - - m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) - return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] - - @njit(fastmath=True) def n_sat_LRG_modified(M_h, logM_cut, M_cut, M_1, sigma, alpha, kappa): """ Standard Zheng et al. (2005) satellite HOD parametrization for LRGs, modified with n_cent_LRG """ - if M_h - kappa*M_cut < 0: + if M_h - kappa * M_cut < 0: return 0 - return ((M_h - kappa*M_cut)/M_1)**alpha*0.5*math.erfc((logM_cut - np.log10(M_h))/(1.41421356*sigma)) + return ( + ((M_h - kappa * M_cut) / M_1) ** alpha + * 0.5 + * math.erfc((logM_cut - np.log10(M_h)) / (1.41421356 * sigma)) + ) @njit(fastmath=True) @@ -231,36 +38,44 @@ def n_cen_LRG(M_h, logM_cut, sigma): """ Standard Zheng et al. (2005) central HOD parametrization for LRGs. """ - return 0.5*math.erfc((logM_cut - np.log10(M_h))/(1.41421356*sigma)) + return 0.5 * math.erfc((logM_cut - np.log10(M_h)) / (1.41421356 * sigma)) + @njit(fastmath=True) -def N_sat_generic(M_h, M_cut, kappa, M_1, alpha, A_s=1.): +def N_sat_generic(M_h, M_cut, kappa, M_1, alpha, A_s=1.0): """ Standard Zheng et al. (2005) satellite HOD parametrization for all tracers with an optional amplitude parameter, A_s. """ - if M_h - kappa*M_cut < 0: + if M_h - kappa * M_cut < 0: return 0 - return A_s*((M_h-kappa*M_cut)/M_1)**alpha + return A_s * ((M_h - kappa * M_cut) / M_1) ** alpha + @njit(fastmath=True) -def N_sat_elg(M_h, M_cut, kappa, M_1, alpha, A_s=1., alpha1 = 0., beta = 0.): +def N_sat_elg(M_h, M_cut, kappa, M_1, alpha, A_s=1.0, alpha1=0.0, beta=0.0): """ Standard power law modulated by an exponential fall off at small M """ # return (M_h/M_1)**alpha/(1+np.exp(-A_s*(np.log10(M_h)-np.log10(kappa*M_cut)))) + beta*(M_h/M_1)**(-alpha1)/100 - if M_h - kappa*M_cut < 0: + if M_h - kappa * M_cut < 0: return 0 - return A_s*((M_h-kappa*M_cut)/M_1)**alpha # + beta*(M_h/M_1)**(-alpha1)/100 + return ( + A_s * ((M_h - kappa * M_cut) / M_1) ** alpha + ) # + beta*(M_h/M_1)**(-alpha1)/100 + @njit(fastmath=True) -def N_cen_ELG_v1(M_h, p_max, Q, logM_cut, sigma, gamma, Anorm = 1): +def N_cen_ELG_v1(M_h, p_max, Q, logM_cut, sigma, gamma, Anorm=1): """ HOD function for ELG centrals taken from arXiv:1910.05095. """ logM_h = np.log10(M_h) phi = phi_fun(logM_h, logM_cut, sigma) Phi = Phi_fun(logM_h, logM_cut, sigma, gamma) - return 2.*(p_max-1./Q)*phi*Phi/Anorm # + 0.5/Q*(1 + math.erf((logM_h-logM_cut-0.8)*3)) + return ( + 2.0 * (p_max - 1.0 / Q) * phi * Phi / Anorm + ) # + 0.5/Q*(1 + math.erf((logM_h-logM_cut-0.8)*3)) + @njit(fastmath=True) def N_cen_ELG_v2(M_h, p_max, logM_cut, sigma, gamma): @@ -269,16 +84,17 @@ def N_cen_ELG_v2(M_h, p_max, logM_cut, sigma, gamma): """ logM_h = np.log10(M_h) if logM_h <= logM_cut: - return p_max*Gaussian_fun(logM_h, logM_cut, sigma) + return p_max * Gaussian_fun(logM_h, logM_cut, sigma) else: - return p_max*(M_h/10**logM_cut)**gamma/(2.5066283*sigma) + return p_max * (M_h / 10**logM_cut) ** gamma / (2.5066283 * sigma) + @njit(fastmath=True) def N_cen_QSO(M_h, logM_cut, sigma): """ HOD function (Zheng et al. (2005) with p_max) for QSO centrals taken from arXiv:2007.09012. """ - return 0.5*(1 + math.erf((np.log10(M_h)-logM_cut)/1.41421356/sigma)) + return 0.5 * (1 + math.erf((np.log10(M_h) - logM_cut) / 1.41421356 / sigma)) @njit(fastmath=True) @@ -289,37 +105,60 @@ def phi_fun(logM_h, logM_cut, sigma): phi = Gaussian_fun(logM_h, logM_cut, sigma) return phi + @njit(fastmath=True) def Phi_fun(logM_h, logM_cut, sigma, gamma): """ Aiding function for N_cen_ELG_v1(). """ - x = gamma*(logM_h-logM_cut)/sigma - Phi = 0.5*(1 + math.erf(x/np.sqrt(2))) + x = gamma * (logM_h - logM_cut) / sigma + Phi = 0.5 * (1 + math.erf(x / np.sqrt(2))) return Phi + @njit(fastmath=True) def Gaussian_fun(x, mean, sigma): """ Gaussian function with centered at `mean' with standard deviation `sigma'. """ - return 0.3989422804014327/sigma*np.exp(-(x - mean)**2/2/sigma**2) + return 0.3989422804014327 / sigma * np.exp(-((x - mean) ** 2) / 2 / sigma**2) @njit(fastmath=True) def wrap(x, L): - '''Fast scalar mod implementation''' - L2 = L/2 + """Fast scalar mod implementation""" + L2 = L / 2 if x >= L2: return x - L elif x < -L2: return x + L return x + @njit(parallel=True, fastmath=True) -def gen_cent(pos, vel, mass, ids, multis, randoms, vdev, deltac, fenv, shear, - LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, CSMF_hod_dict, - rsd, inv_velz2kms, lbox, want_LRG, want_ELG, want_QSO, want_CSMF, Nthread, origin): +def gen_cent( + pos, + vel, + mass, + ids, + multis, + randoms, + vdev, + deltac, + fenv, + shear, + LRG_hod_dict, + ELG_hod_dict, + QSO_hod_dict, + rsd, + inv_velz2kms, + lbox, + want_LRG, + want_ELG, + want_QSO, + Nthread, + origin, +): """ Generate central galaxies in place in memory with a two pass numba parallel implementation. """ @@ -327,30 +166,47 @@ def gen_cent(pos, vel, mass, ids, multis, randoms, vdev, deltac, fenv, shear, if want_LRG: # parse out the hod parameters logM_cut_L, sigma_L = LRG_hod_dict['logM_cut'], LRG_hod_dict['sigma'] - ic_L, alpha_c_L, Ac_L, Bc_L = LRG_hod_dict['ic'], LRG_hod_dict['alpha_c'], LRG_hod_dict['Acent'], LRG_hod_dict['Bcent'] + ic_L, alpha_c_L, Ac_L, Bc_L = ( + LRG_hod_dict['ic'], + LRG_hod_dict['alpha_c'], + LRG_hod_dict['Acent'], + LRG_hod_dict['Bcent'], + ) if want_ELG: - pmax_E, Q_E, logM_cut_E, sigma_E, gamma_E = \ - ELG_hod_dict['p_max'], ELG_hod_dict['Q'], ELG_hod_dict['logM_cut'], ELG_hod_dict['sigma'], ELG_hod_dict['gamma'] - alpha_c_E, Ac_E, Bc_E, Cc_E, ic_E = ELG_hod_dict['alpha_c'], ELG_hod_dict['Acent'], ELG_hod_dict['Bcent'], ELG_hod_dict['Ccent'], ELG_hod_dict['ic'] + pmax_E, Q_E, logM_cut_E, sigma_E, gamma_E = ( + ELG_hod_dict['p_max'], + ELG_hod_dict['Q'], + ELG_hod_dict['logM_cut'], + ELG_hod_dict['sigma'], + ELG_hod_dict['gamma'], + ) + alpha_c_E, Ac_E, Bc_E, Cc_E, ic_E = ( + ELG_hod_dict['alpha_c'], + ELG_hod_dict['Acent'], + ELG_hod_dict['Bcent'], + ELG_hod_dict['Ccent'], + ELG_hod_dict['ic'], + ) if want_QSO: logM_cut_Q, sigma_Q = QSO_hod_dict['logM_cut'], QSO_hod_dict['sigma'] - alpha_c_Q, Ac_Q, Bc_Q, ic_Q = QSO_hod_dict['alpha_c'], QSO_hod_dict['Acent'], QSO_hod_dict['Bcent'], QSO_hod_dict['ic'] - - if want_CSMF: - Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C = CSMF_hod_dict['Mstar_low'], CSMF_hod_dict['Mstar_up'], CSMF_hod_dict['M_1'], CSMF_hod_dict['M_0'], CSMF_hod_dict['gamma_1'], CSMF_hod_dict['gamma_2'], CSMF_hod_dict['sigma_c'] - ic_C, alpha_c_C, Ac_C, Bc_C = CSMF_hod_dict['ic'], CSMF_hod_dict['alpha_c'], CSMF_hod_dict['Acent'], CSMF_hod_dict['Bcent'] - - # print(Ac_C,Bc_C) + alpha_c_Q, Ac_Q, Bc_Q, ic_Q = ( + QSO_hod_dict['alpha_c'], + QSO_hod_dict['Acent'], + QSO_hod_dict['Bcent'], + QSO_hod_dict['ic'], + ) H = len(mass) - + numba.set_num_threads(Nthread) - Nout = np.zeros((Nthread, 4, 8), dtype = np.int64) - hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype(np.int64) # starting index of each thread + Nout = np.zeros((Nthread, 3, 8), dtype=np.int64) + hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype( + np.int64 + ) # starting index of each thread - keep = np.empty(H, dtype = np.int8) # mask array tracking which halos to keep + keep = np.empty(H, dtype=np.int8) # mask array tracking which halos to keep # figuring out the number of halos kept for each thread for tid in numba.prange(Nthread): @@ -360,210 +216,174 @@ def gen_cent(pos, vel, mass, ids, multis, randoms, vdev, deltac, fenv, shear, if want_LRG: # do assembly bias and secondary bias logM_cut_L_temp = logM_cut_L + Ac_L * deltac[i] + Bc_L * fenv[i] - LRG_marker += n_cen_LRG(mass[i], logM_cut_L_temp, sigma_L) * ic_L * multis[i] + LRG_marker += ( + n_cen_LRG(mass[i], logM_cut_L_temp, sigma_L) * ic_L * multis[i] + ) ELG_marker = LRG_marker if want_ELG: - logM_cut_E_temp = logM_cut_E + Ac_E * deltac[i] + Bc_E * fenv[i] + Cc_E * shear[i] - ELG_marker += N_cen_ELG_v1(mass[i], pmax_E, Q_E, logM_cut_E_temp, sigma_E, gamma_E) * ic_E * multis[i] + logM_cut_E_temp = ( + logM_cut_E + Ac_E * deltac[i] + Bc_E * fenv[i] + Cc_E * shear[i] + ) + ELG_marker += ( + N_cen_ELG_v1( + mass[i], pmax_E, Q_E, logM_cut_E_temp, sigma_E, gamma_E + ) + * ic_E + * multis[i] + ) QSO_marker = ELG_marker if want_QSO: logM_cut_Q_temp = logM_cut_Q + Ac_Q * deltac[i] + Bc_Q * fenv[i] - QSO_marker += N_cen_QSO(mass[i], logM_cut_Q_temp, sigma_Q) * ic_Q * multis[i] - CSMF_marker = QSO_marker - if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) - - ncen = n_cen_CSMF(mass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) - - CSMF_marker += ncen * ic_C * multis[i] - + QSO_marker += ( + N_cen_QSO(mass[i], logM_cut_Q_temp, sigma_Q) * ic_Q * multis[i] + ) + if randoms[i] <= LRG_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 elif randoms[i] <= ELG_marker: - Nout[tid, 1, 0] += 1 # counting + Nout[tid, 1, 0] += 1 # counting keep[i] = 2 elif randoms[i] <= QSO_marker: - Nout[tid, 2, 0] += 1 # counting + Nout[tid, 2, 0] += 1 # counting keep[i] = 3 - elif randoms[i] <= CSMF_marker: - Nout[tid, 3, 0] += 1 # counting - keep[i] = 4 else: keep[i] = 0 - + # compose galaxy array, first create array of galaxy starting indices for the threads - gstart = np.empty((Nthread + 1, 4), dtype = np.int64) + gstart = np.empty((Nthread + 1, 3), dtype=np.int64) gstart[0, :] = 0 gstart[1:, 0] = Nout[:, 0, 0].cumsum() gstart[1:, 1] = Nout[:, 1, 0].cumsum() gstart[1:, 2] = Nout[:, 2, 0].cumsum() - gstart[1:, 3] = Nout[:, 3, 0].cumsum() # galaxy arrays N_lrg = gstart[-1, 0] - lrg_x = np.empty(N_lrg, dtype = mass.dtype) - lrg_y = np.empty(N_lrg, dtype = mass.dtype) - lrg_z = np.empty(N_lrg, dtype = mass.dtype) - lrg_vx = np.empty(N_lrg, dtype = mass.dtype) - lrg_vy = np.empty(N_lrg, dtype = mass.dtype) - lrg_vz = np.empty(N_lrg, dtype = mass.dtype) - lrg_mass = np.empty(N_lrg, dtype = mass.dtype) - lrg_id = np.empty(N_lrg, dtype = ids.dtype) + lrg_x = np.empty(N_lrg, dtype=mass.dtype) + lrg_y = np.empty(N_lrg, dtype=mass.dtype) + lrg_z = np.empty(N_lrg, dtype=mass.dtype) + lrg_vx = np.empty(N_lrg, dtype=mass.dtype) + lrg_vy = np.empty(N_lrg, dtype=mass.dtype) + lrg_vz = np.empty(N_lrg, dtype=mass.dtype) + lrg_mass = np.empty(N_lrg, dtype=mass.dtype) + lrg_id = np.empty(N_lrg, dtype=ids.dtype) # galaxy arrays N_elg = gstart[-1, 1] - elg_x = np.empty(N_elg, dtype = mass.dtype) - elg_y = np.empty(N_elg, dtype = mass.dtype) - elg_z = np.empty(N_elg, dtype = mass.dtype) - elg_vx = np.empty(N_elg, dtype = mass.dtype) - elg_vy = np.empty(N_elg, dtype = mass.dtype) - elg_vz = np.empty(N_elg, dtype = mass.dtype) - elg_mass = np.empty(N_elg, dtype = mass.dtype) - elg_id = np.empty(N_elg, dtype = ids.dtype) + elg_x = np.empty(N_elg, dtype=mass.dtype) + elg_y = np.empty(N_elg, dtype=mass.dtype) + elg_z = np.empty(N_elg, dtype=mass.dtype) + elg_vx = np.empty(N_elg, dtype=mass.dtype) + elg_vy = np.empty(N_elg, dtype=mass.dtype) + elg_vz = np.empty(N_elg, dtype=mass.dtype) + elg_mass = np.empty(N_elg, dtype=mass.dtype) + elg_id = np.empty(N_elg, dtype=ids.dtype) # galaxy arrays N_qso = gstart[-1, 2] - qso_x = np.empty(N_qso, dtype = mass.dtype) - qso_y = np.empty(N_qso, dtype = mass.dtype) - qso_z = np.empty(N_qso, dtype = mass.dtype) - qso_vx = np.empty(N_qso, dtype = mass.dtype) - qso_vy = np.empty(N_qso, dtype = mass.dtype) - qso_vz = np.empty(N_qso, dtype = mass.dtype) - qso_mass = np.empty(N_qso, dtype = mass.dtype) - qso_id = np.empty(N_qso, dtype = ids.dtype) - - # galaxy arrays - N_CSMF = gstart[-1, 3] - CSMF_x = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_y = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_z = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_id = np.empty(N_CSMF, dtype = ids.dtype) + qso_x = np.empty(N_qso, dtype=mass.dtype) + qso_y = np.empty(N_qso, dtype=mass.dtype) + qso_z = np.empty(N_qso, dtype=mass.dtype) + qso_vx = np.empty(N_qso, dtype=mass.dtype) + qso_vy = np.empty(N_qso, dtype=mass.dtype) + qso_vz = np.empty(N_qso, dtype=mass.dtype) + qso_mass = np.empty(N_qso, dtype=mass.dtype) + qso_id = np.empty(N_qso, dtype=ids.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): - j1, j2, j3, j4 = gstart[tid] + j1, j2, j3 = gstart[tid] for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: # loop thru three directions to assign galaxy velocities and positions - lrg_x[j1] = pos[i,0] - lrg_vx[j1] = vel[i,0] + alpha_c_L * vdev[i,0] # velocity bias - lrg_y[j1] = pos[i,1] - lrg_vy[j1] = vel[i,1] + alpha_c_L * vdev[i,1] # velocity bias - lrg_z[j1] = pos[i,2] - lrg_vz[j1] = vel[i,2] + alpha_c_L * vdev[i,2] # velocity bias + lrg_x[j1] = pos[i, 0] + lrg_vx[j1] = vel[i, 0] + alpha_c_L * vdev[i, 0] # velocity bias + lrg_y[j1] = pos[i, 1] + lrg_vy[j1] = vel[i, 1] + alpha_c_L * vdev[i, 1] # velocity bias + lrg_z[j1] = pos[i, 2] + lrg_vz[j1] = vel[i, 2] + alpha_c_L * vdev[i, 2] # velocity bias # rsd only applies to the z direction if rsd and origin is not None: nx = lrg_x[j1] - origin[0] ny = lrg_y[j1] - origin[1] nz = lrg_z[j1] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms * (lrg_vx[j1]*nx + lrg_vy[j1]*ny + lrg_vz[j1]*nz) - lrg_x[j1] = lrg_x[j1]+proj*nx - lrg_y[j1] = lrg_y[j1]+proj*ny - lrg_z[j1] = lrg_z[j1]+proj*nz + proj = inv_velz2kms * ( + lrg_vx[j1] * nx + lrg_vy[j1] * ny + lrg_vz[j1] * nz + ) + lrg_x[j1] = lrg_x[j1] + proj * nx + lrg_y[j1] = lrg_y[j1] + proj * ny + lrg_z[j1] = lrg_z[j1] + proj * nz elif rsd: - lrg_z[j1] = wrap(pos[i,2] + lrg_vz[j1] * inv_velz2kms, lbox) + lrg_z[j1] = wrap(pos[i, 2] + lrg_vz[j1] * inv_velz2kms, lbox) lrg_mass[j1] = mass[i] lrg_id[j1] = ids[i] j1 += 1 elif keep[i] == 2: # loop thru three directions to assign galaxy velocities and positions - elg_x[j2] = pos[i,0] - elg_vx[j2] = vel[i,0] + alpha_c_E * vdev[i,0] # velocity bias - elg_y[j2] = pos[i,1] - elg_vy[j2] = vel[i,1] + alpha_c_E * vdev[i,1] # velocity bias - elg_z[j2] = pos[i,2] - elg_vz[j2] = vel[i,2] + alpha_c_E * vdev[i,2] # velocity bias + elg_x[j2] = pos[i, 0] + elg_vx[j2] = vel[i, 0] + alpha_c_E * vdev[i, 0] # velocity bias + elg_y[j2] = pos[i, 1] + elg_vy[j2] = vel[i, 1] + alpha_c_E * vdev[i, 1] # velocity bias + elg_z[j2] = pos[i, 2] + elg_vz[j2] = vel[i, 2] + alpha_c_E * vdev[i, 2] # velocity bias # rsd only applies to the z direction if rsd and origin is not None: nx = elg_x[j2] - origin[0] ny = elg_y[j2] - origin[1] nz = elg_z[j2] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms*(elg_vx[j2]*nx+elg_vy[j2]*ny+elg_vz[j2]*nz) - elg_x[j2] = elg_x[j2]+proj*nx - elg_y[j2] = elg_y[j2]+proj*ny - elg_z[j2] = elg_z[j2]+proj*nz + proj = inv_velz2kms * ( + elg_vx[j2] * nx + elg_vy[j2] * ny + elg_vz[j2] * nz + ) + elg_x[j2] = elg_x[j2] + proj * nx + elg_y[j2] = elg_y[j2] + proj * ny + elg_z[j2] = elg_z[j2] + proj * nz elif rsd: - elg_z[j2] = wrap(pos[i,2] + elg_vz[j2] * inv_velz2kms, lbox) + elg_z[j2] = wrap(pos[i, 2] + elg_vz[j2] * inv_velz2kms, lbox) elg_mass[j2] = mass[i] elg_id[j2] = ids[i] j2 += 1 elif keep[i] == 3: # loop thru three directions to assign galaxy velocities and positions - qso_x[j3] = pos[i,0] - qso_vx[j3] = vel[i,0] + alpha_c_Q * vdev[i,0] # velocity bias - qso_y[j3] = pos[i,1] - qso_vy[j3] = vel[i,1] + alpha_c_Q * vdev[i,1] # velocity bias - qso_z[j3] = pos[i,2] - qso_vz[j3] = vel[i,2] + alpha_c_Q * vdev[i,2] # velocity bias + qso_x[j3] = pos[i, 0] + qso_vx[j3] = vel[i, 0] + alpha_c_Q * vdev[i, 0] # velocity bias + qso_y[j3] = pos[i, 1] + qso_vy[j3] = vel[i, 1] + alpha_c_Q * vdev[i, 1] # velocity bias + qso_z[j3] = pos[i, 2] + qso_vz[j3] = vel[i, 2] + alpha_c_Q * vdev[i, 2] # velocity bias # rsd only applies to the z direction if rsd and origin is not None: nx = qso_x[j3] - origin[0] ny = qso_y[j3] - origin[1] nz = qso_z[j3] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms*(qso_vx[j3]*nx+qso_vy[j3]*ny+qso_vz[j3]*nz) - qso_x[j3] = qso_x[j3]+proj*nx - qso_y[j3] = qso_y[j3]+proj*ny - qso_z[j3] = qso_z[j3]+proj*nz + proj = inv_velz2kms * ( + qso_vx[j3] * nx + qso_vy[j3] * ny + qso_vz[j3] * nz + ) + qso_x[j3] = qso_x[j3] + proj * nx + qso_y[j3] = qso_y[j3] + proj * ny + qso_z[j3] = qso_z[j3] + proj * nz elif rsd: - qso_z[j3] = wrap(pos[i,2] + qso_vz[j3] * inv_velz2kms, lbox) + qso_z[j3] = wrap(pos[i, 2] + qso_vz[j3] * inv_velz2kms, lbox) qso_mass[j3] = mass[i] qso_id[j3] = ids[i] j3 += 1 - elif keep[i] == 4: - # loop thru three directions to assign galaxy velocities and positions - CSMF_x[j4] = pos[i,0] - CSMF_vx[j4] = vel[i,0] + alpha_c_C * vdev[i,0] # velocity bias - CSMF_y[j4] = pos[i,1] - CSMF_vy[j4] = vel[i,1] + alpha_c_C * vdev[i,1] # velocity bias - CSMF_z[j4] = pos[i,2] - CSMF_vz[j4] = vel[i,2] + alpha_c_C * vdev[i,2] # velocity bias - - # rsd only applies to the z direction - if rsd and origin is not None: - nx = CSMF_x[j4] - origin[0] - ny = CSMF_y[j4] - origin[1] - nz = CSMF_z[j4] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) - nx *= inv_norm - ny *= inv_norm - nz *= inv_norm - proj = inv_velz2kms*(CSMF_vx[j4]*nx+CSMF_vy[j4]*ny+CSMF_vz[j4]*nz) - CSMF_x[j4] = CSMF_x[j4]+proj*nx - CSMF_y[j4] = CSMF_y[j4]+proj*ny - CSMF_z[j4] = CSMF_z[j4]+proj*nz - elif rsd: - CSMF_z[j4] = wrap(pos[i,2] + CSMF_vz[j4] * inv_velz2kms, lbox) - - CSMF_mass[j4] = mass[i] - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) - CSMF_stellarmass[j4] = get_random_cen_stellarmass_linearinterpolation(mass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) - CSMF_id[j4] = ids[i] - j4 += 1 # assert j == gstart[tid + 1] - LRG_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) - ELG_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) - QSO_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) - ID_dict = Dict.empty(key_type = types.unicode_type, value_type = int_array) + LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) LRG_dict['x'] = lrg_x LRG_dict['y'] = lrg_y LRG_dict['z'] = lrg_z @@ -590,27 +410,18 @@ def gen_cent(pos, vel, mass, ids, multis, randoms, vdev, deltac, fenv, shear, QSO_dict['vz'] = qso_vz QSO_dict['mass'] = qso_mass ID_dict['QSO'] = qso_id - - CSMF_dict['x'] = CSMF_x - CSMF_dict['y'] = CSMF_y - CSMF_dict['z'] = CSMF_z - CSMF_dict['vx'] = CSMF_vx - CSMF_dict['vy'] = CSMF_vy - CSMF_dict['vz'] = CSMF_vz - CSMF_dict['mass'] = CSMF_mass - CSMF_dict['stellarmass'] = CSMF_stellarmass - ID_dict['CSMF'] = CSMF_id - return LRG_dict, ELG_dict, QSO_dict, CSMF_dict, ID_dict, keep + return LRG_dict, ELG_dict, QSO_dict, ID_dict, keep + @njit(parallel=True, fastmath=True) def getPointsOnSphere(nPoints, Nthread, seed=None): - ''' + """ --- Aiding function for NFW computation, generate random points in a sphere - ''' + """ numba.set_num_threads(Nthread) ind = min(Nthread, nPoints) # starting index of each thread - hstart = np.rint(np.linspace(0, nPoints, ind+1)) + hstart = np.rint(np.linspace(0, nPoints, ind + 1)) ur = np.zeros((nPoints, 3), dtype=np.float64) cmin = -1 cmax = +1 @@ -620,19 +431,38 @@ def getPointsOnSphere(nPoints, Nthread, seed=None): np.random.seed(seed[tid]) for i in range(hstart[tid], hstart[tid + 1]): u1, u2 = np.random.uniform(0, 1), np.random.uniform(0, 1) - ra = 0 + u1*(2*np.pi-0) - dec = np.pi - (np.arccos(cmin+u2*(cmax-cmin))) + ra = 0 + u1 * (2 * np.pi - 0) + dec = np.pi - (np.arccos(cmin + u2 * (cmax - cmin))) ur[i, 0] = np.sin(dec) * np.cos(ra) ur[i, 1] = np.sin(dec) * np.sin(ra) ur[i, 2] = np.cos(dec) return ur -@njit(fastmath=True, parallel=True) # parallel=True, -def compute_fast_NFW(NFW_draw, h_id, x_h, y_h, z_h, vx_h, vy_h, vz_h, vrms_h, c, M, Rvir, - rd_pos, num_sat, f_sigv, vel_sat = 'rd_normal', Nthread = 16, - exp_frac=0, exp_scale=1, nfw_rescale=1): +@njit(fastmath=True, parallel=True) # parallel=True, +def compute_fast_NFW( + NFW_draw, + h_id, + x_h, + y_h, + z_h, + vx_h, + vy_h, + vz_h, + vrms_h, + c, + M, + Rvir, + rd_pos, + num_sat, + f_sigv, + vel_sat='rd_normal', + Nthread=16, + exp_frac=0, + exp_scale=1, + nfw_rescale=1, +): """ --- Compute NFW positions and velocities for satelitte galaxies c: r98/r25 @@ -663,36 +493,57 @@ def compute_fast_NFW(NFW_draw, h_id, x_h, y_h, z_h, vx_h, vy_h, vz_h, vrms_h, c, for tid in numba.prange(Nthread): for i in range(int(hstart[tid]), int(hstart[tid + 1])): ind = i - #while (NFW_draw[ind] > c[i]): + # while (NFW_draw[ind] > c[i]): # ind = np.random.randint(0, len(NFW_draw)) - #etaVir = NFW_draw[ind]/c[i] # =r/rvir - if np.random.uniform(0,1) c[i]): + while NFW_draw[ind] > c[i]: ind = np.random.randint(0, len(NFW_draw)) - etaVir = NFW_draw[ind]/c[i]*nfw_rescale + etaVir = NFW_draw[ind] / c[i] * nfw_rescale p = etaVir * Rvir[i] x_sat[i] = x_h[i] + rd_pos[i, 0] * p y_sat[i] = y_h[i] + rd_pos[i, 1] * p z_sat[i] = z_h[i] + rd_pos[i, 2] * p - # if vel_sat == 'rd_normal': - sig = vrms_h[i]*0.577*f_sigv - vx_sat[i] = np.random.normal(loc=vx_h[i], scale=sig) - vy_sat[i] = np.random.normal(loc=vy_h[i], scale=sig) - vz_sat[i] = np.random.normal(loc=vz_h[i], scale=sig) - # else: - # raise ValueError( - # 'Wrong vel_sat argument only "rd_normal"') + if vel_sat == 'rd_normal': + sig = vrms_h[i] * 0.577 * f_sigv + vx_sat[i] = np.random.normal(loc=vx_h[i], scale=sig) + vy_sat[i] = np.random.normal(loc=vy_h[i], scale=sig) + vz_sat[i] = np.random.normal(loc=vz_h[i], scale=sig) + else: + raise ValueError('Wrong vel_sat argument only "rd_normal"') return h_id, x_sat, y_sat, z_sat, vx_sat, vy_sat, vz_sat, M -@njit(fastmath = True, parallel=True) -def gen_sats_nfw(NFW_draw, hpos, hvel, hmass, hid, hdeltac, hfenv, hshear, hvrms, hc, hrvir, - LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, CSMF_hod_dict, want_LRG, want_ELG, want_QSO, want_CSMF, - rsd, inv_velz2kms, lbox, keep_cent, vel_sat = 'rd_normal', Nthread = 16): + +@njit(fastmath=True, parallel=True) +def gen_sats_nfw( + NFW_draw, + hpos, + hvel, + hmass, + hid, + hdeltac, + hfenv, + hshear, + hvrms, + hc, + hrvir, + LRG_hod_dict, + ELG_hod_dict, + QSO_hod_dict, + want_LRG, + want_ELG, + want_QSO, + rsd, + inv_velz2kms, + lbox, + keep_cent, + vel_sat='rd_normal', + Nthread=16, +): """ Generate satellite galaxies on an NFW profile, with option for an extended profile. See Rocher et al. 2023. @@ -701,130 +552,245 @@ def gen_sats_nfw(NFW_draw, hpos, hvel, hmass, hid, hdeltac, hfenv, hshear, hvrms """ if want_LRG: - logM_cut_L, logM1_L, sigma_L, alpha_L, kappa_L = \ - LRG_hod_dict['logM_cut'], LRG_hod_dict['logM1'], LRG_hod_dict['sigma'], LRG_hod_dict['alpha'], LRG_hod_dict['kappa'] - Ac_L, As_L, Bc_L, Bs_L, ic_L = \ - LRG_hod_dict['Acent'], LRG_hod_dict['Asat'], LRG_hod_dict['Bcent'], \ - LRG_hod_dict['Bsat'], LRG_hod_dict['ic'] + logM_cut_L, logM1_L, sigma_L, alpha_L, kappa_L = ( + LRG_hod_dict['logM_cut'], + LRG_hod_dict['logM1'], + LRG_hod_dict['sigma'], + LRG_hod_dict['alpha'], + LRG_hod_dict['kappa'], + ) + Ac_L, As_L, Bc_L, Bs_L, ic_L = ( + LRG_hod_dict['Acent'], + LRG_hod_dict['Asat'], + LRG_hod_dict['Bcent'], + LRG_hod_dict['Bsat'], + LRG_hod_dict['ic'], + ) f_sigv_L = LRG_hod_dict['f_sigv'] if want_ELG: - logM_cut_E, kappa_E, logM1_E, alpha_E, A_E = \ - ELG_hod_dict['logM_cut'], ELG_hod_dict['kappa'], ELG_hod_dict['logM1'], ELG_hod_dict['alpha'], ELG_hod_dict['A_s'] - Ac_E, As_E, Bc_E, Bs_E, Cc_E, Cs_E, ic_E, = \ - ELG_hod_dict['Acent'], ELG_hod_dict['Asat'], ELG_hod_dict['Bcent'], \ - ELG_hod_dict['Bsat'], ELG_hod_dict['Ccent'], ELG_hod_dict['Csat'], ELG_hod_dict['ic'] - logM1_EE, alpha_EE, logM1_EL, alpha_EL = \ - ELG_hod_dict['logM1_EE'], ELG_hod_dict['alpha_EE'], \ - ELG_hod_dict['logM1_EL'], ELG_hod_dict['alpha_EL'] + logM_cut_E, kappa_E, logM1_E, alpha_E, A_E = ( + ELG_hod_dict['logM_cut'], + ELG_hod_dict['kappa'], + ELG_hod_dict['logM1'], + ELG_hod_dict['alpha'], + ELG_hod_dict['A_s'], + ) + ( + Ac_E, + As_E, + Bc_E, + Bs_E, + Cc_E, + Cs_E, + ic_E, + ) = ( + ELG_hod_dict['Acent'], + ELG_hod_dict['Asat'], + ELG_hod_dict['Bcent'], + ELG_hod_dict['Bsat'], + ELG_hod_dict['Ccent'], + ELG_hod_dict['Csat'], + ELG_hod_dict['ic'], + ) + logM1_EE, alpha_EE, logM1_EL, alpha_EL = ( + ELG_hod_dict['logM1_EE'], + ELG_hod_dict['alpha_EE'], + ELG_hod_dict['logM1_EL'], + ELG_hod_dict['alpha_EL'], + ) f_sigv_E = ELG_hod_dict['f_sigv'] exp_frac = ELG_hod_dict['exp_frac'] exp_scale = ELG_hod_dict['exp_scale'] nfw_rescale = ELG_hod_dict['nfw_rescale'] if want_QSO: - logM_cut_Q, kappa_Q, logM1_Q, alpha_Q = \ - QSO_hod_dict['logM_cut'], QSO_hod_dict['kappa'], QSO_hod_dict['logM1'], QSO_hod_dict['alpha'] - Ac_Q, As_Q, Bc_Q, Bs_Q, ic_Q = \ - QSO_hod_dict['Acent'], QSO_hod_dict['Asat'], QSO_hod_dict['Bcent'], \ - QSO_hod_dict['Bsat'], QSO_hod_dict['ic'] + logM_cut_Q, kappa_Q, logM1_Q, alpha_Q = ( + QSO_hod_dict['logM_cut'], + QSO_hod_dict['kappa'], + QSO_hod_dict['logM1'], + QSO_hod_dict['alpha'], + ) + Ac_Q, As_Q, Bc_Q, Bs_Q, ic_Q = ( + QSO_hod_dict['Acent'], + QSO_hod_dict['Asat'], + QSO_hod_dict['Bcent'], + QSO_hod_dict['Bsat'], + QSO_hod_dict['ic'], + ) f_sigv_Q = QSO_hod_dict['f_sigv'] - - if want_CSMF: - Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = \ - CSMF_hod_dict['Mstar_low'], CSMF_hod_dict['Mstar_up'], CSMF_hod_dict['M_1'], CSMF_hod_dict['M_0'], CSMF_hod_dict['gamma_1'], \ - CSMF_hod_dict['gamma_2'], CSMF_hod_dict['sigma_c'], CSMF_hod_dict['a_1'], CSMF_hod_dict['a_2'], CSMF_hod_dict['M_2'], \ - CSMF_hod_dict['b_0'], CSMF_hod_dict['b_1'], CSMF_hod_dict['b_2'], CSMF_hod_dict['delta_1'], CSMF_hod_dict['delta_2'] - Ac_C, As_C, Bc_C, Bs_C, ic_C = \ - CSMF_hod_dict['Acent'], CSMF_hod_dict['Asat'], CSMF_hod_dict['Bcent'], \ - CSMF_hod_dict['Bsat'], CSMF_hod_dict['ic'] - f_sigv_C = CSMF_hod_dict['f_sigv'] numba.set_num_threads(Nthread) # compute nsate for each halo # figuring out the number of particles kept for each thread - num_sats_L = np.zeros(len(hid), dtype = np.int64) - num_sats_E = np.zeros(len(hid), dtype = np.int64) - num_sats_Q = np.zeros(len(hid), dtype = np.int64) - num_sats_C = np.zeros(len(hid), dtype = np.int64) - stellarmass_C = np.zeros(len(hid), dtype = np.int64) - hstart = np.rint(np.linspace(0, len(hid), Nthread + 1)).astype(np.int64) # starting index of each thread + num_sats_L = np.zeros(len(hid), dtype=np.int64) + num_sats_E = np.zeros(len(hid), dtype=np.int64) + num_sats_Q = np.zeros(len(hid), dtype=np.int64) + hstart = np.rint(np.linspace(0, len(hid), Nthread + 1)).astype( + np.int64 + ) # starting index of each thread for tid in range(Nthread): for i in range(hstart[tid], hstart[tid + 1]): if want_LRG: - M1_L_temp = 10**(logM1_L + As_L * hdeltac[i] + Bs_L * hfenv[i]) + M1_L_temp = 10 ** (logM1_L + As_L * hdeltac[i] + Bs_L * hfenv[i]) logM_cut_L_temp = logM_cut_L + Ac_L * hdeltac[i] + Bc_L * hfenv[i] - base_p_L = n_sat_LRG_modified(hmass[i], logM_cut_L_temp, - 10**logM_cut_L_temp, M1_L_temp, sigma_L, alpha_L, kappa_L) * ic_L + base_p_L = ( + n_sat_LRG_modified( + hmass[i], + logM_cut_L_temp, + 10**logM_cut_L_temp, + M1_L_temp, + sigma_L, + alpha_L, + kappa_L, + ) + * ic_L + ) num_sats_L[i] = np.random.poisson(base_p_L) if want_ELG: - M1_E_temp = 10**(logM1_E + As_E * hdeltac[i] + Bs_E * hfenv[i] + Cs_E * hshear[i]) - logM_cut_E_temp = logM_cut_E + Ac_E * hdeltac[i] + Bc_E * hfenv[i] + Cc_E * hshear[i] - base_p_E = N_sat_elg( - hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_E, A_E) * ic_E + M1_E_temp = 10 ** ( + logM1_E + As_E * hdeltac[i] + Bs_E * hfenv[i] + Cs_E * hshear[i] + ) + logM_cut_E_temp = ( + logM_cut_E + Ac_E * hdeltac[i] + Bc_E * hfenv[i] + Cc_E * hshear[i] + ) + base_p_E = ( + N_sat_elg( + hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_E, A_E + ) + * ic_E + ) # elg conformity if keep_cent[i] == 1: - M1_E_temp = 10**(logM1_EL + As_E * hdeltac[i] + Bs_E * hfenv[i]) - base_p_E = N_sat_elg( - hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_EL, A_E) * ic_E + M1_E_temp = 10 ** (logM1_EL + As_E * hdeltac[i] + Bs_E * hfenv[i]) + base_p_E = ( + N_sat_elg( + hmass[i], + 10**logM_cut_E_temp, + kappa_E, + M1_E_temp, + alpha_EL, + A_E, + ) + * ic_E + ) elif keep_cent[i] == 2: - M1_E_temp = 10**(logM1_EE + As_E * hdeltac[i] + Bs_E * hfenv[i]) # M1_E_temp*10**delta_M1 - base_p_E = N_sat_elg( - hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_EE, A_E) * ic_E + M1_E_temp = 10 ** ( + logM1_EE + As_E * hdeltac[i] + Bs_E * hfenv[i] + ) # M1_E_temp*10**delta_M1 + base_p_E = ( + N_sat_elg( + hmass[i], + 10**logM_cut_E_temp, + kappa_E, + M1_E_temp, + alpha_EE, + A_E, + ) + * ic_E + ) num_sats_E[i] = np.random.poisson(base_p_E) if want_QSO: - M1_Q_temp = 10**(logM1_Q + As_Q * hdeltac[i] + Bs_Q * hfenv[i]) + M1_Q_temp = 10 ** (logM1_Q + As_Q * hdeltac[i] + Bs_Q * hfenv[i]) logM_cut_Q_temp = logM_cut_Q + Ac_Q * hdeltac[i] + Bc_Q * hfenv[i] - base_p_Q = N_sat_generic( - hmass[i], 10**logM_cut_Q_temp, kappa_Q, M1_Q_temp, alpha_Q) * ic_Q + base_p_Q = ( + N_sat_generic( + hmass[i], 10**logM_cut_Q_temp, kappa_Q, M1_Q_temp, alpha_Q + ) + * ic_Q + ) num_sats_Q[i] = np.random.poisson(base_p_Q) - - if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) - a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] - base_p_C = n_sat_CSMF(hmass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C_temp, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) * ic_C - - num_sats_C[i] = np.random.poisson(base_p_C) - + # generate rdpos rd_pos_L = getPointsOnSphere(np.sum(num_sats_L), Nthread) rd_pos_E = getPointsOnSphere(np.sum(num_sats_E), Nthread) rd_pos_Q = getPointsOnSphere(np.sum(num_sats_Q), Nthread) - rd_pos_C = getPointsOnSphere(np.sum(num_sats_C), Nthread) # put satellites on NFW - h_id_L, x_sat_L, y_sat_L, z_sat_L, vx_sat_L, vy_sat_L, vz_sat_L, M_L\ - = compute_fast_NFW(NFW_draw, hid, hpos[:, 0], hpos[:, 1], hpos[:, 2], hvel[:, 0], hvel[:, 1], hvel[:, 2], - hvrms, hc, hmass, hrvir, rd_pos_L, num_sats_L, f_sigv_L, vel_sat, Nthread, - exp_frac, exp_scale, nfw_rescale) - h_id_E, x_sat_E, y_sat_E, z_sat_E, vx_sat_E, vy_sat_E, vz_sat_E, M_E\ - = compute_fast_NFW(NFW_draw, hid, hpos[:, 0], hpos[:, 1], hpos[:, 2], hvel[:, 0], hvel[:, 1], hvel[:, 2], - hvrms, hc, hmass, hrvir, rd_pos_E, num_sats_E, f_sigv_E, vel_sat, Nthread, - exp_frac, exp_scale, nfw_rescale) - h_id_Q, x_sat_Q, y_sat_Q, z_sat_Q, vx_sat_Q, vy_sat_Q, vz_sat_Q, M_Q\ - = compute_fast_NFW(NFW_draw, hid, hpos[:, 0], hpos[:, 1], hpos[:, 2], hvel[:, 0], hvel[:, 1], hvel[:, 2], - hvrms, hc, hmass, hrvir, rd_pos_Q, num_sats_Q, f_sigv_Q, vel_sat, Nthread, - exp_frac, exp_scale, nfw_rescale) - h_id_C, x_sat_C, y_sat_C, z_sat_C, vx_sat_C, vy_sat_C, vz_sat_C, M_C\ - = compute_fast_NFW(NFW_draw, hid, hpos[:, 0], hpos[:, 1], hpos[:, 2], hvel[:, 0], hvel[:, 1], hvel[:, 2], - hvrms, hc, hmass, hrvir, rd_pos_C, num_sats_C, f_sigv_C, vel_sat, Nthread, - exp_frac, exp_scale, nfw_rescale) - - + h_id_L, x_sat_L, y_sat_L, z_sat_L, vx_sat_L, vy_sat_L, vz_sat_L, M_L = ( + compute_fast_NFW( + NFW_draw, + hid, + hpos[:, 0], + hpos[:, 1], + hpos[:, 2], + hvel[:, 0], + hvel[:, 1], + hvel[:, 2], + hvrms, + hc, + hmass, + hrvir, + rd_pos_L, + num_sats_L, + f_sigv_L, + vel_sat, + Nthread, + exp_frac, + exp_scale, + nfw_rescale, + ) + ) + h_id_E, x_sat_E, y_sat_E, z_sat_E, vx_sat_E, vy_sat_E, vz_sat_E, M_E = ( + compute_fast_NFW( + NFW_draw, + hid, + hpos[:, 0], + hpos[:, 1], + hpos[:, 2], + hvel[:, 0], + hvel[:, 1], + hvel[:, 2], + hvrms, + hc, + hmass, + hrvir, + rd_pos_E, + num_sats_E, + f_sigv_E, + vel_sat, + Nthread, + exp_frac, + exp_scale, + nfw_rescale, + ) + ) + h_id_Q, x_sat_Q, y_sat_Q, z_sat_Q, vx_sat_Q, vy_sat_Q, vz_sat_Q, M_Q = ( + compute_fast_NFW( + NFW_draw, + hid, + hpos[:, 0], + hpos[:, 1], + hpos[:, 2], + hvel[:, 0], + hvel[:, 1], + hvel[:, 2], + hvrms, + hc, + hmass, + hrvir, + rd_pos_Q, + num_sats_Q, + f_sigv_Q, + vel_sat, + Nthread, + exp_frac, + exp_scale, + nfw_rescale, + ) + ) # do rsd if rsd: z_sat_L = (z_sat_L + vz_sat_L * inv_velz2kms) % lbox z_sat_E = (z_sat_E + vz_sat_E * inv_velz2kms) % lbox z_sat_Q = (z_sat_Q + vz_sat_Q * inv_velz2kms) % lbox - z_sat_C = (z_sat_C + vz_sat_C * inv_velz2kms) % lbox - LRG_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) - ELG_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) - QSO_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) - ID_dict = Dict.empty(key_type = types.unicode_type, value_type = int_array) + LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) LRG_dict['x'] = x_sat_L LRG_dict['y'] = y_sat_L LRG_dict['z'] = z_sat_L @@ -851,93 +817,170 @@ def gen_sats_nfw(NFW_draw, hpos, hvel, hmass, hid, hdeltac, hfenv, hshear, hvrms QSO_dict['vz'] = vz_sat_Q QSO_dict['mass'] = M_Q ID_dict['QSO'] = h_id_Q - - CSMF_dict['x'] = x_sat_C - CSMF_dict['y'] = y_sat_C - CSMF_dict['z'] = z_sat_C - CSMF_dict['vx'] = vx_sat_C - CSMF_dict['vy'] = vy_sat_C - CSMF_dict['vz'] = vz_sat_C - CSMF_dict['mass'] = M_C - stellarmass_C = np.empty_like(M_C) - ## compute stellarmass of all the satelites - if want_CSMF: - hstart = np.rint(np.linspace(0, num_sats_C.sum(), Nthread + 1)) - for tid in numba.prange(Nthread): - for i in range(int(hstart[tid]), int(hstart[tid + 1])): - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) - a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] - stellarmass_C[i] = get_random_sat_stellarmass_linearinterpolation(M_C[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, a1_C_temp, s, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) - CSMF_dict['stellarmass'] = stellarmass_C - ID_dict['CSMF'] = h_id_C - - return LRG_dict, ELG_dict, QSO_dict, CSMF_dict, ID_dict - -@njit(parallel = True, fastmath = True) -def gen_sats(ppos, pvel, hvel, hmass, hid, weights, randoms, hdeltac, hfenv, hshear, - enable_ranks, ranks, ranksv, ranksp, ranksr, ranksc, - LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, CSMF_hod_dict, - rsd, inv_velz2kms, lbox, Mpart, want_LRG, want_ELG, want_QSO, want_CSMF, Nthread, origin, keep_cent): + return LRG_dict, ELG_dict, QSO_dict, ID_dict + + +@njit(parallel=True, fastmath=True) +def gen_sats( + ppos, + pvel, + hvel, + hmass, + hid, + weights, + randoms, + hdeltac, + hfenv, + hshear, + enable_ranks, + ranks, + ranksv, + ranksp, + ranksr, + ranksc, + LRG_hod_dict, + ELG_hod_dict, + QSO_hod_dict, + rsd, + inv_velz2kms, + lbox, + Mpart, + want_LRG, + want_ELG, + want_QSO, + Nthread, + origin, + keep_cent, +): """ Generate satellite galaxies in place in memory with a two pass numba parallel implementation. """ if want_LRG: - logM_cut_L, logM1_L, sigma_L, alpha_L, kappa_L = \ - LRG_hod_dict['logM_cut'], LRG_hod_dict['logM1'], LRG_hod_dict['sigma'], LRG_hod_dict['alpha'], LRG_hod_dict['kappa'] - alpha_s_L, s_L, s_v_L, s_p_L, s_r_L, Ac_L, As_L, Bc_L, Bs_L, ic_L = \ - LRG_hod_dict['alpha_s'], LRG_hod_dict['s'], LRG_hod_dict['s_v'], LRG_hod_dict['s_p'], \ - LRG_hod_dict['s_r'], LRG_hod_dict['Acent'], LRG_hod_dict['Asat'], LRG_hod_dict['Bcent'], \ - LRG_hod_dict['Bsat'], LRG_hod_dict['ic'] + logM_cut_L, logM1_L, sigma_L, alpha_L, kappa_L = ( + LRG_hod_dict['logM_cut'], + LRG_hod_dict['logM1'], + LRG_hod_dict['sigma'], + LRG_hod_dict['alpha'], + LRG_hod_dict['kappa'], + ) + alpha_s_L, s_L, s_v_L, s_p_L, s_r_L, Ac_L, As_L, Bc_L, Bs_L, ic_L = ( + LRG_hod_dict['alpha_s'], + LRG_hod_dict['s'], + LRG_hod_dict['s_v'], + LRG_hod_dict['s_p'], + LRG_hod_dict['s_r'], + LRG_hod_dict['Acent'], + LRG_hod_dict['Asat'], + LRG_hod_dict['Bcent'], + LRG_hod_dict['Bsat'], + LRG_hod_dict['ic'], + ) if want_ELG: - logM_cut_E, kappa_E, logM1_E, alpha_E, A_E = \ - ELG_hod_dict['logM_cut'], ELG_hod_dict['kappa'], ELG_hod_dict['logM1'], ELG_hod_dict['alpha'], ELG_hod_dict['A_s'] - alpha_s_E, s_E, s_v_E, s_p_E, s_r_E, Ac_E, As_E, Bc_E, Bs_E, Cc_E, Cs_E, ic_E, logM1_EE, alpha_EE, logM1_EL, alpha_EL = \ - ELG_hod_dict['alpha_s'], ELG_hod_dict['s'], ELG_hod_dict['s_v'], ELG_hod_dict['s_p'], \ - ELG_hod_dict['s_r'], ELG_hod_dict['Acent'], ELG_hod_dict['Asat'], ELG_hod_dict['Bcent'], \ - ELG_hod_dict['Bsat'], ELG_hod_dict['Ccent'], ELG_hod_dict['Csat'], ELG_hod_dict['ic'], \ - ELG_hod_dict['logM1_EE'], ELG_hod_dict['alpha_EE'], ELG_hod_dict['logM1_EL'], ELG_hod_dict['alpha_EL'] + logM_cut_E, kappa_E, logM1_E, alpha_E, A_E = ( + ELG_hod_dict['logM_cut'], + ELG_hod_dict['kappa'], + ELG_hod_dict['logM1'], + ELG_hod_dict['alpha'], + ELG_hod_dict['A_s'], + ) + ( + alpha_s_E, + s_E, + s_v_E, + s_p_E, + s_r_E, + Ac_E, + As_E, + Bc_E, + Bs_E, + Cc_E, + Cs_E, + ic_E, + logM1_EE, + alpha_EE, + logM1_EL, + alpha_EL, + ) = ( + ELG_hod_dict['alpha_s'], + ELG_hod_dict['s'], + ELG_hod_dict['s_v'], + ELG_hod_dict['s_p'], + ELG_hod_dict['s_r'], + ELG_hod_dict['Acent'], + ELG_hod_dict['Asat'], + ELG_hod_dict['Bcent'], + ELG_hod_dict['Bsat'], + ELG_hod_dict['Ccent'], + ELG_hod_dict['Csat'], + ELG_hod_dict['ic'], + ELG_hod_dict['logM1_EE'], + ELG_hod_dict['alpha_EE'], + ELG_hod_dict['logM1_EL'], + ELG_hod_dict['alpha_EL'], + ) if want_QSO: - logM_cut_Q, kappa_Q, logM1_Q, alpha_Q = \ - QSO_hod_dict['logM_cut'], QSO_hod_dict['kappa'], QSO_hod_dict['logM1'], QSO_hod_dict['alpha'] - alpha_s_Q, s_Q, s_v_Q, s_p_Q, s_r_Q, Ac_Q, As_Q, Bc_Q, Bs_Q, ic_Q = \ - QSO_hod_dict['alpha_s'], QSO_hod_dict['s'], QSO_hod_dict['s_v'], QSO_hod_dict['s_p'], \ - QSO_hod_dict['s_r'], QSO_hod_dict['Acent'], QSO_hod_dict['Asat'], QSO_hod_dict['Bcent'], \ - QSO_hod_dict['Bsat'], QSO_hod_dict['ic'] - - if want_CSMF: - Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = \ - CSMF_hod_dict['Mstar_low'], CSMF_hod_dict['Mstar_up'], CSMF_hod_dict['M_1'], CSMF_hod_dict['M_0'], CSMF_hod_dict['gamma_1'], \ - CSMF_hod_dict['gamma_2'], CSMF_hod_dict['sigma_c'], CSMF_hod_dict['a_1'], CSMF_hod_dict['a_2'], CSMF_hod_dict['M_2'], \ - CSMF_hod_dict['b_0'], CSMF_hod_dict['b_1'], CSMF_hod_dict['b_2'], CSMF_hod_dict['delta_1'], CSMF_hod_dict['delta_2'] - alpha_s_C, s_C, s_v_C, s_p_C, s_r_C, Ac_C, As_C, Bc_C, Bs_C, ic_C = CSMF_hod_dict['alpha_s'], CSMF_hod_dict['s'],\ - CSMF_hod_dict['s_v'], CSMF_hod_dict['s_p'], CSMF_hod_dict['s_r'], CSMF_hod_dict['Acent'],\ - CSMF_hod_dict['Asat'], CSMF_hod_dict['Bcent'], CSMF_hod_dict['Bsat'], CSMF_hod_dict['ic'] - - # print(Ac_C, As_C, Bc_C, Bs_C) - H = len(hmass) # num of particles + logM_cut_Q, kappa_Q, logM1_Q, alpha_Q = ( + QSO_hod_dict['logM_cut'], + QSO_hod_dict['kappa'], + QSO_hod_dict['logM1'], + QSO_hod_dict['alpha'], + ) + alpha_s_Q, s_Q, s_v_Q, s_p_Q, s_r_Q, Ac_Q, As_Q, Bc_Q, Bs_Q, ic_Q = ( + QSO_hod_dict['alpha_s'], + QSO_hod_dict['s'], + QSO_hod_dict['s_v'], + QSO_hod_dict['s_p'], + QSO_hod_dict['s_r'], + QSO_hod_dict['Acent'], + QSO_hod_dict['Asat'], + QSO_hod_dict['Bcent'], + QSO_hod_dict['Bsat'], + QSO_hod_dict['ic'], + ) + + H = len(hmass) # num of particles numba.set_num_threads(Nthread) - Nout = np.zeros((Nthread, 4, 8), dtype = np.int64) - hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype(np.int64) # starting index of each thread + Nout = np.zeros((Nthread, 3, 8), dtype=np.int64) + hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype( + np.int64 + ) # starting index of each thread - keep = np.empty(H, dtype = np.int8) # mask array tracking which halos to keep + keep = np.empty(H, dtype=np.int8) # mask array tracking which halos to keep # figuring out the number of particles kept for each thread - for tid in numba.prange(Nthread): #numba.prange(Nthread): + for tid in numba.prange(Nthread): # numba.prange(Nthread): for i in range(hstart[tid], hstart[tid + 1]): # print(logM1, As, hdeltac[i], Bs, hfenv[i]) LRG_marker = 0 if want_LRG: - M1_L_temp = 10**(logM1_L + As_L * hdeltac[i] + Bs_L * hfenv[i]) + M1_L_temp = 10 ** (logM1_L + As_L * hdeltac[i] + Bs_L * hfenv[i]) logM_cut_L_temp = logM_cut_L + Ac_L * hdeltac[i] + Bc_L * hfenv[i] - base_p_L = n_sat_LRG_modified(hmass[i], logM_cut_L_temp, - 10**logM_cut_L_temp, M1_L_temp, sigma_L, alpha_L, kappa_L) * weights[i] * ic_L + base_p_L = ( + n_sat_LRG_modified( + hmass[i], + logM_cut_L_temp, + 10**logM_cut_L_temp, + M1_L_temp, + sigma_L, + alpha_L, + kappa_L, + ) + * weights[i] + * ic_L + ) if enable_ranks: - decorator_L = 1 + s_L * ranks[i] + s_v_L * ranksv[i] + s_p_L * ranksp[i] + s_r_L * ranksr[i] + decorator_L = ( + 1 + + s_L * ranks[i] + + s_v_L * ranksv[i] + + s_p_L * ranksp[i] + + s_r_L * ranksr[i] + ) exp_sat = base_p_L * decorator_L else: exp_sat = base_p_L @@ -945,19 +988,50 @@ def gen_sats(ppos, pvel, hvel, hmass, hid, weights, randoms, hdeltac, hfenv, hsh ELG_marker = LRG_marker if want_ELG: - M1_E_temp = 10**(logM1_E + As_E * hdeltac[i] + Bs_E * hfenv[i] + Cs_E * hshear[i]) - logM_cut_E_temp = logM_cut_E + Ac_E * hdeltac[i] + Bc_E * hfenv[i] + Cc_E * hshear[i] - base_p_E = N_sat_elg( - hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_E, A_E) * weights[i] * ic_E + M1_E_temp = 10 ** ( + logM1_E + As_E * hdeltac[i] + Bs_E * hfenv[i] + Cs_E * hshear[i] + ) + logM_cut_E_temp = ( + logM_cut_E + Ac_E * hdeltac[i] + Bc_E * hfenv[i] + Cc_E * hshear[i] + ) + base_p_E = ( + N_sat_elg( + hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_E, A_E + ) + * weights[i] + * ic_E + ) # elg conformity if keep_cent[i] == 1: - M1_E_temp = 10**(logM1_EL + As_E * hdeltac[i] + Bs_E * hfenv[i]) - base_p_E = N_sat_elg( - hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_EL, A_E) * weights[i] * ic_E + M1_E_temp = 10 ** (logM1_EL + As_E * hdeltac[i] + Bs_E * hfenv[i]) + base_p_E = ( + N_sat_elg( + hmass[i], + 10**logM_cut_E_temp, + kappa_E, + M1_E_temp, + alpha_EL, + A_E, + ) + * weights[i] + * ic_E + ) elif keep_cent[i] == 2: - M1_E_temp = 10**(logM1_EE + As_E * hdeltac[i] + Bs_E * hfenv[i]) # M1_E_temp*10**delta_M1 - base_p_E = N_sat_elg( - hmass[i], 10**logM_cut_E_temp, kappa_E, M1_E_temp, alpha_EE, A_E) * weights[i] * ic_E + M1_E_temp = 10 ** ( + logM1_EE + As_E * hdeltac[i] + Bs_E * hfenv[i] + ) # M1_E_temp*10**delta_M1 + base_p_E = ( + N_sat_elg( + hmass[i], + 10**logM_cut_E_temp, + kappa_E, + M1_E_temp, + alpha_EE, + A_E, + ) + * weights[i] + * ic_E + ) # if base_p_E > 1: # print("ExE new p", base_p_E, np.log10(hmass[i]), N_sat_elg( @@ -965,134 +1039,124 @@ def gen_sats(ppos, pvel, hvel, hmass, hid, weights, randoms, hdeltac, hfenv, hsh # rank mods if enable_ranks: - decorator_E = 1 + s_E * ranks[i] + s_v_E * ranksv[i] + s_p_E * ranksp[i] + s_r_E * ranksr[i] + decorator_E = ( + 1 + + s_E * ranks[i] + + s_v_E * ranksv[i] + + s_p_E * ranksp[i] + + s_r_E * ranksr[i] + ) base_p_E = base_p_E * decorator_E ELG_marker += base_p_E QSO_marker = ELG_marker if want_QSO: - M1_Q_temp = 10**(logM1_Q + As_Q * hdeltac[i] + Bs_Q * hfenv[i]) + M1_Q_temp = 10 ** (logM1_Q + As_Q * hdeltac[i] + Bs_Q * hfenv[i]) logM_cut_Q_temp = logM_cut_Q + Ac_Q * hdeltac[i] + Bc_Q * hfenv[i] - base_p_Q = N_sat_generic( - hmass[i], 10**logM_cut_Q_temp, kappa_Q, M1_Q_temp, alpha_Q) * weights[i] * ic_Q + base_p_Q = ( + N_sat_generic( + hmass[i], 10**logM_cut_Q_temp, kappa_Q, M1_Q_temp, alpha_Q + ) + * weights[i] + * ic_Q + ) if enable_ranks: - decorator_Q = 1 + s_Q * ranks[i] + s_v_Q * ranksv[i] + s_p_Q * ranksp[i] + s_r_Q * ranksr[i] + decorator_Q = ( + 1 + + s_Q * ranks[i] + + s_v_Q * ranksv[i] + + s_p_Q * ranksp[i] + + s_r_Q * ranksr[i] + ) exp_sat = base_p_Q * decorator_Q else: exp_sat = base_p_Q QSO_marker += exp_sat - - CSMF_marker = QSO_marker - if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) - a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] - base_p_C = n_sat_CSMF(hmass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C_temp, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) * weights[i] * ic_C - if enable_ranks: - decorator_C = 1 + s_C * ranks[i] + s_v_C * ranksv[i] + s_p_C * ranksp[i] + s_r_C * ranksr[i] - exp_sat = base_p_C * decorator_C - else: - exp_sat = base_p_C - CSMF_marker += exp_sat if randoms[i] <= LRG_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 elif randoms[i] <= ELG_marker: - Nout[tid, 1, 0] += 1 # counting + Nout[tid, 1, 0] += 1 # counting keep[i] = 2 elif randoms[i] <= QSO_marker: - Nout[tid, 2, 0] += 1 # counting + Nout[tid, 2, 0] += 1 # counting keep[i] = 3 - elif randoms[i] <= CSMF_marker: - Nout[tid, 3, 0] += 1 # counting - keep[i] = 4 else: keep[i] = 0 -# if(0 not in keep): -# print("Waring: All particles are used for satelites") -# else: -# unused_haloes = sum(keep == 0) -# print("From "+str(H)+" particles we used:",(H-unused_haloes)/H) - # compose galaxy array, first create array of galaxy starting indices for the threads - gstart = np.empty((Nthread + 1, 4), dtype = np.int64) + gstart = np.empty((Nthread + 1, 3), dtype=np.int64) gstart[0, :] = 0 gstart[1:, 0] = Nout[:, 0, 0].cumsum() gstart[1:, 1] = Nout[:, 1, 0].cumsum() gstart[1:, 2] = Nout[:, 2, 0].cumsum() - gstart[1:, 3] = Nout[:, 3, 0].cumsum() # galaxy arrays N_lrg = gstart[-1, 0] - lrg_x = np.empty(N_lrg, dtype = hmass.dtype) - lrg_y = np.empty(N_lrg, dtype = hmass.dtype) - lrg_z = np.empty(N_lrg, dtype = hmass.dtype) - lrg_vx = np.empty(N_lrg, dtype = hmass.dtype) - lrg_vy = np.empty(N_lrg, dtype = hmass.dtype) - lrg_vz = np.empty(N_lrg, dtype = hmass.dtype) - lrg_mass = np.empty(N_lrg, dtype = hmass.dtype) - lrg_id = np.empty(N_lrg, dtype = hid.dtype) + lrg_x = np.empty(N_lrg, dtype=hmass.dtype) + lrg_y = np.empty(N_lrg, dtype=hmass.dtype) + lrg_z = np.empty(N_lrg, dtype=hmass.dtype) + lrg_vx = np.empty(N_lrg, dtype=hmass.dtype) + lrg_vy = np.empty(N_lrg, dtype=hmass.dtype) + lrg_vz = np.empty(N_lrg, dtype=hmass.dtype) + lrg_mass = np.empty(N_lrg, dtype=hmass.dtype) + lrg_id = np.empty(N_lrg, dtype=hid.dtype) # galaxy arrays N_elg = gstart[-1, 1] - elg_x = np.empty(N_elg, dtype = hmass.dtype) - elg_y = np.empty(N_elg, dtype = hmass.dtype) - elg_z = np.empty(N_elg, dtype = hmass.dtype) - elg_vx = np.empty(N_elg, dtype = hmass.dtype) - elg_vy = np.empty(N_elg, dtype = hmass.dtype) - elg_vz = np.empty(N_elg, dtype = hmass.dtype) - elg_mass = np.empty(N_elg, dtype = hmass.dtype) - elg_id = np.empty(N_elg, dtype = hid.dtype) + elg_x = np.empty(N_elg, dtype=hmass.dtype) + elg_y = np.empty(N_elg, dtype=hmass.dtype) + elg_z = np.empty(N_elg, dtype=hmass.dtype) + elg_vx = np.empty(N_elg, dtype=hmass.dtype) + elg_vy = np.empty(N_elg, dtype=hmass.dtype) + elg_vz = np.empty(N_elg, dtype=hmass.dtype) + elg_mass = np.empty(N_elg, dtype=hmass.dtype) + elg_id = np.empty(N_elg, dtype=hid.dtype) # galaxy arrays N_qso = gstart[-1, 2] - qso_x = np.empty(N_qso, dtype = hmass.dtype) - qso_y = np.empty(N_qso, dtype = hmass.dtype) - qso_z = np.empty(N_qso, dtype = hmass.dtype) - qso_vx = np.empty(N_qso, dtype = hmass.dtype) - qso_vy = np.empty(N_qso, dtype = hmass.dtype) - qso_vz = np.empty(N_qso, dtype = hmass.dtype) - qso_mass = np.empty(N_qso, dtype = hmass.dtype) - qso_id = np.empty(N_qso, dtype = hid.dtype) - - # galaxy arrays - N_CSMF = gstart[-1, 3] - CSMF_x = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_y = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_z = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_id = np.empty(N_CSMF, dtype = hid.dtype) - + qso_x = np.empty(N_qso, dtype=hmass.dtype) + qso_y = np.empty(N_qso, dtype=hmass.dtype) + qso_z = np.empty(N_qso, dtype=hmass.dtype) + qso_vx = np.empty(N_qso, dtype=hmass.dtype) + qso_vy = np.empty(N_qso, dtype=hmass.dtype) + qso_vz = np.empty(N_qso, dtype=hmass.dtype) + qso_mass = np.empty(N_qso, dtype=hmass.dtype) + qso_id = np.empty(N_qso, dtype=hid.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): - j1, j2, j3, j4 = gstart[tid] + j1, j2, j3 = gstart[tid] for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: lrg_x[j1] = ppos[i, 0] - lrg_vx[j1] = hvel[i, 0] + alpha_s_L * (pvel[i, 0] - hvel[i, 0]) # velocity bias + lrg_vx[j1] = hvel[i, 0] + alpha_s_L * ( + pvel[i, 0] - hvel[i, 0] + ) # velocity bias lrg_y[j1] = ppos[i, 1] - lrg_vy[j1] = hvel[i, 1] + alpha_s_L * (pvel[i, 1] - hvel[i, 1]) # velocity bias + lrg_vy[j1] = hvel[i, 1] + alpha_s_L * ( + pvel[i, 1] - hvel[i, 1] + ) # velocity bias lrg_z[j1] = ppos[i, 2] - lrg_vz[j1] = hvel[i, 2] + alpha_s_L * (pvel[i, 2] - hvel[i, 2]) # velocity bias + lrg_vz[j1] = hvel[i, 2] + alpha_s_L * ( + pvel[i, 2] - hvel[i, 2] + ) # velocity bias if rsd and origin is not None: nx = lrg_x[j1] - origin[0] ny = lrg_y[j1] - origin[1] nz = lrg_z[j1] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms*(lrg_vx[j1]*nx+lrg_vy[j1]*ny+lrg_vz[j1]*nz) - lrg_x[j1] = lrg_x[j1]+proj*nx - lrg_y[j1] = lrg_y[j1]+proj*ny - lrg_z[j1] = lrg_z[j1]+proj*nz + proj = inv_velz2kms * ( + lrg_vx[j1] * nx + lrg_vy[j1] * ny + lrg_vz[j1] * nz + ) + lrg_x[j1] = lrg_x[j1] + proj * nx + lrg_y[j1] = lrg_y[j1] + proj * ny + lrg_z[j1] = lrg_z[j1] + proj * nz elif rsd: lrg_z[j1] = wrap(lrg_z[j1] + lrg_vz[j1] * inv_velz2kms, lbox) lrg_mass[j1] = hmass[i] @@ -1100,23 +1164,31 @@ def gen_sats(ppos, pvel, hvel, hmass, hid, weights, randoms, hdeltac, hfenv, hsh j1 += 1 elif keep[i] == 2: elg_x[j2] = ppos[i, 0] - elg_vx[j2] = hvel[i, 0] + alpha_s_E * (pvel[i, 0] - hvel[i, 0]) # velocity bias + elg_vx[j2] = hvel[i, 0] + alpha_s_E * ( + pvel[i, 0] - hvel[i, 0] + ) # velocity bias elg_y[j2] = ppos[i, 1] - elg_vy[j2] = hvel[i, 1] + alpha_s_E * (pvel[i, 1] - hvel[i, 1]) # velocity bias + elg_vy[j2] = hvel[i, 1] + alpha_s_E * ( + pvel[i, 1] - hvel[i, 1] + ) # velocity bias elg_z[j2] = ppos[i, 2] - elg_vz[j2] = hvel[i, 2] + alpha_s_E * (pvel[i, 2] - hvel[i, 2]) # velocity bias + elg_vz[j2] = hvel[i, 2] + alpha_s_E * ( + pvel[i, 2] - hvel[i, 2] + ) # velocity bias if rsd and origin is not None: nx = elg_x[j2] - origin[0] ny = elg_y[j2] - origin[1] nz = elg_z[j2] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms*(elg_vx[j2]*nx+elg_vy[j2]*ny+elg_vz[j2]*nz) - elg_x[j2] = elg_x[j2]+proj*nx - elg_y[j2] = elg_y[j2]+proj*ny - elg_z[j2] = elg_z[j2]+proj*nz + proj = inv_velz2kms * ( + elg_vx[j2] * nx + elg_vy[j2] * ny + elg_vz[j2] * nz + ) + elg_x[j2] = elg_x[j2] + proj * nx + elg_y[j2] = elg_y[j2] + proj * ny + elg_z[j2] = elg_z[j2] + proj * nz elif rsd: elg_z[j2] = wrap(elg_z[j2] + elg_vz[j2] * inv_velz2kms, lbox) elg_mass[j2] = hmass[i] @@ -1124,70 +1196,42 @@ def gen_sats(ppos, pvel, hvel, hmass, hid, weights, randoms, hdeltac, hfenv, hsh j2 += 1 elif keep[i] == 3: qso_x[j3] = ppos[i, 0] - qso_vx[j3] = hvel[i, 0] + alpha_s_Q * (pvel[i, 0] - hvel[i, 0]) # velocity bias + qso_vx[j3] = hvel[i, 0] + alpha_s_Q * ( + pvel[i, 0] - hvel[i, 0] + ) # velocity bias qso_y[j3] = ppos[i, 1] - qso_vy[j3] = hvel[i, 1] + alpha_s_Q * (pvel[i, 1] - hvel[i, 1]) # velocity bias + qso_vy[j3] = hvel[i, 1] + alpha_s_Q * ( + pvel[i, 1] - hvel[i, 1] + ) # velocity bias qso_z[j3] = ppos[i, 2] - qso_vz[j3] = hvel[i, 2] + alpha_s_Q * (pvel[i, 2] - hvel[i, 2]) # velocity bias + qso_vz[j3] = hvel[i, 2] + alpha_s_Q * ( + pvel[i, 2] - hvel[i, 2] + ) # velocity bias if rsd and origin is not None: nx = qso_x[j3] - origin[0] ny = qso_y[j3] - origin[1] nz = qso_z[j3] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms*(qso_vx[j3]*nx+qso_vy[j3]*ny+qso_vz[j3]*nz) - qso_x[j3] = qso_x[j3]+proj*nx - qso_y[j3] = qso_y[j3]+proj*ny - qso_z[j3] = qso_z[j3]+proj*nz + proj = inv_velz2kms * ( + qso_vx[j3] * nx + qso_vy[j3] * ny + qso_vz[j3] * nz + ) + qso_x[j3] = qso_x[j3] + proj * nx + qso_y[j3] = qso_y[j3] + proj * ny + qso_z[j3] = qso_z[j3] + proj * nz elif rsd: qso_z[j3] = wrap(qso_z[j3] + qso_vz[j3] * inv_velz2kms, lbox) qso_mass[j3] = hmass[i] qso_id[j3] = hid[i] j3 += 1 - elif keep[i] == 4: - CSMF_x[j4] = ppos[i, 0] - CSMF_vx[j4] = hvel[i, 0] + alpha_s_C * (pvel[i, 0] - hvel[i, 0]) # velocity bias - CSMF_y[j4] = ppos[i, 1] - CSMF_vy[j4] = hvel[i, 1] + alpha_s_C * (pvel[i, 1] - hvel[i, 1]) # velocity bias - CSMF_z[j4] = ppos[i, 2] - CSMF_vz[j4] = hvel[i, 2] + alpha_s_C * (pvel[i, 2] - hvel[i, 2]) # velocity bias - if rsd and origin is not None: - nx = CSMF_x[j4] - origin[0] - ny = CSMF_y[j4] - origin[1] - nz = CSMF_z[j4] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) - nx *= inv_norm - ny *= inv_norm - nz *= inv_norm - proj = inv_velz2kms*(CSMF_vx[j4]*nx+CSMF_vy[j4]*ny+CSMF_vz[j4]*nz) - CSMF_x[j4] = CSMF_x[j4]+proj*nx - CSMF_y[j4] = CSMF_y[j4]+proj*ny - CSMF_z[j4] = CSMF_z[j4]+proj*nz - elif rsd: - CSMF_z[j4] = wrap(CSMF_z[j4] + CSMF_vz[j4] * inv_velz2kms, lbox) - - - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) - a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] - # stellar_mass = get_random_sat_stellarmass_linearinterpolation(hmass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, a1_C_temp, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) - stellar_mass = get_random_sat_stellarmass(hmass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, a1_C_temp, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) - # if(j4>16000): - - CSMF_stellarmass[j4] = stellar_mass - CSMF_mass[j4] = hmass[i] - # CSMF_stellarmass[j4] = hmass[i] - CSMF_id[j4] = hid[i] - j4 += 1 # assert j == gstart[tid + 1] - LRG_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) - ELG_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) - QSO_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) - ID_dict = Dict.empty(key_type = types.unicode_type, value_type = int_array) - + LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) LRG_dict['x'] = lrg_x LRG_dict['y'] = lrg_y LRG_dict['z'] = lrg_z @@ -1214,22 +1258,12 @@ def gen_sats(ppos, pvel, hvel, hmass, hid, weights, randoms, hdeltac, hfenv, hsh QSO_dict['vz'] = qso_vz QSO_dict['mass'] = qso_mass ID_dict['QSO'] = qso_id - - CSMF_dict['x'] = CSMF_x - CSMF_dict['y'] = CSMF_y - CSMF_dict['z'] = CSMF_z - CSMF_dict['vx'] = CSMF_vx - CSMF_dict['vy'] = CSMF_vy - CSMF_dict['vz'] = CSMF_vz - CSMF_dict['mass'] = CSMF_mass - CSMF_dict['stellarmass'] = CSMF_stellarmass - ID_dict['CSMF'] = CSMF_id - - return LRG_dict, ELG_dict, QSO_dict, CSMF_dict, ID_dict - -@njit(parallel = True, fastmath = True) + return LRG_dict, ELG_dict, QSO_dict, ID_dict + + +@njit(parallel=True, fastmath=True) def fast_concatenate(array1, array2, Nthread): - ''' Fast concatenate with numba parallel''' + """Fast concatenate with numba parallel""" N1 = len(array1) N2 = len(array2) @@ -1238,7 +1272,7 @@ def fast_concatenate(array1, array2, Nthread): elif N2 == 0: return array1 - final_array = np.empty(N1 + N2, dtype = array1.dtype) + final_array = np.empty(N1 + N2, dtype=array1.dtype) # if one thread, then no need to parallel if Nthread == 1: for i in range(N1): @@ -1253,7 +1287,7 @@ def fast_concatenate(array1, array2, Nthread): hstart1 = np.rint(np.linspace(0, N1, Nthread1 + 1)).astype(np.int64) hstart2 = np.rint(np.linspace(0, N2, Nthread2 + 1)).astype(np.int64) + N1 - for tid in numba.prange(Nthread): #numba.prange(Nthread): + for tid in numba.prange(Nthread): # numba.prange(Nthread): if tid < Nthread1: for i in range(hstart1[tid], hstart1[tid + 1]): final_array[i] = array1[i] @@ -1264,7 +1298,18 @@ def fast_concatenate(array1, array2, Nthread): return final_array -def gen_gals(halos_array, subsample, tracers, params, Nthread, enable_ranks, rsd, verbose, nfw, NFW_draw = None): +def gen_gals( + halos_array, + subsample, + tracers, + params, + Nthread, + enable_ranks, + rsd, + verbose, + nfw, + NFW_draw=None, +): """ parse hod parameters, pass them on to central and satellite generators and then format the results @@ -1300,26 +1345,26 @@ def gen_gals(halos_array, subsample, tracers, params, Nthread, enable_ranks, rsd ELG_HOD = tracers[tracer] if tracer == 'QSO': QSO_HOD = tracers[tracer] - if tracer == 'CSMF': - CSMF_HOD = tracers[tracer] - - # print(CSMF_HOD) if 'LRG' in tracers.keys(): want_LRG = True - LRG_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + LRG_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) for key, value in LRG_HOD.items(): LRG_hod_dict[key] = value # LRG design and decorations logM_cut_L, logM1_L = map(LRG_HOD.get, ('logM_cut', 'logM1')) # z-evolving HOD - Delta_a = 1./(1+params['z']) - 1./(1+LRG_HOD.get('z_pivot', params['z'])) + Delta_a = 1.0 / (1 + params['z']) - 1.0 / ( + 1 + LRG_HOD.get('z_pivot', params['z']) + ) logM_cut_pr = LRG_HOD.get('logM_cut_pr', 0.0) logM1_pr = LRG_HOD.get('logM1_pr', 0.0) - logM_cut_L = logM_cut_L + logM_cut_pr*Delta_a - logM1_L = logM1_L + logM1_pr*Delta_a + logM_cut_L = logM_cut_L + logM_cut_pr * Delta_a + logM1_L = logM1_L + logM1_pr * Delta_a LRG_hod_dict['logM_cut'] = logM_cut_L LRG_hod_dict['logM1'] = logM1_L @@ -1334,25 +1379,30 @@ def gen_gals(halos_array, subsample, tracers, params, Nthread, enable_ranks, rsd else: want_LRG = False - LRG_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) - + LRG_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) if 'ELG' in tracers.keys(): # ELG design want_ELG = True - ELG_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + ELG_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) for key, value in ELG_HOD.items(): ELG_hod_dict[key] = value logM_cut_E, logM1_E = map(ELG_HOD.get, ('logM_cut', 'logM1')) # z-evolving HOD - Delta_a = 1./(1+params['z']) - 1./(1+ELG_HOD.get('z_pivot', params['z'])) + Delta_a = 1.0 / (1 + params['z']) - 1.0 / ( + 1 + ELG_HOD.get('z_pivot', params['z']) + ) logM_cut_pr = ELG_HOD.get('logM_cut_pr', 0.0) logM1_pr = ELG_HOD.get('logM1_pr', 0.0) - logM_cut_E = logM_cut_E + logM_cut_pr*Delta_a - logM1_E = logM1_E + logM1_pr*Delta_a + logM_cut_E = logM_cut_E + logM_cut_pr * Delta_a + logM1_E = logM1_E + logM1_pr * Delta_a ELG_hod_dict['logM_cut'] = logM_cut_E ELG_hod_dict['logM1'] = logM1_E @@ -1376,22 +1426,28 @@ def gen_gals(halos_array, subsample, tracers, params, Nthread, enable_ranks, rsd else: want_ELG = False - ELG_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + ELG_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) if 'QSO' in tracers.keys(): # QSO design want_QSO = True - QSO_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + QSO_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) for key, value in QSO_HOD.items(): QSO_hod_dict[key] = value logM_cut_Q, logM1_Q = map(QSO_HOD.get, ('logM_cut', 'logM1')) # z-evolving HOD - Delta_a = 1./(1+params['z']) - 1./(1+QSO_HOD.get('z_pivot', params['z'])) + Delta_a = 1.0 / (1 + params['z']) - 1.0 / ( + 1 + QSO_HOD.get('z_pivot', params['z']) + ) logM_cut_pr = QSO_HOD.get('logM_cut_pr', 0.0) logM1_pr = QSO_HOD.get('logM1_pr', 0.0) - logM_cut_Q = logM_cut_Q + logM_cut_pr*Delta_a - logM1_Q = logM1_Q + logM1_pr*Delta_a + logM_cut_Q = logM_cut_Q + logM_cut_pr * Delta_a + logM1_Q = logM1_Q + logM1_pr * Delta_a QSO_hod_dict['logM_cut'] = logM_cut_Q QSO_hod_dict['logM1'] = logM1_Q @@ -1406,96 +1462,150 @@ def gen_gals(halos_array, subsample, tracers, params, Nthread, enable_ranks, rsd else: want_QSO = False - QSO_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) - - - if 'CSMF' in tracers.keys(): - want_CSMF = True - - CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) - for key, value in CSMF_HOD.items(): - CSMF_hod_dict[key] = value - - CSMF_hod_dict['Acent'] = CSMF_HOD.get('Acent', 0.0) - CSMF_hod_dict['Asat'] = CSMF_HOD.get('Asat', 0.0) - CSMF_hod_dict['Bcent'] = CSMF_HOD.get('Bcent', 0.0) - CSMF_hod_dict['Bsat'] = CSMF_HOD.get('Bsat', 0.0) - CSMF_hod_dict['ic'] = CSMF_HOD.get('ic', 1.0) - CSMF_hod_dict['f_sigv'] = CSMF_HOD.get('f_sigv', 0) + QSO_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) - else: - want_CSMF = False - CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) - - # print(CSMF_hod_dict) start = time.time() velz2kms = params['velz2kms'] - inv_velz2kms = 1/velz2kms + inv_velz2kms = 1 / velz2kms lbox = params['Lbox'] origin = params['origin'] - LRG_dict_cent, ELG_dict_cent, QSO_dict_cent, CSMF_dict_cent, ID_dict_cent, keep_cent = \ - gen_cent(halos_array['hpos'], halos_array['hvel'], halos_array['hmass'], halos_array['hid'], halos_array['hmultis'], - halos_array['hrandoms'], halos_array['hveldev'], - halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), - halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), - halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), - LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, CSMF_hod_dict, rsd, inv_velz2kms, lbox, want_LRG, want_ELG, want_QSO, want_CSMF, Nthread, origin) - + LRG_dict_cent, ELG_dict_cent, QSO_dict_cent, ID_dict_cent, keep_cent = gen_cent( + halos_array['hpos'], + halos_array['hvel'], + halos_array['hmass'], + halos_array['hid'], + halos_array['hmultis'], + halos_array['hrandoms'], + halos_array['hveldev'], + halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), + halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), + halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), + LRG_hod_dict, + ELG_hod_dict, + QSO_hod_dict, + rsd, + inv_velz2kms, + lbox, + want_LRG, + want_ELG, + want_QSO, + Nthread, + origin, + ) if verbose: - print("generating centrals took ", time.time() - start) + print('generating centrals took ', time.time() - start) start = time.time() if nfw: - warnings.warn("NFW profile is unoptimized. It has different velocity bias. It does not support lightcone.") - LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, CSMF_dict_sat, ID_dict_sat = \ - gen_sats_nfw(NFW_draw, halos_array['hpos'], halos_array['hvel'], halos_array['hmass'], halos_array['hid'], - halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), - halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), - halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), - halos_array['hsigma3d'], halos_array['hc'], halos_array['hrvir'], - LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, CSMF_hod_dict, want_LRG, want_ELG, want_QSO, want_CSMF, - rsd, inv_velz2kms, lbox, keep_cent, Nthread = Nthread) + warnings.warn( + 'NFW profile is unoptimized. It has different velocity bias. It does not support lightcone.' + ) + LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, ID_dict_sat = gen_sats_nfw( + NFW_draw, + halos_array['hpos'], + halos_array['hvel'], + halos_array['hmass'], + halos_array['hid'], + halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), + halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), + halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), + halos_array['hsigma3d'], + halos_array['hc'], + halos_array['hrvir'], + LRG_hod_dict, + ELG_hod_dict, + QSO_hod_dict, + want_LRG, + want_ELG, + want_QSO, + rsd, + inv_velz2kms, + lbox, + keep_cent, + Nthread=Nthread, + ) else: - LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, CSMF_dict_sat, ID_dict_sat = \ - gen_sats(subsample['ppos'], subsample['pvel'], subsample['phvel'], subsample['phmass'], subsample['phid'], - subsample['pweights'], subsample['prandoms'], - subsample.get('pdeltac', np.zeros(len(subsample['phid']))), - subsample.get('pfenv', np.zeros(len(subsample['phid']))), - subsample.get('pshear', np.zeros(len(subsample['phid']))), - enable_ranks, subsample['pranks'], subsample['pranksv'], subsample['pranksp'], subsample['pranksr'], subsample['pranksc'], - LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, CSMF_hod_dict, rsd, inv_velz2kms, lbox, params['Mpart'], - want_LRG, want_ELG, want_QSO, want_CSMF, Nthread, origin, keep_cent[subsample['pinds']]) + LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, ID_dict_sat = gen_sats( + subsample['ppos'], + subsample['pvel'], + subsample['phvel'], + subsample['phmass'], + subsample['phid'], + subsample['pweights'], + subsample['prandoms'], + subsample.get('pdeltac', np.zeros(len(subsample['phid']))), + subsample.get('pfenv', np.zeros(len(subsample['phid']))), + subsample.get('pshear', np.zeros(len(subsample['phid']))), + enable_ranks, + subsample['pranks'], + subsample['pranksv'], + subsample['pranksp'], + subsample['pranksr'], + subsample['pranksc'], + LRG_hod_dict, + ELG_hod_dict, + QSO_hod_dict, + rsd, + inv_velz2kms, + lbox, + params['Mpart'], + want_LRG, + want_ELG, + want_QSO, + Nthread, + origin, + keep_cent[subsample['pinds']], + ) if verbose: - print("generating satellites took ", time.time() - start) - - + print('generating satellites took ', time.time() - start) + # B.H. TODO: need a for loop above so we don't need to do this by hand - HOD_dict_sat = {'LRG': LRG_dict_sat, 'ELG': ELG_dict_sat, 'QSO': QSO_dict_sat, 'CSMF': CSMF_dict_sat} - HOD_dict_cent = {'LRG': LRG_dict_cent, 'ELG': ELG_dict_cent, 'QSO': QSO_dict_cent, 'CSMF': CSMF_dict_cent} + HOD_dict_sat = {'LRG': LRG_dict_sat, 'ELG': ELG_dict_sat, 'QSO': QSO_dict_sat} + HOD_dict_cent = {'LRG': LRG_dict_cent, 'ELG': ELG_dict_cent, 'QSO': QSO_dict_cent} - # do a concatenate in numba parallel start = time.time() HOD_dict = {} for tracer in tracers: - tracer_dict = {'Ncent':len(HOD_dict_cent[tracer]['x'])} + tracer_dict = {'Ncent': len(HOD_dict_cent[tracer]['x'])} for k in HOD_dict_cent[tracer]: - tracer_dict[k] = fast_concatenate(HOD_dict_cent[tracer][k], HOD_dict_sat[tracer][k], Nthread) - tracer_dict['id'] = fast_concatenate(ID_dict_cent[tracer], ID_dict_sat[tracer], Nthread) + tracer_dict[k] = fast_concatenate( + HOD_dict_cent[tracer][k], HOD_dict_sat[tracer][k], Nthread + ) + tracer_dict['id'] = fast_concatenate( + ID_dict_cent[tracer], ID_dict_sat[tracer], Nthread + ) if verbose: - print(tracer, "number of galaxies ", len(tracer_dict['x'])) - print("satellite fraction ", len(HOD_dict_sat[tracer]['x'])/len(tracer_dict['x'])) + print(tracer, 'number of galaxies ', len(tracer_dict['x'])) + print( + 'satellite fraction ', + len(HOD_dict_sat[tracer]['x']) / len(tracer_dict['x']), + ) HOD_dict[tracer] = tracer_dict if verbose: - print("organizing outputs took ", time.time() - start) + print('organizing outputs took ', time.time() - start) return HOD_dict -def gen_gal_cat(halo_data, particle_data, tracers, params, Nthread = 16, - enable_ranks = False, rsd = True, nfw = False, NFW_draw = None, - write_to_disk = False, savedir = "./", verbose = False, fn_ext = None): +def gen_gal_cat( + halo_data, + particle_data, + tracers, + params, + Nthread=16, + enable_ranks=False, + rsd=True, + nfw=False, + NFW_draw=None, + write_to_disk=False, + savedir='./', + verbose=False, + fn_ext=None, +): """ pass on inputs to the gen_gals function and takes care of I/O @@ -1545,42 +1655,67 @@ def gen_gal_cat(halo_data, particle_data, tracers, params, Nthread = 16, """ if not isinstance(rsd, bool): - raise ValueError("Error: rsd has to be a boolean") + raise ValueError('Error: rsd has to be a boolean') # find the halos, populate them with galaxies and write them to files - HOD_dict = gen_gals(halo_data, particle_data, tracers, params, Nthread, enable_ranks, rsd, verbose, nfw, NFW_draw) + HOD_dict = gen_gals( + halo_data, + particle_data, + tracers, + params, + Nthread, + enable_ranks, + rsd, + verbose, + nfw, + NFW_draw, + ) # how many galaxies were generated and write them to disk for tracer in tracers.keys(): Ncent = HOD_dict[tracer]['Ncent'] if verbose: - print("generated %ss:"%tracer, len(HOD_dict[tracer]['x']), - "satellite fraction ", 1 - Ncent/len(HOD_dict[tracer]['x'])) + print( + 'generated %ss:' % tracer, + len(HOD_dict[tracer]['x']), + 'satellite fraction ', + 1 - Ncent / len(HOD_dict[tracer]['x']), + ) if write_to_disk: if verbose: - print("outputting galaxies to disk") + print('outputting galaxies to disk') if rsd: - rsd_string = "_rsd" + rsd_string = '_rsd' else: - rsd_string = "" + rsd_string = '' if fn_ext is None: - outdir = (savedir) / ("galaxies"+rsd_string) + outdir = (savedir) / ('galaxies' + rsd_string) else: - outdir = (savedir) / ("galaxies"+rsd_string+fn_ext) + outdir = (savedir) / ('galaxies' + rsd_string + fn_ext) # create directories if not existing - os.makedirs(outdir, exist_ok = True) + os.makedirs(outdir, exist_ok=True) # save to file # outdict = HOD_dict[tracer].pop('Ncent', None) - table = Table(HOD_dict[tracer], meta = {'Ncent': Ncent, 'Gal_type': tracer, **tracers[tracer]}) + table = Table( + HOD_dict[tracer], + meta={'Ncent': Ncent, 'Gal_type': tracer, **tracers[tracer]}, + ) if params['chunk'] == -1: - ascii.write(table, outdir / (f"{tracer}s.dat"), overwrite = True, format = 'ecsv') + ascii.write( + table, outdir / (f'{tracer}s.dat'), overwrite=True, format='ecsv' + ) else: - ascii.write(table, outdir / (f"{tracer}s_chunk{params['chunk']:d}.dat"), overwrite = True, format = 'ecsv') + ascii.write( + table, + outdir / (f"{tracer}s_chunk{params['chunk']:d}.dat"), + overwrite=True, + format='ecsv', + ) return HOD_dict diff --git a/abacusnbody/hod/abacus_hod.py b/abacusnbody/hod/abacus_hod.py index 6d7c615b..0ad22d49 100644 --- a/abacusnbody/hod/abacus_hod.py +++ b/abacusnbody/hod/abacus_hod.py @@ -33,6 +33,8 @@ N_sat_elg, N_cen_QSO, N_sat_generic, + n_cen_CSMF, + n_sat_CSMF ) # TODO B.H.: staging can be shorter and prettier; perhaps asdf for h5 and ecsv? @@ -392,6 +394,7 @@ def staging(self): ('ELG' not in self.tracers.keys()) and ('QSO' not in self.tracers.keys()) and (not self.force_mt) + and ('CSMF' not in self.tracers.keys()) ): halofilename = subsample_dir / ( 'halos_xcom_%d_seed600_abacushod_oldfenv' % eslab From 7b2b4339e9da7f6d3e75c647d7b1a0489670cd39 Mon Sep 17 00:00:00 2001 From: Pierre Burger Date: Fri, 20 Sep 2024 11:21:20 -0400 Subject: [PATCH 03/13] CSMF HOD Addition of a conditional stellar mass function HOD --- abacusnbody/hod/GRAND_HOD.py | 581 +++++++++++++++++++++++++++++++++- abacusnbody/hod/abacus_hod.py | 31 +- 2 files changed, 592 insertions(+), 20 deletions(-) diff --git a/abacusnbody/hod/GRAND_HOD.py b/abacusnbody/hod/GRAND_HOD.py index 739790dd..0723944e 100644 --- a/abacusnbody/hod/GRAND_HOD.py +++ b/abacusnbody/hod/GRAND_HOD.py @@ -19,6 +19,196 @@ G = 4.302e-6 # in kpc/Msol (km.s)^2 + +@njit(fastmath=True) +def n_cen_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): + """ + Standard Cacciato et al. (2008) centrals HOD parametrization for CSMF + """ + M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) + x_low = np.log10(Mstar_low/M_c_value)/(1.41421356*sigma_c) + x_up = np.log10(Mstar_up/M_c_value)/(1.41421356*sigma_c) + + return 0.5*(math.erf(x_up) - math.erf(x_low)) + +@njit(fastmath=True) +def CSMF_centrals(M_h, Mstar, M_1, M_0, gamma1, gamma2, sigma_c): + """ + Eq. (34) from Cacciato et al. (2008) + """ + + M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) + + return 1/(1.41421356*np.sqrt(np.pi)*np.log(10)*sigma_c*Mstar)*np.exp(-(np.log10(Mstar)-np.log10(M_c_value))**2/(2*sigma_c**2)) + +@njit(fastmath=True) +def M_c(M_h, M_1, M_0, gamma1, gamma2): + """ + Eq. (37) from Cacciato et al. (2008) + """ + return M_0 * (M_h/M_1)**gamma1/(1+M_h/M_1)**(gamma1-gamma2) + + + +@njit(fastmath=True) +def get_random_cen_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): + + nbins = 1000 + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + + CSMF_cen = CSMF_centrals(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c) + + cdf = [] + cdf_current_value = 0 + for k in range(len(delta_stellar_masses)): + cdf_current_value = cdf_current_value + CSMF_cen[k]*delta_stellar_masses[k] + cdf.append(cdf_current_value) + cdf=np.array(cdf) + + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin_clostest = (np.abs(cdf - random_rv)).argmin() + + return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) + + + + +@njit(fastmath=True) +def get_random_cen_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): + + nbins = 1000 + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + + CSMF_cen = CSMF_centrals(M_h, stellarmass, M_1, M_0, gamma1, gamma2, sigma_c) + + delta=stellarmass[1:]-stellarmass[:-1] + + cdf = [] + cdf_current_value = 0 + for i in range(len(delta)): + cdf_current_value = cdf_current_value + CSMF_cen[i]*delta[i] + cdf.append(cdf_current_value) + cdf=np.array(cdf) + + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin = np.where(cdf>random_rv)[0][0] + + m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) + return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + + +@njit(fastmath=True) +def n_sat_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c, a1, a2, M2, b0, b1, b2, delta1, delta2): + """ + Standard Cacciato et al. (2008) satellite HOD parametrization for CSMF + """ + nbins = 1000 + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + + CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + + nsat = 0 + for i in range(nbins-1): + nsat += (CSMF_sat[i+1]-CSMF_sat[i])*(stellarmass[i+1]-stellarmass[i])/2 + (stellarmass[i+1]-stellarmass[i])*CSMF_sat[i] + + return nsat#*ncen + + +@njit(fastmath=True) +def CSMF_satelites(M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): + """ + Eq. (36) from Cacciato et al. (2008) + """ + M_s_value = M_s(M_h, M_1, M_0, gamma1, gamma2) + alpha_s_value = alpha_s(M_h, a1, a2, M2) + phi_s_value = phi_s(M_h, b0, b1, b2) + + delta = 10 ** (delta1 + delta2 * (np.log10(M_h) - 12)) + + return phi_s_value/M_s_value * (Mstar/M_s_value)**alpha_s_value * np.exp(-delta*(Mstar/M_s_value)**2) + +@njit(fastmath=True) +def M_s(M_h, M_1, M_0, gamma1, gamma2): + """ + Eq. (38) from Cacciato et al. (2008) + """ + return 0.562 * M_c(M_h, M_1, M_0, gamma1, gamma2) + +@njit(fastmath=True) +def alpha_s(M_h, a1, a2, M2): + """ + Eq. (39) from Cacciato et al. (2008) + """ + return -2.0 + a1 * (1-2/np.pi*np.arctan(a2*np.log10(M_h/M2))) + +@njit(fastmath=True) +def phi_s(M_h, b0, b1, b2): + """ + Eq. (40) from Cacciato et al. (2008) + """ + M12 = M_h/1e12 + log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12)**2 + return 10**log_phi_s + + +@njit(fastmath=True) +def get_random_sat_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): + + nbins = 1000 + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + + CSMF_sat = CSMF_satelites(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + + + cdf = [] + cdf_current_value = 0 + for k in range(len(delta_stellar_masses)): + cdf_current_value = cdf_current_value + CSMF_sat[k]*delta_stellar_masses[k] + cdf.append(cdf_current_value) + cdf=np.array(cdf) + + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin_clostest = (np.abs(cdf - random_rv)).argmin() + + return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) + + +@njit(fastmath=True) +def get_random_sat_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): + + nbins = 100 + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + + CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + + delta=stellarmass[1:]-stellarmass[:-1] + + cdf = [] + cdf_current_value = 0 + for i in range(len(delta)): + cdf_current_value = cdf_current_value + CSMF_sat[i]*delta[i] + cdf.append(cdf_current_value) + cdf=np.array(cdf) + + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin = np.where(cdf>random_rv)[0][0] + + m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) + return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + + @njit(fastmath=True) def n_sat_LRG_modified(M_h, logM_cut, M_cut, M_1, sigma, alpha, kappa): """ @@ -150,12 +340,14 @@ def gen_cent( LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, + CSMF_hod_dict, rsd, inv_velz2kms, lbox, want_LRG, want_ELG, want_QSO, + want_CSMF, Nthread, origin, ): @@ -197,11 +389,28 @@ def gen_cent( QSO_hod_dict['Bcent'], QSO_hod_dict['ic'], ) + if want_CSMF: + Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'] + ) + ic_C, alpha_c_C, Ac_C, Bc_C = ( + CSMF_hod_dict['ic'], + CSMF_hod_dict['alpha_c'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Bcent'] + ) + H = len(mass) numba.set_num_threads(Nthread) - Nout = np.zeros((Nthread, 3, 8), dtype=np.int64) + Nout = np.zeros((Nthread, 4, 8), dtype=np.int64) hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype( np.int64 ) # starting index of each thread @@ -237,6 +446,14 @@ def gen_cent( QSO_marker += ( N_cen_QSO(mass[i], logM_cut_Q_temp, sigma_Q) * ic_Q * multis[i] ) + + CSMF_marker = QSO_marker + if want_CSMF: + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + + ncen = n_cen_CSMF(mass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) + + CSMF_marker += ncen * ic_C * multis[i] if randoms[i] <= LRG_marker: Nout[tid, 0, 0] += 1 # counting @@ -247,15 +464,19 @@ def gen_cent( elif randoms[i] <= QSO_marker: Nout[tid, 2, 0] += 1 # counting keep[i] = 3 + elif randoms[i] <= CSMF_marker: + Nout[tid, 3, 0] += 1 # counting + keep[i] = 4 else: keep[i] = 0 # compose galaxy array, first create array of galaxy starting indices for the threads - gstart = np.empty((Nthread + 1, 3), dtype=np.int64) + gstart = np.empty((Nthread + 1, 4), dtype=np.int64) gstart[0, :] = 0 gstart[1:, 0] = Nout[:, 0, 0].cumsum() gstart[1:, 1] = Nout[:, 1, 0].cumsum() gstart[1:, 2] = Nout[:, 2, 0].cumsum() + gstart[1:, 3] = Nout[:, 3, 0].cumsum() # galaxy arrays N_lrg = gstart[-1, 0] @@ -289,10 +510,22 @@ def gen_cent( qso_vz = np.empty(N_qso, dtype=mass.dtype) qso_mass = np.empty(N_qso, dtype=mass.dtype) qso_id = np.empty(N_qso, dtype=ids.dtype) + + # galaxy arrays + N_CSMF = gstart[-1, 3] + CSMF_x = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_y = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_z = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_id = np.empty(N_CSMF, dtype = ids.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): - j1, j2, j3 = gstart[tid] + j1, j2, j3, j4 = gstart[tid] for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: # loop thru three directions to assign galaxy velocities and positions @@ -378,11 +611,44 @@ def gen_cent( qso_mass[j3] = mass[i] qso_id[j3] = ids[i] j3 += 1 + elif keep[i] == 4: + # loop thru three directions to assign galaxy velocities and positions + CSMF_x[j4] = pos[i,0] + CSMF_vx[j4] = vel[i,0] + alpha_c_C * vdev[i,0] # velocity bias + CSMF_y[j4] = pos[i,1] + CSMF_vy[j4] = vel[i,1] + alpha_c_C * vdev[i,1] # velocity bias + CSMF_z[j4] = pos[i,2] + CSMF_vz[j4] = vel[i,2] + alpha_c_C * vdev[i,2] # velocity bias + + # rsd only applies to the z direction + if rsd and origin is not None: + nx = CSMF_x[j4] - origin[0] + ny = CSMF_y[j4] - origin[1] + nz = CSMF_z[j4] - origin[2] + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + nx *= inv_norm + ny *= inv_norm + nz *= inv_norm + proj = inv_velz2kms*(CSMF_vx[j4]*nx+CSMF_vy[j4]*ny+CSMF_vz[j4]*nz) + CSMF_x[j4] = CSMF_x[j4]+proj*nx + CSMF_y[j4] = CSMF_y[j4]+proj*ny + CSMF_z[j4] = CSMF_z[j4]+proj*nz + elif rsd: + CSMF_z[j4] = wrap(pos[i,2] + CSMF_vz[j4] * inv_velz2kms, lbox) + + CSMF_mass[j4] = mass[i] + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + CSMF_stellarmass[j4] = get_random_cen_stellarmass_linearinterpolation( + mass[i], Mstar_low_C, Mstar_up_C, + M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) + CSMF_id[j4] = ids[i] + j4 += 1 # assert j == gstart[tid + 1] LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) LRG_dict['x'] = lrg_x LRG_dict['y'] = lrg_y @@ -410,7 +676,17 @@ def gen_cent( QSO_dict['vz'] = qso_vz QSO_dict['mass'] = qso_mass ID_dict['QSO'] = qso_id - return LRG_dict, ELG_dict, QSO_dict, ID_dict, keep + + CSMF_dict['x'] = CSMF_x + CSMF_dict['y'] = CSMF_y + CSMF_dict['z'] = CSMF_z + CSMF_dict['vx'] = CSMF_vx + CSMF_dict['vy'] = CSMF_vy + CSMF_dict['vz'] = CSMF_vz + CSMF_dict['mass'] = CSMF_mass + CSMF_dict['stellarmass'] = CSMF_stellarmass + ID_dict['CSMF'] = CSMF_id + return LRG_dict, ELG_dict, QSO_dict, CSMF_dict, ID_dict, keep @njit(parallel=True, fastmath=True) @@ -534,9 +810,11 @@ def gen_sats_nfw( LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, + CSMF_hod_dict, want_LRG, want_ELG, want_QSO, + want_CSMF, rsd, inv_velz2kms, lbox, @@ -619,6 +897,33 @@ def gen_sats_nfw( QSO_hod_dict['ic'], ) f_sigv_Q = QSO_hod_dict['f_sigv'] + + if want_CSMF: + Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], + CSMF_hod_dict['M_2'], + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'] + ) + Ac_C, As_C, Bc_C, Bs_C, ic_C = ( + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Asat'], + CSMF_hod_dict['Bcent'], + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'] + ) + f_sigv_C = CSMF_hod_dict['f_sigv'] numba.set_num_threads(Nthread) @@ -627,6 +932,8 @@ def gen_sats_nfw( num_sats_L = np.zeros(len(hid), dtype=np.int64) num_sats_E = np.zeros(len(hid), dtype=np.int64) num_sats_Q = np.zeros(len(hid), dtype=np.int64) + num_sats_C = np.zeros(len(hid), dtype = np.int64) + stellarmass_C = np.zeros(len(hid), dtype = np.int64) hstart = np.rint(np.linspace(0, len(hid), Nthread + 1)).astype( np.int64 ) # starting index of each thread @@ -702,11 +1009,38 @@ def gen_sats_nfw( * ic_Q ) num_sats_Q[i] = np.random.poisson(base_p_Q) + + if want_CSMF: + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] + base_p_C = ( + n_sat_CSMF( + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C) + * ic_C + ) + + num_sats_C[i] = np.random.poisson(base_p_C) # generate rdpos rd_pos_L = getPointsOnSphere(np.sum(num_sats_L), Nthread) rd_pos_E = getPointsOnSphere(np.sum(num_sats_E), Nthread) rd_pos_Q = getPointsOnSphere(np.sum(num_sats_Q), Nthread) + rd_pos_C = getPointsOnSphere(np.sum(num_sats_C), Nthread) # put satellites on NFW h_id_L, x_sat_L, y_sat_L, z_sat_L, vx_sat_L, vy_sat_L, vz_sat_L, M_L = ( @@ -781,15 +1115,44 @@ def gen_sats_nfw( nfw_rescale, ) ) + + h_id_C, x_sat_C, y_sat_C, z_sat_C, vx_sat_C, vy_sat_C, vz_sat_C, M_C = ( + compute_fast_NFW( + NFW_draw, + hid, + hpos[:, 0], + hpos[:, 1], + hpos[:, 2], + hvel[:, 0], + hvel[:, 1], + hvel[:, 2], + hvrms, + hc, + hmass, + hrvir, + rd_pos_C, + num_sats_C, + f_sigv_C, + vel_sat, + Nthread, + exp_frac, + exp_scale, + nfw_rescale + ) + ) + + # do rsd if rsd: z_sat_L = (z_sat_L + vz_sat_L * inv_velz2kms) % lbox z_sat_E = (z_sat_E + vz_sat_E * inv_velz2kms) % lbox z_sat_Q = (z_sat_Q + vz_sat_Q * inv_velz2kms) % lbox + z_sat_C = (z_sat_C + vz_sat_C * inv_velz2kms) % lbox LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) LRG_dict['x'] = x_sat_L LRG_dict['y'] = y_sat_L @@ -817,8 +1180,30 @@ def gen_sats_nfw( QSO_dict['vz'] = vz_sat_Q QSO_dict['mass'] = M_Q ID_dict['QSO'] = h_id_Q - - return LRG_dict, ELG_dict, QSO_dict, ID_dict + + CSMF_dict['x'] = x_sat_C + CSMF_dict['y'] = y_sat_C + CSMF_dict['z'] = z_sat_C + CSMF_dict['vx'] = vx_sat_C + CSMF_dict['vy'] = vy_sat_C + CSMF_dict['vz'] = vz_sat_C + CSMF_dict['mass'] = M_C + stellarmass_C = np.empty_like(M_C) + ## compute stellarmass of all the satelites + if want_CSMF: + hstart = np.rint(np.linspace(0, num_sats_C.sum(), Nthread + 1)) + for tid in numba.prange(Nthread): + for i in range(int(hstart[tid]), int(hstart[tid + 1])): + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] + stellarmass_C[i] = get_random_sat_stellarmass_linearinterpolation( + M_C[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, + a1_C_temp, s, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) + + CSMF_dict['stellarmass'] = stellarmass_C + ID_dict['CSMF'] = h_id_C + + return LRG_dict, ELG_dict, QSO_dict, CSMF_dict, ID_dict @njit(parallel=True, fastmath=True) @@ -842,6 +1227,7 @@ def gen_sats( LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, + CSMF_hod_dict, rsd, inv_velz2kms, lbox, @@ -849,6 +1235,7 @@ def gen_sats( want_LRG, want_ELG, want_QSO, + want_CSMF, Nthread, origin, keep_cent, @@ -941,11 +1328,43 @@ def gen_sats( QSO_hod_dict['Bsat'], QSO_hod_dict['ic'], ) + + if want_CSMF: + Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], + CSMF_hod_dict['M_2'], + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'] + ) + + alpha_s_C, s_C, s_v_C, s_p_C, s_r_C, Ac_C, As_C, Bc_C, Bs_C, ic_C = ( + CSMF_hod_dict['alpha_s'], + CSMF_hod_dict['s'], + CSMF_hod_dict['s_v'], + CSMF_hod_dict['s_p'], + CSMF_hod_dict['s_r'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Asat'], + CSMF_hod_dict['Bcent'], + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'] + ) H = len(hmass) # num of particles numba.set_num_threads(Nthread) - Nout = np.zeros((Nthread, 3, 8), dtype=np.int64) + Nout = np.zeros((Nthread, 4, 8), dtype=np.int64) hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype( np.int64 ) # starting index of each thread @@ -1073,6 +1492,39 @@ def gen_sats( else: exp_sat = base_p_Q QSO_marker += exp_sat + + + CSMF_marker = QSO_marker + if want_CSMF: + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] + base_p_C = ( + n_sat_CSMF( + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C + ) + * weights[i] + * ic_C) + if enable_ranks: + decorator_C = 1 + s_C * ranks[i] + s_v_C * ranksv[i] + s_p_C * ranksp[i] + s_r_C * ranksr[i] + exp_sat = base_p_C * decorator_C + else: + exp_sat = base_p_C + CSMF_marker += exp_sat if randoms[i] <= LRG_marker: Nout[tid, 0, 0] += 1 # counting @@ -1083,15 +1535,19 @@ def gen_sats( elif randoms[i] <= QSO_marker: Nout[tid, 2, 0] += 1 # counting keep[i] = 3 + elif randoms[i] <= CSMF_marker: + Nout[tid, 3, 0] += 1 # counting + keep[i] = 4 else: keep[i] = 0 # compose galaxy array, first create array of galaxy starting indices for the threads - gstart = np.empty((Nthread + 1, 3), dtype=np.int64) + gstart = np.empty((Nthread + 1, 4), dtype=np.int64) gstart[0, :] = 0 gstart[1:, 0] = Nout[:, 0, 0].cumsum() gstart[1:, 1] = Nout[:, 1, 0].cumsum() gstart[1:, 2] = Nout[:, 2, 0].cumsum() + gstart[1:, 3] = Nout[:, 3, 0].cumsum() # galaxy arrays N_lrg = gstart[-1, 0] @@ -1125,10 +1581,22 @@ def gen_sats( qso_vz = np.empty(N_qso, dtype=hmass.dtype) qso_mass = np.empty(N_qso, dtype=hmass.dtype) qso_id = np.empty(N_qso, dtype=hid.dtype) + + # galaxy arrays + N_CSMF = gstart[-1, 3] + CSMF_x = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_y = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_z = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_id = np.empty(N_CSMF, dtype = hid.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): - j1, j2, j3 = gstart[tid] + j1, j2, j3, j4 = gstart[tid] for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: lrg_x[j1] = ppos[i, 0] @@ -1226,11 +1694,57 @@ def gen_sats( qso_mass[j3] = hmass[i] qso_id[j3] = hid[i] j3 += 1 + elif keep[i] == 4: + CSMF_x[j4] = ppos[i, 0] + CSMF_vx[j4] = hvel[i, 0] + alpha_s_C * (pvel[i, 0] - hvel[i, 0]) # velocity bias + CSMF_y[j4] = ppos[i, 1] + CSMF_vy[j4] = hvel[i, 1] + alpha_s_C * (pvel[i, 1] - hvel[i, 1]) # velocity bias + CSMF_z[j4] = ppos[i, 2] + CSMF_vz[j4] = hvel[i, 2] + alpha_s_C * (pvel[i, 2] - hvel[i, 2]) # velocity bias + if rsd and origin is not None: + nx = CSMF_x[j4] - origin[0] + ny = CSMF_y[j4] - origin[1] + nz = CSMF_z[j4] - origin[2] + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + nx *= inv_norm + ny *= inv_norm + nz *= inv_norm + proj = inv_velz2kms*(CSMF_vx[j4]*nx+CSMF_vy[j4]*ny+CSMF_vz[j4]*nz) + CSMF_x[j4] = CSMF_x[j4]+proj*nx + CSMF_y[j4] = CSMF_y[j4]+proj*ny + CSMF_z[j4] = CSMF_z[j4]+proj*nz + elif rsd: + CSMF_z[j4] = wrap(CSMF_z[j4] + CSMF_vz[j4] * inv_velz2kms, lbox) + + + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] + CSMF_stellarmass[j4] = get_random_sat_stellarmass( + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C + ) + CSMF_mass[j4] = hmass[i] + CSMF_id[j4] = hid[i] + j4 += 1 # assert j == gstart[tid + 1] LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) LRG_dict['x'] = lrg_x LRG_dict['y'] = lrg_y @@ -1258,7 +1772,18 @@ def gen_sats( QSO_dict['vz'] = qso_vz QSO_dict['mass'] = qso_mass ID_dict['QSO'] = qso_id - return LRG_dict, ELG_dict, QSO_dict, ID_dict + + CSMF_dict['x'] = CSMF_x + CSMF_dict['y'] = CSMF_y + CSMF_dict['z'] = CSMF_z + CSMF_dict['vx'] = CSMF_vx + CSMF_dict['vy'] = CSMF_vy + CSMF_dict['vz'] = CSMF_vz + CSMF_dict['mass'] = CSMF_mass + CSMF_dict['stellarmass'] = CSMF_stellarmass + ID_dict['CSMF'] = CSMF_id + + return LRG_dict, ELG_dict, QSO_dict, CSMF_dict, ID_dict @njit(parallel=True, fastmath=True) @@ -1345,6 +1870,8 @@ def gen_gals( ELG_HOD = tracers[tracer] if tracer == 'QSO': QSO_HOD = tracers[tracer] + if tracer == 'CSMF': + CSMF_HOD = tracers[tracer] if 'LRG' in tracers.keys(): want_LRG = True @@ -1465,6 +1992,24 @@ def gen_gals( QSO_hod_dict = nb.typed.Dict.empty( key_type=nb.types.unicode_type, value_type=nb.types.float64 ) + + if 'CSMF' in tracers.keys(): + want_CSMF = True + + CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + for key, value in CSMF_HOD.items(): + CSMF_hod_dict[key] = value + + CSMF_hod_dict['Acent'] = CSMF_HOD.get('Acent', 0.0) + CSMF_hod_dict['Asat'] = CSMF_HOD.get('Asat', 0.0) + CSMF_hod_dict['Bcent'] = CSMF_HOD.get('Bcent', 0.0) + CSMF_hod_dict['Bsat'] = CSMF_HOD.get('Bsat', 0.0) + CSMF_hod_dict['ic'] = CSMF_HOD.get('ic', 1.0) + CSMF_hod_dict['f_sigv'] = CSMF_HOD.get('f_sigv', 0) + + else: + want_CSMF = False + CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) start = time.time() @@ -1473,7 +2018,7 @@ def gen_gals( lbox = params['Lbox'] origin = params['origin'] - LRG_dict_cent, ELG_dict_cent, QSO_dict_cent, ID_dict_cent, keep_cent = gen_cent( + LRG_dict_cent, ELG_dict_cent, QSO_dict_cent, CSMF_dict_cent, ID_dict_cent, keep_cent = gen_cent( halos_array['hpos'], halos_array['hvel'], halos_array['hmass'], @@ -1487,12 +2032,14 @@ def gen_gals( LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, + CSMF_hod_dict, rsd, inv_velz2kms, lbox, want_LRG, want_ELG, want_QSO, + want_CSMF, Nthread, origin, ) @@ -1504,7 +2051,7 @@ def gen_gals( warnings.warn( 'NFW profile is unoptimized. It has different velocity bias. It does not support lightcone.' ) - LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, ID_dict_sat = gen_sats_nfw( + LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, CSMF_dict_sat, ID_dict_sat = gen_sats_nfw( NFW_draw, halos_array['hpos'], halos_array['hvel'], @@ -1519,9 +2066,11 @@ def gen_gals( LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, + CSMF_hod_dict, want_LRG, want_ELG, want_QSO, + want_CSMF, rsd, inv_velz2kms, lbox, @@ -1529,7 +2078,7 @@ def gen_gals( Nthread=Nthread, ) else: - LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, ID_dict_sat = gen_sats( + LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, CSMF_dict_sat, ID_dict_sat = gen_sats( subsample['ppos'], subsample['pvel'], subsample['phvel'], @@ -1549,6 +2098,7 @@ def gen_gals( LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, + CSMF_hod_dict, rsd, inv_velz2kms, lbox, @@ -1556,6 +2106,7 @@ def gen_gals( want_LRG, want_ELG, want_QSO, + want_CSMF, Nthread, origin, keep_cent[subsample['pinds']], @@ -1564,8 +2115,8 @@ def gen_gals( print('generating satellites took ', time.time() - start) # B.H. TODO: need a for loop above so we don't need to do this by hand - HOD_dict_sat = {'LRG': LRG_dict_sat, 'ELG': ELG_dict_sat, 'QSO': QSO_dict_sat} - HOD_dict_cent = {'LRG': LRG_dict_cent, 'ELG': ELG_dict_cent, 'QSO': QSO_dict_cent} + HOD_dict_sat = {'LRG': LRG_dict_sat, 'ELG': ELG_dict_sat, 'QSO': QSO_dict_sat, 'CSMF': CSMF_dict_sat} + HOD_dict_cent = {'LRG': LRG_dict_cent, 'ELG': ELG_dict_cent, 'QSO': QSO_dict_cent, 'CSMF': CSMF_dict_cent} # do a concatenate in numba parallel start = time.time() diff --git a/abacusnbody/hod/abacus_hod.py b/abacusnbody/hod/abacus_hod.py index 0ad22d49..cdd94383 100644 --- a/abacusnbody/hod/abacus_hod.py +++ b/abacusnbody/hod/abacus_hod.py @@ -900,11 +900,32 @@ def compute_ngal(self, tracers=None, Nthread=16): elif etracer == 'CSMF': newngal = AbacusHOD._compute_ngal_CSMF( - self.logMbins, self.deltacbins, self.fenvbins, self.halo_mass_func, - tracer_hod['Mstar_low'], tracer_hod['Mstar_up'], tracer_hod['M_1'], tracer_hod['M_0'], tracer_hod['gamma_1'], tracer_hod['gamma_2'], - tracer_hod['sigma_c'], tracer_hod['a_1'], tracer_hod['a_2'], tracer_hod['M_2'], tracer_hod['b_0'], tracer_hod['b_1'], - tracer_hod['b_2'], tracer_hod.get('delta_1', 0), tracer_hod.get('delta_2', 0), tracer_hod.get('Acent', 0), - tracer_hod.get('Asat', 0), tracer_hod.get('Bcent', 0), tracer_hod.get('Bsat', 0), tracer_hod.get('ic', 1), Nthread) + self.logMbins, + self.deltacbins, + self.fenvbins, + self.halo_mass_func, + tracer_hod['Mstar_low'], + tracer_hod['Mstar_up'], + tracer_hod['M_1'], + tracer_hod['M_0'], + tracer_hod['gamma_1'], + tracer_hod['gamma_2'], + tracer_hod['sigma_c'], + tracer_hod['a_1'], + tracer_hod['a_2'], + tracer_hod['M_2'], + tracer_hod['b_0'], + tracer_hod['b_1'], + tracer_hod['b_2'], + tracer_hod.get('delta_1', 0), + tracer_hod.get('delta_2', 0), + tracer_hod.get('Acent', 0), + tracer_hod.get('Asat', 0), + tracer_hod.get('Bcent', 0), + tracer_hod.get('Bsat', 0), + tracer_hod.get('ic', 1), + Nthread + ) ngal_dict[etracer] = newngal[0] + newngal[1] fsat_dict[etracer] = newngal[1] / (newngal[0] + newngal[1]) return ngal_dict, fsat_dict From 3f623443558adedd2faf53823bd1270ba21f2d50 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 20 Sep 2024 15:36:12 +0000 Subject: [PATCH 04/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- abacusnbody/hod/GRAND_HOD.py | 887 +++++++++++++++++++++------------- abacusnbody/hod/abacus_hod.py | 137 ++++-- 2 files changed, 647 insertions(+), 377 deletions(-) diff --git a/abacusnbody/hod/GRAND_HOD.py b/abacusnbody/hod/GRAND_HOD.py index 0723944e..9c9afbfb 100644 --- a/abacusnbody/hod/GRAND_HOD.py +++ b/abacusnbody/hod/GRAND_HOD.py @@ -19,119 +19,167 @@ G = 4.302e-6 # in kpc/Msol (km.s)^2 - @njit(fastmath=True) def n_cen_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): """ Standard Cacciato et al. (2008) centrals HOD parametrization for CSMF """ M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - x_low = np.log10(Mstar_low/M_c_value)/(1.41421356*sigma_c) - x_up = np.log10(Mstar_up/M_c_value)/(1.41421356*sigma_c) - - return 0.5*(math.erf(x_up) - math.erf(x_low)) + x_low = np.log10(Mstar_low / M_c_value) / (1.41421356 * sigma_c) + x_up = np.log10(Mstar_up / M_c_value) / (1.41421356 * sigma_c) + + return 0.5 * (math.erf(x_up) - math.erf(x_low)) + @njit(fastmath=True) def CSMF_centrals(M_h, Mstar, M_1, M_0, gamma1, gamma2, sigma_c): """ Eq. (34) from Cacciato et al. (2008) """ - + M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - - return 1/(1.41421356*np.sqrt(np.pi)*np.log(10)*sigma_c*Mstar)*np.exp(-(np.log10(Mstar)-np.log10(M_c_value))**2/(2*sigma_c**2)) + + return ( + 1 + / (1.41421356 * np.sqrt(np.pi) * np.log(10) * sigma_c * Mstar) + * np.exp(-((np.log10(Mstar) - np.log10(M_c_value)) ** 2) / (2 * sigma_c**2)) + ) + @njit(fastmath=True) def M_c(M_h, M_1, M_0, gamma1, gamma2): """ Eq. (37) from Cacciato et al. (2008) """ - return M_0 * (M_h/M_1)**gamma1/(1+M_h/M_1)**(gamma1-gamma2) - + return M_0 * (M_h / M_1) ** gamma1 / (1 + M_h / M_1) ** (gamma1 - gamma2) @njit(fastmath=True) -def get_random_cen_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): - +def get_random_cen_stellarmass( + M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c +): nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + + CSMF_cen = CSMF_centrals( + M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c + ) - CSMF_cen = CSMF_centrals(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c) - - cdf = [] + cdf = [] cdf_current_value = 0 for k in range(len(delta_stellar_masses)): - cdf_current_value = cdf_current_value + CSMF_cen[k]*delta_stellar_masses[k] + cdf_current_value = cdf_current_value + CSMF_cen[k] * delta_stellar_masses[k] cdf.append(cdf_current_value) - cdf=np.array(cdf) - - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) + cdf = np.array(cdf) + + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) - + + return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) - @njit(fastmath=True) -def get_random_cen_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): - +def get_random_cen_stellarmass_linearinterpolation( + M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c +): nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) CSMF_cen = CSMF_centrals(M_h, stellarmass, M_1, M_0, gamma1, gamma2, sigma_c) - - delta=stellarmass[1:]-stellarmass[:-1] - - cdf = [] + + delta = stellarmass[1:] - stellarmass[:-1] + + cdf = [] cdf_current_value = 0 for i in range(len(delta)): - cdf_current_value = cdf_current_value + CSMF_cen[i]*delta[i] + cdf_current_value = cdf_current_value + CSMF_cen[i] * delta[i] cdf.append(cdf_current_value) - cdf=np.array(cdf) - - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) - bin = np.where(cdf>random_rv)[0][0] - - m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) - return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + cdf = np.array(cdf) + + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) + bin = np.where(cdf > random_rv)[0][0] + + m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) + return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] @njit(fastmath=True) -def n_sat_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c, a1, a2, M2, b0, b1, b2, delta1, delta2): +def n_sat_CSMF( + M_h, + Mstar_low, + Mstar_up, + M_1, + M_0, + gamma1, + gamma2, + sigma_c, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, +): """ Standard Cacciato et al. (2008) satellite HOD parametrization for CSMF """ nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - - CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + + CSMF_sat = CSMF_satelites( + M_h, + stellarmass, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, + ) nsat = 0 - for i in range(nbins-1): - nsat += (CSMF_sat[i+1]-CSMF_sat[i])*(stellarmass[i+1]-stellarmass[i])/2 + (stellarmass[i+1]-stellarmass[i])*CSMF_sat[i] - - return nsat#*ncen + for i in range(nbins - 1): + nsat += (CSMF_sat[i + 1] - CSMF_sat[i]) * ( + stellarmass[i + 1] - stellarmass[i] + ) / 2 + (stellarmass[i + 1] - stellarmass[i]) * CSMF_sat[i] + + return nsat # *ncen @njit(fastmath=True) -def CSMF_satelites(M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): +def CSMF_satelites( + M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2 +): """ Eq. (36) from Cacciato et al. (2008) """ M_s_value = M_s(M_h, M_1, M_0, gamma1, gamma2) alpha_s_value = alpha_s(M_h, a1, a2, M2) phi_s_value = phi_s(M_h, b0, b1, b2) - + delta = 10 ** (delta1 + delta2 * (np.log10(M_h) - 12)) - - return phi_s_value/M_s_value * (Mstar/M_s_value)**alpha_s_value * np.exp(-delta*(Mstar/M_s_value)**2) + + return ( + phi_s_value + / M_s_value + * (Mstar / M_s_value) ** alpha_s_value + * np.exp(-delta * (Mstar / M_s_value) ** 2) + ) + @njit(fastmath=True) def M_s(M_h, M_1, M_0, gamma1, gamma2): @@ -140,73 +188,134 @@ def M_s(M_h, M_1, M_0, gamma1, gamma2): """ return 0.562 * M_c(M_h, M_1, M_0, gamma1, gamma2) + @njit(fastmath=True) def alpha_s(M_h, a1, a2, M2): """ Eq. (39) from Cacciato et al. (2008) """ - return -2.0 + a1 * (1-2/np.pi*np.arctan(a2*np.log10(M_h/M2))) + return -2.0 + a1 * (1 - 2 / np.pi * np.arctan(a2 * np.log10(M_h / M2))) + @njit(fastmath=True) def phi_s(M_h, b0, b1, b2): """ Eq. (40) from Cacciato et al. (2008) """ - M12 = M_h/1e12 - log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12)**2 + M12 = M_h / 1e12 + log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12) ** 2 return 10**log_phi_s @njit(fastmath=True) -def get_random_sat_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): - +def get_random_sat_stellarmass( + M_h, + Mstar_low, + Mstar_up, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, +): nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - - CSMF_sat = CSMF_satelites(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) - - - cdf = [] + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + + CSMF_sat = CSMF_satelites( + M_h, + stellarmass_centers, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, + ) + + cdf = [] cdf_current_value = 0 for k in range(len(delta_stellar_masses)): - cdf_current_value = cdf_current_value + CSMF_sat[k]*delta_stellar_masses[k] + cdf_current_value = cdf_current_value + CSMF_sat[k] * delta_stellar_masses[k] cdf.append(cdf_current_value) - cdf=np.array(cdf) - - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) + cdf = np.array(cdf) + + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) - + + return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) + @njit(fastmath=True) -def get_random_sat_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): - +def get_random_sat_stellarmass_linearinterpolation( + M_h, + Mstar_low, + Mstar_up, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, +): nbins = 100 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - - CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) - - delta=stellarmass[1:]-stellarmass[:-1] - - cdf = [] + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + + CSMF_sat = CSMF_satelites( + M_h, + stellarmass, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, + ) + + delta = stellarmass[1:] - stellarmass[:-1] + + cdf = [] cdf_current_value = 0 for i in range(len(delta)): - cdf_current_value = cdf_current_value + CSMF_sat[i]*delta[i] + cdf_current_value = cdf_current_value + CSMF_sat[i] * delta[i] cdf.append(cdf_current_value) - cdf=np.array(cdf) - - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) - bin = np.where(cdf>random_rv)[0][0] - - m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) - return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + cdf = np.array(cdf) + + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) + bin = np.where(cdf > random_rv)[0][0] + + m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) + return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] @njit(fastmath=True) @@ -391,21 +500,20 @@ def gen_cent( ) if want_CSMF: Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'] + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], ) ic_C, alpha_c_C, Ac_C, Bc_C = ( - CSMF_hod_dict['ic'], - CSMF_hod_dict['alpha_c'], - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Bcent'] + CSMF_hod_dict['ic'], + CSMF_hod_dict['alpha_c'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Bcent'], ) - H = len(mass) @@ -446,13 +554,22 @@ def gen_cent( QSO_marker += ( N_cen_QSO(mass[i], logM_cut_Q_temp, sigma_Q) * ic_Q * multis[i] ) - + CSMF_marker = QSO_marker if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) - - ncen = n_cen_CSMF(mass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) - + M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + + ncen = n_cen_CSMF( + mass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + ) + CSMF_marker += ncen * ic_C * multis[i] if randoms[i] <= LRG_marker: @@ -465,7 +582,7 @@ def gen_cent( Nout[tid, 2, 0] += 1 # counting keep[i] = 3 elif randoms[i] <= CSMF_marker: - Nout[tid, 3, 0] += 1 # counting + Nout[tid, 3, 0] += 1 # counting keep[i] = 4 else: keep[i] = 0 @@ -510,18 +627,18 @@ def gen_cent( qso_vz = np.empty(N_qso, dtype=mass.dtype) qso_mass = np.empty(N_qso, dtype=mass.dtype) qso_id = np.empty(N_qso, dtype=ids.dtype) - + # galaxy arrays N_CSMF = gstart[-1, 3] - CSMF_x = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_y = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_z = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_id = np.empty(N_CSMF, dtype = ids.dtype) + CSMF_x = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_y = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_z = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_id = np.empty(N_CSMF, dtype=ids.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): @@ -613,34 +730,43 @@ def gen_cent( j3 += 1 elif keep[i] == 4: # loop thru three directions to assign galaxy velocities and positions - CSMF_x[j4] = pos[i,0] - CSMF_vx[j4] = vel[i,0] + alpha_c_C * vdev[i,0] # velocity bias - CSMF_y[j4] = pos[i,1] - CSMF_vy[j4] = vel[i,1] + alpha_c_C * vdev[i,1] # velocity bias - CSMF_z[j4] = pos[i,2] - CSMF_vz[j4] = vel[i,2] + alpha_c_C * vdev[i,2] # velocity bias - + CSMF_x[j4] = pos[i, 0] + CSMF_vx[j4] = vel[i, 0] + alpha_c_C * vdev[i, 0] # velocity bias + CSMF_y[j4] = pos[i, 1] + CSMF_vy[j4] = vel[i, 1] + alpha_c_C * vdev[i, 1] # velocity bias + CSMF_z[j4] = pos[i, 2] + CSMF_vz[j4] = vel[i, 2] + alpha_c_C * vdev[i, 2] # velocity bias + # rsd only applies to the z direction if rsd and origin is not None: nx = CSMF_x[j4] - origin[0] ny = CSMF_y[j4] - origin[1] nz = CSMF_z[j4] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms*(CSMF_vx[j4]*nx+CSMF_vy[j4]*ny+CSMF_vz[j4]*nz) - CSMF_x[j4] = CSMF_x[j4]+proj*nx - CSMF_y[j4] = CSMF_y[j4]+proj*ny - CSMF_z[j4] = CSMF_z[j4]+proj*nz + proj = inv_velz2kms * ( + CSMF_vx[j4] * nx + CSMF_vy[j4] * ny + CSMF_vz[j4] * nz + ) + CSMF_x[j4] = CSMF_x[j4] + proj * nx + CSMF_y[j4] = CSMF_y[j4] + proj * ny + CSMF_z[j4] = CSMF_z[j4] + proj * nz elif rsd: - CSMF_z[j4] = wrap(pos[i,2] + CSMF_vz[j4] * inv_velz2kms, lbox) - + CSMF_z[j4] = wrap(pos[i, 2] + CSMF_vz[j4] * inv_velz2kms, lbox) + CSMF_mass[j4] = mass[i] - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) CSMF_stellarmass[j4] = get_random_cen_stellarmass_linearinterpolation( - mass[i], Mstar_low_C, Mstar_up_C, - M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) + mass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + ) CSMF_id[j4] = ids[i] j4 += 1 # assert j == gstart[tid + 1] @@ -648,7 +774,7 @@ def gen_cent( LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) LRG_dict['x'] = lrg_x LRG_dict['y'] = lrg_y @@ -676,7 +802,7 @@ def gen_cent( QSO_dict['vz'] = qso_vz QSO_dict['mass'] = qso_mass ID_dict['QSO'] = qso_id - + CSMF_dict['x'] = CSMF_x CSMF_dict['y'] = CSMF_y CSMF_dict['z'] = CSMF_z @@ -897,31 +1023,47 @@ def gen_sats_nfw( QSO_hod_dict['ic'], ) f_sigv_Q = QSO_hod_dict['f_sigv'] - + if want_CSMF: - Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], - CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'] + ( + Mstar_low_C, + Mstar_up_C, + M_1_C, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], + CSMF_hod_dict['M_2'], + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'], ) Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Asat'], CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'] + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'], ) f_sigv_C = CSMF_hod_dict['f_sigv'] @@ -932,8 +1074,8 @@ def gen_sats_nfw( num_sats_L = np.zeros(len(hid), dtype=np.int64) num_sats_E = np.zeros(len(hid), dtype=np.int64) num_sats_Q = np.zeros(len(hid), dtype=np.int64) - num_sats_C = np.zeros(len(hid), dtype = np.int64) - stellarmass_C = np.zeros(len(hid), dtype = np.int64) + num_sats_C = np.zeros(len(hid), dtype=np.int64) + stellarmass_C = np.zeros(len(hid), dtype=np.int64) hstart = np.rint(np.linspace(0, len(hid), Nthread + 1)).astype( np.int64 ) # starting index of each thread @@ -1009,32 +1151,35 @@ def gen_sats_nfw( * ic_Q ) num_sats_Q[i] = np.random.poisson(base_p_Q) - + if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] base_p_C = ( n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) * ic_C ) - - num_sats_C[i] = np.random.poisson(base_p_C) + + num_sats_C[i] = np.random.poisson(base_p_C) # generate rdpos rd_pos_L = getPointsOnSphere(np.sum(num_sats_L), Nthread) @@ -1115,33 +1260,32 @@ def gen_sats_nfw( nfw_rescale, ) ) - + h_id_C, x_sat_C, y_sat_C, z_sat_C, vx_sat_C, vy_sat_C, vz_sat_C, M_C = ( compute_fast_NFW( - NFW_draw, - hid, - hpos[:, 0], - hpos[:, 1], - hpos[:, 2], - hvel[:, 0], - hvel[:, 1], + NFW_draw, + hid, + hpos[:, 0], + hpos[:, 1], + hpos[:, 2], + hvel[:, 0], + hvel[:, 1], hvel[:, 2], - hvrms, - hc, - hmass, - hrvir, - rd_pos_C, - num_sats_C, - f_sigv_C, - vel_sat, + hvrms, + hc, + hmass, + hrvir, + rd_pos_C, + num_sats_C, + f_sigv_C, + vel_sat, Nthread, - exp_frac, - exp_scale, - nfw_rescale + exp_frac, + exp_scale, + nfw_rescale, ) ) - - + # do rsd if rsd: z_sat_L = (z_sat_L + vz_sat_L * inv_velz2kms) % lbox @@ -1152,7 +1296,7 @@ def gen_sats_nfw( LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) LRG_dict['x'] = x_sat_L LRG_dict['y'] = y_sat_L @@ -1180,7 +1324,7 @@ def gen_sats_nfw( QSO_dict['vz'] = vz_sat_Q QSO_dict['mass'] = M_Q ID_dict['QSO'] = h_id_Q - + CSMF_dict['x'] = x_sat_C CSMF_dict['y'] = y_sat_C CSMF_dict['z'] = z_sat_C @@ -1194,12 +1338,29 @@ def gen_sats_nfw( hstart = np.rint(np.linspace(0, num_sats_C.sum(), Nthread + 1)) for tid in numba.prange(Nthread): for i in range(int(hstart[tid]), int(hstart[tid + 1])): - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] stellarmass_C[i] = get_random_sat_stellarmass_linearinterpolation( - M_C[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, - a1_C_temp, s, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) - + M_C[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + a1_C_temp, + s, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) + CSMF_dict['stellarmass'] = stellarmass_C ID_dict['CSMF'] = h_id_C @@ -1328,37 +1489,53 @@ def gen_sats( QSO_hod_dict['Bsat'], QSO_hod_dict['ic'], ) - + if want_CSMF: - Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], + ( + Mstar_low_C, + Mstar_up_C, + M_1_C, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'] + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'], ) - + alpha_s_C, s_C, s_v_C, s_p_C, s_r_C, Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['alpha_s'], + CSMF_hod_dict['alpha_s'], CSMF_hod_dict['s'], - CSMF_hod_dict['s_v'], - CSMF_hod_dict['s_p'], - CSMF_hod_dict['s_r'], + CSMF_hod_dict['s_v'], + CSMF_hod_dict['s_p'], + CSMF_hod_dict['s_r'], CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], - CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'] + CSMF_hod_dict['Asat'], + CSMF_hod_dict['Bcent'], + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'], ) H = len(hmass) # num of particles @@ -1492,35 +1669,43 @@ def gen_sats( else: exp_sat = base_p_Q QSO_marker += exp_sat - - + CSMF_marker = QSO_marker if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] base_p_C = ( n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C - ) - * weights[i] - * ic_C) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) + * weights[i] + * ic_C + ) if enable_ranks: - decorator_C = 1 + s_C * ranks[i] + s_v_C * ranksv[i] + s_p_C * ranksp[i] + s_r_C * ranksr[i] + decorator_C = ( + 1 + + s_C * ranks[i] + + s_v_C * ranksv[i] + + s_p_C * ranksp[i] + + s_r_C * ranksr[i] + ) exp_sat = base_p_C * decorator_C else: exp_sat = base_p_C @@ -1536,7 +1721,7 @@ def gen_sats( Nout[tid, 2, 0] += 1 # counting keep[i] = 3 elif randoms[i] <= CSMF_marker: - Nout[tid, 3, 0] += 1 # counting + Nout[tid, 3, 0] += 1 # counting keep[i] = 4 else: keep[i] = 0 @@ -1581,18 +1766,18 @@ def gen_sats( qso_vz = np.empty(N_qso, dtype=hmass.dtype) qso_mass = np.empty(N_qso, dtype=hmass.dtype) qso_id = np.empty(N_qso, dtype=hid.dtype) - + # galaxy arrays N_CSMF = gstart[-1, 3] - CSMF_x = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_y = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_z = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_id = np.empty(N_CSMF, dtype = hid.dtype) + CSMF_x = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_y = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_z = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_id = np.empty(N_CSMF, dtype=hid.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): @@ -1696,46 +1881,55 @@ def gen_sats( j3 += 1 elif keep[i] == 4: CSMF_x[j4] = ppos[i, 0] - CSMF_vx[j4] = hvel[i, 0] + alpha_s_C * (pvel[i, 0] - hvel[i, 0]) # velocity bias + CSMF_vx[j4] = hvel[i, 0] + alpha_s_C * ( + pvel[i, 0] - hvel[i, 0] + ) # velocity bias CSMF_y[j4] = ppos[i, 1] - CSMF_vy[j4] = hvel[i, 1] + alpha_s_C * (pvel[i, 1] - hvel[i, 1]) # velocity bias + CSMF_vy[j4] = hvel[i, 1] + alpha_s_C * ( + pvel[i, 1] - hvel[i, 1] + ) # velocity bias CSMF_z[j4] = ppos[i, 2] - CSMF_vz[j4] = hvel[i, 2] + alpha_s_C * (pvel[i, 2] - hvel[i, 2]) # velocity bias + CSMF_vz[j4] = hvel[i, 2] + alpha_s_C * ( + pvel[i, 2] - hvel[i, 2] + ) # velocity bias if rsd and origin is not None: nx = CSMF_x[j4] - origin[0] ny = CSMF_y[j4] - origin[1] nz = CSMF_z[j4] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms*(CSMF_vx[j4]*nx+CSMF_vy[j4]*ny+CSMF_vz[j4]*nz) - CSMF_x[j4] = CSMF_x[j4]+proj*nx - CSMF_y[j4] = CSMF_y[j4]+proj*ny - CSMF_z[j4] = CSMF_z[j4]+proj*nz + proj = inv_velz2kms * ( + CSMF_vx[j4] * nx + CSMF_vy[j4] * ny + CSMF_vz[j4] * nz + ) + CSMF_x[j4] = CSMF_x[j4] + proj * nx + CSMF_y[j4] = CSMF_y[j4] + proj * ny + CSMF_z[j4] = CSMF_z[j4] + proj * nz elif rsd: - CSMF_z[j4] = wrap(CSMF_z[j4] + CSMF_vz[j4] * inv_velz2kms, lbox) - - - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + CSMF_z[j4] = wrap(CSMF_z[j4] + CSMF_vz[j4] * inv_velz2kms, lbox) + + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] CSMF_stellarmass[j4] = get_random_sat_stellarmass( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C - ) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) CSMF_mass[j4] = hmass[i] CSMF_id[j4] = hid[i] j4 += 1 @@ -1744,7 +1938,7 @@ def gen_sats( LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) LRG_dict['x'] = lrg_x LRG_dict['y'] = lrg_y @@ -1772,7 +1966,7 @@ def gen_sats( QSO_dict['vz'] = qso_vz QSO_dict['mass'] = qso_mass ID_dict['QSO'] = qso_id - + CSMF_dict['x'] = CSMF_x CSMF_dict['y'] = CSMF_y CSMF_dict['z'] = CSMF_z @@ -1782,7 +1976,7 @@ def gen_sats( CSMF_dict['mass'] = CSMF_mass CSMF_dict['stellarmass'] = CSMF_stellarmass ID_dict['CSMF'] = CSMF_id - + return LRG_dict, ELG_dict, QSO_dict, CSMF_dict, ID_dict @@ -1992,11 +2186,13 @@ def gen_gals( QSO_hod_dict = nb.typed.Dict.empty( key_type=nb.types.unicode_type, value_type=nb.types.float64 ) - + if 'CSMF' in tracers.keys(): want_CSMF = True - CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + CSMF_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) for key, value in CSMF_HOD.items(): CSMF_hod_dict[key] = value @@ -2009,7 +2205,9 @@ def gen_gals( else: want_CSMF = False - CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + CSMF_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) start = time.time() @@ -2018,7 +2216,14 @@ def gen_gals( lbox = params['Lbox'] origin = params['origin'] - LRG_dict_cent, ELG_dict_cent, QSO_dict_cent, CSMF_dict_cent, ID_dict_cent, keep_cent = gen_cent( + ( + LRG_dict_cent, + ELG_dict_cent, + QSO_dict_cent, + CSMF_dict_cent, + ID_dict_cent, + keep_cent, + ) = gen_cent( halos_array['hpos'], halos_array['hvel'], halos_array['hmass'], @@ -2051,31 +2256,33 @@ def gen_gals( warnings.warn( 'NFW profile is unoptimized. It has different velocity bias. It does not support lightcone.' ) - LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, CSMF_dict_sat, ID_dict_sat = gen_sats_nfw( - NFW_draw, - halos_array['hpos'], - halos_array['hvel'], - halos_array['hmass'], - halos_array['hid'], - halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), - halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), - halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), - halos_array['hsigma3d'], - halos_array['hc'], - halos_array['hrvir'], - LRG_hod_dict, - ELG_hod_dict, - QSO_hod_dict, - CSMF_hod_dict, - want_LRG, - want_ELG, - want_QSO, - want_CSMF, - rsd, - inv_velz2kms, - lbox, - keep_cent, - Nthread=Nthread, + LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, CSMF_dict_sat, ID_dict_sat = ( + gen_sats_nfw( + NFW_draw, + halos_array['hpos'], + halos_array['hvel'], + halos_array['hmass'], + halos_array['hid'], + halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), + halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), + halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), + halos_array['hsigma3d'], + halos_array['hc'], + halos_array['hrvir'], + LRG_hod_dict, + ELG_hod_dict, + QSO_hod_dict, + CSMF_hod_dict, + want_LRG, + want_ELG, + want_QSO, + want_CSMF, + rsd, + inv_velz2kms, + lbox, + keep_cent, + Nthread=Nthread, + ) ) else: LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, CSMF_dict_sat, ID_dict_sat = gen_sats( @@ -2115,8 +2322,18 @@ def gen_gals( print('generating satellites took ', time.time() - start) # B.H. TODO: need a for loop above so we don't need to do this by hand - HOD_dict_sat = {'LRG': LRG_dict_sat, 'ELG': ELG_dict_sat, 'QSO': QSO_dict_sat, 'CSMF': CSMF_dict_sat} - HOD_dict_cent = {'LRG': LRG_dict_cent, 'ELG': ELG_dict_cent, 'QSO': QSO_dict_cent, 'CSMF': CSMF_dict_cent} + HOD_dict_sat = { + 'LRG': LRG_dict_sat, + 'ELG': ELG_dict_sat, + 'QSO': QSO_dict_sat, + 'CSMF': CSMF_dict_sat, + } + HOD_dict_cent = { + 'LRG': LRG_dict_cent, + 'ELG': ELG_dict_cent, + 'QSO': QSO_dict_cent, + 'CSMF': CSMF_dict_cent, + } # do a concatenate in numba parallel start = time.time() diff --git a/abacusnbody/hod/abacus_hod.py b/abacusnbody/hod/abacus_hod.py index cdd94383..9d1962bc 100644 --- a/abacusnbody/hod/abacus_hod.py +++ b/abacusnbody/hod/abacus_hod.py @@ -34,7 +34,7 @@ N_cen_QSO, N_sat_generic, n_cen_CSMF, - n_sat_CSMF + n_sat_CSMF, ) # TODO B.H.: staging can be shorter and prettier; perhaps asdf for h5 and ecsv? @@ -312,7 +312,7 @@ def staging(self): if ( ('ELG' not in self.tracers.keys()) and ('QSO' not in self.tracers.keys()) - and (not self.force_mt) + and (not self.force_mt) and ('CSMF' not in self.tracers.keys()) ): halofilename = subsample_dir / ( @@ -896,35 +896,34 @@ def compute_ngal(self, tracers=None, Nthread=16): ) ngal_dict[etracer] = newngal[0] + newngal[1] fsat_dict[etracer] = newngal[1] / (newngal[0] + newngal[1]) - - + elif etracer == 'CSMF': newngal = AbacusHOD._compute_ngal_CSMF( - self.logMbins, - self.deltacbins, - self.fenvbins, + self.logMbins, + self.deltacbins, + self.fenvbins, self.halo_mass_func, - tracer_hod['Mstar_low'], - tracer_hod['Mstar_up'], - tracer_hod['M_1'], - tracer_hod['M_0'], - tracer_hod['gamma_1'], - tracer_hod['gamma_2'], - tracer_hod['sigma_c'], - tracer_hod['a_1'], - tracer_hod['a_2'], - tracer_hod['M_2'], - tracer_hod['b_0'], - tracer_hod['b_1'], - tracer_hod['b_2'], - tracer_hod.get('delta_1', 0), - tracer_hod.get('delta_2', 0), + tracer_hod['Mstar_low'], + tracer_hod['Mstar_up'], + tracer_hod['M_1'], + tracer_hod['M_0'], + tracer_hod['gamma_1'], + tracer_hod['gamma_2'], + tracer_hod['sigma_c'], + tracer_hod['a_1'], + tracer_hod['a_2'], + tracer_hod['M_2'], + tracer_hod['b_0'], + tracer_hod['b_1'], + tracer_hod['b_2'], + tracer_hod.get('delta_1', 0), + tracer_hod.get('delta_2', 0), tracer_hod.get('Acent', 0), - tracer_hod.get('Asat', 0), - tracer_hod.get('Bcent', 0), - tracer_hod.get('Bsat', 0), - tracer_hod.get('ic', 1), - Nthread + tracer_hod.get('Asat', 0), + tracer_hod.get('Bcent', 0), + tracer_hod.get('Bsat', 0), + tracer_hod.get('ic', 1), + Nthread, ) ngal_dict[etracer] = newngal[0] + newngal[1] fsat_dict[etracer] = newngal[1] / (newngal[0] + newngal[1]) @@ -1130,35 +1129,89 @@ def _compute_ngal_qso( ngal_cent += halo_mass_func[i, j, k] * ncent_temp * ic ngal_sat += halo_mass_func[i, j, k] * nsat_temp * ic return ngal_cent, ngal_sat - - + @staticmethod - @njit(fastmath = True, parallel = True) - def _compute_ngal_CSMF(logMbins, deltacbins, fenvbins, halo_mass_func, - Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c, a1, a2, M2, b0, b1, b2, delta1, delta2, Acent, Asat, Bcent, Bsat, ic, Nthread): + @njit(fastmath=True, parallel=True) + def _compute_ngal_CSMF( + logMbins, + deltacbins, + fenvbins, + halo_mass_func, + Mstar_low, + Mstar_up, + M_1, + M_0, + gamma1, + gamma2, + sigma_c, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, + Acent, + Asat, + Bcent, + Bsat, + ic, + Nthread, + ): """ internal helper to compute number of CSMFs """ numba.set_num_threads(Nthread) - logMs = 0.5*(logMbins[1:] + logMbins[:-1]) - deltacs = 0.5*(deltacbins[1:] + deltacbins[:-1]) - fenvs = 0.5*(fenvbins[1:] + fenvbins[:-1]) + logMs = 0.5 * (logMbins[1:] + logMbins[:-1]) + deltacs = 0.5 * (deltacbins[1:] + deltacbins[:-1]) + fenvs = 0.5 * (fenvbins[1:] + fenvbins[:-1]) ngal_cent = 0 ngal_sat = 0 for i in numba.prange(len(logMbins) - 1): for j in range(len(deltacbins) - 1): for k in range(len(fenvbins) - 1): - Mh_temp = 10**logMs[i] - M_1_temp = 10**(np.log10(M_1)) #+ Acent * deltacs[j] + Bcent * fenvs[k]) - M2_temp = 10**(np.log10(M2)) #+ Asat * deltacs[j] + Bsat * fenvs[k]) - - ncent_temp = n_cen_CSMF(Mh_temp, Mstar_low, Mstar_up, M_1_temp, M_0, gamma1, gamma2, sigma_c) - nsat_temp = n_sat_CSMF(Mh_temp, Mstar_low, Mstar_up, M_1_temp, M_0, gamma1, gamma2, sigma_c, a1, a2, M2_temp, b0, b1, b2, delta1, delta2) + Mh_temp = 10 ** logMs[i] + M_1_temp = 10 ** ( + np.log10(M_1) + ) # + Acent * deltacs[j] + Bcent * fenvs[k]) + M2_temp = 10 ** ( + np.log10(M2) + ) # + Asat * deltacs[j] + Bsat * fenvs[k]) + + ncent_temp = n_cen_CSMF( + Mh_temp, + Mstar_low, + Mstar_up, + M_1_temp, + M_0, + gamma1, + gamma2, + sigma_c, + ) + nsat_temp = n_sat_CSMF( + Mh_temp, + Mstar_low, + Mstar_up, + M_1_temp, + M_0, + gamma1, + gamma2, + sigma_c, + a1, + a2, + M2_temp, + b0, + b1, + b2, + delta1, + delta2, + ) ngal_cent += halo_mass_func[i, j, k] * ncent_temp * ic ngal_sat += halo_mass_func[i, j, k] * nsat_temp * ic - + return ngal_cent, ngal_sat def compute_clustering(self, mock_dict, *args, **kwargs): From 4ea04d42d6a509c34a0570211dbd0afdfa15c3d9 Mon Sep 17 00:00:00 2001 From: Pierre Burger Date: Wed, 9 Oct 2024 15:56:28 -0400 Subject: [PATCH 05/13] Separate CSMF python file Separation of GRAND_HOD to GRAND_HOD and CSMF_HOD, also removed the for loops --- .DS_Store | Bin 0 -> 6148 bytes abacusnbody/.DS_Store | Bin 0 -> 6148 bytes abacusnbody/hod/CSMF_HOD.py | 1153 +++++++++++++++++++++++++++++++++ abacusnbody/hod/GRAND_HOD.py | 842 ++---------------------- abacusnbody/hod/abacus_hod.py | 53 +- 5 files changed, 1228 insertions(+), 820 deletions(-) create mode 100644 .DS_Store create mode 100644 abacusnbody/.DS_Store create mode 100644 abacusnbody/hod/CSMF_HOD.py diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5e75159d3da651e8cdbbeda424a7a9c34fd09479 GIT binary patch literal 6148 zcmeHK%Wl&^6upzC#32P`0TL{bEU}G(QwdRtO`0ZxN`*y@U;!xDaex%p6U7cu6-CM( zzJXt0&5!Uetl-?4sbuG|MTpLo?wontW&;{f{`AP@D8eWjBIM1hQ4+~nM^2L{Ra{>WI8~?Guiu@|_YV#lZfoCLG~9X1d+0UX zM_y~Os5-arKX`oF{}i7m@{3A<1Qso2Hw<3FS5&r$^4Xsxu}nT7g4s1XqNkANIX$5s z`P8HdmcVR9y6=e0OIOD7C?R)Rj1%`(;wp4P$C#mq>-*GOiRLWwSFO^RM(`o^@i0@& z@qUHzYlU$W(HHOu?&b`<0cFS_L+hO4BDcgD7iT-geH~Jlo&lR+L`8;rVZP;Fm_1^~ ziDLQ4i#6NCW3}?U3UV3X9yzycqh|T&UBnFe8LUOJUinBgueO=j->F?~8Ep1e0jt0U z1$cdM;f!sAOO0yjK%tHRz%sg(q0T=G%yA924K6ie1ST{UsHwsnF@&b0-!*Ykfj>6% B)8zmF literal 0 HcmV?d00001 diff --git a/abacusnbody/.DS_Store b/abacusnbody/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..d15d1b7b523a7922dd16e014a18f6dfe0b569e05 GIT binary patch literal 6148 zcmeHK&ubGw82u)R?SvrYpr8j~z-y?Rf>`jfjrHKwh#pkpt{dp$b|-WvY4{%<;IZquG0Ul>%j0+Ps<5Ba!lq#G>v4=c`0WBUekBqp5-R9| zDk!Th{`W0@hcd8hN;Y}w(H5^wr(`MHVtu`ul`g7TxW^yqtjMc6ihinAyK`r6U+?Rl z{y09hi*cRQ^J0;8q_rjKSl9$tSC&;rYC?DSCWo8#6J1?`! zX~(`pcXfBtf4E#89vuzL@Gx2p%yJk#jRxjfG+eE^`u^i5&tFfz6lbO7TSqATMO@qs zxQ1`2XgZ@Qn{$mxzreX{{#q+24u}Kdz>ORT6O!)TNG&39KpeQu4)FY7kc@G_#-p1$ zP`E7su!Ck7sPo4nYfQj6VB--z2<4?fFI8cSp}ZXJCh+2bjYlsh6*eC#Tv>%3igH)S z`HiEKiakmx4u}Jr4z&GYpV$AF*YE!~MRFt#hy%CE0o6GfpN#NGVQt-cI9_WF@)MGS od5y=v5>(t$3@$Il50Earandom_rv)[0][0] + + m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) + return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + + +@njit(fastmath=True) +def n_sat_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c, a1, a2, M2, b0, b1, b2, delta1, delta2): + """ + Standard Cacciato et al. (2008) satellite HOD parametrization for CSMF + """ + nbins = 1000 + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + + CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + + nsat = 0 + for i in range(nbins-1): + nsat += (CSMF_sat[i+1]-CSMF_sat[i])*(stellarmass[i+1]-stellarmass[i])/2 + (stellarmass[i+1]-stellarmass[i])*CSMF_sat[i] + + return nsat#*ncen + + +@njit(fastmath=True) +def CSMF_satelites(M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): + """ + Eq. (36) from Cacciato et al. (2008) + """ + M_s_value = M_s(M_h, M_1, M_0, gamma1, gamma2) + alpha_s_value = alpha_s(M_h, a1, a2, M2) + phi_s_value = phi_s(M_h, b0, b1, b2) + + delta = 10 ** (delta1 + delta2 * (np.log10(M_h) - 12)) + + return phi_s_value/M_s_value * (Mstar/M_s_value)**alpha_s_value * np.exp(-delta*(Mstar/M_s_value)**2) + +@njit(fastmath=True) +def M_s(M_h, M_1, M_0, gamma1, gamma2): + """ + Eq. (38) from Cacciato et al. (2008) + """ + return 0.562 * M_c(M_h, M_1, M_0, gamma1, gamma2) + +@njit(fastmath=True) +def alpha_s(M_h, a1, a2, M2): + """ + Eq. (39) from Cacciato et al. (2008) + """ + return -2.0 + a1 * (1-2/np.pi*np.arctan(a2*np.log10(M_h/M2))) + +@njit(fastmath=True) +def phi_s(M_h, b0, b1, b2): + """ + Eq. (40) from Cacciato et al. (2008) + """ + M12 = M_h/1e12 + log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12)**2 + return 10**log_phi_s + + +@njit(fastmath=True) +def get_random_sat_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): + + nbins = 1000 + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + + CSMF_sat = CSMF_satelites(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + + cdf = np.cumsum(CSMF_sat*delta_stellar_masses) + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin_clostest = (np.abs(cdf - random_rv)).argmin() + + return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) + + +@njit(fastmath=True) +def get_random_sat_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): + + nbins = 100 + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + + CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + + cdf = np.cumsum(CSMF_sat[:-1]*delta_stellar_masses) + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin = np.where(cdf>random_rv)[0][0] + + m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) + return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + + +@njit(fastmath=True) +def Gaussian_fun(x, mean, sigma): + """ + Gaussian function with centered at `mean' with standard deviation `sigma'. + """ + return 0.3989422804014327 / sigma * np.exp(-((x - mean) ** 2) / 2 / sigma**2) + + +@njit(fastmath=True) +def wrap(x, L): + """Fast scalar mod implementation""" + L2 = L / 2 + if x >= L2: + return x - L + elif x < -L2: + return x + L + return x + + +@njit(parallel=True, fastmath=True) +def gen_cent( + pos, + vel, + mass, + ids, + multis, + randoms, + vdev, + deltac, + fenv, + shear, + CSMF_hod_dict, + rsd, + inv_velz2kms, + lbox, + want_CSMF, + Nthread, + origin, +): + """ + Generate central galaxies in place in memory with a two pass numba parallel implementation. + """ + if want_CSMF: + Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'] + ) + ic_C, alpha_c_C, Ac_C, Bc_C = ( + CSMF_hod_dict['ic'], + CSMF_hod_dict['alpha_c'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Bcent'] + ) + + + H = len(mass) + + numba.set_num_threads(Nthread) + Nout = np.zeros((Nthread, 1, 8), dtype=np.int64) + hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype( + np.int64 + ) # starting index of each thread + + keep = np.empty(H, dtype=np.int8) # mask array tracking which halos to keep + + # figuring out the number of halos kept for each thread + for tid in numba.prange(Nthread): + for i in range(hstart[tid], hstart[tid + 1]): + # first create the markers between 0 and 1 for different tracers + CSMF_marker = 0 + if want_CSMF: + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + + ncen = n_cen_CSMF(mass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) + + CSMF_marker += ncen * ic_C * multis[i] + + if randoms[i] <= CSMF_marker: + Nout[tid, 0, 0] += 1 # counting + keep[i] = 1 + else: + keep[i] = 0 + + # compose galaxy array, first create array of galaxy starting indices for the threads + gstart = np.empty((Nthread + 1, 1), dtype=np.int64) + gstart[0, :] = 0 + gstart[1:, 0] = Nout[:, 0, 0].cumsum() + + + # galaxy arrays + N_CSMF = gstart[-1, 0] + CSMF_x = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_y = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_z = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_id = np.empty(N_CSMF, dtype = ids.dtype) + + # fill in the galaxy arrays + for tid in numba.prange(Nthread): + j1 = gstart[tid] + for i in range(hstart[tid], hstart[tid + 1]): + if keep[i] == 1: + # loop thru three directions to assign galaxy velocities and positions + CSMF_x[j1] = pos[i,0] + CSMF_vx[j1] = vel[i,0] + alpha_c_C * vdev[i,0] # velocity bias + CSMF_y[j1] = pos[i,1] + CSMF_vy[j1] = vel[i,1] + alpha_c_C * vdev[i,1] # velocity bias + CSMF_z[j1] = pos[i,2] + CSMF_vz[j1] = vel[i,2] + alpha_c_C * vdev[i,2] # velocity bias + + # rsd only applies to the z direction + if rsd and origin is not None: + nx = CSMF_x[j1] - origin[0] + ny = CSMF_y[j1] - origin[1] + nz = CSMF_z[j1] - origin[2] + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + nx *= inv_norm + ny *= inv_norm + nz *= inv_norm + proj = inv_velz2kms*(CSMF_vx[j1]*nx+CSMF_vy[j1]*ny+CSMF_vz[j1]*nz) + CSMF_x[j1] = CSMF_x[j1]+proj*nx + CSMF_y[j1] = CSMF_y[j1]+proj*ny + CSMF_z[j1] = CSMF_z[j1]+proj*nz + elif rsd: + CSMF_z[j1] = wrap(pos[i,2] + CSMF_vz[j1] * inv_velz2kms, lbox) + + CSMF_mass[j1] = mass[i] + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + CSMF_stellarmass[j1] = get_random_cen_stellarmass_linearinterpolation( + mass[i], Mstar_low_C, Mstar_up_C, + M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) + CSMF_id[j1] = ids[i] + j1 += 1 + # assert j == gstart[tid + 1] + + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) + + + CSMF_dict['x'] = CSMF_x + CSMF_dict['y'] = CSMF_y + CSMF_dict['z'] = CSMF_z + CSMF_dict['vx'] = CSMF_vx + CSMF_dict['vy'] = CSMF_vy + CSMF_dict['vz'] = CSMF_vz + CSMF_dict['mass'] = CSMF_mass + CSMF_dict['stellarmass'] = CSMF_stellarmass + ID_dict['CSMF'] = CSMF_id + return CSMF_dict, ID_dict, keep + + +@njit(parallel=True, fastmath=True) +def getPointsOnSphere(nPoints, Nthread, seed=None): + """ + --- Aiding function for NFW computation, generate random points in a sphere + """ + numba.set_num_threads(Nthread) + ind = min(Nthread, nPoints) + # starting index of each thread + hstart = np.rint(np.linspace(0, nPoints, ind + 1)) + ur = np.zeros((nPoints, 3), dtype=np.float64) + cmin = -1 + cmax = +1 + + for tid in numba.prange(Nthread): + if seed is not None: + np.random.seed(seed[tid]) + for i in range(hstart[tid], hstart[tid + 1]): + u1, u2 = np.random.uniform(0, 1), np.random.uniform(0, 1) + ra = 0 + u1 * (2 * np.pi - 0) + dec = np.pi - (np.arccos(cmin + u2 * (cmax - cmin))) + + ur[i, 0] = np.sin(dec) * np.cos(ra) + ur[i, 1] = np.sin(dec) * np.sin(ra) + ur[i, 2] = np.cos(dec) + return ur + + +@njit(fastmath=True, parallel=True) # parallel=True, +def compute_fast_NFW( + NFW_draw, + h_id, + x_h, + y_h, + z_h, + vx_h, + vy_h, + vz_h, + vrms_h, + c, + M, + Rvir, + rd_pos, + num_sat, + f_sigv, + vel_sat='rd_normal', + Nthread=16, + exp_frac=0, + exp_scale=1, + nfw_rescale=1, +): + """ + --- Compute NFW positions and velocities for satelitte galaxies + c: r98/r25 + vrms_h: 'sigmav3d_L2com' + """ + # numba.set_num_threads(Nthread) + # figuring out the number of halos kept for each thread + h_id = np.repeat(h_id, num_sat) + M = np.repeat(M, num_sat) + c = np.repeat(c, num_sat) + Rvir = np.repeat(Rvir, num_sat) + x_h = np.repeat(x_h, num_sat) + y_h = np.repeat(y_h, num_sat) + z_h = np.repeat(z_h, num_sat) + vx_h = np.repeat(vx_h, num_sat) + vy_h = np.repeat(vy_h, num_sat) + vz_h = np.repeat(vz_h, num_sat) + vrms_h = np.repeat(vrms_h, num_sat) + x_sat = np.empty_like(x_h) + y_sat = np.empty_like(y_h) + z_sat = np.empty_like(z_h) + vx_sat = np.empty_like(vx_h) + vy_sat = np.empty_like(vy_h) + vz_sat = np.empty_like(vz_h) + + # starting index of each thread + hstart = np.rint(np.linspace(0, num_sat.sum(), Nthread + 1)) + for tid in numba.prange(Nthread): + for i in range(int(hstart[tid]), int(hstart[tid + 1])): + ind = i + # while (NFW_draw[ind] > c[i]): + # ind = np.random.randint(0, len(NFW_draw)) + # etaVir = NFW_draw[ind]/c[i] # =r/rvir + if np.random.uniform(0, 1) < exp_frac: + tt = np.random.exponential(exp_scale, size=1)[0] + etaVir = tt / c[i] + else: + while NFW_draw[ind] > c[i]: + ind = np.random.randint(0, len(NFW_draw)) + etaVir = NFW_draw[ind] / c[i] * nfw_rescale + + p = etaVir * Rvir[i] + x_sat[i] = x_h[i] + rd_pos[i, 0] * p + y_sat[i] = y_h[i] + rd_pos[i, 1] * p + z_sat[i] = z_h[i] + rd_pos[i, 2] * p + if vel_sat == 'rd_normal': + sig = vrms_h[i] * 0.577 * f_sigv + vx_sat[i] = np.random.normal(loc=vx_h[i], scale=sig) + vy_sat[i] = np.random.normal(loc=vy_h[i], scale=sig) + vz_sat[i] = np.random.normal(loc=vz_h[i], scale=sig) + else: + raise ValueError('Wrong vel_sat argument only "rd_normal"') + return h_id, x_sat, y_sat, z_sat, vx_sat, vy_sat, vz_sat, M + + +@njit(fastmath=True, parallel=True) +def gen_sats_nfw( + NFW_draw, + hpos, + hvel, + hmass, + hid, + hdeltac, + hfenv, + hshear, + hvrms, + hc, + hrvir, + CSMF_hod_dict, + want_CSMF, + rsd, + inv_velz2kms, + lbox, + keep_cent, + vel_sat='rd_normal', + Nthread=16, +): + """ + Generate satellite galaxies on an NFW profile, with option for an extended profile. See Rocher et al. 2023. + + Not yet on lightcone!! Different velocity bias treatment!! Not built for performance!! + + """ + if want_CSMF: + Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], + CSMF_hod_dict['M_2'], + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'] + ) + Ac_C, As_C, Bc_C, Bs_C, ic_C = ( + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Asat'], + CSMF_hod_dict['Bcent'], + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'] + ) + f_sigv_C = CSMF_hod_dict['f_sigv'] + + numba.set_num_threads(Nthread) + + # compute nsate for each halo + # figuring out the number of particles kept for each thread + num_sats_C = np.zeros(len(hid), dtype = np.int64) + stellarmass_C = np.zeros(len(hid), dtype = np.int64) + hstart = np.rint(np.linspace(0, len(hid), Nthread + 1)).astype( + np.int64 + ) # starting index of each thread + for tid in range(Nthread): + for i in range(hstart[tid], hstart[tid + 1]): + if want_CSMF: + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] + base_p_C = ( + n_sat_CSMF( + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C) + * ic_C + ) + num_sats_C[i] = np.random.poisson(base_p_C) + + # generate rdpos + rd_pos_C = getPointsOnSphere(np.sum(num_sats_C), Nthread) + + # put satellites on NFW + h_id_C, x_sat_C, y_sat_C, z_sat_C, vx_sat_C, vy_sat_C, vz_sat_C, M_C = ( + compute_fast_NFW( + NFW_draw, + hid, + hpos[:, 0], + hpos[:, 1], + hpos[:, 2], + hvel[:, 0], + hvel[:, 1], + hvel[:, 2], + hvrms, + hc, + hmass, + hrvir, + rd_pos_C, + num_sats_C, + f_sigv_C, + vel_sat, + Nthread, + exp_frac, + exp_scale, + nfw_rescale + ) + ) + + + # do rsd + if rsd: + z_sat_C = (z_sat_C + vz_sat_C * inv_velz2kms) % lbox + + + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) + + CSMF_dict['x'] = x_sat_C + CSMF_dict['y'] = y_sat_C + CSMF_dict['z'] = z_sat_C + CSMF_dict['vx'] = vx_sat_C + CSMF_dict['vy'] = vy_sat_C + CSMF_dict['vz'] = vz_sat_C + CSMF_dict['mass'] = M_C + stellarmass_C = np.empty_like(M_C) + ## compute stellarmass of all the satelites + if want_CSMF: + hstart = np.rint(np.linspace(0, num_sats_C.sum(), Nthread + 1)) + for tid in numba.prange(Nthread): + for i in range(int(hstart[tid]), int(hstart[tid + 1])): + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] + stellarmass_C[i] = get_random_sat_stellarmass_linearinterpolation( + M_C[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, + a1_C_temp, s, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) + + CSMF_dict['stellarmass'] = stellarmass_C + ID_dict['CSMF'] = h_id_C + + return CSMF_dict, ID_dict + + +@njit(parallel=True, fastmath=True) +def gen_sats( + ppos, + pvel, + hvel, + hmass, + hid, + weights, + randoms, + hdeltac, + hfenv, + hshear, + enable_ranks, + ranks, + ranksv, + ranksp, + ranksr, + ranksc, + CSMF_hod_dict, + rsd, + inv_velz2kms, + lbox, + Mpart, + want_CSMF, + Nthread, + origin, + keep_cent, +): + """ + Generate satellite galaxies in place in memory with a two pass numba parallel implementation. + """ + + + + if want_CSMF: + Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], + CSMF_hod_dict['M_2'], + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'] + ) + + alpha_s_C, s_C, s_v_C, s_p_C, s_r_C, Ac_C, As_C, Bc_C, Bs_C, ic_C = ( + CSMF_hod_dict['alpha_s'], + CSMF_hod_dict['s'], + CSMF_hod_dict['s_v'], + CSMF_hod_dict['s_p'], + CSMF_hod_dict['s_r'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Asat'], + CSMF_hod_dict['Bcent'], + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'] + ) + + H = len(hmass) # num of particles + + numba.set_num_threads(Nthread) + Nout = np.zeros((Nthread, 1, 8), dtype=np.int64) + hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype( + np.int64 + ) # starting index of each thread + + keep = np.empty(H, dtype=np.int8) # mask array tracking which halos to keep + + # figuring out the number of particles kept for each thread + for tid in numba.prange(Nthread): # numba.prange(Nthread): + for i in range(hstart[tid], hstart[tid + 1]): + # print(logM1, As, hdeltac[i], Bs, hfenv[i]) + CSMF_marker = 0 + if want_CSMF: + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] + base_p_C = ( + n_sat_CSMF( + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C + ) + * weights[i] + * ic_C) + if enable_ranks: + decorator_C = 1 + s_C * ranks[i] + s_v_C * ranksv[i] + s_p_C * ranksp[i] + s_r_C * ranksr[i] + exp_sat = base_p_C * decorator_C + else: + exp_sat = base_p_C + CSMF_marker += exp_sat + + + if randoms[i] <= CSMF_marker: + Nout[tid, 0, 0] += 1 # counting + keep[i] = 1 + else: + keep[i] = 0 + + # compose galaxy array, first create array of galaxy starting indices for the threads + gstart = np.empty((Nthread + 1, 1), dtype=np.int64) + gstart[0, :] = 0 + gstart[1:, 0] = Nout[:, 0, 0].cumsum() + + # galaxy arrays + N_CSMF = gstart[-1, 0] + CSMF_x = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_y = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_z = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_id = np.empty(N_CSMF, dtype = hid.dtype) + + # fill in the galaxy arrays + for tid in numba.prange(Nthread): + j1 = gstart[tid] + for i in range(hstart[tid], hstart[tid + 1]): + if keep[i] == 1: + + CSMF_x[j1] = ppos[i, 0] + CSMF_vx[j1] = hvel[i, 0] + alpha_s_C * (pvel[i, 0] - hvel[i, 0]) # velocity bias + CSMF_y[j1] = ppos[i, 1] + CSMF_vy[j1] = hvel[i, 1] + alpha_s_C * (pvel[i, 1] - hvel[i, 1]) # velocity bias + CSMF_z[j1] = ppos[i, 2] + CSMF_vz[j1] = hvel[i, 2] + alpha_s_C * (pvel[i, 2] - hvel[i, 2]) # velocity bias + if rsd and origin is not None: + nx = CSMF_x[j1] - origin[0] + ny = CSMF_y[j1] - origin[1] + nz = CSMF_z[j1] - origin[2] + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + nx *= inv_norm + ny *= inv_norm + nz *= inv_norm + proj = inv_velz2kms*(CSMF_vx[j1]*nx+CSMF_vy[j1]*ny+CSMF_vz[j1]*nz) + CSMF_x[j1] = CSMF_x[j1]+proj*nx + CSMF_y[j1] = CSMF_y[j1]+proj*ny + CSMF_z[j1] = CSMF_z[j1]+proj*nz + elif rsd: + CSMF_z[j1] = wrap(CSMF_z[j1] + CSMF_vz[j1] * inv_velz2kms, lbox) + + + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] + CSMF_stellarmass[j1] = get_random_sat_stellarmass( + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C + ) + CSMF_mass[j1] = hmass[i] + CSMF_id[j1] = hid[i] + j1 += 1 + # assert j == gstart[tid + 1] + + + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) + + + CSMF_dict['x'] = CSMF_x + CSMF_dict['y'] = CSMF_y + CSMF_dict['z'] = CSMF_z + CSMF_dict['vx'] = CSMF_vx + CSMF_dict['vy'] = CSMF_vy + CSMF_dict['vz'] = CSMF_vz + CSMF_dict['mass'] = CSMF_mass + CSMF_dict['stellarmass'] = CSMF_stellarmass + ID_dict['CSMF'] = CSMF_id + + return CSMF_dict, ID_dict + + +@njit(parallel=True, fastmath=True) +def fast_concatenate(array1, array2, Nthread): + """Fast concatenate with numba parallel""" + + N1 = len(array1) + N2 = len(array2) + if N1 == 0: + return array2 + elif N2 == 0: + return array1 + + final_array = np.empty(N1 + N2, dtype=array1.dtype) + # if one thread, then no need to parallel + if Nthread == 1: + for i in range(N1): + final_array[i] = array1[i] + for j in range(N2): + final_array[j + N1] = array2[j] + return final_array + + numba.set_num_threads(Nthread) + Nthread1 = max(1, int(np.floor(Nthread * N1 / (N1 + N2)))) + Nthread2 = Nthread - Nthread1 + hstart1 = np.rint(np.linspace(0, N1, Nthread1 + 1)).astype(np.int64) + hstart2 = np.rint(np.linspace(0, N2, Nthread2 + 1)).astype(np.int64) + N1 + + for tid in numba.prange(Nthread): # numba.prange(Nthread): + if tid < Nthread1: + for i in range(hstart1[tid], hstart1[tid + 1]): + final_array[i] = array1[i] + else: + for i in range(hstart2[tid - Nthread1], hstart2[tid + 1 - Nthread1]): + final_array[i] = array2[i - N1] + # final_array = np.concatenate((array1, array2)) + return final_array + + +def gen_gals( + halos_array, + subsample, + tracers, + params, + Nthread, + enable_ranks, + rsd, + verbose, + nfw, + NFW_draw=None, +): + """ + parse hod parameters, pass them on to central and satellite generators + and then format the results + + Parameters + ---------- + + halos_array : dictionary of arrays + a dictionary of halo properties (pos, vel, mass, id, randoms, ...) + + subsample : dictionary of arrays + a dictionary of particle propoerties (pos, vel, hmass, hid, Np, subsampling, randoms, ...) + + tracers : dictionary of dictionaries + Dictionary of multi-tracer HODs + + enable_ranks : boolean + Flag of whether to implement particle ranks. + + rsd : boolean + Flag of whether to implement RSD. + + params : dict + Dictionary of various simulation parameters. + + """ + + # B.H. TODO: pass as dictionary; make what's below more succinct + for tracer in tracers.keys(): + if tracer == 'CSMF': + CSMF_HOD = tracers[tracer] + + + if 'CSMF' in tracers.keys(): + want_CSMF = True + + CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + for key, value in CSMF_HOD.items(): + CSMF_hod_dict[key] = value + + CSMF_hod_dict['Acent'] = CSMF_HOD.get('Acent', 0.0) + CSMF_hod_dict['Asat'] = CSMF_HOD.get('Asat', 0.0) + CSMF_hod_dict['Bcent'] = CSMF_HOD.get('Bcent', 0.0) + CSMF_hod_dict['Bsat'] = CSMF_HOD.get('Bsat', 0.0) + CSMF_hod_dict['ic'] = CSMF_HOD.get('ic', 1.0) + CSMF_hod_dict['f_sigv'] = CSMF_HOD.get('f_sigv', 0) + + else: + want_CSMF = False + CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + + start = time.time() + + velz2kms = params['velz2kms'] + inv_velz2kms = 1 / velz2kms + lbox = params['Lbox'] + origin = params['origin'] + + CSMF_dict_cent, ID_dict_cent, keep_cent = gen_cent( + halos_array['hpos'], + halos_array['hvel'], + halos_array['hmass'], + halos_array['hid'], + halos_array['hmultis'], + halos_array['hrandoms'], + halos_array['hveldev'], + halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), + halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), + halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), + CSMF_hod_dict, + rsd, + inv_velz2kms, + lbox, + want_CSMF, + Nthread, + origin, + ) + if verbose: + print('generating centrals took ', time.time() - start) + + start = time.time() + if nfw: + warnings.warn( + 'NFW profile is unoptimized. It has different velocity bias. It does not support lightcone.' + ) + CSMF_dict_sat, ID_dict_sat = gen_sats_nfw( + NFW_draw, + halos_array['hpos'], + halos_array['hvel'], + halos_array['hmass'], + halos_array['hid'], + halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), + halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), + halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), + halos_array['hsigma3d'], + halos_array['hc'], + halos_array['hrvir'], + CSMF_hod_dict, + want_CSMF, + rsd, + inv_velz2kms, + lbox, + keep_cent, + Nthread=Nthread, + ) + else: + CSMF_dict_sat, ID_dict_sat = gen_sats( + subsample['ppos'], + subsample['pvel'], + subsample['phvel'], + subsample['phmass'], + subsample['phid'], + subsample['pweights'], + subsample['prandoms'], + subsample.get('pdeltac', np.zeros(len(subsample['phid']))), + subsample.get('pfenv', np.zeros(len(subsample['phid']))), + subsample.get('pshear', np.zeros(len(subsample['phid']))), + enable_ranks, + subsample['pranks'], + subsample['pranksv'], + subsample['pranksp'], + subsample['pranksr'], + subsample['pranksc'], + CSMF_hod_dict, + rsd, + inv_velz2kms, + lbox, + params['Mpart'], + want_CSMF, + Nthread, + origin, + keep_cent[subsample['pinds']], + ) + if verbose: + print('generating satellites took ', time.time() - start) + + # B.H. TODO: need a for loop above so we don't need to do this by hand + HOD_dict_sat = {'CSMF': CSMF_dict_sat} + HOD_dict_cent = {'CSMF': CSMF_dict_cent} + + # do a concatenate in numba parallel + start = time.time() + HOD_dict = {} + for tracer in tracers: + tracer_dict = {'Ncent': len(HOD_dict_cent[tracer]['x'])} + for k in HOD_dict_cent[tracer]: + tracer_dict[k] = fast_concatenate( + HOD_dict_cent[tracer][k], HOD_dict_sat[tracer][k], Nthread + ) + tracer_dict['id'] = fast_concatenate( + ID_dict_cent[tracer], ID_dict_sat[tracer], Nthread + ) + if verbose: + print(tracer, 'number of galaxies ', len(tracer_dict['x'])) + print( + 'satellite fraction ', + len(HOD_dict_sat[tracer]['x']) / len(tracer_dict['x']), + ) + HOD_dict[tracer] = tracer_dict + if verbose: + print('organizing outputs took ', time.time() - start) + return HOD_dict + + +def gen_gal_cat_CSMF( + halo_data, + particle_data, + tracers, + params, + Nthread=16, + enable_ranks=False, + rsd=True, + nfw=False, + NFW_draw=None, + write_to_disk=False, + savedir='./', + verbose=False, + fn_ext=None, +): + """ + pass on inputs to the gen_gals function and takes care of I/O + + Parameters + ---------- + + halos_data : dictionary of arrays + a dictionary of halo properties (pos, vel, mass, id, randoms, ...) + + particle_data : dictionary of arrays + a dictionary of particle propoerties (pos, vel, hmass, hid, Np, subsampling, randoms, ...) + + tracers : dictionary of dictionaries + Dictionary of multi-tracer HODs + + enable_ranks : boolean + Flag of whether to implement particle ranks. + + rsd : boolean + Flag of whether to implement RSD. + + nfw : boolean + Flag of whether to generate satellites from an NFW profile. + + write_to_disk : boolean + Flag of whether to output to disk. + + verbose : boolean + Whether to output detailed outputs. + + savedir : str + where to save the output if write_to_disk == True. + + params : dict + Dictionary of various simulation parameters. + + fn_ext: str + filename extension for saved files. Only relevant when ``write_to_disk = True``. + + Output + ------ + + HOD_dict : dictionary of dictionaries + Dictionary of the format: {tracer1_dict, tracer2_dict, ...}, + where tracer1_dict = {x, y, z, vx, vy, vz, mass, id} + + """ + + if not isinstance(rsd, bool): + raise ValueError('Error: rsd has to be a boolean') + + # find the halos, populate them with galaxies and write them to files + HOD_dict = gen_gals( + halo_data, + particle_data, + tracers, + params, + Nthread, + enable_ranks, + rsd, + verbose, + nfw, + NFW_draw, + ) + + # how many galaxies were generated and write them to disk + for tracer in tracers.keys(): + Ncent = HOD_dict[tracer]['Ncent'] + if verbose: + print( + 'generated %ss:' % tracer, + len(HOD_dict[tracer]['x']), + 'satellite fraction ', + 1 - Ncent / len(HOD_dict[tracer]['x']), + ) + + if write_to_disk: + if verbose: + print('outputting galaxies to disk') + + if rsd: + rsd_string = '_rsd' + else: + rsd_string = '' + + if fn_ext is None: + outdir = (savedir) / ('galaxies' + rsd_string) + else: + outdir = (savedir) / ('galaxies' + rsd_string + fn_ext) + + # create directories if not existing + os.makedirs(outdir, exist_ok=True) + + # save to file + # outdict = + HOD_dict[tracer].pop('Ncent', None) + table = Table( + HOD_dict[tracer], + meta={'Ncent': Ncent, 'Gal_type': tracer, **tracers[tracer]}, + ) + if params['chunk'] == -1: + ascii.write( + table, outdir / (f'{tracer}s.dat'), overwrite=True, format='ecsv' + ) + else: + ascii.write( + table, + outdir / (f"{tracer}s_chunk{params['chunk']:d}.dat"), + overwrite=True, + format='ecsv', + ) + + return HOD_dict diff --git a/abacusnbody/hod/GRAND_HOD.py b/abacusnbody/hod/GRAND_HOD.py index 9c9afbfb..0437a4db 100644 --- a/abacusnbody/hod/GRAND_HOD.py +++ b/abacusnbody/hod/GRAND_HOD.py @@ -19,305 +19,6 @@ G = 4.302e-6 # in kpc/Msol (km.s)^2 -@njit(fastmath=True) -def n_cen_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): - """ - Standard Cacciato et al. (2008) centrals HOD parametrization for CSMF - """ - M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - x_low = np.log10(Mstar_low / M_c_value) / (1.41421356 * sigma_c) - x_up = np.log10(Mstar_up / M_c_value) / (1.41421356 * sigma_c) - - return 0.5 * (math.erf(x_up) - math.erf(x_low)) - - -@njit(fastmath=True) -def CSMF_centrals(M_h, Mstar, M_1, M_0, gamma1, gamma2, sigma_c): - """ - Eq. (34) from Cacciato et al. (2008) - """ - - M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - - return ( - 1 - / (1.41421356 * np.sqrt(np.pi) * np.log(10) * sigma_c * Mstar) - * np.exp(-((np.log10(Mstar) - np.log10(M_c_value)) ** 2) / (2 * sigma_c**2)) - ) - - -@njit(fastmath=True) -def M_c(M_h, M_1, M_0, gamma1, gamma2): - """ - Eq. (37) from Cacciato et al. (2008) - """ - return M_0 * (M_h / M_1) ** gamma1 / (1 + M_h / M_1) ** (gamma1 - gamma2) - - -@njit(fastmath=True) -def get_random_cen_stellarmass( - M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c -): - nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) - stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 - delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] - - CSMF_cen = CSMF_centrals( - M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c - ) - - cdf = [] - cdf_current_value = 0 - for k in range(len(delta_stellar_masses)): - cdf_current_value = cdf_current_value + CSMF_cen[k] * delta_stellar_masses[k] - cdf.append(cdf_current_value) - cdf = np.array(cdf) - - cdf = cdf / cdf[-1] - - random_rv = np.random.uniform(cdf.min(), cdf.max()) - bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) - - -@njit(fastmath=True) -def get_random_cen_stellarmass_linearinterpolation( - M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c -): - nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) - - CSMF_cen = CSMF_centrals(M_h, stellarmass, M_1, M_0, gamma1, gamma2, sigma_c) - - delta = stellarmass[1:] - stellarmass[:-1] - - cdf = [] - cdf_current_value = 0 - for i in range(len(delta)): - cdf_current_value = cdf_current_value + CSMF_cen[i] * delta[i] - cdf.append(cdf_current_value) - cdf = np.array(cdf) - - cdf = cdf / cdf[-1] - - random_rv = np.random.uniform(cdf.min(), cdf.max()) - bin = np.where(cdf > random_rv)[0][0] - - m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) - return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] - - -@njit(fastmath=True) -def n_sat_CSMF( - M_h, - Mstar_low, - Mstar_up, - M_1, - M_0, - gamma1, - gamma2, - sigma_c, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, -): - """ - Standard Cacciato et al. (2008) satellite HOD parametrization for CSMF - """ - nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) - - CSMF_sat = CSMF_satelites( - M_h, - stellarmass, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, - ) - - nsat = 0 - for i in range(nbins - 1): - nsat += (CSMF_sat[i + 1] - CSMF_sat[i]) * ( - stellarmass[i + 1] - stellarmass[i] - ) / 2 + (stellarmass[i + 1] - stellarmass[i]) * CSMF_sat[i] - - return nsat # *ncen - - -@njit(fastmath=True) -def CSMF_satelites( - M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2 -): - """ - Eq. (36) from Cacciato et al. (2008) - """ - M_s_value = M_s(M_h, M_1, M_0, gamma1, gamma2) - alpha_s_value = alpha_s(M_h, a1, a2, M2) - phi_s_value = phi_s(M_h, b0, b1, b2) - - delta = 10 ** (delta1 + delta2 * (np.log10(M_h) - 12)) - - return ( - phi_s_value - / M_s_value - * (Mstar / M_s_value) ** alpha_s_value - * np.exp(-delta * (Mstar / M_s_value) ** 2) - ) - - -@njit(fastmath=True) -def M_s(M_h, M_1, M_0, gamma1, gamma2): - """ - Eq. (38) from Cacciato et al. (2008) - """ - return 0.562 * M_c(M_h, M_1, M_0, gamma1, gamma2) - - -@njit(fastmath=True) -def alpha_s(M_h, a1, a2, M2): - """ - Eq. (39) from Cacciato et al. (2008) - """ - return -2.0 + a1 * (1 - 2 / np.pi * np.arctan(a2 * np.log10(M_h / M2))) - - -@njit(fastmath=True) -def phi_s(M_h, b0, b1, b2): - """ - Eq. (40) from Cacciato et al. (2008) - """ - M12 = M_h / 1e12 - log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12) ** 2 - return 10**log_phi_s - - -@njit(fastmath=True) -def get_random_sat_stellarmass( - M_h, - Mstar_low, - Mstar_up, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, -): - nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) - stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 - delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] - - CSMF_sat = CSMF_satelites( - M_h, - stellarmass_centers, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, - ) - - cdf = [] - cdf_current_value = 0 - for k in range(len(delta_stellar_masses)): - cdf_current_value = cdf_current_value + CSMF_sat[k] * delta_stellar_masses[k] - cdf.append(cdf_current_value) - cdf = np.array(cdf) - - cdf = cdf / cdf[-1] - - random_rv = np.random.uniform(cdf.min(), cdf.max()) - bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) - - -@njit(fastmath=True) -def get_random_sat_stellarmass_linearinterpolation( - M_h, - Mstar_low, - Mstar_up, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, -): - nbins = 100 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) - - CSMF_sat = CSMF_satelites( - M_h, - stellarmass, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, - ) - - delta = stellarmass[1:] - stellarmass[:-1] - - cdf = [] - cdf_current_value = 0 - for i in range(len(delta)): - cdf_current_value = cdf_current_value + CSMF_sat[i] * delta[i] - cdf.append(cdf_current_value) - cdf = np.array(cdf) - - cdf = cdf / cdf[-1] - - random_rv = np.random.uniform(cdf.min(), cdf.max()) - bin = np.where(cdf > random_rv)[0][0] - - m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) - return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] - - @njit(fastmath=True) def n_sat_LRG_modified(M_h, logM_cut, M_cut, M_1, sigma, alpha, kappa): """ @@ -449,14 +150,12 @@ def gen_cent( LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, - CSMF_hod_dict, rsd, inv_velz2kms, lbox, want_LRG, want_ELG, want_QSO, - want_CSMF, Nthread, origin, ): @@ -498,27 +197,11 @@ def gen_cent( QSO_hod_dict['Bcent'], QSO_hod_dict['ic'], ) - if want_CSMF: - Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - ) - ic_C, alpha_c_C, Ac_C, Bc_C = ( - CSMF_hod_dict['ic'], - CSMF_hod_dict['alpha_c'], - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Bcent'], - ) H = len(mass) numba.set_num_threads(Nthread) - Nout = np.zeros((Nthread, 4, 8), dtype=np.int64) + Nout = np.zeros((Nthread, 3, 8), dtype=np.int64) hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype( np.int64 ) # starting index of each thread @@ -555,23 +238,6 @@ def gen_cent( N_cen_QSO(mass[i], logM_cut_Q_temp, sigma_Q) * ic_Q * multis[i] ) - CSMF_marker = QSO_marker - if want_CSMF: - M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) - - ncen = n_cen_CSMF( - mass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - ) - - CSMF_marker += ncen * ic_C * multis[i] - if randoms[i] <= LRG_marker: Nout[tid, 0, 0] += 1 # counting keep[i] = 1 @@ -581,19 +247,15 @@ def gen_cent( elif randoms[i] <= QSO_marker: Nout[tid, 2, 0] += 1 # counting keep[i] = 3 - elif randoms[i] <= CSMF_marker: - Nout[tid, 3, 0] += 1 # counting - keep[i] = 4 else: keep[i] = 0 # compose galaxy array, first create array of galaxy starting indices for the threads - gstart = np.empty((Nthread + 1, 4), dtype=np.int64) + gstart = np.empty((Nthread + 1, 3), dtype=np.int64) gstart[0, :] = 0 gstart[1:, 0] = Nout[:, 0, 0].cumsum() gstart[1:, 1] = Nout[:, 1, 0].cumsum() gstart[1:, 2] = Nout[:, 2, 0].cumsum() - gstart[1:, 3] = Nout[:, 3, 0].cumsum() # galaxy arrays N_lrg = gstart[-1, 0] @@ -628,21 +290,9 @@ def gen_cent( qso_mass = np.empty(N_qso, dtype=mass.dtype) qso_id = np.empty(N_qso, dtype=ids.dtype) - # galaxy arrays - N_CSMF = gstart[-1, 3] - CSMF_x = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_y = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_z = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_id = np.empty(N_CSMF, dtype=ids.dtype) - # fill in the galaxy arrays for tid in numba.prange(Nthread): - j1, j2, j3, j4 = gstart[tid] + j1, j2, j3 = gstart[tid] for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: # loop thru three directions to assign galaxy velocities and positions @@ -728,53 +378,11 @@ def gen_cent( qso_mass[j3] = mass[i] qso_id[j3] = ids[i] j3 += 1 - elif keep[i] == 4: - # loop thru three directions to assign galaxy velocities and positions - CSMF_x[j4] = pos[i, 0] - CSMF_vx[j4] = vel[i, 0] + alpha_c_C * vdev[i, 0] # velocity bias - CSMF_y[j4] = pos[i, 1] - CSMF_vy[j4] = vel[i, 1] + alpha_c_C * vdev[i, 1] # velocity bias - CSMF_z[j4] = pos[i, 2] - CSMF_vz[j4] = vel[i, 2] + alpha_c_C * vdev[i, 2] # velocity bias - - # rsd only applies to the z direction - if rsd and origin is not None: - nx = CSMF_x[j4] - origin[0] - ny = CSMF_y[j4] - origin[1] - nz = CSMF_z[j4] - origin[2] - inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) - nx *= inv_norm - ny *= inv_norm - nz *= inv_norm - proj = inv_velz2kms * ( - CSMF_vx[j4] * nx + CSMF_vy[j4] * ny + CSMF_vz[j4] * nz - ) - CSMF_x[j4] = CSMF_x[j4] + proj * nx - CSMF_y[j4] = CSMF_y[j4] + proj * ny - CSMF_z[j4] = CSMF_z[j4] + proj * nz - elif rsd: - CSMF_z[j4] = wrap(pos[i, 2] + CSMF_vz[j4] * inv_velz2kms, lbox) - - CSMF_mass[j4] = mass[i] - M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) - CSMF_stellarmass[j4] = get_random_cen_stellarmass_linearinterpolation( - mass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - ) - CSMF_id[j4] = ids[i] - j4 += 1 # assert j == gstart[tid + 1] LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) LRG_dict['x'] = lrg_x LRG_dict['y'] = lrg_y @@ -802,17 +410,7 @@ def gen_cent( QSO_dict['vz'] = qso_vz QSO_dict['mass'] = qso_mass ID_dict['QSO'] = qso_id - - CSMF_dict['x'] = CSMF_x - CSMF_dict['y'] = CSMF_y - CSMF_dict['z'] = CSMF_z - CSMF_dict['vx'] = CSMF_vx - CSMF_dict['vy'] = CSMF_vy - CSMF_dict['vz'] = CSMF_vz - CSMF_dict['mass'] = CSMF_mass - CSMF_dict['stellarmass'] = CSMF_stellarmass - ID_dict['CSMF'] = CSMF_id - return LRG_dict, ELG_dict, QSO_dict, CSMF_dict, ID_dict, keep + return LRG_dict, ELG_dict, QSO_dict, ID_dict, keep @njit(parallel=True, fastmath=True) @@ -936,11 +534,9 @@ def gen_sats_nfw( LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, - CSMF_hod_dict, want_LRG, want_ELG, want_QSO, - want_CSMF, rsd, inv_velz2kms, lbox, @@ -1024,49 +620,6 @@ def gen_sats_nfw( ) f_sigv_Q = QSO_hod_dict['f_sigv'] - if want_CSMF: - ( - Mstar_low_C, - Mstar_up_C, - M_1_C, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], - CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'], - ) - Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], - CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'], - ) - f_sigv_C = CSMF_hod_dict['f_sigv'] - numba.set_num_threads(Nthread) # compute nsate for each halo @@ -1074,8 +627,6 @@ def gen_sats_nfw( num_sats_L = np.zeros(len(hid), dtype=np.int64) num_sats_E = np.zeros(len(hid), dtype=np.int64) num_sats_Q = np.zeros(len(hid), dtype=np.int64) - num_sats_C = np.zeros(len(hid), dtype=np.int64) - stellarmass_C = np.zeros(len(hid), dtype=np.int64) hstart = np.rint(np.linspace(0, len(hid), Nthread + 1)).astype( np.int64 ) # starting index of each thread @@ -1152,40 +703,10 @@ def gen_sats_nfw( ) num_sats_Q[i] = np.random.poisson(base_p_Q) - if want_CSMF: - M_1_C_temp = 10 ** ( - np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] - ) - a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] - base_p_C = ( - n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) - * ic_C - ) - - num_sats_C[i] = np.random.poisson(base_p_C) - # generate rdpos rd_pos_L = getPointsOnSphere(np.sum(num_sats_L), Nthread) rd_pos_E = getPointsOnSphere(np.sum(num_sats_E), Nthread) rd_pos_Q = getPointsOnSphere(np.sum(num_sats_Q), Nthread) - rd_pos_C = getPointsOnSphere(np.sum(num_sats_C), Nthread) # put satellites on NFW h_id_L, x_sat_L, y_sat_L, z_sat_L, vx_sat_L, vy_sat_L, vz_sat_L, M_L = ( @@ -1260,43 +781,15 @@ def gen_sats_nfw( nfw_rescale, ) ) - - h_id_C, x_sat_C, y_sat_C, z_sat_C, vx_sat_C, vy_sat_C, vz_sat_C, M_C = ( - compute_fast_NFW( - NFW_draw, - hid, - hpos[:, 0], - hpos[:, 1], - hpos[:, 2], - hvel[:, 0], - hvel[:, 1], - hvel[:, 2], - hvrms, - hc, - hmass, - hrvir, - rd_pos_C, - num_sats_C, - f_sigv_C, - vel_sat, - Nthread, - exp_frac, - exp_scale, - nfw_rescale, - ) - ) - # do rsd if rsd: z_sat_L = (z_sat_L + vz_sat_L * inv_velz2kms) % lbox z_sat_E = (z_sat_E + vz_sat_E * inv_velz2kms) % lbox z_sat_Q = (z_sat_Q + vz_sat_Q * inv_velz2kms) % lbox - z_sat_C = (z_sat_C + vz_sat_C * inv_velz2kms) % lbox LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) LRG_dict['x'] = x_sat_L LRG_dict['y'] = y_sat_L @@ -1325,46 +818,7 @@ def gen_sats_nfw( QSO_dict['mass'] = M_Q ID_dict['QSO'] = h_id_Q - CSMF_dict['x'] = x_sat_C - CSMF_dict['y'] = y_sat_C - CSMF_dict['z'] = z_sat_C - CSMF_dict['vx'] = vx_sat_C - CSMF_dict['vy'] = vy_sat_C - CSMF_dict['vz'] = vz_sat_C - CSMF_dict['mass'] = M_C - stellarmass_C = np.empty_like(M_C) - ## compute stellarmass of all the satelites - if want_CSMF: - hstart = np.rint(np.linspace(0, num_sats_C.sum(), Nthread + 1)) - for tid in numba.prange(Nthread): - for i in range(int(hstart[tid]), int(hstart[tid + 1])): - M_1_C_temp = 10 ** ( - np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] - ) - a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] - stellarmass_C[i] = get_random_sat_stellarmass_linearinterpolation( - M_C[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - a1_C_temp, - s, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) - - CSMF_dict['stellarmass'] = stellarmass_C - ID_dict['CSMF'] = h_id_C - - return LRG_dict, ELG_dict, QSO_dict, CSMF_dict, ID_dict + return LRG_dict, ELG_dict, QSO_dict, ID_dict @njit(parallel=True, fastmath=True) @@ -1388,7 +842,6 @@ def gen_sats( LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, - CSMF_hod_dict, rsd, inv_velz2kms, lbox, @@ -1396,7 +849,6 @@ def gen_sats( want_LRG, want_ELG, want_QSO, - want_CSMF, Nthread, origin, keep_cent, @@ -1490,58 +942,10 @@ def gen_sats( QSO_hod_dict['ic'], ) - if want_CSMF: - ( - Mstar_low_C, - Mstar_up_C, - M_1_C, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], - CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'], - ) - - alpha_s_C, s_C, s_v_C, s_p_C, s_r_C, Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['alpha_s'], - CSMF_hod_dict['s'], - CSMF_hod_dict['s_v'], - CSMF_hod_dict['s_p'], - CSMF_hod_dict['s_r'], - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], - CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'], - ) - H = len(hmass) # num of particles numba.set_num_threads(Nthread) - Nout = np.zeros((Nthread, 4, 8), dtype=np.int64) + Nout = np.zeros((Nthread, 3, 8), dtype=np.int64) hstart = np.rint(np.linspace(0, H, Nthread + 1)).astype( np.int64 ) # starting index of each thread @@ -1670,47 +1074,6 @@ def gen_sats( exp_sat = base_p_Q QSO_marker += exp_sat - CSMF_marker = QSO_marker - if want_CSMF: - M_1_C_temp = 10 ** ( - np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] - ) - a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] - base_p_C = ( - n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) - * weights[i] - * ic_C - ) - if enable_ranks: - decorator_C = ( - 1 - + s_C * ranks[i] - + s_v_C * ranksv[i] - + s_p_C * ranksp[i] - + s_r_C * ranksr[i] - ) - exp_sat = base_p_C * decorator_C - else: - exp_sat = base_p_C - CSMF_marker += exp_sat - if randoms[i] <= LRG_marker: Nout[tid, 0, 0] += 1 # counting keep[i] = 1 @@ -1720,19 +1083,15 @@ def gen_sats( elif randoms[i] <= QSO_marker: Nout[tid, 2, 0] += 1 # counting keep[i] = 3 - elif randoms[i] <= CSMF_marker: - Nout[tid, 3, 0] += 1 # counting - keep[i] = 4 else: keep[i] = 0 # compose galaxy array, first create array of galaxy starting indices for the threads - gstart = np.empty((Nthread + 1, 4), dtype=np.int64) + gstart = np.empty((Nthread + 1, 3), dtype=np.int64) gstart[0, :] = 0 gstart[1:, 0] = Nout[:, 0, 0].cumsum() gstart[1:, 1] = Nout[:, 1, 0].cumsum() gstart[1:, 2] = Nout[:, 2, 0].cumsum() - gstart[1:, 3] = Nout[:, 3, 0].cumsum() # galaxy arrays N_lrg = gstart[-1, 0] @@ -1767,21 +1126,9 @@ def gen_sats( qso_mass = np.empty(N_qso, dtype=hmass.dtype) qso_id = np.empty(N_qso, dtype=hid.dtype) - # galaxy arrays - N_CSMF = gstart[-1, 3] - CSMF_x = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_y = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_z = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_id = np.empty(N_CSMF, dtype=hid.dtype) - # fill in the galaxy arrays for tid in numba.prange(Nthread): - j1, j2, j3, j4 = gstart[tid] + j1, j2, j3 = gstart[tid] for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: lrg_x[j1] = ppos[i, 0] @@ -1879,66 +1226,11 @@ def gen_sats( qso_mass[j3] = hmass[i] qso_id[j3] = hid[i] j3 += 1 - elif keep[i] == 4: - CSMF_x[j4] = ppos[i, 0] - CSMF_vx[j4] = hvel[i, 0] + alpha_s_C * ( - pvel[i, 0] - hvel[i, 0] - ) # velocity bias - CSMF_y[j4] = ppos[i, 1] - CSMF_vy[j4] = hvel[i, 1] + alpha_s_C * ( - pvel[i, 1] - hvel[i, 1] - ) # velocity bias - CSMF_z[j4] = ppos[i, 2] - CSMF_vz[j4] = hvel[i, 2] + alpha_s_C * ( - pvel[i, 2] - hvel[i, 2] - ) # velocity bias - if rsd and origin is not None: - nx = CSMF_x[j4] - origin[0] - ny = CSMF_y[j4] - origin[1] - nz = CSMF_z[j4] - origin[2] - inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) - nx *= inv_norm - ny *= inv_norm - nz *= inv_norm - proj = inv_velz2kms * ( - CSMF_vx[j4] * nx + CSMF_vy[j4] * ny + CSMF_vz[j4] * nz - ) - CSMF_x[j4] = CSMF_x[j4] + proj * nx - CSMF_y[j4] = CSMF_y[j4] + proj * ny - CSMF_z[j4] = CSMF_z[j4] + proj * nz - elif rsd: - CSMF_z[j4] = wrap(CSMF_z[j4] + CSMF_vz[j4] * inv_velz2kms, lbox) - - M_1_C_temp = 10 ** ( - np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] - ) - a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] - CSMF_stellarmass[j4] = get_random_sat_stellarmass( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) - CSMF_mass[j4] = hmass[i] - CSMF_id[j4] = hid[i] - j4 += 1 # assert j == gstart[tid + 1] LRG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ELG_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) QSO_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) LRG_dict['x'] = lrg_x LRG_dict['y'] = lrg_y @@ -1966,18 +1258,7 @@ def gen_sats( QSO_dict['vz'] = qso_vz QSO_dict['mass'] = qso_mass ID_dict['QSO'] = qso_id - - CSMF_dict['x'] = CSMF_x - CSMF_dict['y'] = CSMF_y - CSMF_dict['z'] = CSMF_z - CSMF_dict['vx'] = CSMF_vx - CSMF_dict['vy'] = CSMF_vy - CSMF_dict['vz'] = CSMF_vz - CSMF_dict['mass'] = CSMF_mass - CSMF_dict['stellarmass'] = CSMF_stellarmass - ID_dict['CSMF'] = CSMF_id - - return LRG_dict, ELG_dict, QSO_dict, CSMF_dict, ID_dict + return LRG_dict, ELG_dict, QSO_dict, ID_dict @njit(parallel=True, fastmath=True) @@ -2064,8 +1345,6 @@ def gen_gals( ELG_HOD = tracers[tracer] if tracer == 'QSO': QSO_HOD = tracers[tracer] - if tracer == 'CSMF': - CSMF_HOD = tracers[tracer] if 'LRG' in tracers.keys(): want_LRG = True @@ -2187,28 +1466,6 @@ def gen_gals( key_type=nb.types.unicode_type, value_type=nb.types.float64 ) - if 'CSMF' in tracers.keys(): - want_CSMF = True - - CSMF_hod_dict = nb.typed.Dict.empty( - key_type=nb.types.unicode_type, value_type=nb.types.float64 - ) - for key, value in CSMF_HOD.items(): - CSMF_hod_dict[key] = value - - CSMF_hod_dict['Acent'] = CSMF_HOD.get('Acent', 0.0) - CSMF_hod_dict['Asat'] = CSMF_HOD.get('Asat', 0.0) - CSMF_hod_dict['Bcent'] = CSMF_HOD.get('Bcent', 0.0) - CSMF_hod_dict['Bsat'] = CSMF_HOD.get('Bsat', 0.0) - CSMF_hod_dict['ic'] = CSMF_HOD.get('ic', 1.0) - CSMF_hod_dict['f_sigv'] = CSMF_HOD.get('f_sigv', 0) - - else: - want_CSMF = False - CSMF_hod_dict = nb.typed.Dict.empty( - key_type=nb.types.unicode_type, value_type=nb.types.float64 - ) - start = time.time() velz2kms = params['velz2kms'] @@ -2216,14 +1473,7 @@ def gen_gals( lbox = params['Lbox'] origin = params['origin'] - ( - LRG_dict_cent, - ELG_dict_cent, - QSO_dict_cent, - CSMF_dict_cent, - ID_dict_cent, - keep_cent, - ) = gen_cent( + LRG_dict_cent, ELG_dict_cent, QSO_dict_cent, ID_dict_cent, keep_cent = gen_cent( halos_array['hpos'], halos_array['hvel'], halos_array['hmass'], @@ -2237,14 +1487,12 @@ def gen_gals( LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, - CSMF_hod_dict, rsd, inv_velz2kms, lbox, want_LRG, want_ELG, want_QSO, - want_CSMF, Nthread, origin, ) @@ -2256,36 +1504,32 @@ def gen_gals( warnings.warn( 'NFW profile is unoptimized. It has different velocity bias. It does not support lightcone.' ) - LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, CSMF_dict_sat, ID_dict_sat = ( - gen_sats_nfw( - NFW_draw, - halos_array['hpos'], - halos_array['hvel'], - halos_array['hmass'], - halos_array['hid'], - halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), - halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), - halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), - halos_array['hsigma3d'], - halos_array['hc'], - halos_array['hrvir'], - LRG_hod_dict, - ELG_hod_dict, - QSO_hod_dict, - CSMF_hod_dict, - want_LRG, - want_ELG, - want_QSO, - want_CSMF, - rsd, - inv_velz2kms, - lbox, - keep_cent, - Nthread=Nthread, - ) + LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, ID_dict_sat = gen_sats_nfw( + NFW_draw, + halos_array['hpos'], + halos_array['hvel'], + halos_array['hmass'], + halos_array['hid'], + halos_array.get('hdeltac', np.zeros(len(halos_array['hmass']))), + halos_array.get('hfenv', np.zeros(len(halos_array['hmass']))), + halos_array.get('hshear', np.zeros(len(halos_array['hmass']))), + halos_array['hsigma3d'], + halos_array['hc'], + halos_array['hrvir'], + LRG_hod_dict, + ELG_hod_dict, + QSO_hod_dict, + want_LRG, + want_ELG, + want_QSO, + rsd, + inv_velz2kms, + lbox, + keep_cent, + Nthread=Nthread, ) else: - LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, CSMF_dict_sat, ID_dict_sat = gen_sats( + LRG_dict_sat, ELG_dict_sat, QSO_dict_sat, ID_dict_sat = gen_sats( subsample['ppos'], subsample['pvel'], subsample['phvel'], @@ -2305,7 +1549,6 @@ def gen_gals( LRG_hod_dict, ELG_hod_dict, QSO_hod_dict, - CSMF_hod_dict, rsd, inv_velz2kms, lbox, @@ -2313,7 +1556,6 @@ def gen_gals( want_LRG, want_ELG, want_QSO, - want_CSMF, Nthread, origin, keep_cent[subsample['pinds']], @@ -2322,18 +1564,8 @@ def gen_gals( print('generating satellites took ', time.time() - start) # B.H. TODO: need a for loop above so we don't need to do this by hand - HOD_dict_sat = { - 'LRG': LRG_dict_sat, - 'ELG': ELG_dict_sat, - 'QSO': QSO_dict_sat, - 'CSMF': CSMF_dict_sat, - } - HOD_dict_cent = { - 'LRG': LRG_dict_cent, - 'ELG': ELG_dict_cent, - 'QSO': QSO_dict_cent, - 'CSMF': CSMF_dict_cent, - } + HOD_dict_sat = {'LRG': LRG_dict_sat, 'ELG': ELG_dict_sat, 'QSO': QSO_dict_sat} + HOD_dict_cent = {'LRG': LRG_dict_cent, 'ELG': ELG_dict_cent, 'QSO': QSO_dict_cent} # do a concatenate in numba parallel start = time.time() @@ -2486,4 +1718,4 @@ def gen_gal_cat( format='ecsv', ) - return HOD_dict + return HOD_dict \ No newline at end of file diff --git a/abacusnbody/hod/abacus_hod.py b/abacusnbody/hod/abacus_hod.py index 9d1962bc..7663c5f2 100644 --- a/abacusnbody/hod/abacus_hod.py +++ b/abacusnbody/hod/abacus_hod.py @@ -33,10 +33,15 @@ N_sat_elg, N_cen_QSO, N_sat_generic, +) + +from .CSMF_HOD import ( + gen_gal_cat_CSMF, n_cen_CSMF, n_sat_CSMF, ) + # TODO B.H.: staging can be shorter and prettier; perhaps asdf for h5 and ecsv? @@ -760,21 +765,39 @@ def run_hod( ) start = time.time() - mock_dict = gen_gal_cat( - self.halo_data, - self.particle_data, - tracers, - self.params, - Nthread, - enable_ranks=self.want_ranks, - rsd=want_rsd, - nfw=want_nfw, - NFW_draw=NFW_draw, - write_to_disk=write_to_disk, - savedir=self.mock_dir, - verbose=verbose, - fn_ext=fn_ext, - ) + if 'CSMF' in tracers.keys(): + mock_dict = gen_gal_cat_CSMF( + self.halo_data, + self.particle_data, + tracers, + self.params, + Nthread, + enable_ranks=self.want_ranks, + rsd=want_rsd, + nfw=want_nfw, + NFW_draw=NFW_draw, + write_to_disk=write_to_disk, + savedir=self.mock_dir, + verbose=verbose, + fn_ext=fn_ext, + ) + else: + mock_dict = gen_gal_cat( + self.halo_data, + self.particle_data, + tracers, + self.params, + Nthread, + enable_ranks=self.want_ranks, + rsd=want_rsd, + nfw=want_nfw, + NFW_draw=NFW_draw, + write_to_disk=write_to_disk, + savedir=self.mock_dir, + verbose=verbose, + fn_ext=fn_ext, + ) + self.logger.info(f'HOD generated in elapsed time {time.time() - start:.2f} s.') return mock_dict From 3c77aa44a64b16d4843b860f5dcab4560ebe2150 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 9 Oct 2024 19:57:17 +0000 Subject: [PATCH 06/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- abacusnbody/hod/CSMF_HOD.py | 781 +++++++++++++++++++++------------- abacusnbody/hod/GRAND_HOD.py | 2 +- abacusnbody/hod/abacus_hod.py | 2 +- 3 files changed, 487 insertions(+), 298 deletions(-) diff --git a/abacusnbody/hod/CSMF_HOD.py b/abacusnbody/hod/CSMF_HOD.py index 0963ef11..e577ce4c 100644 --- a/abacusnbody/hod/CSMF_HOD.py +++ b/abacusnbody/hod/CSMF_HOD.py @@ -19,108 +19,155 @@ G = 4.302e-6 # in kpc/Msol (km.s)^2 - @njit(fastmath=True) def n_cen_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): """ Standard Cacciato et al. (2008) centrals HOD parametrization for CSMF """ M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - x_low = np.log10(Mstar_low/M_c_value)/(1.41421356*sigma_c) - x_up = np.log10(Mstar_up/M_c_value)/(1.41421356*sigma_c) - - return 0.5*(math.erf(x_up) - math.erf(x_low)) + x_low = np.log10(Mstar_low / M_c_value) / (1.41421356 * sigma_c) + x_up = np.log10(Mstar_up / M_c_value) / (1.41421356 * sigma_c) + + return 0.5 * (math.erf(x_up) - math.erf(x_low)) + @njit(fastmath=True) def CSMF_centrals(M_h, Mstar, M_1, M_0, gamma1, gamma2, sigma_c): """ Eq. (34) from Cacciato et al. (2008) """ - + M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - - return 1/(1.41421356*np.sqrt(np.pi)*np.log(10)*sigma_c*Mstar)*np.exp(-(np.log10(Mstar)-np.log10(M_c_value))**2/(2*sigma_c**2)) + + return ( + 1 + / (1.41421356 * np.sqrt(np.pi) * np.log(10) * sigma_c * Mstar) + * np.exp(-((np.log10(Mstar) - np.log10(M_c_value)) ** 2) / (2 * sigma_c**2)) + ) + @njit(fastmath=True) def M_c(M_h, M_1, M_0, gamma1, gamma2): """ Eq. (37) from Cacciato et al. (2008) """ - return M_0 * (M_h/M_1)**gamma1/(1+M_h/M_1)**(gamma1-gamma2) - + return M_0 * (M_h / M_1) ** gamma1 / (1 + M_h / M_1) ** (gamma1 - gamma2) @njit(fastmath=True) -def get_random_cen_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): - +def get_random_cen_stellarmass( + M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c +): nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - - CSMF_cen = CSMF_centrals(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c) - - - cdf = np.cumsum(CSMF_cen*delta_stellar_masses) - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + + CSMF_cen = CSMF_centrals( + M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c + ) + + cdf = np.cumsum(CSMF_cen * delta_stellar_masses) + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) - + + return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) - @njit(fastmath=True) -def get_random_cen_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): - +def get_random_cen_stellarmass_linearinterpolation( + M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c +): nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + CSMF_cen = CSMF_centrals(M_h, stellarmass, M_1, M_0, gamma1, gamma2, sigma_c) - - cdf = np.cumsum(CSMF_cen[:-1]*delta_stellar_masses) - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) - bin = np.where(cdf>random_rv)[0][0] - - m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) - return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + + cdf = np.cumsum(CSMF_cen[:-1] * delta_stellar_masses) + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) + bin = np.where(cdf > random_rv)[0][0] + + m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) + return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] @njit(fastmath=True) -def n_sat_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c, a1, a2, M2, b0, b1, b2, delta1, delta2): +def n_sat_CSMF( + M_h, + Mstar_low, + Mstar_up, + M_1, + M_0, + gamma1, + gamma2, + sigma_c, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, +): """ Standard Cacciato et al. (2008) satellite HOD parametrization for CSMF """ nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - - CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + + CSMF_sat = CSMF_satelites( + M_h, + stellarmass, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, + ) nsat = 0 - for i in range(nbins-1): - nsat += (CSMF_sat[i+1]-CSMF_sat[i])*(stellarmass[i+1]-stellarmass[i])/2 + (stellarmass[i+1]-stellarmass[i])*CSMF_sat[i] - - return nsat#*ncen + for i in range(nbins - 1): + nsat += (CSMF_sat[i + 1] - CSMF_sat[i]) * ( + stellarmass[i + 1] - stellarmass[i] + ) / 2 + (stellarmass[i + 1] - stellarmass[i]) * CSMF_sat[i] + + return nsat # *ncen @njit(fastmath=True) -def CSMF_satelites(M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): +def CSMF_satelites( + M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2 +): """ Eq. (36) from Cacciato et al. (2008) """ M_s_value = M_s(M_h, M_1, M_0, gamma1, gamma2) alpha_s_value = alpha_s(M_h, a1, a2, M2) phi_s_value = phi_s(M_h, b0, b1, b2) - + delta = 10 ** (delta1 + delta2 * (np.log10(M_h) - 12)) - - return phi_s_value/M_s_value * (Mstar/M_s_value)**alpha_s_value * np.exp(-delta*(Mstar/M_s_value)**2) + + return ( + phi_s_value + / M_s_value + * (Mstar / M_s_value) ** alpha_s_value + * np.exp(-delta * (Mstar / M_s_value) ** 2) + ) + @njit(fastmath=True) def M_s(M_h, M_1, M_0, gamma1, gamma2): @@ -129,60 +176,122 @@ def M_s(M_h, M_1, M_0, gamma1, gamma2): """ return 0.562 * M_c(M_h, M_1, M_0, gamma1, gamma2) + @njit(fastmath=True) def alpha_s(M_h, a1, a2, M2): """ Eq. (39) from Cacciato et al. (2008) """ - return -2.0 + a1 * (1-2/np.pi*np.arctan(a2*np.log10(M_h/M2))) + return -2.0 + a1 * (1 - 2 / np.pi * np.arctan(a2 * np.log10(M_h / M2))) + @njit(fastmath=True) def phi_s(M_h, b0, b1, b2): """ Eq. (40) from Cacciato et al. (2008) """ - M12 = M_h/1e12 - log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12)**2 + M12 = M_h / 1e12 + log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12) ** 2 return 10**log_phi_s @njit(fastmath=True) -def get_random_sat_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): - +def get_random_sat_stellarmass( + M_h, + Mstar_low, + Mstar_up, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, +): nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - - CSMF_sat = CSMF_satelites(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) - - cdf = np.cumsum(CSMF_sat*delta_stellar_masses) - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + + CSMF_sat = CSMF_satelites( + M_h, + stellarmass_centers, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, + ) + + cdf = np.cumsum(CSMF_sat * delta_stellar_masses) + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) - + + return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) + @njit(fastmath=True) -def get_random_sat_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): - +def get_random_sat_stellarmass_linearinterpolation( + M_h, + Mstar_low, + Mstar_up, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, +): nbins = 100 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - - CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) - - cdf = np.cumsum(CSMF_sat[:-1]*delta_stellar_masses) - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) - bin = np.where(cdf>random_rv)[0][0] - - m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) - return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + + CSMF_sat = CSMF_satelites( + M_h, + stellarmass, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, + ) + + cdf = np.cumsum(CSMF_sat[:-1] * delta_stellar_masses) + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) + bin = np.where(cdf > random_rv)[0][0] + + m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) + return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] @njit(fastmath=True) @@ -226,24 +335,23 @@ def gen_cent( ): """ Generate central galaxies in place in memory with a two pass numba parallel implementation. - """ + """ if want_CSMF: Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'] + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], ) ic_C, alpha_c_C, Ac_C, Bc_C = ( - CSMF_hod_dict['ic'], - CSMF_hod_dict['alpha_c'], - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Bcent'] + CSMF_hod_dict['ic'], + CSMF_hod_dict['alpha_c'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Bcent'], ) - H = len(mass) @@ -261,14 +369,23 @@ def gen_cent( # first create the markers between 0 and 1 for different tracers CSMF_marker = 0 if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) - - ncen = n_cen_CSMF(mass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) - + M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + + ncen = n_cen_CSMF( + mass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + ) + CSMF_marker += ncen * ic_C * multis[i] if randoms[i] <= CSMF_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 else: keep[i] = 0 @@ -277,19 +394,18 @@ def gen_cent( gstart = np.empty((Nthread + 1, 1), dtype=np.int64) gstart[0, :] = 0 gstart[1:, 0] = Nout[:, 0, 0].cumsum() - # galaxy arrays N_CSMF = gstart[-1, 0] - CSMF_x = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_y = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_z = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_id = np.empty(N_CSMF, dtype = ids.dtype) + CSMF_x = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_y = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_z = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_id = np.empty(N_CSMF, dtype=ids.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): @@ -297,42 +413,50 @@ def gen_cent( for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: # loop thru three directions to assign galaxy velocities and positions - CSMF_x[j1] = pos[i,0] - CSMF_vx[j1] = vel[i,0] + alpha_c_C * vdev[i,0] # velocity bias - CSMF_y[j1] = pos[i,1] - CSMF_vy[j1] = vel[i,1] + alpha_c_C * vdev[i,1] # velocity bias - CSMF_z[j1] = pos[i,2] - CSMF_vz[j1] = vel[i,2] + alpha_c_C * vdev[i,2] # velocity bias - + CSMF_x[j1] = pos[i, 0] + CSMF_vx[j1] = vel[i, 0] + alpha_c_C * vdev[i, 0] # velocity bias + CSMF_y[j1] = pos[i, 1] + CSMF_vy[j1] = vel[i, 1] + alpha_c_C * vdev[i, 1] # velocity bias + CSMF_z[j1] = pos[i, 2] + CSMF_vz[j1] = vel[i, 2] + alpha_c_C * vdev[i, 2] # velocity bias + # rsd only applies to the z direction if rsd and origin is not None: nx = CSMF_x[j1] - origin[0] ny = CSMF_y[j1] - origin[1] nz = CSMF_z[j1] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms*(CSMF_vx[j1]*nx+CSMF_vy[j1]*ny+CSMF_vz[j1]*nz) - CSMF_x[j1] = CSMF_x[j1]+proj*nx - CSMF_y[j1] = CSMF_y[j1]+proj*ny - CSMF_z[j1] = CSMF_z[j1]+proj*nz + proj = inv_velz2kms * ( + CSMF_vx[j1] * nx + CSMF_vy[j1] * ny + CSMF_vz[j1] * nz + ) + CSMF_x[j1] = CSMF_x[j1] + proj * nx + CSMF_y[j1] = CSMF_y[j1] + proj * ny + CSMF_z[j1] = CSMF_z[j1] + proj * nz elif rsd: - CSMF_z[j1] = wrap(pos[i,2] + CSMF_vz[j1] * inv_velz2kms, lbox) - + CSMF_z[j1] = wrap(pos[i, 2] + CSMF_vz[j1] * inv_velz2kms, lbox) + CSMF_mass[j1] = mass[i] - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) CSMF_stellarmass[j1] = get_random_cen_stellarmass_linearinterpolation( - mass[i], Mstar_low_C, Mstar_up_C, - M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) + mass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + ) CSMF_id[j1] = ids[i] j1 += 1 # assert j == gstart[tid + 1] - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) - - + CSMF_dict['x'] = CSMF_x CSMF_dict['y'] = CSMF_y CSMF_dict['z'] = CSMF_z @@ -479,29 +603,45 @@ def gen_sats_nfw( """ if want_CSMF: - Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], - CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'] + ( + Mstar_low_C, + Mstar_up_C, + M_1_C, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], + CSMF_hod_dict['M_2'], + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'], ) Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Asat'], CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'] + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'], ) f_sigv_C = CSMF_hod_dict['f_sigv'] @@ -509,37 +649,40 @@ def gen_sats_nfw( # compute nsate for each halo # figuring out the number of particles kept for each thread - num_sats_C = np.zeros(len(hid), dtype = np.int64) - stellarmass_C = np.zeros(len(hid), dtype = np.int64) + num_sats_C = np.zeros(len(hid), dtype=np.int64) + stellarmass_C = np.zeros(len(hid), dtype=np.int64) hstart = np.rint(np.linspace(0, len(hid), Nthread + 1)).astype( np.int64 ) # starting index of each thread for tid in range(Nthread): for i in range(hstart[tid], hstart[tid + 1]): if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] base_p_C = ( n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) * ic_C ) - num_sats_C[i] = np.random.poisson(base_p_C) + num_sats_C[i] = np.random.poisson(base_p_C) # generate rdpos rd_pos_C = getPointsOnSphere(np.sum(num_sats_C), Nthread) @@ -547,38 +690,36 @@ def gen_sats_nfw( # put satellites on NFW h_id_C, x_sat_C, y_sat_C, z_sat_C, vx_sat_C, vy_sat_C, vz_sat_C, M_C = ( compute_fast_NFW( - NFW_draw, - hid, - hpos[:, 0], - hpos[:, 1], - hpos[:, 2], - hvel[:, 0], - hvel[:, 1], + NFW_draw, + hid, + hpos[:, 0], + hpos[:, 1], + hpos[:, 2], + hvel[:, 0], + hvel[:, 1], hvel[:, 2], - hvrms, - hc, - hmass, - hrvir, - rd_pos_C, - num_sats_C, - f_sigv_C, - vel_sat, + hvrms, + hc, + hmass, + hrvir, + rd_pos_C, + num_sats_C, + f_sigv_C, + vel_sat, Nthread, - exp_frac, - exp_scale, - nfw_rescale + exp_frac, + exp_scale, + nfw_rescale, ) ) - - + # do rsd if rsd: z_sat_C = (z_sat_C + vz_sat_C * inv_velz2kms) % lbox - - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) - + CSMF_dict['x'] = x_sat_C CSMF_dict['y'] = y_sat_C CSMF_dict['z'] = z_sat_C @@ -592,12 +733,29 @@ def gen_sats_nfw( hstart = np.rint(np.linspace(0, num_sats_C.sum(), Nthread + 1)) for tid in numba.prange(Nthread): for i in range(int(hstart[tid]), int(hstart[tid + 1])): - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] stellarmass_C[i] = get_random_sat_stellarmass_linearinterpolation( - M_C[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, - a1_C_temp, s, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) - + M_C[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + a1_C_temp, + s, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) + CSMF_dict['stellarmass'] = stellarmass_C ID_dict['CSMF'] = h_id_C @@ -636,38 +794,52 @@ def gen_sats( Generate satellite galaxies in place in memory with a two pass numba parallel implementation. """ - - if want_CSMF: - Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], + ( + Mstar_low_C, + Mstar_up_C, + M_1_C, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'] + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'], ) - + alpha_s_C, s_C, s_v_C, s_p_C, s_r_C, Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['alpha_s'], + CSMF_hod_dict['alpha_s'], CSMF_hod_dict['s'], - CSMF_hod_dict['s_v'], - CSMF_hod_dict['s_p'], - CSMF_hod_dict['s_r'], + CSMF_hod_dict['s_v'], + CSMF_hod_dict['s_p'], + CSMF_hod_dict['s_r'], CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], - CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'] + CSMF_hod_dict['Asat'], + CSMF_hod_dict['Bcent'], + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'], ) H = len(hmass) # num of particles @@ -686,39 +858,47 @@ def gen_sats( # print(logM1, As, hdeltac[i], Bs, hfenv[i]) CSMF_marker = 0 if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] base_p_C = ( n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C - ) - * weights[i] - * ic_C) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) + * weights[i] + * ic_C + ) if enable_ranks: - decorator_C = 1 + s_C * ranks[i] + s_v_C * ranksv[i] + s_p_C * ranksp[i] + s_r_C * ranksr[i] + decorator_C = ( + 1 + + s_C * ranks[i] + + s_v_C * ranksv[i] + + s_p_C * ranksp[i] + + s_r_C * ranksr[i] + ) exp_sat = base_p_C * decorator_C else: exp_sat = base_p_C CSMF_marker += exp_sat - if randoms[i] <= CSMF_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 else: keep[i] = 0 @@ -730,73 +910,79 @@ def gen_sats( # galaxy arrays N_CSMF = gstart[-1, 0] - CSMF_x = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_y = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_z = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_id = np.empty(N_CSMF, dtype = hid.dtype) + CSMF_x = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_y = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_z = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_id = np.empty(N_CSMF, dtype=hid.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): j1 = gstart[tid] for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: - CSMF_x[j1] = ppos[i, 0] - CSMF_vx[j1] = hvel[i, 0] + alpha_s_C * (pvel[i, 0] - hvel[i, 0]) # velocity bias + CSMF_vx[j1] = hvel[i, 0] + alpha_s_C * ( + pvel[i, 0] - hvel[i, 0] + ) # velocity bias CSMF_y[j1] = ppos[i, 1] - CSMF_vy[j1] = hvel[i, 1] + alpha_s_C * (pvel[i, 1] - hvel[i, 1]) # velocity bias + CSMF_vy[j1] = hvel[i, 1] + alpha_s_C * ( + pvel[i, 1] - hvel[i, 1] + ) # velocity bias CSMF_z[j1] = ppos[i, 2] - CSMF_vz[j1] = hvel[i, 2] + alpha_s_C * (pvel[i, 2] - hvel[i, 2]) # velocity bias + CSMF_vz[j1] = hvel[i, 2] + alpha_s_C * ( + pvel[i, 2] - hvel[i, 2] + ) # velocity bias if rsd and origin is not None: nx = CSMF_x[j1] - origin[0] ny = CSMF_y[j1] - origin[1] nz = CSMF_z[j1] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms*(CSMF_vx[j1]*nx+CSMF_vy[j1]*ny+CSMF_vz[j1]*nz) - CSMF_x[j1] = CSMF_x[j1]+proj*nx - CSMF_y[j1] = CSMF_y[j1]+proj*ny - CSMF_z[j1] = CSMF_z[j1]+proj*nz + proj = inv_velz2kms * ( + CSMF_vx[j1] * nx + CSMF_vy[j1] * ny + CSMF_vz[j1] * nz + ) + CSMF_x[j1] = CSMF_x[j1] + proj * nx + CSMF_y[j1] = CSMF_y[j1] + proj * ny + CSMF_z[j1] = CSMF_z[j1] + proj * nz elif rsd: - CSMF_z[j1] = wrap(CSMF_z[j1] + CSMF_vz[j1] * inv_velz2kms, lbox) - - - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + CSMF_z[j1] = wrap(CSMF_z[j1] + CSMF_vz[j1] * inv_velz2kms, lbox) + + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] CSMF_stellarmass[j1] = get_random_sat_stellarmass( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C - ) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) CSMF_mass[j1] = hmass[i] CSMF_id[j1] = hid[i] j1 += 1 # assert j == gstart[tid + 1] - - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) - CSMF_dict['x'] = CSMF_x CSMF_dict['y'] = CSMF_y CSMF_dict['z'] = CSMF_z @@ -806,7 +992,7 @@ def gen_sats( CSMF_dict['mass'] = CSMF_mass CSMF_dict['stellarmass'] = CSMF_stellarmass ID_dict['CSMF'] = CSMF_id - + return CSMF_dict, ID_dict @@ -891,11 +1077,12 @@ def gen_gals( if tracer == 'CSMF': CSMF_HOD = tracers[tracer] - if 'CSMF' in tracers.keys(): want_CSMF = True - CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + CSMF_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) for key, value in CSMF_HOD.items(): CSMF_hod_dict[key] = value @@ -908,7 +1095,9 @@ def gen_gals( else: want_CSMF = False - CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + CSMF_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) start = time.time() diff --git a/abacusnbody/hod/GRAND_HOD.py b/abacusnbody/hod/GRAND_HOD.py index 0437a4db..739790dd 100644 --- a/abacusnbody/hod/GRAND_HOD.py +++ b/abacusnbody/hod/GRAND_HOD.py @@ -1718,4 +1718,4 @@ def gen_gal_cat( format='ecsv', ) - return HOD_dict \ No newline at end of file + return HOD_dict diff --git a/abacusnbody/hod/abacus_hod.py b/abacusnbody/hod/abacus_hod.py index 7663c5f2..4de43ebf 100644 --- a/abacusnbody/hod/abacus_hod.py +++ b/abacusnbody/hod/abacus_hod.py @@ -797,7 +797,7 @@ def run_hod( verbose=verbose, fn_ext=fn_ext, ) - + self.logger.info(f'HOD generated in elapsed time {time.time() - start:.2f} s.') return mock_dict From 3601cad3e592801f762c22da87607d3a276cc100 Mon Sep 17 00:00:00 2001 From: Pierre Burger Date: Wed, 9 Oct 2024 16:06:53 -0400 Subject: [PATCH 07/13] Update CSMF_hod.py removed the unused variables --- abacusnbody/hod/CSMF_HOD.py | 778 ++++++++++++++---------------------- 1 file changed, 293 insertions(+), 485 deletions(-) diff --git a/abacusnbody/hod/CSMF_HOD.py b/abacusnbody/hod/CSMF_HOD.py index e577ce4c..a3f4d883 100644 --- a/abacusnbody/hod/CSMF_HOD.py +++ b/abacusnbody/hod/CSMF_HOD.py @@ -19,155 +19,108 @@ G = 4.302e-6 # in kpc/Msol (km.s)^2 + @njit(fastmath=True) def n_cen_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): """ Standard Cacciato et al. (2008) centrals HOD parametrization for CSMF """ M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - x_low = np.log10(Mstar_low / M_c_value) / (1.41421356 * sigma_c) - x_up = np.log10(Mstar_up / M_c_value) / (1.41421356 * sigma_c) - - return 0.5 * (math.erf(x_up) - math.erf(x_low)) - + x_low = np.log10(Mstar_low/M_c_value)/(1.41421356*sigma_c) + x_up = np.log10(Mstar_up/M_c_value)/(1.41421356*sigma_c) + + return 0.5*(math.erf(x_up) - math.erf(x_low)) @njit(fastmath=True) def CSMF_centrals(M_h, Mstar, M_1, M_0, gamma1, gamma2, sigma_c): """ Eq. (34) from Cacciato et al. (2008) """ - + M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - - return ( - 1 - / (1.41421356 * np.sqrt(np.pi) * np.log(10) * sigma_c * Mstar) - * np.exp(-((np.log10(Mstar) - np.log10(M_c_value)) ** 2) / (2 * sigma_c**2)) - ) - + + return 1/(1.41421356*np.sqrt(np.pi)*np.log(10)*sigma_c*Mstar)*np.exp(-(np.log10(Mstar)-np.log10(M_c_value))**2/(2*sigma_c**2)) @njit(fastmath=True) def M_c(M_h, M_1, M_0, gamma1, gamma2): """ Eq. (37) from Cacciato et al. (2008) """ - return M_0 * (M_h / M_1) ** gamma1 / (1 + M_h / M_1) ** (gamma1 - gamma2) + return M_0 * (M_h/M_1)**gamma1/(1+M_h/M_1)**(gamma1-gamma2) + @njit(fastmath=True) -def get_random_cen_stellarmass( - M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c -): +def get_random_cen_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): + nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) - stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 - delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] - - CSMF_cen = CSMF_centrals( - M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c - ) - - cdf = np.cumsum(CSMF_cen * delta_stellar_masses) - cdf = cdf / cdf[-1] - - random_rv = np.random.uniform(cdf.min(), cdf.max()) + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + + CSMF_cen = CSMF_centrals(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c) + + + cdf = np.cumsum(CSMF_cen*delta_stellar_masses) + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) + + return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) + + @njit(fastmath=True) -def get_random_cen_stellarmass_linearinterpolation( - M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c -): +def get_random_cen_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): + nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) - stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 - delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] - + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + # stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + CSMF_cen = CSMF_centrals(M_h, stellarmass, M_1, M_0, gamma1, gamma2, sigma_c) - - cdf = np.cumsum(CSMF_cen[:-1] * delta_stellar_masses) - cdf = cdf / cdf[-1] - - random_rv = np.random.uniform(cdf.min(), cdf.max()) - bin = np.where(cdf > random_rv)[0][0] - - m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) - return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] + + cdf = np.cumsum(CSMF_cen[:-1]*delta_stellar_masses) + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin = np.where(cdf>random_rv)[0][0] + + m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) + return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] @njit(fastmath=True) -def n_sat_CSMF( - M_h, - Mstar_low, - Mstar_up, - M_1, - M_0, - gamma1, - gamma2, - sigma_c, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, -): +def n_sat_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c, a1, a2, M2, b0, b1, b2, delta1, delta2): """ Standard Cacciato et al. (2008) satellite HOD parametrization for CSMF """ nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) - - CSMF_sat = CSMF_satelites( - M_h, - stellarmass, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, - ) + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + + CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) nsat = 0 - for i in range(nbins - 1): - nsat += (CSMF_sat[i + 1] - CSMF_sat[i]) * ( - stellarmass[i + 1] - stellarmass[i] - ) / 2 + (stellarmass[i + 1] - stellarmass[i]) * CSMF_sat[i] - - return nsat # *ncen + for i in range(nbins-1): + nsat += (CSMF_sat[i+1]-CSMF_sat[i])*(stellarmass[i+1]-stellarmass[i])/2 + (stellarmass[i+1]-stellarmass[i])*CSMF_sat[i] + + return nsat#*ncen @njit(fastmath=True) -def CSMF_satelites( - M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2 -): +def CSMF_satelites(M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): """ Eq. (36) from Cacciato et al. (2008) """ M_s_value = M_s(M_h, M_1, M_0, gamma1, gamma2) alpha_s_value = alpha_s(M_h, a1, a2, M2) phi_s_value = phi_s(M_h, b0, b1, b2) - + delta = 10 ** (delta1 + delta2 * (np.log10(M_h) - 12)) - - return ( - phi_s_value - / M_s_value - * (Mstar / M_s_value) ** alpha_s_value - * np.exp(-delta * (Mstar / M_s_value) ** 2) - ) - + + return phi_s_value/M_s_value * (Mstar/M_s_value)**alpha_s_value * np.exp(-delta*(Mstar/M_s_value)**2) @njit(fastmath=True) def M_s(M_h, M_1, M_0, gamma1, gamma2): @@ -176,122 +129,60 @@ def M_s(M_h, M_1, M_0, gamma1, gamma2): """ return 0.562 * M_c(M_h, M_1, M_0, gamma1, gamma2) - @njit(fastmath=True) def alpha_s(M_h, a1, a2, M2): """ Eq. (39) from Cacciato et al. (2008) """ - return -2.0 + a1 * (1 - 2 / np.pi * np.arctan(a2 * np.log10(M_h / M2))) - + return -2.0 + a1 * (1-2/np.pi*np.arctan(a2*np.log10(M_h/M2))) @njit(fastmath=True) def phi_s(M_h, b0, b1, b2): """ Eq. (40) from Cacciato et al. (2008) """ - M12 = M_h / 1e12 - log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12) ** 2 + M12 = M_h/1e12 + log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12)**2 return 10**log_phi_s @njit(fastmath=True) -def get_random_sat_stellarmass( - M_h, - Mstar_low, - Mstar_up, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, -): +def get_random_sat_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): + nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) - stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 - delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] - - CSMF_sat = CSMF_satelites( - M_h, - stellarmass_centers, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, - ) - - cdf = np.cumsum(CSMF_sat * delta_stellar_masses) - cdf = cdf / cdf[-1] - - random_rv = np.random.uniform(cdf.min(), cdf.max()) + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + + CSMF_sat = CSMF_satelites(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + + cdf = np.cumsum(CSMF_sat*delta_stellar_masses) + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) - + + return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) + @njit(fastmath=True) -def get_random_sat_stellarmass_linearinterpolation( - M_h, - Mstar_low, - Mstar_up, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, -): +def get_random_sat_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): + nbins = 100 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) - stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 - delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] - - CSMF_sat = CSMF_satelites( - M_h, - stellarmass, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, - ) - - cdf = np.cumsum(CSMF_sat[:-1] * delta_stellar_masses) - cdf = cdf / cdf[-1] - - random_rv = np.random.uniform(cdf.min(), cdf.max()) - bin = np.where(cdf > random_rv)[0][0] - - m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) - return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + # stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + + CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + + cdf = np.cumsum(CSMF_sat[:-1]*delta_stellar_masses) + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin = np.where(cdf>random_rv)[0][0] + + m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) + return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] @njit(fastmath=True) @@ -335,23 +226,24 @@ def gen_cent( ): """ Generate central galaxies in place in memory with a two pass numba parallel implementation. - """ + """ if want_CSMF: Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'] ) ic_C, alpha_c_C, Ac_C, Bc_C = ( - CSMF_hod_dict['ic'], - CSMF_hod_dict['alpha_c'], - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Bcent'], + CSMF_hod_dict['ic'], + CSMF_hod_dict['alpha_c'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Bcent'] ) + H = len(mass) @@ -369,23 +261,14 @@ def gen_cent( # first create the markers between 0 and 1 for different tracers CSMF_marker = 0 if want_CSMF: - M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) - - ncen = n_cen_CSMF( - mass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - ) - + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + + ncen = n_cen_CSMF(mass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) + CSMF_marker += ncen * ic_C * multis[i] if randoms[i] <= CSMF_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 else: keep[i] = 0 @@ -394,18 +277,19 @@ def gen_cent( gstart = np.empty((Nthread + 1, 1), dtype=np.int64) gstart[0, :] = 0 gstart[1:, 0] = Nout[:, 0, 0].cumsum() + # galaxy arrays N_CSMF = gstart[-1, 0] - CSMF_x = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_y = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_z = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_id = np.empty(N_CSMF, dtype=ids.dtype) + CSMF_x = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_y = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_z = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_id = np.empty(N_CSMF, dtype = ids.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): @@ -413,50 +297,42 @@ def gen_cent( for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: # loop thru three directions to assign galaxy velocities and positions - CSMF_x[j1] = pos[i, 0] - CSMF_vx[j1] = vel[i, 0] + alpha_c_C * vdev[i, 0] # velocity bias - CSMF_y[j1] = pos[i, 1] - CSMF_vy[j1] = vel[i, 1] + alpha_c_C * vdev[i, 1] # velocity bias - CSMF_z[j1] = pos[i, 2] - CSMF_vz[j1] = vel[i, 2] + alpha_c_C * vdev[i, 2] # velocity bias - + CSMF_x[j1] = pos[i,0] + CSMF_vx[j1] = vel[i,0] + alpha_c_C * vdev[i,0] # velocity bias + CSMF_y[j1] = pos[i,1] + CSMF_vy[j1] = vel[i,1] + alpha_c_C * vdev[i,1] # velocity bias + CSMF_z[j1] = pos[i,2] + CSMF_vz[j1] = vel[i,2] + alpha_c_C * vdev[i,2] # velocity bias + # rsd only applies to the z direction if rsd and origin is not None: nx = CSMF_x[j1] - origin[0] ny = CSMF_y[j1] - origin[1] nz = CSMF_z[j1] - origin[2] - inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms * ( - CSMF_vx[j1] * nx + CSMF_vy[j1] * ny + CSMF_vz[j1] * nz - ) - CSMF_x[j1] = CSMF_x[j1] + proj * nx - CSMF_y[j1] = CSMF_y[j1] + proj * ny - CSMF_z[j1] = CSMF_z[j1] + proj * nz + proj = inv_velz2kms*(CSMF_vx[j1]*nx+CSMF_vy[j1]*ny+CSMF_vz[j1]*nz) + CSMF_x[j1] = CSMF_x[j1]+proj*nx + CSMF_y[j1] = CSMF_y[j1]+proj*ny + CSMF_z[j1] = CSMF_z[j1]+proj*nz elif rsd: - CSMF_z[j1] = wrap(pos[i, 2] + CSMF_vz[j1] * inv_velz2kms, lbox) - + CSMF_z[j1] = wrap(pos[i,2] + CSMF_vz[j1] * inv_velz2kms, lbox) + CSMF_mass[j1] = mass[i] - M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) CSMF_stellarmass[j1] = get_random_cen_stellarmass_linearinterpolation( - mass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - ) + mass[i], Mstar_low_C, Mstar_up_C, + M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) CSMF_id[j1] = ids[i] j1 += 1 # assert j == gstart[tid + 1] - CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) - + + CSMF_dict['x'] = CSMF_x CSMF_dict['y'] = CSMF_y CSMF_dict['z'] = CSMF_z @@ -603,45 +479,29 @@ def gen_sats_nfw( """ if want_CSMF: - ( - Mstar_low_C, - Mstar_up_C, - M_1_C, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], - CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'], + Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], + CSMF_hod_dict['M_2'], + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'] ) Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Asat'], CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'], + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'] ) f_sigv_C = CSMF_hod_dict['f_sigv'] @@ -649,40 +509,37 @@ def gen_sats_nfw( # compute nsate for each halo # figuring out the number of particles kept for each thread - num_sats_C = np.zeros(len(hid), dtype=np.int64) - stellarmass_C = np.zeros(len(hid), dtype=np.int64) + num_sats_C = np.zeros(len(hid), dtype = np.int64) + stellarmass_C = np.zeros(len(hid), dtype = np.int64) hstart = np.rint(np.linspace(0, len(hid), Nthread + 1)).astype( np.int64 ) # starting index of each thread for tid in range(Nthread): for i in range(hstart[tid], hstart[tid + 1]): if want_CSMF: - M_1_C_temp = 10 ** ( - np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] - ) + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] base_p_C = ( n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C) * ic_C ) - num_sats_C[i] = np.random.poisson(base_p_C) + num_sats_C[i] = np.random.poisson(base_p_C) # generate rdpos rd_pos_C = getPointsOnSphere(np.sum(num_sats_C), Nthread) @@ -690,36 +547,35 @@ def gen_sats_nfw( # put satellites on NFW h_id_C, x_sat_C, y_sat_C, z_sat_C, vx_sat_C, vy_sat_C, vz_sat_C, M_C = ( compute_fast_NFW( - NFW_draw, - hid, - hpos[:, 0], - hpos[:, 1], - hpos[:, 2], - hvel[:, 0], - hvel[:, 1], + NFW_draw, + hid, + hpos[:, 0], + hpos[:, 1], + hpos[:, 2], + hvel[:, 0], + hvel[:, 1], hvel[:, 2], - hvrms, - hc, - hmass, - hrvir, - rd_pos_C, - num_sats_C, - f_sigv_C, - vel_sat, + hvrms, + hc, + hmass, + hrvir, + rd_pos_C, + num_sats_C, + f_sigv_C, + vel_sat, Nthread, - exp_frac, - exp_scale, - nfw_rescale, ) ) - + + # do rsd if rsd: z_sat_C = (z_sat_C + vz_sat_C * inv_velz2kms) % lbox - CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) + CSMF_dict['x'] = x_sat_C CSMF_dict['y'] = y_sat_C CSMF_dict['z'] = z_sat_C @@ -733,29 +589,12 @@ def gen_sats_nfw( hstart = np.rint(np.linspace(0, num_sats_C.sum(), Nthread + 1)) for tid in numba.prange(Nthread): for i in range(int(hstart[tid]), int(hstart[tid + 1])): - M_1_C_temp = 10 ** ( - np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] - ) + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] stellarmass_C[i] = get_random_sat_stellarmass_linearinterpolation( - M_C[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - a1_C_temp, - s, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) - + M_C[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, + a1_C_temp, s, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) + CSMF_dict['stellarmass'] = stellarmass_C ID_dict['CSMF'] = h_id_C @@ -794,52 +633,38 @@ def gen_sats( Generate satellite galaxies in place in memory with a two pass numba parallel implementation. """ + + if want_CSMF: - ( - Mstar_low_C, - Mstar_up_C, - M_1_C, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], + Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'], + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'] ) - + alpha_s_C, s_C, s_v_C, s_p_C, s_r_C, Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['alpha_s'], + CSMF_hod_dict['alpha_s'], CSMF_hod_dict['s'], - CSMF_hod_dict['s_v'], - CSMF_hod_dict['s_p'], - CSMF_hod_dict['s_r'], + CSMF_hod_dict['s_v'], + CSMF_hod_dict['s_p'], + CSMF_hod_dict['s_r'], CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], - CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'], + CSMF_hod_dict['Asat'], + CSMF_hod_dict['Bcent'], + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'] ) H = len(hmass) # num of particles @@ -858,47 +683,39 @@ def gen_sats( # print(logM1, As, hdeltac[i], Bs, hfenv[i]) CSMF_marker = 0 if want_CSMF: - M_1_C_temp = 10 ** ( - np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] - ) + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] base_p_C = ( n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) - * weights[i] - * ic_C - ) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C + ) + * weights[i] + * ic_C) if enable_ranks: - decorator_C = ( - 1 - + s_C * ranks[i] - + s_v_C * ranksv[i] - + s_p_C * ranksp[i] - + s_r_C * ranksr[i] - ) + decorator_C = 1 + s_C * ranks[i] + s_v_C * ranksv[i] + s_p_C * ranksp[i] + s_r_C * ranksr[i] exp_sat = base_p_C * decorator_C else: exp_sat = base_p_C CSMF_marker += exp_sat + if randoms[i] <= CSMF_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 else: keep[i] = 0 @@ -910,79 +727,73 @@ def gen_sats( # galaxy arrays N_CSMF = gstart[-1, 0] - CSMF_x = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_y = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_z = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_id = np.empty(N_CSMF, dtype=hid.dtype) + CSMF_x = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_y = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_z = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_id = np.empty(N_CSMF, dtype = hid.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): j1 = gstart[tid] for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: + CSMF_x[j1] = ppos[i, 0] - CSMF_vx[j1] = hvel[i, 0] + alpha_s_C * ( - pvel[i, 0] - hvel[i, 0] - ) # velocity bias + CSMF_vx[j1] = hvel[i, 0] + alpha_s_C * (pvel[i, 0] - hvel[i, 0]) # velocity bias CSMF_y[j1] = ppos[i, 1] - CSMF_vy[j1] = hvel[i, 1] + alpha_s_C * ( - pvel[i, 1] - hvel[i, 1] - ) # velocity bias + CSMF_vy[j1] = hvel[i, 1] + alpha_s_C * (pvel[i, 1] - hvel[i, 1]) # velocity bias CSMF_z[j1] = ppos[i, 2] - CSMF_vz[j1] = hvel[i, 2] + alpha_s_C * ( - pvel[i, 2] - hvel[i, 2] - ) # velocity bias + CSMF_vz[j1] = hvel[i, 2] + alpha_s_C * (pvel[i, 2] - hvel[i, 2]) # velocity bias if rsd and origin is not None: nx = CSMF_x[j1] - origin[0] ny = CSMF_y[j1] - origin[1] nz = CSMF_z[j1] - origin[2] - inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms * ( - CSMF_vx[j1] * nx + CSMF_vy[j1] * ny + CSMF_vz[j1] * nz - ) - CSMF_x[j1] = CSMF_x[j1] + proj * nx - CSMF_y[j1] = CSMF_y[j1] + proj * ny - CSMF_z[j1] = CSMF_z[j1] + proj * nz + proj = inv_velz2kms*(CSMF_vx[j1]*nx+CSMF_vy[j1]*ny+CSMF_vz[j1]*nz) + CSMF_x[j1] = CSMF_x[j1]+proj*nx + CSMF_y[j1] = CSMF_y[j1]+proj*ny + CSMF_z[j1] = CSMF_z[j1]+proj*nz elif rsd: - CSMF_z[j1] = wrap(CSMF_z[j1] + CSMF_vz[j1] * inv_velz2kms, lbox) - - M_1_C_temp = 10 ** ( - np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] - ) + CSMF_z[j1] = wrap(CSMF_z[j1] + CSMF_vz[j1] * inv_velz2kms, lbox) + + + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] CSMF_stellarmass[j1] = get_random_sat_stellarmass( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C + ) CSMF_mass[j1] = hmass[i] CSMF_id[j1] = hid[i] j1 += 1 # assert j == gstart[tid + 1] - CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) + CSMF_dict['x'] = CSMF_x CSMF_dict['y'] = CSMF_y CSMF_dict['z'] = CSMF_z @@ -992,7 +803,7 @@ def gen_sats( CSMF_dict['mass'] = CSMF_mass CSMF_dict['stellarmass'] = CSMF_stellarmass ID_dict['CSMF'] = CSMF_id - + return CSMF_dict, ID_dict @@ -1077,12 +888,11 @@ def gen_gals( if tracer == 'CSMF': CSMF_HOD = tracers[tracer] + if 'CSMF' in tracers.keys(): want_CSMF = True - CSMF_hod_dict = nb.typed.Dict.empty( - key_type=nb.types.unicode_type, value_type=nb.types.float64 - ) + CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) for key, value in CSMF_HOD.items(): CSMF_hod_dict[key] = value @@ -1095,9 +905,7 @@ def gen_gals( else: want_CSMF = False - CSMF_hod_dict = nb.typed.Dict.empty( - key_type=nb.types.unicode_type, value_type=nb.types.float64 - ) + CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) start = time.time() From d907717e96d487e11f795cfe1185d84334cf407a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 9 Oct 2024 20:07:09 +0000 Subject: [PATCH 08/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- abacusnbody/hod/CSMF_HOD.py | 771 ++++++++++++++++++++++-------------- 1 file changed, 480 insertions(+), 291 deletions(-) diff --git a/abacusnbody/hod/CSMF_HOD.py b/abacusnbody/hod/CSMF_HOD.py index a3f4d883..277fe40b 100644 --- a/abacusnbody/hod/CSMF_HOD.py +++ b/abacusnbody/hod/CSMF_HOD.py @@ -19,108 +19,155 @@ G = 4.302e-6 # in kpc/Msol (km.s)^2 - @njit(fastmath=True) def n_cen_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): """ Standard Cacciato et al. (2008) centrals HOD parametrization for CSMF """ M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - x_low = np.log10(Mstar_low/M_c_value)/(1.41421356*sigma_c) - x_up = np.log10(Mstar_up/M_c_value)/(1.41421356*sigma_c) - - return 0.5*(math.erf(x_up) - math.erf(x_low)) + x_low = np.log10(Mstar_low / M_c_value) / (1.41421356 * sigma_c) + x_up = np.log10(Mstar_up / M_c_value) / (1.41421356 * sigma_c) + + return 0.5 * (math.erf(x_up) - math.erf(x_low)) + @njit(fastmath=True) def CSMF_centrals(M_h, Mstar, M_1, M_0, gamma1, gamma2, sigma_c): """ Eq. (34) from Cacciato et al. (2008) """ - + M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - - return 1/(1.41421356*np.sqrt(np.pi)*np.log(10)*sigma_c*Mstar)*np.exp(-(np.log10(Mstar)-np.log10(M_c_value))**2/(2*sigma_c**2)) + + return ( + 1 + / (1.41421356 * np.sqrt(np.pi) * np.log(10) * sigma_c * Mstar) + * np.exp(-((np.log10(Mstar) - np.log10(M_c_value)) ** 2) / (2 * sigma_c**2)) + ) + @njit(fastmath=True) def M_c(M_h, M_1, M_0, gamma1, gamma2): """ Eq. (37) from Cacciato et al. (2008) """ - return M_0 * (M_h/M_1)**gamma1/(1+M_h/M_1)**(gamma1-gamma2) - + return M_0 * (M_h / M_1) ** gamma1 / (1 + M_h / M_1) ** (gamma1 - gamma2) @njit(fastmath=True) -def get_random_cen_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): - +def get_random_cen_stellarmass( + M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c +): nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - - CSMF_cen = CSMF_centrals(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c) - - - cdf = np.cumsum(CSMF_cen*delta_stellar_masses) - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + + CSMF_cen = CSMF_centrals( + M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c + ) + + cdf = np.cumsum(CSMF_cen * delta_stellar_masses) + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) - + + return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) - @njit(fastmath=True) -def get_random_cen_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): - +def get_random_cen_stellarmass_linearinterpolation( + M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c +): nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) # stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + CSMF_cen = CSMF_centrals(M_h, stellarmass, M_1, M_0, gamma1, gamma2, sigma_c) - - cdf = np.cumsum(CSMF_cen[:-1]*delta_stellar_masses) - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) - bin = np.where(cdf>random_rv)[0][0] - - m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) - return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + + cdf = np.cumsum(CSMF_cen[:-1] * delta_stellar_masses) + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) + bin = np.where(cdf > random_rv)[0][0] + + m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) + return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] @njit(fastmath=True) -def n_sat_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c, a1, a2, M2, b0, b1, b2, delta1, delta2): +def n_sat_CSMF( + M_h, + Mstar_low, + Mstar_up, + M_1, + M_0, + gamma1, + gamma2, + sigma_c, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, +): """ Standard Cacciato et al. (2008) satellite HOD parametrization for CSMF """ nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - - CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + + CSMF_sat = CSMF_satelites( + M_h, + stellarmass, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, + ) nsat = 0 - for i in range(nbins-1): - nsat += (CSMF_sat[i+1]-CSMF_sat[i])*(stellarmass[i+1]-stellarmass[i])/2 + (stellarmass[i+1]-stellarmass[i])*CSMF_sat[i] - - return nsat#*ncen + for i in range(nbins - 1): + nsat += (CSMF_sat[i + 1] - CSMF_sat[i]) * ( + stellarmass[i + 1] - stellarmass[i] + ) / 2 + (stellarmass[i + 1] - stellarmass[i]) * CSMF_sat[i] + + return nsat # *ncen @njit(fastmath=True) -def CSMF_satelites(M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): +def CSMF_satelites( + M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2 +): """ Eq. (36) from Cacciato et al. (2008) """ M_s_value = M_s(M_h, M_1, M_0, gamma1, gamma2) alpha_s_value = alpha_s(M_h, a1, a2, M2) phi_s_value = phi_s(M_h, b0, b1, b2) - + delta = 10 ** (delta1 + delta2 * (np.log10(M_h) - 12)) - - return phi_s_value/M_s_value * (Mstar/M_s_value)**alpha_s_value * np.exp(-delta*(Mstar/M_s_value)**2) + + return ( + phi_s_value + / M_s_value + * (Mstar / M_s_value) ** alpha_s_value + * np.exp(-delta * (Mstar / M_s_value) ** 2) + ) + @njit(fastmath=True) def M_s(M_h, M_1, M_0, gamma1, gamma2): @@ -129,60 +176,122 @@ def M_s(M_h, M_1, M_0, gamma1, gamma2): """ return 0.562 * M_c(M_h, M_1, M_0, gamma1, gamma2) + @njit(fastmath=True) def alpha_s(M_h, a1, a2, M2): """ Eq. (39) from Cacciato et al. (2008) """ - return -2.0 + a1 * (1-2/np.pi*np.arctan(a2*np.log10(M_h/M2))) + return -2.0 + a1 * (1 - 2 / np.pi * np.arctan(a2 * np.log10(M_h / M2))) + @njit(fastmath=True) def phi_s(M_h, b0, b1, b2): """ Eq. (40) from Cacciato et al. (2008) """ - M12 = M_h/1e12 - log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12)**2 + M12 = M_h / 1e12 + log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12) ** 2 return 10**log_phi_s @njit(fastmath=True) -def get_random_sat_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): - +def get_random_sat_stellarmass( + M_h, + Mstar_low, + Mstar_up, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, +): nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - - CSMF_sat = CSMF_satelites(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) - - cdf = np.cumsum(CSMF_sat*delta_stellar_masses) - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + + CSMF_sat = CSMF_satelites( + M_h, + stellarmass_centers, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, + ) + + cdf = np.cumsum(CSMF_sat * delta_stellar_masses) + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) - + + return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) + @njit(fastmath=True) -def get_random_sat_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): - +def get_random_sat_stellarmass_linearinterpolation( + M_h, + Mstar_low, + Mstar_up, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, +): nbins = 100 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) # stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - - CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) - - cdf = np.cumsum(CSMF_sat[:-1]*delta_stellar_masses) - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) - bin = np.where(cdf>random_rv)[0][0] - - m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) - return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + + CSMF_sat = CSMF_satelites( + M_h, + stellarmass, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, + ) + + cdf = np.cumsum(CSMF_sat[:-1] * delta_stellar_masses) + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) + bin = np.where(cdf > random_rv)[0][0] + + m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) + return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] @njit(fastmath=True) @@ -226,24 +335,23 @@ def gen_cent( ): """ Generate central galaxies in place in memory with a two pass numba parallel implementation. - """ + """ if want_CSMF: Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'] + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], ) ic_C, alpha_c_C, Ac_C, Bc_C = ( - CSMF_hod_dict['ic'], - CSMF_hod_dict['alpha_c'], - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Bcent'] + CSMF_hod_dict['ic'], + CSMF_hod_dict['alpha_c'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Bcent'], ) - H = len(mass) @@ -261,14 +369,23 @@ def gen_cent( # first create the markers between 0 and 1 for different tracers CSMF_marker = 0 if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) - - ncen = n_cen_CSMF(mass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) - + M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + + ncen = n_cen_CSMF( + mass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + ) + CSMF_marker += ncen * ic_C * multis[i] if randoms[i] <= CSMF_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 else: keep[i] = 0 @@ -277,19 +394,18 @@ def gen_cent( gstart = np.empty((Nthread + 1, 1), dtype=np.int64) gstart[0, :] = 0 gstart[1:, 0] = Nout[:, 0, 0].cumsum() - # galaxy arrays N_CSMF = gstart[-1, 0] - CSMF_x = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_y = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_z = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_id = np.empty(N_CSMF, dtype = ids.dtype) + CSMF_x = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_y = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_z = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_id = np.empty(N_CSMF, dtype=ids.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): @@ -297,42 +413,50 @@ def gen_cent( for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: # loop thru three directions to assign galaxy velocities and positions - CSMF_x[j1] = pos[i,0] - CSMF_vx[j1] = vel[i,0] + alpha_c_C * vdev[i,0] # velocity bias - CSMF_y[j1] = pos[i,1] - CSMF_vy[j1] = vel[i,1] + alpha_c_C * vdev[i,1] # velocity bias - CSMF_z[j1] = pos[i,2] - CSMF_vz[j1] = vel[i,2] + alpha_c_C * vdev[i,2] # velocity bias - + CSMF_x[j1] = pos[i, 0] + CSMF_vx[j1] = vel[i, 0] + alpha_c_C * vdev[i, 0] # velocity bias + CSMF_y[j1] = pos[i, 1] + CSMF_vy[j1] = vel[i, 1] + alpha_c_C * vdev[i, 1] # velocity bias + CSMF_z[j1] = pos[i, 2] + CSMF_vz[j1] = vel[i, 2] + alpha_c_C * vdev[i, 2] # velocity bias + # rsd only applies to the z direction if rsd and origin is not None: nx = CSMF_x[j1] - origin[0] ny = CSMF_y[j1] - origin[1] nz = CSMF_z[j1] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms*(CSMF_vx[j1]*nx+CSMF_vy[j1]*ny+CSMF_vz[j1]*nz) - CSMF_x[j1] = CSMF_x[j1]+proj*nx - CSMF_y[j1] = CSMF_y[j1]+proj*ny - CSMF_z[j1] = CSMF_z[j1]+proj*nz + proj = inv_velz2kms * ( + CSMF_vx[j1] * nx + CSMF_vy[j1] * ny + CSMF_vz[j1] * nz + ) + CSMF_x[j1] = CSMF_x[j1] + proj * nx + CSMF_y[j1] = CSMF_y[j1] + proj * ny + CSMF_z[j1] = CSMF_z[j1] + proj * nz elif rsd: - CSMF_z[j1] = wrap(pos[i,2] + CSMF_vz[j1] * inv_velz2kms, lbox) - + CSMF_z[j1] = wrap(pos[i, 2] + CSMF_vz[j1] * inv_velz2kms, lbox) + CSMF_mass[j1] = mass[i] - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) CSMF_stellarmass[j1] = get_random_cen_stellarmass_linearinterpolation( - mass[i], Mstar_low_C, Mstar_up_C, - M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) + mass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + ) CSMF_id[j1] = ids[i] j1 += 1 # assert j == gstart[tid + 1] - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) - - + CSMF_dict['x'] = CSMF_x CSMF_dict['y'] = CSMF_y CSMF_dict['z'] = CSMF_z @@ -479,29 +603,45 @@ def gen_sats_nfw( """ if want_CSMF: - Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], - CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'] + ( + Mstar_low_C, + Mstar_up_C, + M_1_C, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], + CSMF_hod_dict['M_2'], + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'], ) Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Asat'], CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'] + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'], ) f_sigv_C = CSMF_hod_dict['f_sigv'] @@ -509,37 +649,40 @@ def gen_sats_nfw( # compute nsate for each halo # figuring out the number of particles kept for each thread - num_sats_C = np.zeros(len(hid), dtype = np.int64) - stellarmass_C = np.zeros(len(hid), dtype = np.int64) + num_sats_C = np.zeros(len(hid), dtype=np.int64) + stellarmass_C = np.zeros(len(hid), dtype=np.int64) hstart = np.rint(np.linspace(0, len(hid), Nthread + 1)).astype( np.int64 ) # starting index of each thread for tid in range(Nthread): for i in range(hstart[tid], hstart[tid + 1]): if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] base_p_C = ( n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) * ic_C ) - num_sats_C[i] = np.random.poisson(base_p_C) + num_sats_C[i] = np.random.poisson(base_p_C) # generate rdpos rd_pos_C = getPointsOnSphere(np.sum(num_sats_C), Nthread) @@ -547,35 +690,33 @@ def gen_sats_nfw( # put satellites on NFW h_id_C, x_sat_C, y_sat_C, z_sat_C, vx_sat_C, vy_sat_C, vz_sat_C, M_C = ( compute_fast_NFW( - NFW_draw, - hid, - hpos[:, 0], - hpos[:, 1], - hpos[:, 2], - hvel[:, 0], - hvel[:, 1], + NFW_draw, + hid, + hpos[:, 0], + hpos[:, 1], + hpos[:, 2], + hvel[:, 0], + hvel[:, 1], hvel[:, 2], - hvrms, - hc, - hmass, - hrvir, - rd_pos_C, - num_sats_C, - f_sigv_C, - vel_sat, + hvrms, + hc, + hmass, + hrvir, + rd_pos_C, + num_sats_C, + f_sigv_C, + vel_sat, Nthread, ) ) - - + # do rsd if rsd: z_sat_C = (z_sat_C + vz_sat_C * inv_velz2kms) % lbox - - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) - + CSMF_dict['x'] = x_sat_C CSMF_dict['y'] = y_sat_C CSMF_dict['z'] = z_sat_C @@ -589,12 +730,29 @@ def gen_sats_nfw( hstart = np.rint(np.linspace(0, num_sats_C.sum(), Nthread + 1)) for tid in numba.prange(Nthread): for i in range(int(hstart[tid]), int(hstart[tid + 1])): - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] stellarmass_C[i] = get_random_sat_stellarmass_linearinterpolation( - M_C[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, - a1_C_temp, s, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) - + M_C[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + a1_C_temp, + s, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) + CSMF_dict['stellarmass'] = stellarmass_C ID_dict['CSMF'] = h_id_C @@ -633,38 +791,52 @@ def gen_sats( Generate satellite galaxies in place in memory with a two pass numba parallel implementation. """ - - if want_CSMF: - Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], + ( + Mstar_low_C, + Mstar_up_C, + M_1_C, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'] + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'], ) - + alpha_s_C, s_C, s_v_C, s_p_C, s_r_C, Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['alpha_s'], + CSMF_hod_dict['alpha_s'], CSMF_hod_dict['s'], - CSMF_hod_dict['s_v'], - CSMF_hod_dict['s_p'], - CSMF_hod_dict['s_r'], + CSMF_hod_dict['s_v'], + CSMF_hod_dict['s_p'], + CSMF_hod_dict['s_r'], CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], - CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'] + CSMF_hod_dict['Asat'], + CSMF_hod_dict['Bcent'], + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'], ) H = len(hmass) # num of particles @@ -683,39 +855,47 @@ def gen_sats( # print(logM1, As, hdeltac[i], Bs, hfenv[i]) CSMF_marker = 0 if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] base_p_C = ( n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C - ) - * weights[i] - * ic_C) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) + * weights[i] + * ic_C + ) if enable_ranks: - decorator_C = 1 + s_C * ranks[i] + s_v_C * ranksv[i] + s_p_C * ranksp[i] + s_r_C * ranksr[i] + decorator_C = ( + 1 + + s_C * ranks[i] + + s_v_C * ranksv[i] + + s_p_C * ranksp[i] + + s_r_C * ranksr[i] + ) exp_sat = base_p_C * decorator_C else: exp_sat = base_p_C CSMF_marker += exp_sat - if randoms[i] <= CSMF_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 else: keep[i] = 0 @@ -727,73 +907,79 @@ def gen_sats( # galaxy arrays N_CSMF = gstart[-1, 0] - CSMF_x = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_y = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_z = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_id = np.empty(N_CSMF, dtype = hid.dtype) + CSMF_x = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_y = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_z = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_id = np.empty(N_CSMF, dtype=hid.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): j1 = gstart[tid] for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: - CSMF_x[j1] = ppos[i, 0] - CSMF_vx[j1] = hvel[i, 0] + alpha_s_C * (pvel[i, 0] - hvel[i, 0]) # velocity bias + CSMF_vx[j1] = hvel[i, 0] + alpha_s_C * ( + pvel[i, 0] - hvel[i, 0] + ) # velocity bias CSMF_y[j1] = ppos[i, 1] - CSMF_vy[j1] = hvel[i, 1] + alpha_s_C * (pvel[i, 1] - hvel[i, 1]) # velocity bias + CSMF_vy[j1] = hvel[i, 1] + alpha_s_C * ( + pvel[i, 1] - hvel[i, 1] + ) # velocity bias CSMF_z[j1] = ppos[i, 2] - CSMF_vz[j1] = hvel[i, 2] + alpha_s_C * (pvel[i, 2] - hvel[i, 2]) # velocity bias + CSMF_vz[j1] = hvel[i, 2] + alpha_s_C * ( + pvel[i, 2] - hvel[i, 2] + ) # velocity bias if rsd and origin is not None: nx = CSMF_x[j1] - origin[0] ny = CSMF_y[j1] - origin[1] nz = CSMF_z[j1] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms*(CSMF_vx[j1]*nx+CSMF_vy[j1]*ny+CSMF_vz[j1]*nz) - CSMF_x[j1] = CSMF_x[j1]+proj*nx - CSMF_y[j1] = CSMF_y[j1]+proj*ny - CSMF_z[j1] = CSMF_z[j1]+proj*nz + proj = inv_velz2kms * ( + CSMF_vx[j1] * nx + CSMF_vy[j1] * ny + CSMF_vz[j1] * nz + ) + CSMF_x[j1] = CSMF_x[j1] + proj * nx + CSMF_y[j1] = CSMF_y[j1] + proj * ny + CSMF_z[j1] = CSMF_z[j1] + proj * nz elif rsd: - CSMF_z[j1] = wrap(CSMF_z[j1] + CSMF_vz[j1] * inv_velz2kms, lbox) - - - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + CSMF_z[j1] = wrap(CSMF_z[j1] + CSMF_vz[j1] * inv_velz2kms, lbox) + + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] CSMF_stellarmass[j1] = get_random_sat_stellarmass( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C - ) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) CSMF_mass[j1] = hmass[i] CSMF_id[j1] = hid[i] j1 += 1 # assert j == gstart[tid + 1] - - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) - CSMF_dict['x'] = CSMF_x CSMF_dict['y'] = CSMF_y CSMF_dict['z'] = CSMF_z @@ -803,7 +989,7 @@ def gen_sats( CSMF_dict['mass'] = CSMF_mass CSMF_dict['stellarmass'] = CSMF_stellarmass ID_dict['CSMF'] = CSMF_id - + return CSMF_dict, ID_dict @@ -888,11 +1074,12 @@ def gen_gals( if tracer == 'CSMF': CSMF_HOD = tracers[tracer] - if 'CSMF' in tracers.keys(): want_CSMF = True - CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + CSMF_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) for key, value in CSMF_HOD.items(): CSMF_hod_dict[key] = value @@ -905,7 +1092,9 @@ def gen_gals( else: want_CSMF = False - CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + CSMF_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) start = time.time() From 90838e20b307b860fabffd11b425a7abb745cfee Mon Sep 17 00:00:00 2001 From: Pierre Burger Date: Wed, 9 Oct 2024 17:11:18 -0400 Subject: [PATCH 09/13] Corrected undefined and unused variables Corrected undefined and unused variables --- abacusnbody/hod/CSMF_HOD.py | 770 +++++++++++++--------------------- abacusnbody/hod/abacus_hod.py | 6 +- 2 files changed, 293 insertions(+), 483 deletions(-) diff --git a/abacusnbody/hod/CSMF_HOD.py b/abacusnbody/hod/CSMF_HOD.py index 277fe40b..f81895b9 100644 --- a/abacusnbody/hod/CSMF_HOD.py +++ b/abacusnbody/hod/CSMF_HOD.py @@ -19,155 +19,108 @@ G = 4.302e-6 # in kpc/Msol (km.s)^2 + @njit(fastmath=True) def n_cen_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): """ Standard Cacciato et al. (2008) centrals HOD parametrization for CSMF """ M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - x_low = np.log10(Mstar_low / M_c_value) / (1.41421356 * sigma_c) - x_up = np.log10(Mstar_up / M_c_value) / (1.41421356 * sigma_c) - - return 0.5 * (math.erf(x_up) - math.erf(x_low)) - + x_low = np.log10(Mstar_low/M_c_value)/(1.41421356*sigma_c) + x_up = np.log10(Mstar_up/M_c_value)/(1.41421356*sigma_c) + + return 0.5*(math.erf(x_up) - math.erf(x_low)) @njit(fastmath=True) def CSMF_centrals(M_h, Mstar, M_1, M_0, gamma1, gamma2, sigma_c): """ Eq. (34) from Cacciato et al. (2008) """ - + M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - - return ( - 1 - / (1.41421356 * np.sqrt(np.pi) * np.log(10) * sigma_c * Mstar) - * np.exp(-((np.log10(Mstar) - np.log10(M_c_value)) ** 2) / (2 * sigma_c**2)) - ) - + + return 1/(1.41421356*np.sqrt(np.pi)*np.log(10)*sigma_c*Mstar)*np.exp(-(np.log10(Mstar)-np.log10(M_c_value))**2/(2*sigma_c**2)) @njit(fastmath=True) def M_c(M_h, M_1, M_0, gamma1, gamma2): """ Eq. (37) from Cacciato et al. (2008) """ - return M_0 * (M_h / M_1) ** gamma1 / (1 + M_h / M_1) ** (gamma1 - gamma2) + return M_0 * (M_h/M_1)**gamma1/(1+M_h/M_1)**(gamma1-gamma2) + @njit(fastmath=True) -def get_random_cen_stellarmass( - M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c -): +def get_random_cen_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): + nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) - stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 - delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] - - CSMF_cen = CSMF_centrals( - M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c - ) - - cdf = np.cumsum(CSMF_cen * delta_stellar_masses) - cdf = cdf / cdf[-1] - - random_rv = np.random.uniform(cdf.min(), cdf.max()) + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + + CSMF_cen = CSMF_centrals(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c) + + + cdf = np.cumsum(CSMF_cen*delta_stellar_masses) + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) + + return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) + + @njit(fastmath=True) -def get_random_cen_stellarmass_linearinterpolation( - M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c -): +def get_random_cen_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): + nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) # stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] - + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + CSMF_cen = CSMF_centrals(M_h, stellarmass, M_1, M_0, gamma1, gamma2, sigma_c) - - cdf = np.cumsum(CSMF_cen[:-1] * delta_stellar_masses) - cdf = cdf / cdf[-1] - - random_rv = np.random.uniform(cdf.min(), cdf.max()) - bin = np.where(cdf > random_rv)[0][0] - - m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) - return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] + + cdf = np.cumsum(CSMF_cen[:-1]*delta_stellar_masses) + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin = np.where(cdf>random_rv)[0][0] + + m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) + return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] @njit(fastmath=True) -def n_sat_CSMF( - M_h, - Mstar_low, - Mstar_up, - M_1, - M_0, - gamma1, - gamma2, - sigma_c, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, -): +def n_sat_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c, a1, a2, M2, b0, b1, b2, delta1, delta2): """ Standard Cacciato et al. (2008) satellite HOD parametrization for CSMF """ nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) - - CSMF_sat = CSMF_satelites( - M_h, - stellarmass, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, - ) + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + + CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) nsat = 0 - for i in range(nbins - 1): - nsat += (CSMF_sat[i + 1] - CSMF_sat[i]) * ( - stellarmass[i + 1] - stellarmass[i] - ) / 2 + (stellarmass[i + 1] - stellarmass[i]) * CSMF_sat[i] - - return nsat # *ncen + for i in range(nbins-1): + nsat += (CSMF_sat[i+1]-CSMF_sat[i])*(stellarmass[i+1]-stellarmass[i])/2 + (stellarmass[i+1]-stellarmass[i])*CSMF_sat[i] + + return nsat#*ncen @njit(fastmath=True) -def CSMF_satelites( - M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2 -): +def CSMF_satelites(M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): """ Eq. (36) from Cacciato et al. (2008) """ M_s_value = M_s(M_h, M_1, M_0, gamma1, gamma2) alpha_s_value = alpha_s(M_h, a1, a2, M2) phi_s_value = phi_s(M_h, b0, b1, b2) - + delta = 10 ** (delta1 + delta2 * (np.log10(M_h) - 12)) - - return ( - phi_s_value - / M_s_value - * (Mstar / M_s_value) ** alpha_s_value - * np.exp(-delta * (Mstar / M_s_value) ** 2) - ) - + + return phi_s_value/M_s_value * (Mstar/M_s_value)**alpha_s_value * np.exp(-delta*(Mstar/M_s_value)**2) @njit(fastmath=True) def M_s(M_h, M_1, M_0, gamma1, gamma2): @@ -176,122 +129,60 @@ def M_s(M_h, M_1, M_0, gamma1, gamma2): """ return 0.562 * M_c(M_h, M_1, M_0, gamma1, gamma2) - @njit(fastmath=True) def alpha_s(M_h, a1, a2, M2): """ Eq. (39) from Cacciato et al. (2008) """ - return -2.0 + a1 * (1 - 2 / np.pi * np.arctan(a2 * np.log10(M_h / M2))) - + return -2.0 + a1 * (1-2/np.pi*np.arctan(a2*np.log10(M_h/M2))) @njit(fastmath=True) def phi_s(M_h, b0, b1, b2): """ Eq. (40) from Cacciato et al. (2008) """ - M12 = M_h / 1e12 - log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12) ** 2 + M12 = M_h/1e12 + log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12)**2 return 10**log_phi_s @njit(fastmath=True) -def get_random_sat_stellarmass( - M_h, - Mstar_low, - Mstar_up, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, -): +def get_random_sat_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): + nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) - stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 - delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] - - CSMF_sat = CSMF_satelites( - M_h, - stellarmass_centers, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, - ) - - cdf = np.cumsum(CSMF_sat * delta_stellar_masses) - cdf = cdf / cdf[-1] - - random_rv = np.random.uniform(cdf.min(), cdf.max()) + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + + CSMF_sat = CSMF_satelites(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + + cdf = np.cumsum(CSMF_sat*delta_stellar_masses) + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) - + + return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) + @njit(fastmath=True) -def get_random_sat_stellarmass_linearinterpolation( - M_h, - Mstar_low, - Mstar_up, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, -): +def get_random_sat_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): + nbins = 100 - stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) # stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] - - CSMF_sat = CSMF_satelites( - M_h, - stellarmass, - M_1, - M_0, - gamma1, - gamma2, - a1, - a2, - M2, - b0, - b1, - b2, - delta1, - delta2, - ) - - cdf = np.cumsum(CSMF_sat[:-1] * delta_stellar_masses) - cdf = cdf / cdf[-1] - - random_rv = np.random.uniform(cdf.min(), cdf.max()) - bin = np.where(cdf > random_rv)[0][0] - - m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) - return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] + delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] + + CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + + cdf = np.cumsum(CSMF_sat[:-1]*delta_stellar_masses) + cdf=cdf/cdf[-1] + + random_rv = np.random.uniform(cdf.min(),cdf.max()) + bin = np.where(cdf>random_rv)[0][0] + + m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) + return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] @njit(fastmath=True) @@ -335,23 +226,24 @@ def gen_cent( ): """ Generate central galaxies in place in memory with a two pass numba parallel implementation. - """ + """ if want_CSMF: Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'] ) ic_C, alpha_c_C, Ac_C, Bc_C = ( - CSMF_hod_dict['ic'], - CSMF_hod_dict['alpha_c'], - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Bcent'], + CSMF_hod_dict['ic'], + CSMF_hod_dict['alpha_c'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Bcent'] ) + H = len(mass) @@ -369,23 +261,14 @@ def gen_cent( # first create the markers between 0 and 1 for different tracers CSMF_marker = 0 if want_CSMF: - M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) - - ncen = n_cen_CSMF( - mass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - ) - + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + + ncen = n_cen_CSMF(mass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) + CSMF_marker += ncen * ic_C * multis[i] if randoms[i] <= CSMF_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 else: keep[i] = 0 @@ -394,18 +277,19 @@ def gen_cent( gstart = np.empty((Nthread + 1, 1), dtype=np.int64) gstart[0, :] = 0 gstart[1:, 0] = Nout[:, 0, 0].cumsum() + # galaxy arrays N_CSMF = gstart[-1, 0] - CSMF_x = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_y = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_z = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype=mass.dtype) - CSMF_id = np.empty(N_CSMF, dtype=ids.dtype) + CSMF_x = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_y = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_z = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype = mass.dtype) + CSMF_id = np.empty(N_CSMF, dtype = ids.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): @@ -413,50 +297,42 @@ def gen_cent( for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: # loop thru three directions to assign galaxy velocities and positions - CSMF_x[j1] = pos[i, 0] - CSMF_vx[j1] = vel[i, 0] + alpha_c_C * vdev[i, 0] # velocity bias - CSMF_y[j1] = pos[i, 1] - CSMF_vy[j1] = vel[i, 1] + alpha_c_C * vdev[i, 1] # velocity bias - CSMF_z[j1] = pos[i, 2] - CSMF_vz[j1] = vel[i, 2] + alpha_c_C * vdev[i, 2] # velocity bias - + CSMF_x[j1] = pos[i,0] + CSMF_vx[j1] = vel[i,0] + alpha_c_C * vdev[i,0] # velocity bias + CSMF_y[j1] = pos[i,1] + CSMF_vy[j1] = vel[i,1] + alpha_c_C * vdev[i,1] # velocity bias + CSMF_z[j1] = pos[i,2] + CSMF_vz[j1] = vel[i,2] + alpha_c_C * vdev[i,2] # velocity bias + # rsd only applies to the z direction if rsd and origin is not None: nx = CSMF_x[j1] - origin[0] ny = CSMF_y[j1] - origin[1] nz = CSMF_z[j1] - origin[2] - inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms * ( - CSMF_vx[j1] * nx + CSMF_vy[j1] * ny + CSMF_vz[j1] * nz - ) - CSMF_x[j1] = CSMF_x[j1] + proj * nx - CSMF_y[j1] = CSMF_y[j1] + proj * ny - CSMF_z[j1] = CSMF_z[j1] + proj * nz + proj = inv_velz2kms*(CSMF_vx[j1]*nx+CSMF_vy[j1]*ny+CSMF_vz[j1]*nz) + CSMF_x[j1] = CSMF_x[j1]+proj*nx + CSMF_y[j1] = CSMF_y[j1]+proj*ny + CSMF_z[j1] = CSMF_z[j1]+proj*nz elif rsd: - CSMF_z[j1] = wrap(pos[i, 2] + CSMF_vz[j1] * inv_velz2kms, lbox) - + CSMF_z[j1] = wrap(pos[i,2] + CSMF_vz[j1] * inv_velz2kms, lbox) + CSMF_mass[j1] = mass[i] - M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) CSMF_stellarmass[j1] = get_random_cen_stellarmass_linearinterpolation( - mass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - ) + mass[i], Mstar_low_C, Mstar_up_C, + M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) CSMF_id[j1] = ids[i] j1 += 1 # assert j == gstart[tid + 1] - CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) - + + CSMF_dict['x'] = CSMF_x CSMF_dict['y'] = CSMF_y CSMF_dict['z'] = CSMF_z @@ -603,45 +479,29 @@ def gen_sats_nfw( """ if want_CSMF: - ( - Mstar_low_C, - Mstar_up_C, - M_1_C, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], - CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'], + Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], + CSMF_hod_dict['M_2'], + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'] ) Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Asat'], CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'], + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'] ) f_sigv_C = CSMF_hod_dict['f_sigv'] @@ -649,40 +509,37 @@ def gen_sats_nfw( # compute nsate for each halo # figuring out the number of particles kept for each thread - num_sats_C = np.zeros(len(hid), dtype=np.int64) - stellarmass_C = np.zeros(len(hid), dtype=np.int64) + num_sats_C = np.zeros(len(hid), dtype = np.int64) + stellarmass_C = np.zeros(len(hid), dtype = np.int64) hstart = np.rint(np.linspace(0, len(hid), Nthread + 1)).astype( np.int64 ) # starting index of each thread for tid in range(Nthread): for i in range(hstart[tid], hstart[tid + 1]): if want_CSMF: - M_1_C_temp = 10 ** ( - np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] - ) + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] base_p_C = ( n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C) * ic_C ) - num_sats_C[i] = np.random.poisson(base_p_C) + num_sats_C[i] = np.random.poisson(base_p_C) # generate rdpos rd_pos_C = getPointsOnSphere(np.sum(num_sats_C), Nthread) @@ -690,33 +547,35 @@ def gen_sats_nfw( # put satellites on NFW h_id_C, x_sat_C, y_sat_C, z_sat_C, vx_sat_C, vy_sat_C, vz_sat_C, M_C = ( compute_fast_NFW( - NFW_draw, - hid, - hpos[:, 0], - hpos[:, 1], - hpos[:, 2], - hvel[:, 0], - hvel[:, 1], + NFW_draw, + hid, + hpos[:, 0], + hpos[:, 1], + hpos[:, 2], + hvel[:, 0], + hvel[:, 1], hvel[:, 2], - hvrms, - hc, - hmass, - hrvir, - rd_pos_C, - num_sats_C, - f_sigv_C, - vel_sat, + hvrms, + hc, + hmass, + hrvir, + rd_pos_C, + num_sats_C, + f_sigv_C, + vel_sat, Nthread, ) ) - + + # do rsd if rsd: z_sat_C = (z_sat_C + vz_sat_C * inv_velz2kms) % lbox - CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) - ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) + CSMF_dict['x'] = x_sat_C CSMF_dict['y'] = y_sat_C CSMF_dict['z'] = z_sat_C @@ -730,29 +589,12 @@ def gen_sats_nfw( hstart = np.rint(np.linspace(0, num_sats_C.sum(), Nthread + 1)) for tid in numba.prange(Nthread): for i in range(int(hstart[tid]), int(hstart[tid + 1])): - M_1_C_temp = 10 ** ( - np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] - ) + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] stellarmass_C[i] = get_random_sat_stellarmass_linearinterpolation( - M_C[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - a1_C_temp, - s, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) - + M_C[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, + a1_C_temp, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) + CSMF_dict['stellarmass'] = stellarmass_C ID_dict['CSMF'] = h_id_C @@ -791,52 +633,38 @@ def gen_sats( Generate satellite galaxies in place in memory with a two pass numba parallel implementation. """ + + if want_CSMF: - ( - Mstar_low_C, - Mstar_up_C, - M_1_C, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], + Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'], + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'] ) - + alpha_s_C, s_C, s_v_C, s_p_C, s_r_C, Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['alpha_s'], + CSMF_hod_dict['alpha_s'], CSMF_hod_dict['s'], - CSMF_hod_dict['s_v'], - CSMF_hod_dict['s_p'], - CSMF_hod_dict['s_r'], + CSMF_hod_dict['s_v'], + CSMF_hod_dict['s_p'], + CSMF_hod_dict['s_r'], CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], - CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'], + CSMF_hod_dict['Asat'], + CSMF_hod_dict['Bcent'], + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'] ) H = len(hmass) # num of particles @@ -855,47 +683,39 @@ def gen_sats( # print(logM1, As, hdeltac[i], Bs, hfenv[i]) CSMF_marker = 0 if want_CSMF: - M_1_C_temp = 10 ** ( - np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] - ) + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] base_p_C = ( n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) - * weights[i] - * ic_C - ) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C + ) + * weights[i] + * ic_C) if enable_ranks: - decorator_C = ( - 1 - + s_C * ranks[i] - + s_v_C * ranksv[i] - + s_p_C * ranksp[i] - + s_r_C * ranksr[i] - ) + decorator_C = 1 + s_C * ranks[i] + s_v_C * ranksv[i] + s_p_C * ranksp[i] + s_r_C * ranksr[i] exp_sat = base_p_C * decorator_C else: exp_sat = base_p_C CSMF_marker += exp_sat + if randoms[i] <= CSMF_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 else: keep[i] = 0 @@ -907,15 +727,15 @@ def gen_sats( # galaxy arrays N_CSMF = gstart[-1, 0] - CSMF_x = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_y = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_z = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype=hmass.dtype) - CSMF_id = np.empty(N_CSMF, dtype=hid.dtype) + CSMF_x = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_y = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_z = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype = hmass.dtype) + CSMF_id = np.empty(N_CSMF, dtype = hid.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): @@ -923,63 +743,56 @@ def gen_sats( for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: CSMF_x[j1] = ppos[i, 0] - CSMF_vx[j1] = hvel[i, 0] + alpha_s_C * ( - pvel[i, 0] - hvel[i, 0] - ) # velocity bias + CSMF_vx[j1] = hvel[i, 0] + alpha_s_C * (pvel[i, 0] - hvel[i, 0]) # velocity bias CSMF_y[j1] = ppos[i, 1] - CSMF_vy[j1] = hvel[i, 1] + alpha_s_C * ( - pvel[i, 1] - hvel[i, 1] - ) # velocity bias + CSMF_vy[j1] = hvel[i, 1] + alpha_s_C * (pvel[i, 1] - hvel[i, 1]) # velocity bias CSMF_z[j1] = ppos[i, 2] - CSMF_vz[j1] = hvel[i, 2] + alpha_s_C * ( - pvel[i, 2] - hvel[i, 2] - ) # velocity bias + CSMF_vz[j1] = hvel[i, 2] + alpha_s_C * (pvel[i, 2] - hvel[i, 2]) # velocity bias if rsd and origin is not None: nx = CSMF_x[j1] - origin[0] ny = CSMF_y[j1] - origin[1] nz = CSMF_z[j1] - origin[2] - inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) + inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms * ( - CSMF_vx[j1] * nx + CSMF_vy[j1] * ny + CSMF_vz[j1] * nz - ) - CSMF_x[j1] = CSMF_x[j1] + proj * nx - CSMF_y[j1] = CSMF_y[j1] + proj * ny - CSMF_z[j1] = CSMF_z[j1] + proj * nz + proj = inv_velz2kms*(CSMF_vx[j1]*nx+CSMF_vy[j1]*ny+CSMF_vz[j1]*nz) + CSMF_x[j1] = CSMF_x[j1]+proj*nx + CSMF_y[j1] = CSMF_y[j1]+proj*ny + CSMF_z[j1] = CSMF_z[j1]+proj*nz elif rsd: - CSMF_z[j1] = wrap(CSMF_z[j1] + CSMF_vz[j1] * inv_velz2kms, lbox) - - M_1_C_temp = 10 ** ( - np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] - ) + CSMF_z[j1] = wrap(CSMF_z[j1] + CSMF_vz[j1] * inv_velz2kms, lbox) + + + M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] CSMF_stellarmass[j1] = get_random_sat_stellarmass( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C, - ) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C + ) CSMF_mass[j1] = hmass[i] CSMF_id[j1] = hid[i] j1 += 1 # assert j == gstart[tid + 1] - CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) + + CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) + CSMF_dict['x'] = CSMF_x CSMF_dict['y'] = CSMF_y CSMF_dict['z'] = CSMF_z @@ -989,7 +802,7 @@ def gen_sats( CSMF_dict['mass'] = CSMF_mass CSMF_dict['stellarmass'] = CSMF_stellarmass ID_dict['CSMF'] = CSMF_id - + return CSMF_dict, ID_dict @@ -1074,12 +887,11 @@ def gen_gals( if tracer == 'CSMF': CSMF_HOD = tracers[tracer] + if 'CSMF' in tracers.keys(): want_CSMF = True - CSMF_hod_dict = nb.typed.Dict.empty( - key_type=nb.types.unicode_type, value_type=nb.types.float64 - ) + CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) for key, value in CSMF_HOD.items(): CSMF_hod_dict[key] = value @@ -1092,9 +904,7 @@ def gen_gals( else: want_CSMF = False - CSMF_hod_dict = nb.typed.Dict.empty( - key_type=nb.types.unicode_type, value_type=nb.types.float64 - ) + CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) start = time.time() diff --git a/abacusnbody/hod/abacus_hod.py b/abacusnbody/hod/abacus_hod.py index 4de43ebf..7f64b14a 100644 --- a/abacusnbody/hod/abacus_hod.py +++ b/abacusnbody/hod/abacus_hod.py @@ -797,7 +797,7 @@ def run_hod( verbose=verbose, fn_ext=fn_ext, ) - + self.logger.info(f'HOD generated in elapsed time {time.time() - start:.2f} s.') return mock_dict @@ -1188,8 +1188,8 @@ def _compute_ngal_CSMF( numba.set_num_threads(Nthread) logMs = 0.5 * (logMbins[1:] + logMbins[:-1]) - deltacs = 0.5 * (deltacbins[1:] + deltacbins[:-1]) - fenvs = 0.5 * (fenvbins[1:] + fenvbins[:-1]) + # deltacs = 0.5 * (deltacbins[1:] + deltacbins[:-1]) + # fenvs = 0.5 * (fenvbins[1:] + fenvbins[:-1]) ngal_cent = 0 ngal_sat = 0 From 5468ef825539bc39ee05a255391fcd252ffb2bfd Mon Sep 17 00:00:00 2001 From: Pierre Burger Date: Wed, 9 Oct 2024 17:11:35 -0400 Subject: [PATCH 10/13] ds ds --- .DS_Store | Bin 6148 -> 6148 bytes .gitignore | 1 + abacusnbody/.DS_Store | Bin 6148 -> 8196 bytes abacusnbody/hod/.DS_Store | Bin 0 -> 6148 bytes 4 files changed, 1 insertion(+) create mode 100644 abacusnbody/hod/.DS_Store diff --git a/.DS_Store b/.DS_Store index 5e75159d3da651e8cdbbeda424a7a9c34fd09479..8b526270aca7bdc631fc4806cbdafe5d48ec3890 100644 GIT binary patch delta 19 acmZoMXffDukA=-rN5Rm@aPt!u1z`X{xCQ6{ delta 19 acmZoMXffDukA=-bN5Rm@VDl3e1z`X{rv>H! diff --git a/.gitignore b/.gitignore index 38af2255..0acbd9ee 100644 --- a/.gitignore +++ b/.gitignore @@ -14,3 +14,4 @@ issues/ abacusnbody/metadata/abacussummit_headers.asdf *-jvsc-*.ipynb *~ +abacusnbody/hod/CSMF_HOD.py diff --git a/abacusnbody/.DS_Store b/abacusnbody/.DS_Store index d15d1b7b523a7922dd16e014a18f6dfe0b569e05..299dab670e7e0c4a83548d28f538e273fed55c07 100644 GIT binary patch literal 8196 zcmeHMU2GIp6h3EKU}j3`Ewr@V5wg+*kS=vwp_YQW{R0XjKifYFxX$hjbi#C|?96Vd zgjhAc2pUZ^K8ZikM={1|q7i}#55xx_OeLU+51RO*55}m8i3y%Ncec<49(+(poZH-c z?z!ij`*ZGh=I*@(0PHCkZ2*k`AXDa)RZ($^#N|BSl%%kQ3z6&rEbt*if2}ZilXlb) zc_8vYZhi6 z?&VC!u$&`9Qos2gEiWH>Yzu{XhGh1YOxt5QMT;O!+b)nTH>cK#jS1Y^oJtB#p&)Id$G8-1hl60! zquex{hS%W?yayk{dAI~$!d3VNzJ>4L2lyF&g}<;8tFZW}af z`~ufdCV0j$j;jNJw%~f~zz4AtyKoEk<5t{;0|eDU9KvB7#W4b^K~Oz{7Ut1IpWylw z&fwGdEIvnIeFA%0c{!yP363X125$}mA;MQ7S=&bB{E=f*R0 zeiwI|s}K`t1jseM;?eg@Gp7*EBn%YJv2b+#4R6--?ihEpPZ6<}d>Q*{ByD diff --git a/abacusnbody/hod/.DS_Store b/abacusnbody/hod/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..666cc43bb60580b19a5fecc448b2c6948bad4a44 GIT binary patch literal 6148 zcmeHKOHRWu5Pg$YsHg~a$ucLX#0^3TzXc%C0~A^is#GWfiUh0dISi-aFs$Is*s7>g zSRz0((fHYud3N$_#W4Ws&CWN04uBS2Fj9AcO3eH^2-Sx;Fn32O~|IALE$ z%>RZb&k1jxE9w>(EYA!FxZ~4aEgMtb5^f7KuFEkdno zY`VDuu7E4>M-||nt=4J>^wJe@1zdr-0{VT(>54_fDxiHjSlAJO*kss@b^TdXP9Cv{ zSOw%Angl8_P$eEQB+wbpV_p%l3K-~+c=(VwvcwaL*y!v(w{%Dm&`Vdq6 Date: Wed, 9 Oct 2024 21:12:01 +0000 Subject: [PATCH 11/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- abacusnbody/hod/CSMF_HOD.py | 769 +++++++++++++++++++++------------- abacusnbody/hod/abacus_hod.py | 2 +- 2 files changed, 480 insertions(+), 291 deletions(-) diff --git a/abacusnbody/hod/CSMF_HOD.py b/abacusnbody/hod/CSMF_HOD.py index f81895b9..a20777c8 100644 --- a/abacusnbody/hod/CSMF_HOD.py +++ b/abacusnbody/hod/CSMF_HOD.py @@ -19,108 +19,155 @@ G = 4.302e-6 # in kpc/Msol (km.s)^2 - @njit(fastmath=True) def n_cen_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): """ Standard Cacciato et al. (2008) centrals HOD parametrization for CSMF """ M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - x_low = np.log10(Mstar_low/M_c_value)/(1.41421356*sigma_c) - x_up = np.log10(Mstar_up/M_c_value)/(1.41421356*sigma_c) - - return 0.5*(math.erf(x_up) - math.erf(x_low)) + x_low = np.log10(Mstar_low / M_c_value) / (1.41421356 * sigma_c) + x_up = np.log10(Mstar_up / M_c_value) / (1.41421356 * sigma_c) + + return 0.5 * (math.erf(x_up) - math.erf(x_low)) + @njit(fastmath=True) def CSMF_centrals(M_h, Mstar, M_1, M_0, gamma1, gamma2, sigma_c): """ Eq. (34) from Cacciato et al. (2008) """ - + M_c_value = M_c(M_h, M_1, M_0, gamma1, gamma2) - - return 1/(1.41421356*np.sqrt(np.pi)*np.log(10)*sigma_c*Mstar)*np.exp(-(np.log10(Mstar)-np.log10(M_c_value))**2/(2*sigma_c**2)) + + return ( + 1 + / (1.41421356 * np.sqrt(np.pi) * np.log(10) * sigma_c * Mstar) + * np.exp(-((np.log10(Mstar) - np.log10(M_c_value)) ** 2) / (2 * sigma_c**2)) + ) + @njit(fastmath=True) def M_c(M_h, M_1, M_0, gamma1, gamma2): """ Eq. (37) from Cacciato et al. (2008) """ - return M_0 * (M_h/M_1)**gamma1/(1+M_h/M_1)**(gamma1-gamma2) - + return M_0 * (M_h / M_1) ** gamma1 / (1 + M_h / M_1) ** (gamma1 - gamma2) @njit(fastmath=True) -def get_random_cen_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): - +def get_random_cen_stellarmass( + M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c +): nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - - CSMF_cen = CSMF_centrals(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c) - - - cdf = np.cumsum(CSMF_cen*delta_stellar_masses) - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + + CSMF_cen = CSMF_centrals( + M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, sigma_c + ) + + cdf = np.cumsum(CSMF_cen * delta_stellar_masses) + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) - + + return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) - @njit(fastmath=True) -def get_random_cen_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c): - +def get_random_cen_stellarmass_linearinterpolation( + M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c +): nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) # stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + CSMF_cen = CSMF_centrals(M_h, stellarmass, M_1, M_0, gamma1, gamma2, sigma_c) - - cdf = np.cumsum(CSMF_cen[:-1]*delta_stellar_masses) - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) - bin = np.where(cdf>random_rv)[0][0] - - m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) - return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + + cdf = np.cumsum(CSMF_cen[:-1] * delta_stellar_masses) + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) + bin = np.where(cdf > random_rv)[0][0] + + m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) + return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] @njit(fastmath=True) -def n_sat_CSMF(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, sigma_c, a1, a2, M2, b0, b1, b2, delta1, delta2): +def n_sat_CSMF( + M_h, + Mstar_low, + Mstar_up, + M_1, + M_0, + gamma1, + gamma2, + sigma_c, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, +): """ Standard Cacciato et al. (2008) satellite HOD parametrization for CSMF """ nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - - CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + + CSMF_sat = CSMF_satelites( + M_h, + stellarmass, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, + ) nsat = 0 - for i in range(nbins-1): - nsat += (CSMF_sat[i+1]-CSMF_sat[i])*(stellarmass[i+1]-stellarmass[i])/2 + (stellarmass[i+1]-stellarmass[i])*CSMF_sat[i] - - return nsat#*ncen + for i in range(nbins - 1): + nsat += (CSMF_sat[i + 1] - CSMF_sat[i]) * ( + stellarmass[i + 1] - stellarmass[i] + ) / 2 + (stellarmass[i + 1] - stellarmass[i]) * CSMF_sat[i] + + return nsat # *ncen @njit(fastmath=True) -def CSMF_satelites(M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): +def CSMF_satelites( + M_h, Mstar, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2 +): """ Eq. (36) from Cacciato et al. (2008) """ M_s_value = M_s(M_h, M_1, M_0, gamma1, gamma2) alpha_s_value = alpha_s(M_h, a1, a2, M2) phi_s_value = phi_s(M_h, b0, b1, b2) - + delta = 10 ** (delta1 + delta2 * (np.log10(M_h) - 12)) - - return phi_s_value/M_s_value * (Mstar/M_s_value)**alpha_s_value * np.exp(-delta*(Mstar/M_s_value)**2) + + return ( + phi_s_value + / M_s_value + * (Mstar / M_s_value) ** alpha_s_value + * np.exp(-delta * (Mstar / M_s_value) ** 2) + ) + @njit(fastmath=True) def M_s(M_h, M_1, M_0, gamma1, gamma2): @@ -129,60 +176,122 @@ def M_s(M_h, M_1, M_0, gamma1, gamma2): """ return 0.562 * M_c(M_h, M_1, M_0, gamma1, gamma2) + @njit(fastmath=True) def alpha_s(M_h, a1, a2, M2): """ Eq. (39) from Cacciato et al. (2008) """ - return -2.0 + a1 * (1-2/np.pi*np.arctan(a2*np.log10(M_h/M2))) + return -2.0 + a1 * (1 - 2 / np.pi * np.arctan(a2 * np.log10(M_h / M2))) + @njit(fastmath=True) def phi_s(M_h, b0, b1, b2): """ Eq. (40) from Cacciato et al. (2008) """ - M12 = M_h/1e12 - log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12)**2 + M12 = M_h / 1e12 + log_phi_s = b0 + b1 * np.log10(M12) + b2 * np.log10(M12) ** 2 return 10**log_phi_s @njit(fastmath=True) -def get_random_sat_stellarmass(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): - +def get_random_sat_stellarmass( + M_h, + Mstar_low, + Mstar_up, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, +): nbins = 1000 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) - stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - - CSMF_sat = CSMF_satelites(M_h, stellarmass_centers, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) - - cdf = np.cumsum(CSMF_sat*delta_stellar_masses) - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) + stellarmass_centers = stellarmass[1:] / 2 + stellarmass[:-1] / 2 + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + + CSMF_sat = CSMF_satelites( + M_h, + stellarmass_centers, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, + ) + + cdf = np.cumsum(CSMF_sat * delta_stellar_masses) + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) bin_clostest = (np.abs(cdf - random_rv)).argmin() - - return np.random.uniform(stellarmass[bin_clostest],stellarmass[bin_clostest+1]) - + + return np.random.uniform(stellarmass[bin_clostest], stellarmass[bin_clostest + 1]) + @njit(fastmath=True) -def get_random_sat_stellarmass_linearinterpolation(M_h, Mstar_low, Mstar_up, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2): - +def get_random_sat_stellarmass_linearinterpolation( + M_h, + Mstar_low, + Mstar_up, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, +): nbins = 100 - stellarmass = np.logspace(np.log10(Mstar_low),np.log10(Mstar_up),nbins) + stellarmass = np.logspace(np.log10(Mstar_low), np.log10(Mstar_up), nbins) # stellarmass_centers = stellarmass[1:]/2+stellarmass[:-1]/2 - delta_stellar_masses = stellarmass[1:]-stellarmass[:-1] - - CSMF_sat = CSMF_satelites(M_h, stellarmass, M_1, M_0, gamma1, gamma2, a1, a2, M2, b0, b1, b2, delta1, delta2) - - cdf = np.cumsum(CSMF_sat[:-1]*delta_stellar_masses) - cdf=cdf/cdf[-1] - - random_rv = np.random.uniform(cdf.min(),cdf.max()) - bin = np.where(cdf>random_rv)[0][0] - - m = (stellarmass[bin]-stellarmass[bin-1])/(cdf[bin]-cdf[bin-1]) - return m * (random_rv - cdf[bin-1]) + stellarmass[bin-1] + delta_stellar_masses = stellarmass[1:] - stellarmass[:-1] + + CSMF_sat = CSMF_satelites( + M_h, + stellarmass, + M_1, + M_0, + gamma1, + gamma2, + a1, + a2, + M2, + b0, + b1, + b2, + delta1, + delta2, + ) + + cdf = np.cumsum(CSMF_sat[:-1] * delta_stellar_masses) + cdf = cdf / cdf[-1] + + random_rv = np.random.uniform(cdf.min(), cdf.max()) + bin = np.where(cdf > random_rv)[0][0] + + m = (stellarmass[bin] - stellarmass[bin - 1]) / (cdf[bin] - cdf[bin - 1]) + return m * (random_rv - cdf[bin - 1]) + stellarmass[bin - 1] @njit(fastmath=True) @@ -226,24 +335,23 @@ def gen_cent( ): """ Generate central galaxies in place in memory with a two pass numba parallel implementation. - """ + """ if want_CSMF: Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'] + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], ) ic_C, alpha_c_C, Ac_C, Bc_C = ( - CSMF_hod_dict['ic'], - CSMF_hod_dict['alpha_c'], - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Bcent'] + CSMF_hod_dict['ic'], + CSMF_hod_dict['alpha_c'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Bcent'], ) - H = len(mass) @@ -261,14 +369,23 @@ def gen_cent( # first create the markers between 0 and 1 for different tracers CSMF_marker = 0 if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) - - ncen = n_cen_CSMF(mass[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) - + M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + + ncen = n_cen_CSMF( + mass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + ) + CSMF_marker += ncen * ic_C * multis[i] if randoms[i] <= CSMF_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 else: keep[i] = 0 @@ -277,19 +394,18 @@ def gen_cent( gstart = np.empty((Nthread + 1, 1), dtype=np.int64) gstart[0, :] = 0 gstart[1:, 0] = Nout[:, 0, 0].cumsum() - # galaxy arrays N_CSMF = gstart[-1, 0] - CSMF_x = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_y = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_z = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype = mass.dtype) - CSMF_id = np.empty(N_CSMF, dtype = ids.dtype) + CSMF_x = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_y = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_z = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype=mass.dtype) + CSMF_id = np.empty(N_CSMF, dtype=ids.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): @@ -297,42 +413,50 @@ def gen_cent( for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: # loop thru three directions to assign galaxy velocities and positions - CSMF_x[j1] = pos[i,0] - CSMF_vx[j1] = vel[i,0] + alpha_c_C * vdev[i,0] # velocity bias - CSMF_y[j1] = pos[i,1] - CSMF_vy[j1] = vel[i,1] + alpha_c_C * vdev[i,1] # velocity bias - CSMF_z[j1] = pos[i,2] - CSMF_vz[j1] = vel[i,2] + alpha_c_C * vdev[i,2] # velocity bias - + CSMF_x[j1] = pos[i, 0] + CSMF_vx[j1] = vel[i, 0] + alpha_c_C * vdev[i, 0] # velocity bias + CSMF_y[j1] = pos[i, 1] + CSMF_vy[j1] = vel[i, 1] + alpha_c_C * vdev[i, 1] # velocity bias + CSMF_z[j1] = pos[i, 2] + CSMF_vz[j1] = vel[i, 2] + alpha_c_C * vdev[i, 2] # velocity bias + # rsd only applies to the z direction if rsd and origin is not None: nx = CSMF_x[j1] - origin[0] ny = CSMF_y[j1] - origin[1] nz = CSMF_z[j1] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms*(CSMF_vx[j1]*nx+CSMF_vy[j1]*ny+CSMF_vz[j1]*nz) - CSMF_x[j1] = CSMF_x[j1]+proj*nx - CSMF_y[j1] = CSMF_y[j1]+proj*ny - CSMF_z[j1] = CSMF_z[j1]+proj*nz + proj = inv_velz2kms * ( + CSMF_vx[j1] * nx + CSMF_vy[j1] * ny + CSMF_vz[j1] * nz + ) + CSMF_x[j1] = CSMF_x[j1] + proj * nx + CSMF_y[j1] = CSMF_y[j1] + proj * ny + CSMF_z[j1] = CSMF_z[j1] + proj * nz elif rsd: - CSMF_z[j1] = wrap(pos[i,2] + CSMF_vz[j1] * inv_velz2kms, lbox) - + CSMF_z[j1] = wrap(pos[i, 2] + CSMF_vz[j1] * inv_velz2kms, lbox) + CSMF_mass[j1] = mass[i] - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) + M_1_C_temp = 10 ** (np.log10(M_1_C) + Ac_C * deltac[i] + Bc_C * fenv[i]) CSMF_stellarmass[j1] = get_random_cen_stellarmass_linearinterpolation( - mass[i], Mstar_low_C, Mstar_up_C, - M_1_C_temp, M_0_C, gamma1_C, gamma2_C, sigma_c_C) + mass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + ) CSMF_id[j1] = ids[i] j1 += 1 # assert j == gstart[tid + 1] - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) - - + CSMF_dict['x'] = CSMF_x CSMF_dict['y'] = CSMF_y CSMF_dict['z'] = CSMF_z @@ -479,29 +603,45 @@ def gen_sats_nfw( """ if want_CSMF: - Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], - CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], - CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'] + ( + Mstar_low_C, + Mstar_up_C, + M_1_C, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], + CSMF_hod_dict['gamma_1'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], + CSMF_hod_dict['M_2'], + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'], ) Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], + CSMF_hod_dict['Acent'], + CSMF_hod_dict['Asat'], CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'] + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'], ) f_sigv_C = CSMF_hod_dict['f_sigv'] @@ -509,37 +649,40 @@ def gen_sats_nfw( # compute nsate for each halo # figuring out the number of particles kept for each thread - num_sats_C = np.zeros(len(hid), dtype = np.int64) - stellarmass_C = np.zeros(len(hid), dtype = np.int64) + num_sats_C = np.zeros(len(hid), dtype=np.int64) + stellarmass_C = np.zeros(len(hid), dtype=np.int64) hstart = np.rint(np.linspace(0, len(hid), Nthread + 1)).astype( np.int64 ) # starting index of each thread for tid in range(Nthread): for i in range(hstart[tid], hstart[tid + 1]): if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] base_p_C = ( n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) * ic_C ) - num_sats_C[i] = np.random.poisson(base_p_C) + num_sats_C[i] = np.random.poisson(base_p_C) # generate rdpos rd_pos_C = getPointsOnSphere(np.sum(num_sats_C), Nthread) @@ -547,35 +690,33 @@ def gen_sats_nfw( # put satellites on NFW h_id_C, x_sat_C, y_sat_C, z_sat_C, vx_sat_C, vy_sat_C, vz_sat_C, M_C = ( compute_fast_NFW( - NFW_draw, - hid, - hpos[:, 0], - hpos[:, 1], - hpos[:, 2], - hvel[:, 0], - hvel[:, 1], + NFW_draw, + hid, + hpos[:, 0], + hpos[:, 1], + hpos[:, 2], + hvel[:, 0], + hvel[:, 1], hvel[:, 2], - hvrms, - hc, - hmass, - hrvir, - rd_pos_C, - num_sats_C, - f_sigv_C, - vel_sat, + hvrms, + hc, + hmass, + hrvir, + rd_pos_C, + num_sats_C, + f_sigv_C, + vel_sat, Nthread, ) ) - - + # do rsd if rsd: z_sat_C = (z_sat_C + vz_sat_C * inv_velz2kms) % lbox - - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) - + CSMF_dict['x'] = x_sat_C CSMF_dict['y'] = y_sat_C CSMF_dict['z'] = z_sat_C @@ -589,12 +730,28 @@ def gen_sats_nfw( hstart = np.rint(np.linspace(0, num_sats_C.sum(), Nthread + 1)) for tid in numba.prange(Nthread): for i in range(int(hstart[tid]), int(hstart[tid + 1])): - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] stellarmass_C[i] = get_random_sat_stellarmass_linearinterpolation( - M_C[i], Mstar_low_C, Mstar_up_C, M_1_C_temp, M_0_C, gamma1_C, gamma2_C, - a1_C_temp, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C) - + M_C[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) + CSMF_dict['stellarmass'] = stellarmass_C ID_dict['CSMF'] = h_id_C @@ -633,38 +790,52 @@ def gen_sats( Generate satellite galaxies in place in memory with a two pass numba parallel implementation. """ - - if want_CSMF: - Mstar_low_C, Mstar_up_C, M_1_C, M_0_C, gamma1_C, gamma2_C, sigma_c_C, a1_C, a2_C, M2_C, b0_C, b1_C, b2_C, delta1_C, delta2_C = ( - CSMF_hod_dict['Mstar_low'], - CSMF_hod_dict['Mstar_up'], - CSMF_hod_dict['M_1'], - CSMF_hod_dict['M_0'], + ( + Mstar_low_C, + Mstar_up_C, + M_1_C, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) = ( + CSMF_hod_dict['Mstar_low'], + CSMF_hod_dict['Mstar_up'], + CSMF_hod_dict['M_1'], + CSMF_hod_dict['M_0'], CSMF_hod_dict['gamma_1'], - CSMF_hod_dict['gamma_2'], - CSMF_hod_dict['sigma_c'], - CSMF_hod_dict['a_1'], - CSMF_hod_dict['a_2'], + CSMF_hod_dict['gamma_2'], + CSMF_hod_dict['sigma_c'], + CSMF_hod_dict['a_1'], + CSMF_hod_dict['a_2'], CSMF_hod_dict['M_2'], - CSMF_hod_dict['b_0'], - CSMF_hod_dict['b_1'], - CSMF_hod_dict['b_2'], - CSMF_hod_dict['delta_1'], - CSMF_hod_dict['delta_2'] + CSMF_hod_dict['b_0'], + CSMF_hod_dict['b_1'], + CSMF_hod_dict['b_2'], + CSMF_hod_dict['delta_1'], + CSMF_hod_dict['delta_2'], ) - + alpha_s_C, s_C, s_v_C, s_p_C, s_r_C, Ac_C, As_C, Bc_C, Bs_C, ic_C = ( - CSMF_hod_dict['alpha_s'], + CSMF_hod_dict['alpha_s'], CSMF_hod_dict['s'], - CSMF_hod_dict['s_v'], - CSMF_hod_dict['s_p'], - CSMF_hod_dict['s_r'], + CSMF_hod_dict['s_v'], + CSMF_hod_dict['s_p'], + CSMF_hod_dict['s_r'], CSMF_hod_dict['Acent'], - CSMF_hod_dict['Asat'], - CSMF_hod_dict['Bcent'], - CSMF_hod_dict['Bsat'], - CSMF_hod_dict['ic'] + CSMF_hod_dict['Asat'], + CSMF_hod_dict['Bcent'], + CSMF_hod_dict['Bsat'], + CSMF_hod_dict['ic'], ) H = len(hmass) # num of particles @@ -683,39 +854,47 @@ def gen_sats( # print(logM1, As, hdeltac[i], Bs, hfenv[i]) CSMF_marker = 0 if want_CSMF: - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] base_p_C = ( n_sat_CSMF( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - sigma_c_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C - ) - * weights[i] - * ic_C) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + sigma_c_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) + * weights[i] + * ic_C + ) if enable_ranks: - decorator_C = 1 + s_C * ranks[i] + s_v_C * ranksv[i] + s_p_C * ranksp[i] + s_r_C * ranksr[i] + decorator_C = ( + 1 + + s_C * ranks[i] + + s_v_C * ranksv[i] + + s_p_C * ranksp[i] + + s_r_C * ranksr[i] + ) exp_sat = base_p_C * decorator_C else: exp_sat = base_p_C CSMF_marker += exp_sat - if randoms[i] <= CSMF_marker: - Nout[tid, 0, 0] += 1 # counting + Nout[tid, 0, 0] += 1 # counting keep[i] = 1 else: keep[i] = 0 @@ -727,15 +906,15 @@ def gen_sats( # galaxy arrays N_CSMF = gstart[-1, 0] - CSMF_x = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_y = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_z = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vx = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vy = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_vz = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_mass = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_stellarmass = np.empty(N_CSMF, dtype = hmass.dtype) - CSMF_id = np.empty(N_CSMF, dtype = hid.dtype) + CSMF_x = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_y = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_z = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_vx = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_vy = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_vz = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_mass = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_stellarmass = np.empty(N_CSMF, dtype=hmass.dtype) + CSMF_id = np.empty(N_CSMF, dtype=hid.dtype) # fill in the galaxy arrays for tid in numba.prange(Nthread): @@ -743,56 +922,63 @@ def gen_sats( for i in range(hstart[tid], hstart[tid + 1]): if keep[i] == 1: CSMF_x[j1] = ppos[i, 0] - CSMF_vx[j1] = hvel[i, 0] + alpha_s_C * (pvel[i, 0] - hvel[i, 0]) # velocity bias + CSMF_vx[j1] = hvel[i, 0] + alpha_s_C * ( + pvel[i, 0] - hvel[i, 0] + ) # velocity bias CSMF_y[j1] = ppos[i, 1] - CSMF_vy[j1] = hvel[i, 1] + alpha_s_C * (pvel[i, 1] - hvel[i, 1]) # velocity bias + CSMF_vy[j1] = hvel[i, 1] + alpha_s_C * ( + pvel[i, 1] - hvel[i, 1] + ) # velocity bias CSMF_z[j1] = ppos[i, 2] - CSMF_vz[j1] = hvel[i, 2] + alpha_s_C * (pvel[i, 2] - hvel[i, 2]) # velocity bias + CSMF_vz[j1] = hvel[i, 2] + alpha_s_C * ( + pvel[i, 2] - hvel[i, 2] + ) # velocity bias if rsd and origin is not None: nx = CSMF_x[j1] - origin[0] ny = CSMF_y[j1] - origin[1] nz = CSMF_z[j1] - origin[2] - inv_norm = 1./np.sqrt(nx*nx + ny*ny + nz*nz) + inv_norm = 1.0 / np.sqrt(nx * nx + ny * ny + nz * nz) nx *= inv_norm ny *= inv_norm nz *= inv_norm - proj = inv_velz2kms*(CSMF_vx[j1]*nx+CSMF_vy[j1]*ny+CSMF_vz[j1]*nz) - CSMF_x[j1] = CSMF_x[j1]+proj*nx - CSMF_y[j1] = CSMF_y[j1]+proj*ny - CSMF_z[j1] = CSMF_z[j1]+proj*nz + proj = inv_velz2kms * ( + CSMF_vx[j1] * nx + CSMF_vy[j1] * ny + CSMF_vz[j1] * nz + ) + CSMF_x[j1] = CSMF_x[j1] + proj * nx + CSMF_y[j1] = CSMF_y[j1] + proj * ny + CSMF_z[j1] = CSMF_z[j1] + proj * nz elif rsd: - CSMF_z[j1] = wrap(CSMF_z[j1] + CSMF_vz[j1] * inv_velz2kms, lbox) - - - M_1_C_temp = 10**(np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i]) + CSMF_z[j1] = wrap(CSMF_z[j1] + CSMF_vz[j1] * inv_velz2kms, lbox) + + M_1_C_temp = 10 ** ( + np.log10(M_1_C) + Ac_C * hdeltac[i] + Bc_C * hfenv[i] + ) a1_C_temp = a1_C + As_C * hdeltac[i] + Bs_C * hfenv[i] CSMF_stellarmass[j1] = get_random_sat_stellarmass( - hmass[i], - Mstar_low_C, - Mstar_up_C, - M_1_C_temp, - M_0_C, - gamma1_C, - gamma2_C, - a1_C_temp, - a2_C, - M2_C, - b0_C, - b1_C, - b2_C, - delta1_C, - delta2_C - ) + hmass[i], + Mstar_low_C, + Mstar_up_C, + M_1_C_temp, + M_0_C, + gamma1_C, + gamma2_C, + a1_C_temp, + a2_C, + M2_C, + b0_C, + b1_C, + b2_C, + delta1_C, + delta2_C, + ) CSMF_mass[j1] = hmass[i] CSMF_id[j1] = hid[i] j1 += 1 # assert j == gstart[tid + 1] - - CSMF_dict = Dict.empty(key_type = types.unicode_type, value_type = float_array) + CSMF_dict = Dict.empty(key_type=types.unicode_type, value_type=float_array) ID_dict = Dict.empty(key_type=types.unicode_type, value_type=int_array) - CSMF_dict['x'] = CSMF_x CSMF_dict['y'] = CSMF_y CSMF_dict['z'] = CSMF_z @@ -802,7 +988,7 @@ def gen_sats( CSMF_dict['mass'] = CSMF_mass CSMF_dict['stellarmass'] = CSMF_stellarmass ID_dict['CSMF'] = CSMF_id - + return CSMF_dict, ID_dict @@ -887,11 +1073,12 @@ def gen_gals( if tracer == 'CSMF': CSMF_HOD = tracers[tracer] - if 'CSMF' in tracers.keys(): want_CSMF = True - CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + CSMF_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) for key, value in CSMF_HOD.items(): CSMF_hod_dict[key] = value @@ -904,7 +1091,9 @@ def gen_gals( else: want_CSMF = False - CSMF_hod_dict = nb.typed.Dict.empty(key_type=nb.types.unicode_type, value_type= nb.types.float64) + CSMF_hod_dict = nb.typed.Dict.empty( + key_type=nb.types.unicode_type, value_type=nb.types.float64 + ) start = time.time() diff --git a/abacusnbody/hod/abacus_hod.py b/abacusnbody/hod/abacus_hod.py index 7f64b14a..feabe014 100644 --- a/abacusnbody/hod/abacus_hod.py +++ b/abacusnbody/hod/abacus_hod.py @@ -797,7 +797,7 @@ def run_hod( verbose=verbose, fn_ext=fn_ext, ) - + self.logger.info(f'HOD generated in elapsed time {time.time() - start:.2f} s.') return mock_dict From a20048c3e655a31382f5f4ca18e05123f7997dc6 Mon Sep 17 00:00:00 2001 From: Pierre Burger Date: Tue, 15 Oct 2024 10:54:38 -0400 Subject: [PATCH 12/13] test notebook I created a test notebook showing how to create galaxy catalogues. I added Google Drive links to download the reference Haloes and Galaxy catalogues --- .DS_Store | Bin 6148 -> 6148 bytes abacusnbody/.DS_Store | Bin 8196 -> 8196 bytes abacusnbody/hod/.DS_Store | Bin 6148 -> 6148 bytes tests/.DS_Store | Bin 0 -> 8196 bytes tests/CSMF_test.ipynb | 473 +++++++++++++++++++++++++++ tests/data_CSMF/.DS_Store | Bin 0 -> 6148 bytes tests/data_CSMF/HOD_parameters.fits | Bin 0 -> 8640 bytes tests/data_CSMF/config_CSMF.yaml | 91 ++++++ tests/data_CSMF/halos_data/.DS_Store | Bin 0 -> 6148 bytes 9 files changed, 564 insertions(+) create mode 100644 tests/.DS_Store create mode 100644 tests/CSMF_test.ipynb create mode 100644 tests/data_CSMF/.DS_Store create mode 100644 tests/data_CSMF/HOD_parameters.fits create mode 100644 tests/data_CSMF/config_CSMF.yaml create mode 100644 tests/data_CSMF/halos_data/.DS_Store diff --git a/.DS_Store b/.DS_Store index 8b526270aca7bdc631fc4806cbdafe5d48ec3890..4d78fc8fcc8bfa950713ab81552e650dc1cf4199 100644 GIT binary patch delta 98 zcmZoMXfc@J&&azmU^g=(?_@<*sm*FE511z3VdLgWDlaZb%E?b+U|`rWIfJE8U81_$ z$jnqn!N{^!N1@u#$jDen!OYmGww9AaR9W9TC_XzUH!r_yG9#-zWB0^Dqs{CbfB6BA C&>W8d delta 32 ocmZoMXfc@J&&abeU^g=(&tyebsm*FE512NVF*0su=lIJH0HJRQssI20 diff --git a/abacusnbody/.DS_Store b/abacusnbody/.DS_Store index 299dab670e7e0c4a83548d28f538e273fed55c07..810d60f9639cfaa336b17380f3dd7d9f40e73672 100644 GIT binary patch delta 32 ncmZp1XmQx!BPeX4qhM%cP^+U*ZD?R>prc@JV7NJ6a48=Ek%I_? delta 32 ncmZp1XmQx!BPeXCqhMrUQmdm-ZD?R%q@!SFXt6n7a48=Ek=F=} diff --git a/abacusnbody/hod/.DS_Store b/abacusnbody/hod/.DS_Store index 666cc43bb60580b19a5fecc448b2c6948bad4a44..be3feb435b266f1da08c706b0e2f2a0472df3d06 100644 GIT binary patch delta 57 zcmZoMXfc@J&&a+pU^gQp`(z#_OHp2iM1~|FPG%@&C}xOf$Y984NSUm_B*MffyV;Cs LI_qY3j=%f>iIWXj delta 31 ncmZoMXfc@J&&akhU^gQp+hiUl%gyyn9jp@@tTwZA{N)D#nfwXI diff --git a/tests/.DS_Store b/tests/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..edd31f3eb8794fa576ccc1e828208875187a8bc2 GIT binary patch literal 8196 zcmeHM&2G~`5S~p6bx3~#0#qR(S>jrSHib%vOGpb65^x|04uFE)q^VWMj%tUfRFQIq zci!e9M65OHZW65#k%!7oNn1V93#b6SS=RQ}5Y$`dG zQo@NzI5Fj!sT~Rv-huOlIe_V0B%M&E$q=W)6h#%Tym zKDnNt*8hrvE20|~I)5daYYt?bMr3|0yd;W!&A9)E_z zILh*kjo++PUbu2~(OdMY-pBSqI&SAdKFV6b(R+UPCQak;^g0ONM!j*Tc5^39@*qlj zM*H#&#dxvlFj&h*7Onx)MnP^_IJ zr}r#vgg2JCKO0REcOrhem#T)7r+Wo+uEe9LqgFJOxhl-fR}n28tB-VmG4r`s z;FCwmn5ow_pF_lI2y6$~4D9d2LPiHBM}&A87M2YQD!+?-^T-)Uoh1>CLKGrspg@`{ zj)4ncpe%i^RUx0Aeg1y|BXcet1CD|7Vt_5|ws%|bcKWl=cFgzMKI$h_Ubx;=N&!J5 mw&RduI}SPf!w}a#R5_-p= args_dict['zmin']) & (redshift <= args_dict['zmax']) & (data['stellarmass']>0)\n", + " central_id = np.zeros_like(data['stellarmass'])\n", + " central_id[:data['Ncent']] = 1\n", + "\n", + " data_positions_sky.append(np.c_[ra[mask],dec[mask],dist[mask],redshift[mask],data['stellarmass'][mask],data['mass'][mask],central_id[mask]])\n", + "data_positions_sky = np.concatenate(data_positions_sky, axis=0)\n", + "\n", + "HOD_table = Table()\n", + "HOD_table['ra'] = data_positions_sky[:,0].astype(np.float32)\n", + "HOD_table['dec'] = data_positions_sky[:,1].astype(np.float32)\n", + "HOD_table['comoving_distance_overh'] = data_positions_sky[:,2].astype(np.float32)*u.Mpc\n", + "HOD_table['redshift'] = data_positions_sky[:,3].astype(np.float32)\n", + "HOD_table['log10Mstar_overhsquare'] = np.log10(data_positions_sky[:,4]).astype(np.float64)*u.Msun\n", + "HOD_table['log10Mhalo_overh'] = np.log10(data_positions_sky[:,5]).astype(np.float64)*u.Msun\n", + "HOD_table['central flag'] = data_positions_sky[:,6].astype(np.int32)\n", + "\n", + "HOD_table.write(f'data_CSMF/galaxy_catalogue_cos{cosmo:03}_ph{phase:03}_hod{hod:03}_seed{seed:03}.fits',overwrite=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "464b26e4-8916-4844-a27d-0b87a4a58b64", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/html": [ + "
Table length=53476\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
radeccomoving_distance_overhredshiftlog10Mstar_overhsquarelog10Mhalo_overhcentral flag
MpcsolMasssolMass
float32float32float32float32float64float64int32
89.8014561.005356313.052550.1071883310.66969980354864413.158514070143971
89.9681315.081102314.773860.1077940710.66142834259780312.5673795125458821
89.6028477.63843297.20590.1016207511.0449160596814812.7873889764244381
89.9094879.4423362.139130.1245379810.45763573074627812.5010523470493451
89.01124679.60034358.288850.1231713811.10461325073985813.2197368712865161
89.3260175.91486308.977870.1057551910.8864025102048713.1957830321479521
89.7543373.1615343.472560.11792166511.38814491116055212.8208846821624791
89.6906662.50087302.82540.1035932711.1117148633010613.3626737924240151
89.6663764.31412341.108830.1170854811.29916501245077713.097221290495771
.....................
3.02285461.523003363.07550.1248704910.69783597047026212.7178440068105160
3.0290871.5405174362.9810.1248369410.43369322370287812.7178440068105160
3.04004071.5345762363.18280.12490860410.51711550816800212.7178440068105160
2.33026340.028755056365.990420.1259059510.53792253062864113.2524380017272990
2.32547240.055217672365.898320.1258732210.42413881871390113.2524380017272990
2.31868840.051970188366.020660.1259166910.53982087771925413.2524380017272990
2.32047770.04352397365.932830.1258854910.7197391644448613.2524380017272990
2.33236310.010766512365.902370.1258746710.44755170619331813.2524380017272990
2.34656450.03120295366.135220.125957410.66761811117959613.2524380017272990
2.34314130.0038444304365.97320.1258998410.50368217174970213.2524380017272990
" + ], + "text/plain": [ + "\n", + " ra dec ... log10Mhalo_overh central flag\n", + " ... solMass \n", + " float32 float32 ... float64 int32 \n", + "--------- ------------ ... ------------------ ------------\n", + " 89.80145 61.005356 ... 13.15851407014397 1\n", + " 89.96813 15.081102 ... 12.567379512545882 1\n", + " 89.60284 77.63843 ... 12.787388976424438 1\n", + " 89.90948 79.4423 ... 12.501052347049345 1\n", + "89.011246 79.60034 ... 13.219736871286516 1\n", + " 89.32601 75.91486 ... 13.195783032147952 1\n", + " 89.75433 73.1615 ... 12.820884682162479 1\n", + " 89.69066 62.50087 ... 13.362673792424015 1\n", + " 89.66637 64.31412 ... 13.09722129049577 1\n", + " ... ... ... ... ...\n", + "3.0228546 1.523003 ... 12.717844006810516 0\n", + " 3.029087 1.5405174 ... 12.717844006810516 0\n", + "3.0400407 1.5345762 ... 12.717844006810516 0\n", + "2.3302634 0.028755056 ... 13.252438001727299 0\n", + "2.3254724 0.055217672 ... 13.252438001727299 0\n", + "2.3186884 0.051970188 ... 13.252438001727299 0\n", + "2.3204777 0.04352397 ... 13.252438001727299 0\n", + "2.3323631 0.010766512 ... 13.252438001727299 0\n", + "2.3465645 0.03120295 ... 13.252438001727299 0\n", + "2.3431413 0.0038444304 ... 13.252438001727299 0" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "HOD_table" + ] + }, + { + "cell_type": "markdown", + "id": "a27302e4-b366-432e-a883-c16fd7f0806d", + "metadata": {}, + "source": [ + "## Download reference catalogue and compare it with the new generated output. Note that stellar masses are different for each run by construction." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "469d1e69-940a-4a02-9af0-0942ee49cdb7", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Downloading...\n", + "From: https://drive.google.com/uc?id=1JB5UvLM13wj6IdRvF37-Dh8kRjJWzWh0\n", + "To: /global/u2/p/pierre/abacusutils_modified/tests/data_CSMF/reference_galaxy_catalogue_cos000_ph000_hod000_seed000.fits\n", + "100%|██████████| 1.93M/1.93M [00:00<00:00, 16.1MB/s]\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\u001b[92m \"ra\" is close to the reference catalogue\n", + "\u001b[92m \"dec\" is close to the reference catalogue\n", + "\u001b[92m \"comoving_distance_overh\" is close to the reference catalogue\n", + "\u001b[92m \"redshift\" is close to the reference catalogue\n", + "\u001b[91m \"log10Mstar_overhsquare\" is different to the reference catalogue\n", + "\u001b[92m \"log10Mhalo_overh\" is close to the reference catalogue\n", + "\u001b[92m \"central flag\" is close to the reference catalogue\n" + ] + } + ], + "source": [ + "from common import check_close\n", + "\n", + "gdown.download(id='1JB5UvLM13wj6IdRvF37-Dh8kRjJWzWh0',output=f'data_CSMF/reference_galaxy_catalogue_cos{cosmo:03}_ph{phase:03}_hod{hod:03}_seed{seed:03}.fits')\n", + "HOD_table_ref = Table.read(f'data_CSMF/reference_galaxy_catalogue_cos{cosmo:03}_ph{phase:03}_hod{hod:03}_seed{seed:03}.fits')\n", + "\n", + "\n", + "for name in HOD_table_ref.colnames:\n", + " if(check_close(HOD_table_ref[name],HOD_table[name])):\n", + " print('\\033[92m \"'+name + r'\" is close to the reference catalogue')\n", + " else:\n", + " print('\\033[91m \"'+name + r'\" is different to the reference catalogue')" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "AbacusHOD_env", + "language": "python", + "name": "abacushod_env" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.8" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/data_CSMF/.DS_Store b/tests/data_CSMF/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..d2784d5e0298bc863c021734671ba4515ca3d670 GIT binary patch literal 6148 zcmeHKU5nE|6uq-schg;jfGg~a5b(9swu=k0FR@!iA4IUC4=OQfq79~*(xg-=g*@wz z@zr1AfALAr+?kf`)&*f;MJ#kW40Gr2O=g`A{;Y71yY*zXs8fR%ywpu^M)^_d2&aSiT>^txJLpkf0 zVL2_j;pB}fy_7PGmiaJxl_s;mJvfzl8K!wMVL}=w2zmQD&0{(1%4r@K%WushvWz9+8i_3P?;1a$ebBmtP zbLv5hF|?S#M^LqdPoU66sw(E$eLX>Rf*nw-1kc9+X>jFx{p;MUmg}^5Sq01!)jD~ zdnV2|SZUOB66W$D%*evrPy~;T{yh#S;cK+1Rlq8+uE4e~y6XIYbn*Rvon%W^0jt2Z zQb5#t{azPKGH2_?;?!B|!QaD`skllbQ;?YBm^11qK7_NO%~KhmZ?Mvc8kqefz%tmx JDsZI=`~kxYuXg|d literal 0 HcmV?d00001 diff --git a/tests/data_CSMF/HOD_parameters.fits b/tests/data_CSMF/HOD_parameters.fits new file mode 100644 index 0000000000000000000000000000000000000000..00885bb700674817f4200e94532ca811656a4363 GIT binary patch literal 8640 zcmeI0eM}Q)9LGVYpaM?J=G1JS(rZ;vC?cR1p9e~rMJRPGFH&MDw8d&^u`SGsIz)mF zG$K=|ZXyO@zRg7=5UblUqO~r7AhJY61Z9gb=MXBsbm*?F229Gr;$?};J^%FS<;ka? z`#!(R{hspJk=zItf&?LU1|RW8l9dXnQkAMkG)g3l&F3L%jYuIDsl;~Hyf+K5=P8`A4k2I`2U;8||rN6hw!`F{AxXh?%4&KL);=GWeQl_UNX(E*NTClh zy#9uMc$71KSkz~cHXeRTvW~aoS^TW~`^t|_fS-*=oZli#z0gubsi`6%5q>rvaejWo z{Co+>&&DIpZ?V-+Ela`kOD4e2#v{(p-|A=CTg!O6dbT_tHXd<)0am{x{MjYO&&DIp zFYuKgQTf?;#Q80;)t@yzBKFtDBhC+hky-kKWQjsUNd4J(1p3jferl0LkAUX`f05z$ zONbx&|Jtjn{-1t8^r2y2OHMeXns@o$*dm0>YCUfBzbph-Sb=AsTrq;?`RjEZCINJu zSNdi4ic-+3Kgl~A-~-y77FS1qbr~BlR8ZRLL$N#U@ujg_PJvFxhl}z9PC-i9_X*S8 z9UwZ#SU6`*0t{ZY8Vkyg2aQ#3lxePGz_rc7)=RurI6-*Q^q}WNjf0*s-?Q@|Ff-cQ zBbmk6AakIS`5*%v*LrQ}=uZUq&&_J;n$iW)_iypR#|l^J9J#~q|6tB)jt zYxl~yW&B=-6@h*qH5jUg@NQtWZ? zV6OLd6GZnvDjcK=Aez}&T^)1;2BUE+bZv*gi9Ng8GM;yUUoQNrT~v|@y{dDz`pY9~ zoam}$`rKonsTkQCnX3ak?&>BuvK_FC=t|+fgB!6n7j-=Pd?Q{5g7SG^OomjqKK3}? z3`lP>_?UyTKvS-Mxi;MsT;e?x&HmU8&g}W^FTswxVEQK7$Lro4etDt z48O!K4pzA?hrMWyt&?wZ5qTacdg%b-!6C%&Gx5l|1(atK^_Y z5xuaIJ0DJKTl!;H#RkxAeo15S?qSzixBDt~KLuy9FWgXX>%e*^*X7ab z6p*6dG||&-1w_lEi~>d_T#{XHTD3bAG$(Xq-E{H;b?54GW;XNSlzj04$G#Zsy7Sn) zb+d9nrw%m4?A!~UtkwCIuBKqM6{e}*`X2?ovZ8hCpBWs;K_);ZKqf#YKqf#Y@Xjaj EH}o3yO#lD@ literal 0 HcmV?d00001 diff --git a/tests/data_CSMF/config_CSMF.yaml b/tests/data_CSMF/config_CSMF.yaml new file mode 100644 index 00000000..fad2a64b --- /dev/null +++ b/tests/data_CSMF/config_CSMF.yaml @@ -0,0 +1,91 @@ +# Comment out the probes you don't need +# CAREFUL: yaml is super sensitive to spaces - don't use tabs + +# Simulation parameters +sim_params: + sim_name: 'AbacusSummit_base_c000_ph000' # which simulation + sim_dir: '/global/cfs/cdirs/desi/public/cosmosim/AbacusSummit/halo_light_cones/' # where is the simulation + subsample_dir: 'data_CSMF/halos_data/' + cleaned_halos: 'data_CSMF/halos_data/' + + # subsample_dir: '/pscratch/sd/p/pierre/AbacusSummit/subsamples/' + # cleaned_halos: '/pscratch/sd/p/pierre/AbacusSummit/subsamples/' + + Nthread_load: 10 + z_mock: 0.100 + halo_lc: True + + # subsample_dir: ‘/pscratch/sd/b/boryanah/AbacusHOD_scratch/mocks_lc_subsample/’ + +prepare_sim: + Nparallel_load: 10 # best for cori haswell node + +# HOD parameters +HOD_params: + use_particles: True + want_ranks: True # want to implement satellite profile flexiblities? + want_AB: True # want to implement assembly bias? + density_sigma: 3 # scale radius in Mpc / h for local density calculation + tracer_flags: # which tracer do we want? + LRG: False + ELG: False + QSO: False + CSMF: True + want_rsd: True # want RSD? + write_to_disk: False + + CSMF_params: + Mstar_low: 10.4 + Mstar_up: 12.1 + M_0: 9.935 + M_1: 11.07 + gamma_1: 3.273 + gamma_2: 0.255 + sigma_c: 0.143 + a_1: 0.501 + a_2: 0 + M_2: 14.28 + b_0: -0.766 + b_1: 1.008 + b_2: 0 + delta_1: 0.0 + delta_2: 0.0 + alpha_c: 0.0 + alpha_s: 0.0 + s: 0 + s_v: 0 + s_p: 0 + s_r: 0 + Acent: 0 + Asat: 0 + Bcent: 0 + Bsat: 0 + ic: 1.0 + +fit_params: + M_0: [0, 'CSMF'] + M_1: [1, 'CSMF'] + gamma_1: [2, 'CSMF'] + gamma_2: [3, 'CSMF'] + sigma_c: [4, 'CSMF'] + a_1: [5, 'CSMF'] + b_0: [6, 'CSMF'] + b_1: [7, 'CSMF'] + s: [8, 'CSMF'] + Bcent: [9, 'CSMF'] + Bsat: [10, 'CSMF'] + + +# =========================================================================================== +# The following dictionaries are for using the provided emcee/likelihood templates. They +# are not needed for using the AbacusHOD class and running HODs. + +# parameters that indicate where the observed data that you would like to fit live. +data_params: + + + +# Configuration parameters for emcee +ch_config_params: + + diff --git a/tests/data_CSMF/halos_data/.DS_Store b/tests/data_CSMF/halos_data/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6 GIT binary patch literal 6148 zcmeH~Jr2S!425mzP>H1@V-^m;4Wg<&0T*E43hX&L&p$$qDprKhvt+--jT7}7np#A3 zem<@ulZcFPQ@L2!n>{z**++&mCkOWA81W14cNZlEfg7;MkzE(HCqgga^y>{tEnwC%0;vJ&^%eQ zLs35+`xjp>T0 Date: Tue, 15 Oct 2024 14:54:50 +0000 Subject: [PATCH 13/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/CSMF_test.ipynb | 182 +++++++++++++++++++++---------- tests/data_CSMF/config_CSMF.yaml | 26 ++--- 2 files changed, 135 insertions(+), 73 deletions(-) diff --git a/tests/CSMF_test.ipynb b/tests/CSMF_test.ipynb index aa960f4a..a33d3d46 100644 --- a/tests/CSMF_test.ipynb +++ b/tests/CSMF_test.ipynb @@ -16,22 +16,19 @@ "metadata": {}, "outputs": [], "source": [ - "import time\n", "import yaml\n", "import numpy as np\n", "from abacusnbody.hod.abacus_hod import AbacusHOD\n", - "from astropy.table import Table, vstack\n", + "from astropy.table import Table\n", "from cosmoprimo.utils import DistanceToRedshift\n", "from pyrecon import utils\n", "from cosmoprimo.fiducial import AbacusSummit\n", - "from scipy.interpolate import InterpolatedUnivariateSpline\n", "import logging\n", "import warnings\n", - "from tqdm import tqdm \n", "from astropy import units as u\n", "\n", - "warnings.filterwarnings(\"ignore\", category=np.VisibleDeprecationWarning)\n", - "logger = logging.getLogger(\"ds_abacus_lightcone\")" + "warnings.filterwarnings('ignore', category=np.VisibleDeprecationWarning)\n", + "logger = logging.getLogger('ds_abacus_lightcone')" ] }, { @@ -56,18 +53,26 @@ " for key in param_mapping.keys():\n", " mapping_idx = param_mapping[key]\n", " tracer_type = param_tracer[key]\n", - " if key == \"sigma\" and tracer_type == \"LRG\":\n", + " if key == 'sigma' and tracer_type == 'LRG':\n", " Ball.tracers[tracer_type][key] = 10 ** p[mapping_idx]\n", - " elif key == \"M_0\" and tracer_type == \"CSMF\":\n", - " Ball.tracers[tracer_type][key] = 10 ** p[mapping_idx] \n", - " elif key == \"M_1\" and tracer_type == \"CSMF\":\n", - " Ball.tracers[tracer_type][key] = 10 ** p[mapping_idx] \n", - " elif key == \"M_2\" and tracer_type == \"CSMF\":\n", + " elif key == 'M_0' and tracer_type == 'CSMF':\n", + " Ball.tracers[tracer_type][key] = 10 ** p[mapping_idx]\n", + " elif key == 'M_1' and tracer_type == 'CSMF':\n", + " Ball.tracers[tracer_type][key] = 10 ** p[mapping_idx]\n", + " elif key == 'M_2' and tracer_type == 'CSMF':\n", " Ball.tracers[tracer_type][key] = 10 ** p[mapping_idx]\n", " else:\n", " Ball.tracers[tracer_type][key] = p[mapping_idx]\n", - " \n", - " mock_dict = Ball.run_hod(tracers=Ball.tracers, want_rsd=Ball.want_rsd, want_nfw = False, NFW_draw=None, Nthread=nthread, reseed=seed, verbose=True)\n", + "\n", + " mock_dict = Ball.run_hod(\n", + " tracers=Ball.tracers,\n", + " want_rsd=Ball.want_rsd,\n", + " want_nfw=False,\n", + " NFW_draw=None,\n", + " Nthread=nthread,\n", + " reseed=seed,\n", + " verbose=True,\n", + " )\n", " return mock_dict" ] }, @@ -88,19 +93,23 @@ "source": [ "def setup_hod(config):\n", " print(f\"Processing {config['sim_params']['sim_name']}\")\n", - " sim_params = config[\"sim_params\"]\n", - " HOD_params = config[\"HOD_params\"]\n", - " fit_params = config[\"fit_params\"]\n", + " sim_params = config['sim_params']\n", + " HOD_params = config['HOD_params']\n", + " fit_params = config['fit_params']\n", "\n", - " if(HOD_params['CSMF_params']['Mstar_low']<100):\n", - " HOD_params['CSMF_params']['Mstar_low'] = 10**HOD_params['CSMF_params']['Mstar_low']\n", - " HOD_params['CSMF_params']['Mstar_up'] = 10**HOD_params['CSMF_params']['Mstar_up']\n", + " if HOD_params['CSMF_params']['Mstar_low'] < 100:\n", + " HOD_params['CSMF_params']['Mstar_low'] = (\n", + " 10 ** HOD_params['CSMF_params']['Mstar_low']\n", + " )\n", + " HOD_params['CSMF_params']['Mstar_up'] = (\n", + " 10 ** HOD_params['CSMF_params']['Mstar_up']\n", + " )\n", "\n", " Balls = []\n", " for ez in zranges:\n", - " sim_params[\"z_mock\"] = ez\n", + " sim_params['z_mock'] = ez\n", " Balls += [AbacusHOD(sim_params, HOD_params)]\n", - " \n", + "\n", " param_mapping = {}\n", " param_tracer = {}\n", " for key in fit_params.keys():\n", @@ -108,7 +117,7 @@ " tracer_type = fit_params[key][-1]\n", " param_mapping[key] = mapping_idx\n", " param_tracer[key] = tracer_type\n", - " \n", + "\n", " return Balls, param_mapping, param_tracer" ] }, @@ -128,19 +137,19 @@ "outputs": [], "source": [ "args_dict = {\n", - " \"start_hod\":0,\n", - " \"n_hod\":1,\n", - " \"start_cosmo\":0,\n", - " \"n_cosmo\":1,\n", - " \"start_phase\":0,\n", - " \"n_phase\":1,\n", - " \"survey\":\"DESI\",\n", - " \"tracer\":\"CSMF\",\n", - " \"zmin\":0.1,\n", - " \"zmax\":0.4,\n", - " \"start_seed\":0,\n", - " \"n_seed\":1,\n", - " \"nthreads\":5\n", + " 'start_hod': 0,\n", + " 'n_hod': 1,\n", + " 'start_cosmo': 0,\n", + " 'n_cosmo': 1,\n", + " 'start_phase': 0,\n", + " 'n_phase': 1,\n", + " 'survey': 'DESI',\n", + " 'tracer': 'CSMF',\n", + " 'zmin': 0.1,\n", + " 'zmax': 0.4,\n", + " 'start_seed': 0,\n", + " 'n_seed': 1,\n", + " 'nthreads': 5,\n", "}" ] }, @@ -159,7 +168,7 @@ "metadata": {}, "outputs": [], "source": [ - "config = yaml.safe_load(open(\"data_CSMF/config_CSMF.yaml\"))\n", + "config = yaml.safe_load(open('data_CSMF/config_CSMF.yaml'))\n", "zranges = [0.100]" ] }, @@ -211,8 +220,15 @@ ], "source": [ "import gdown\n", - "gdown.download(id='1VCBXlkYHChn1HCwRPnq9iDSgzILfXfwL',output='data_CSMF/halos_data/AbacusSummit_base_c000_ph000/z0.100/particles_xcom_0_seed600_abacushod_oldfenv_MT_withranks_new.h5')\n", - "gdown.download(id='1xXlbXr8h8AiRFc2p5knYSYloJ9jfcAY9',output='data_CSMF/halos_data/AbacusSummit_base_c000_ph000/z0.100/halos_xcom_0_seed600_abacushod_oldfenv_MT_new.h5')" + "\n", + "gdown.download(\n", + " id='1VCBXlkYHChn1HCwRPnq9iDSgzILfXfwL',\n", + " output='data_CSMF/halos_data/AbacusSummit_base_c000_ph000/z0.100/particles_xcom_0_seed600_abacushod_oldfenv_MT_withranks_new.h5',\n", + ")\n", + "gdown.download(\n", + " id='1xXlbXr8h8AiRFc2p5knYSYloJ9jfcAY9',\n", + " output='data_CSMF/halos_data/AbacusSummit_base_c000_ph000/z0.100/halos_xcom_0_seed600_abacushod_oldfenv_MT_new.h5',\n", + ")" ] }, { @@ -247,7 +263,19 @@ "hod = 0\n", "\n", "HOD_params = Table.read('data_CSMF/HOD_parameters.fits')\n", - "names = ['log_M_0', 'log_M_1', 'gamma_1', 'gamma_2', 'sigma_c','a_1', 'b_0', 'b_1', 's', 'Bcent', 'Bsat']\n", + "names = [\n", + " 'log_M_0',\n", + " 'log_M_1',\n", + " 'gamma_1',\n", + " 'gamma_2',\n", + " 'sigma_c',\n", + " 'a_1',\n", + " 'b_0',\n", + " 'b_1',\n", + " 's',\n", + " 'Bcent',\n", + " 'Bsat',\n", + "]\n", "\n", "HOD_params_np = []\n", "for name in names:\n", @@ -292,31 +320,62 @@ "\n", "data_positions_sky = []\n", "for i, newBall in enumerate(Balls):\n", - "\n", - " hod_dict = get_hod(HOD_params_np[hod],param_mapping,param_tracer,newBall,args_dict['nthreads'],seed)\n", + " hod_dict = get_hod(\n", + " HOD_params_np[hod],\n", + " param_mapping,\n", + " param_tracer,\n", + " newBall,\n", + " args_dict['nthreads'],\n", + " seed,\n", + " )\n", " data = hod_dict[args_dict['tracer']]\n", "\n", - " dist, ra, dec = utils.cartesian_to_sky(np.c_[data['x'] + 990, data['y'] + 990, data['z'] + 990])\n", + " dist, ra, dec = utils.cartesian_to_sky(\n", + " np.c_[data['x'] + 990, data['y'] + 990, data['z'] + 990]\n", + " )\n", " d2z = DistanceToRedshift(mock_cosmo.comoving_radial_distance)\n", " redshift = d2z(dist)\n", "\n", - " mask = (redshift >= args_dict['zmin']) & (redshift <= args_dict['zmax']) & (data['stellarmass']>0)\n", + " mask = (\n", + " (redshift >= args_dict['zmin'])\n", + " & (redshift <= args_dict['zmax'])\n", + " & (data['stellarmass'] > 0)\n", + " )\n", " central_id = np.zeros_like(data['stellarmass'])\n", - " central_id[:data['Ncent']] = 1\n", + " central_id[: data['Ncent']] = 1\n", "\n", - " data_positions_sky.append(np.c_[ra[mask],dec[mask],dist[mask],redshift[mask],data['stellarmass'][mask],data['mass'][mask],central_id[mask]])\n", + " data_positions_sky.append(\n", + " np.c_[\n", + " ra[mask],\n", + " dec[mask],\n", + " dist[mask],\n", + " redshift[mask],\n", + " data['stellarmass'][mask],\n", + " data['mass'][mask],\n", + " central_id[mask],\n", + " ]\n", + " )\n", "data_positions_sky = np.concatenate(data_positions_sky, axis=0)\n", "\n", "HOD_table = Table()\n", - "HOD_table['ra'] = data_positions_sky[:,0].astype(np.float32)\n", - "HOD_table['dec'] = data_positions_sky[:,1].astype(np.float32)\n", - "HOD_table['comoving_distance_overh'] = data_positions_sky[:,2].astype(np.float32)*u.Mpc\n", - "HOD_table['redshift'] = data_positions_sky[:,3].astype(np.float32)\n", - "HOD_table['log10Mstar_overhsquare'] = np.log10(data_positions_sky[:,4]).astype(np.float64)*u.Msun\n", - "HOD_table['log10Mhalo_overh'] = np.log10(data_positions_sky[:,5]).astype(np.float64)*u.Msun\n", - "HOD_table['central flag'] = data_positions_sky[:,6].astype(np.int32)\n", + "HOD_table['ra'] = data_positions_sky[:, 0].astype(np.float32)\n", + "HOD_table['dec'] = data_positions_sky[:, 1].astype(np.float32)\n", + "HOD_table['comoving_distance_overh'] = (\n", + " data_positions_sky[:, 2].astype(np.float32) * u.Mpc\n", + ")\n", + "HOD_table['redshift'] = data_positions_sky[:, 3].astype(np.float32)\n", + "HOD_table['log10Mstar_overhsquare'] = (\n", + " np.log10(data_positions_sky[:, 4]).astype(np.float64) * u.Msun\n", + ")\n", + "HOD_table['log10Mhalo_overh'] = (\n", + " np.log10(data_positions_sky[:, 5]).astype(np.float64) * u.Msun\n", + ")\n", + "HOD_table['central flag'] = data_positions_sky[:, 6].astype(np.int32)\n", "\n", - "HOD_table.write(f'data_CSMF/galaxy_catalogue_cos{cosmo:03}_ph{phase:03}_hod{hod:03}_seed{seed:03}.fits',overwrite=True)" + "HOD_table.write(\n", + " f'data_CSMF/galaxy_catalogue_cos{cosmo:03}_ph{phase:03}_hod{hod:03}_seed{seed:03}.fits',\n", + " overwrite=True,\n", + ")" ] }, { @@ -437,15 +496,20 @@ "source": [ "from common import check_close\n", "\n", - "gdown.download(id='1JB5UvLM13wj6IdRvF37-Dh8kRjJWzWh0',output=f'data_CSMF/reference_galaxy_catalogue_cos{cosmo:03}_ph{phase:03}_hod{hod:03}_seed{seed:03}.fits')\n", - "HOD_table_ref = Table.read(f'data_CSMF/reference_galaxy_catalogue_cos{cosmo:03}_ph{phase:03}_hod{hod:03}_seed{seed:03}.fits')\n", + "gdown.download(\n", + " id='1JB5UvLM13wj6IdRvF37-Dh8kRjJWzWh0',\n", + " output=f'data_CSMF/reference_galaxy_catalogue_cos{cosmo:03}_ph{phase:03}_hod{hod:03}_seed{seed:03}.fits',\n", + ")\n", + "HOD_table_ref = Table.read(\n", + " f'data_CSMF/reference_galaxy_catalogue_cos{cosmo:03}_ph{phase:03}_hod{hod:03}_seed{seed:03}.fits'\n", + ")\n", "\n", "\n", "for name in HOD_table_ref.colnames:\n", - " if(check_close(HOD_table_ref[name],HOD_table[name])):\n", - " print('\\033[92m \"'+name + r'\" is close to the reference catalogue')\n", + " if check_close(HOD_table_ref[name], HOD_table[name]):\n", + " print('\\033[92m \"' + name + r'\" is close to the reference catalogue')\n", " else:\n", - " print('\\033[91m \"'+name + r'\" is different to the reference catalogue')" + " print('\\033[91m \"' + name + r'\" is different to the reference catalogue')" ] } ], diff --git a/tests/data_CSMF/config_CSMF.yaml b/tests/data_CSMF/config_CSMF.yaml index fad2a64b..a2f1a9ad 100644 --- a/tests/data_CSMF/config_CSMF.yaml +++ b/tests/data_CSMF/config_CSMF.yaml @@ -7,31 +7,31 @@ sim_params: sim_dir: '/global/cfs/cdirs/desi/public/cosmosim/AbacusSummit/halo_light_cones/' # where is the simulation subsample_dir: 'data_CSMF/halos_data/' cleaned_halos: 'data_CSMF/halos_data/' - + # subsample_dir: '/pscratch/sd/p/pierre/AbacusSummit/subsamples/' # cleaned_halos: '/pscratch/sd/p/pierre/AbacusSummit/subsamples/' - + Nthread_load: 10 - z_mock: 0.100 + z_mock: 0.100 halo_lc: True # subsample_dir: ‘/pscratch/sd/b/boryanah/AbacusHOD_scratch/mocks_lc_subsample/’ prepare_sim: - Nparallel_load: 10 # best for cori haswell node + Nparallel_load: 10 # best for cori haswell node # HOD parameters HOD_params: use_particles: True want_ranks: True # want to implement satellite profile flexiblities? - want_AB: True # want to implement assembly bias? + want_AB: True # want to implement assembly bias? density_sigma: 3 # scale radius in Mpc / h for local density calculation tracer_flags: # which tracer do we want? - LRG: False + LRG: False ELG: False QSO: False CSMF: True - want_rsd: True # want RSD? + want_rsd: True # want RSD? write_to_disk: False CSMF_params: @@ -75,17 +75,15 @@ fit_params: Bcent: [9, 'CSMF'] Bsat: [10, 'CSMF'] - + # =========================================================================================== -# The following dictionaries are for using the provided emcee/likelihood templates. They -# are not needed for using the AbacusHOD class and running HODs. +# The following dictionaries are for using the provided emcee/likelihood templates. They +# are not needed for using the AbacusHOD class and running HODs. -# parameters that indicate where the observed data that you would like to fit live. +# parameters that indicate where the observed data that you would like to fit live. data_params: - + # Configuration parameters for emcee ch_config_params: - -