import warnings
import numpy as np
from scipy.interpolate import InterpolatedUnivariateSpline as IUS
from astropy import units
from astropy import constants
from astropy.cosmology import z_at_value
from astropy.table import Table
from frb.halos.models import ModifiedNFW, ICM
from frb.defs import frb_cosmo as cosmo
from IPython import embed
[docs]
def init_hmf():
"""
Initialize the Aemulus Halo Mass Function
WARNING: This uses the original version which codes Tinker+2008
We may refactor to use the more accurate, new version
Returns:
hmfe (hmf_emulator.hmf_emulator): An Aemulus halo mass function emulator.
"""
# Hidden here to avoid it becoming a dependency
import hmf_emulator
# Setup HMF
# https://github.com/astropy/astropy/blob/master/astropy/cosmology/parameters.py
#sigma8 = 0.8159
ns = 0.9667
Neff = 3.046
#cosmo_dict = {"om":cosmo.Om0,"ob":cosmo.Ob0,"ol":1.-cosmo.Om0,"ok":0.0,
# "h":cosmo.h,"s8":sigma8,"ns":ns,"w0":-1.0,"Neff":Neff} # "wa":0.0 is assumed internally
cosmo_dict = {"omega_cdm":(cosmo.Om0-cosmo.Ob0)*cosmo.h**2,
"omega_b":cosmo.Ob0*cosmo.h**2,"ok":0.0,
"ln10As": 3.098, # THIS REPLACES sigma8
"H0":cosmo.H0.to('km/(s*Mpc)').value,
"n_s":ns,"w0":-1.0,"N_eff":Neff} # "wa":0.0 is assumed internally
hmfe = hmf_emulator.hmf_emulator()
hmfe.set_cosmology(cosmo_dict)
# Return
return hmfe
# Storing for use
# Needs to come after def init_hmf()
try:
import hmf_emulator
except:
warnings.warn("hmf_emulator not imported. Hope you are not intending to use the hmf.py module..")
else:
hmfe = init_hmf()
[docs]
def frac_in_halos(zvals, Mlow, Mhigh, rmax=1.):
"""
Calculate the fraction of matter in collapsed halos
over a mass range and at a given redshift
Note that the fraction of DM associated with these halos
will be scaled down by an additional factor of f_diffuse
Requires Aemulus HMF to be installed
Args:
zvals: ndarray
Mlow: float
In h^-1 units already so this will be applied for the halo mass function
Mhigh: float
In h^-1 units already
rmax: float
Extent of the halo in units of rvir
WARNING: This calculation assumes a single concentration for all halos
Returns:
ratios: ndarray
rho_halo / rho_m
"""
# Deal with scalar input
if np.isscalar(zvals):
zvals = np.atleast_1d([zvals])
# Cheeky edge case
if np.isclose(Mlow, Mhigh, rtol=1e-05):
return np.zeros(len(zvals))
M = np.logspace(np.log10(Mlow*cosmo.h), np.log10(Mhigh*cosmo.h), num=1000)
lM = np.log(M)
ratios = []
for z in zvals:
# Setup
#dndlM = np.array([hmfe.dndlnM(Mi, a)[0] for Mi in M])
dndlM = M*hmfe.dndM(M, z)
M_spl = IUS(lM, M * dndlM)
# Integrate
rho_tot = M_spl.integral(np.log(Mlow*cosmo.h), np.log(Mhigh*cosmo.h)) * units.M_sun / units.Mpc ** 3
# Cosmology
rho_M = cosmo.critical_density(z) * cosmo.Om(z)/(1+z)**3 # Tinker calculations are all mass
ratio = (rho_tot*cosmo.h**2 / rho_M).decompose()
#
ratios.append(ratio)
ratios = np.array(ratios)
# Boost halos if extend beyond rvir (homologous in mass, but constant concentration is an approx)
if rmax != 1.:
#from pyigm.cgm.models import ModifiedNFW
c = 7.7
nfw = ModifiedNFW(c=c)
M_ratio = nfw.fy_dm(rmax * nfw.c) / nfw.fy_dm(nfw.c)
ratios *= M_ratio
# Return
return np.array(ratios)
[docs]
def halo_incidence(Mlow, zFRB, radius=None, hmfe=None,
Mhigh=1e16, nsample=20, cumul=False):
"""
Calculate the (approximate) average number of
intersections to halos of a
given minimum mass to a given zFRB.
Requires Aemulus HMF to be installed
Args:
Mlow: float
Mass of minimum halo in Solar masses
The code deals with h^-1 factors so that you do not
The minimum value is 2e10
zFRB: float
Redshift of the FRB
radius: Quantity, optional
Physical separation from the sightline for the calculation.
The calculation will specify this radius as rvir derived from
Mlow unless this is specified. And this rvir *will* vary with redshift
hmfe (hmf.hmf_emulator, optional): Halo mass function emulator from Aeumulus
Mhigh: float, optional
Mass of maximum halo in Solar masses
nsammple: int, optional
Number of samplings in redshift
20 should be enough
cumul: bool, optional
Return the cumulative quantities instead
Returns:
If cumul is False
Navg: float
Number of average intersections
elif cumul is True
zeval: ndarray
Ncumul: ndarray
"""
# Mlow limit
if Mlow < 2e10:
raise IOError("Calculations are limited to Mlow > 2e10")
# HMF
if hmfe is None:
hmfe = init_hmf()
#
zs = np.linspace(0., zFRB, nsample)
# Mean density
ns = []
for iz in zs:
ins = hmfe.n_in_bins((Mlow * cosmo.h, Mhigh * cosmo.h), iz)
ns.append(ins*cosmo.h**3) # * units.Mpc**-3
# Interpolate
ns = units.Quantity(ns*units.Mpc**-3)
# Radii
if radius is None:
rhoc = cosmo.critical_density(zs)
#https://arxiv.org/pdf/1312.4629.pdf eq5
q = cosmo.Ode0/(cosmo.Ode0+cosmo.Om0*(1+zs)**3)
rhovir = (18*np.pi**2-82*q-39*q**2)*rhoc
r200 = (((3*Mlow*constants.M_sun.cgs) / (4*np.pi*rhovir))**(1/3)).to('kpc')
else:
r200 = np.ones_like(zs) * radius
# Ap
Ap = np.pi * r200**2
# l(X)
loX = ((constants.c/cosmo.H0) * ns * Ap).decompose().value
# dX
X = cosmo.absorption_distance(zs)
dX = X - np.roll(X,1)
dX[0] = 0.
# Finish
if cumul:
Navg = np.cumsum(loX * dX)
return zs, Navg
else:
Navg = np.sum(loX * dX)
return Navg
[docs]
def build_grid(z_FRB=1., ntrial=10, seed=12345, Mlow=1e10, r_max=2., outfile=None, dz_box=0.1,
dz_grid=0.01, f_hot=0.75, verbose=True):
"""
Generate a universe of dark matter halos with DM measurements
Mainly an internal function for generating useful output grids.
Requires the Aemulus Halo Mass function
Args:
z_FRB: float, optional
ntrial: int, optional
seed: int, optional
Mlow: float, optional
h^-1 mass
r_max: float, optional
Extent of the halo in units of rvir
outfile: str, optional
Write
dz_box: float, optional
Size of the slice of the universe for each sub-calculation
dz_grid: float, optional
redshift spacing in the DM grid
f_hot: float
Fraction of the cosmic fraction of matter in diffuse gas (for DM)
Returns:
DM_grid: ndarray (ntrial, nz)
halo_tbl: Table
Table of all the halos intersected
"""
Mhigh = 1e16 # Msun
# mNFW
y0 = 2.
alpha = 2.
warnings.warn("Ought to do concentration properly someday!")
cgm = ModifiedNFW(alpha=alpha, y0=y0, f_hot=f_hot)
icm = ICM()
# Random numbers
rstate = np.random.RandomState(seed)
# Init HMF
hmfe = init_hmf()
# Boxes
nbox = int(z_FRB / dz_box)
nz = int(z_FRB / dz_grid)
dX = int(np.sqrt(ntrial))+1
#
npad = 6 # Mpc
base_l = 2*dX + npad
print('L_base = {} cMpc'.format(base_l))
warnings.warn("Worry about being big enough given cMpc vs pMpc")
DM_grid = np.zeros((ntrial,nz))
# Spline distance to z
D_max = cosmo.comoving_distance(z_FRB)
D_val = np.linspace(1e-3,D_max.value,200) # IS THIS FINE ENOUGH?
z_val = np.array([z_at_value(cosmo.comoving_distance, iz) for iz in D_val*units.Mpc])
D_to_z = IUS(D_val, z_val)
# Save halo info
#halos = [[] for i in range(ntrial)]
halo_i, M_i, R_i, DM_i, z_i = [], [], [], [], []
# Loop me
prev_zbox = 0.
#for ss in range(nbox):
#for ss in [0]:
for ss in [5]:
zbox = ss*dz_box + dz_box/2.
print('zbox = {}'.format(zbox))
a = 1./(1.0 + zbox) # Scale factor
# Mass function
M = np.logspace(np.log10(Mlow*cosmo.h), np.log10(Mhigh*cosmo.h), num=1000)
lM = np.log(M)
dndlM = np.array([hmf.dndlM(Mi, a) for Mi in M])
n_spl = IUS(lM, dndlM)
cum_n = np.array([n_spl.integral(np.log(Mlow*cosmo.h), ilM) for ilM in lM])
ncum_n = cum_n/cum_n[-1]
# As z increases, we have numerical issues at the high mass end (they are too rare)
try:
mhalo_spl = IUS(ncum_n, lM)
except ValueError:
# Kludge me
print("REDUCING Mhigh by 2x")
Mhigh /= 2.
M = np.logspace(np.log10(Mlow*cosmo.h), np.log10(Mhigh*cosmo.h), num=1000)
lM = np.log(M)
dndlM = np.array([hmf.dndlM(Mi, a) for Mi in M])
n_spl = IUS(lM, dndlM)
cum_n = np.array([n_spl.integral(np.log(Mlow*cosmo.h), ilM) for ilM in lM])
ncum_n = cum_n/cum_n[-1]
#
mhalo_spl = IUS(ncum_n, lM)
# Volume -- Box with base l = 2Mpc
D_zn = cosmo.comoving_distance(zbox + dz_box/2.) # Full box
D_zp = cosmo.comoving_distance(ss*dz_box) # Previous
D_z = D_zn - D_zp
V = D_z * (base_l*units.Mpc)**2
# Average N_halo
avg_n = hmf.n_bin(Mlow*cosmo.h, Mhigh*cosmo.h, a) * cosmo.h**3 * units.Mpc**-3
avg_N = (V * avg_n).value
# Assume Gaussian stats for number of halos
N_halo = int(np.round(avg_N + np.sqrt(avg_N)*rstate.randn(1)))
# Random masses
randM = rstate.random_sample(N_halo)
rM = np.exp(mhalo_spl(randM)) / cosmo.h
# r200
r200 = (((3*rM*units.M_sun.cgs) / (4*np.pi*200*cosmo.critical_density(zbox)))**(1/3)).to('kpc')
# Random locations (X,Y,Z)
X_c = rstate.random_sample(N_halo)*base_l # Mpc
Y_c = rstate.random_sample(N_halo)*base_l # Mpc
Z_c = (rstate.random_sample(N_halo)*D_z.to('Mpc') + D_zp).value
# Check mass fraction
if verbose:
Mtot = np.log10(np.sum(rM))
M_m = (cosmo.critical_density(zbox)*cosmo.Om(zbox) * V/(1+zbox)**3).to('M_sun')
#print("N_halo: {} avg_N: {}".format(N_halo, avg_N))
print("z: {} Mhalo/M_m = {}".format(zbox, 10**Mtot/M_m.value))
print(frac_in_halos([zbox], Mlow, Mhigh))
# Redshifts
z_ran = D_to_z(Z_c)
# Loop on trials
all_DMs = []
all_nhalo = []
all_r200 = []
for itrial in range(ntrial):
# X,Y trial
X_trial = npad//2 + (2*itrial%dX) # Step by 2Mpc
Y_trial = npad//2 + 2*itrial // dX
# Impact parameters
try:
R_com = np.sqrt((X_c-X_trial)**2 + (Y_c-Y_trial)**2) # Mpc
except:
embed()
R_phys = R_com * 1000. / (1+z_ran) * units.kpc
# Cut
intersect = R_phys < r_max*r200
print("We hit {} halos".format(np.sum(intersect)))
all_nhalo.append(np.sum(intersect))
if not np.any(intersect):
all_DMs.append(0.)
continue
# Loop -- FIND A WAY TO SPEED THIS UP!
DMs = []
for iobj in np.where(intersect)[0]:
# Init
if rM[iobj] > 1e14: # Use ICM model
model = icm
else:
model = cgm
model.log_Mhalo=np.log10(rM[iobj])
model.M_halo = 10.**model.log_Mhalo * constants.M_sun.cgs
model.z = zbox # To be consistent with above; should be close enough
model.setup_param(cosmo=cosmo)
# DM
DM = model.Ne_Rperp(R_phys[iobj], rmax=r_max, add_units=False)/(1+model.z)
DMs.append(DM)
# Save halo info
halo_i.append(itrial)
M_i.append(model.M_halo.value)
R_i.append(R_phys[iobj].value)
DM_i.append(DM)
z_i.append(z_ran[iobj])
all_r200.append(cgm.r200.value)
# Save em
iz = (z_ran[intersect]/dz_grid).astype(int)
DM_grid[itrial,iz] += DMs
all_DMs.append(np.sum(DMs))
#print(DMs, np.log10(rM[intersect]), R_phys[intersect])
if (itrial % 100) == 0:
embed()
# Table the halos
halo_tbl = Table()
halo_tbl['trial'] = halo_i
halo_tbl['M'] = M_i
halo_tbl['R'] = R_i
halo_tbl['DM'] = DM_i
halo_tbl['z'] = z_i
# Write
if outfile is not None:
print("Writing to {}".format(outfile))
np.save(outfile, DM_grid, allow_pickle=False)
halo_tbl.write(outfile+'.fits', overwrite=True)
return DM_grid, halo_tbl