import numpy as np
import jax.numpy as jnp
from jax import config, devices, device_put
from jax.scipy.ndimage import map_coordinates
from .cosmology import TCMB, me, mu_e, c, hbar, kB, not4
config.update("jax_enable_x64", True)
import os
file_dir = os.path.dirname(__file__)
#Tabulated values of effective recombination coefficients to interpolate.
R_tab = jnp.array(np.loadtxt(file_dir+"/tabs/R_inf.dat"))
#Tabulated values of 2s-2p transition rates to interpolate.
alpha_tab = jnp.array(np.loadtxt(file_dir+"/tabs/Alpha_inf.dat"))
try:
gpus = devices('gpu')
R_tab = device_put(
R_tab, device=gpus[0])
alpha_tab = device_put(
alpha_tab, device=gpus[0])
except:
pass
# File handling and interpolating related constants.
# Do not change these unless something about the tabulated files have changed.
TR_MIN = 0.004 # Lower bound of radiation temperature axis, in eV.
TR_MAX = 0.4 # Upper bound of radiation temperature axis, in eV.
NTR = 100 # Number of points in radiation temperature axis.
T_RATIO_MIN = 0.1 # Lower bound of temperature ratio axis, defined as min(Tm/Tr, Tr/Tm).
T_RATIO_MAX = 1.0 # Upper bound of temperature ratio axis, defined as min(Tm/Tr, Tr/Tm).
NTM = 40 # Number of points in temperature ratio axis
TR_axis = jnp.linspace(jnp.log(TR_MIN), jnp.log(TR_MAX), NTR)
T_RATIO_axis = jnp.linspace(T_RATIO_MIN, T_RATIO_MAX, NTM)
A2s_table = alpha_tab.reshape((NTR, NTM, 4))
A2p_table = alpha_tab.reshape((NTR, NTM, 4))
E21 = jnp.float64(10.198714553953742) # Energy difference in n=1, 2 for hydrogen, in eV.
E31 = jnp.float64(12.087365397278509) # Energy difference in n=1, 3 for hydrogen, in eV.
E41 = jnp.float64(12.748393192442178) # Energy difference in n=1, 4 for hydrogen, in eV.
E32 = jnp.float64(1.8886508433247664) # Energy difference in n=2, 3 for hydrogen, in eV.
E42 = jnp.float64(2.5496786384884356) # Energy difference in n=2, 4 for hydrogen, in eV.
rydberg = jnp.float64(13.598286071938324) # Ionization energy of hydrogen, in eV
lya_eng = rydberg*3./4. # Lyman-alpha transition energy, in eV
lya_freq = lya_eng / (2.*jnp.pi*hbar) # Lyman-alpha transition frequency, in s^{-1}
thomson_xsec = jnp.float64(6.652458734e-25) # Thomson cross section, in cm^2
stef_bolt = jnp.pi**2 / (60.*hbar**3*c**2) # Stefan-Boltzmann constant, eV^{-3} cm^{-2} s^{-1}
R2s1s = jnp.float64(8.2206) # 2s to 1s transition rate, in s^{-1}
[docs]
def alpha_H(Tm):
"""
Case-B recombination coefficient for hydrogen.
Dimensions: cm^3 s^{-1}
WARNING: Assumes input Tm to be in eV.
Parameters:
-----------
Tm : float
Matter temperature.
Returns:
--------
alpha_H : float
Case B recombination coefficient for hydrogen at given matter temperature.
"""
fudge_factor = 1.0 # Recommended by arXiv:1011.3758
T_in_K = Tm / kB # Convert matter temperature to Kelvin
# Compute the coefficient in cm^3/s
alpha_H = fudge_factor * 4.309e-13 * (1.0e-4 * T_in_K)**(-0.6166) \
/ (1.+0.6703*(1.0e-4*T_in_K)**0.5300)
return alpha_H
[docs]
def beta_H(TCMB):
"""
Case-B photoionization coefficient for hydrogen.
Dimensions: s^{-1}
WARNING: Assumes input TCMB to be in eV.
Parameters:
-----------
TCMB : float
CMB temperature.
Returns:
--------
beta_H : float
Case B photoionization coefficient for hydrogen at given CMB (radiation) temperature.
"""
# Electron de Broglie wavelength at given CMB temperature, in cm (keep in mind h=2pi*hbar)
lambda_dB = 2.*jnp.pi*hbar*c / jnp.sqrt(2.*jnp.pi*mu_e*TCMB)
beta_H = (1./lambda_dB)**3 * jnp.exp(-rydberg/4./TCMB) * alpha_H(TCMB) / 4.
return beta_H
[docs]
def peebles_C(z, xHII, H, nH):
"""
Peebles C factor, probability for an n=2 hydrogen atom to reach the ground state before becoming photoionized, a function of redshift and ionized proton fraction.
Dimensions: None
Parameters:
-----------
z : float/jnp.array
Requested redshift(s).
xHII : float
Ionized hydrogen fraction.
H : float/jnp.array
Value(s) of Hubble at given redshift(s).
nH : float/jnp.array
Value(s) of hydrogen number density at given redshift(s)
Returns:
--------
C : float/jnp.array
Peebles C factor.
"""
# (2p to 1s rate) x (1 - xHII), in s^{-1}
rate_2p1s_times_x1s = 8*jnp.pi*H / (3 * (nH*(c/lya_freq)**3))
# s^{-1}
rate_exc = 3. * rate_2p1s_times_x1s/4. + (1.-xHII) * R2s1s/4.
# Ionization rate, in s^{-1}
rate_ion = (1-xHII) * beta_H(TCMB(z))
return rate_exc / (rate_exc + rate_ion)
[docs]
def Gamma_compton(xe, TCMB, YHe):
"""
Computes the Compton scattering rate. See Eq.(2) of 1904.09296
Dimensions: s^{-1}
Parameters:
-----------
xe : float
Free electron fraction.
TCMB : float
CMB temperature, in eV.
YHe : float
Helium mass fraction.
Returns:
--------
GammaC : float
The Compton scattering rate
"""
# Helium to Hydrogen number density ratio. Here we assumed that all Helium are Helium-4.
# 1) n_H = (1-YHe)*n_b
# 2) n_He = (YHe/4)*n_b
FHe = YHe/not4/(1-YHe)
GammaC = xe / (1.+FHe+xe) * 8.*thomson_xsec*(4.*stef_bolt)*TCMB**4 / (3.*me)
return GammaC
[docs]
def xe_Saha(TCMB, nH):
"""
Computes the free electron fraction in the Saha equilibrium approximation, valid
at early times when matter and photons are in Chemical Equilibrium.
Parameters:
-----------
TCMB : float
CMB temperature, in eV
nH : float
Current Hydrogen number density, in cm^{-3}
Returns:
--------
xe_saha : float
Saha equilibrium prediction of free electron fraction.
s : float
xe_saha^2/(1-xe_saha)
"""
ge = (2.*jnp.pi*mu_e*TCMB)**(3./2.) / (2.*jnp.pi*hbar*c)**3. / nH
s = ge * jnp.exp(-rydberg/TCMB)
xe_saha = (jnp.sqrt(s**2.+4.*s) - s) / 2.
return xe_saha, s
[docs]
def effective_coefficients(TCMB, Tm, H, nH, x1s):
"""
Computes the effective coefficients for, Case-B recombination (A2s, A2p), Case-B photoionization (B2s, B2p),
and Generalized Peebles factors (C2s, C2p) in the effetive four-level-atom strategy outlined by
arXiv:1011.3758.
Dimensions:
A2s, A2p : cm^3 s^{-1}
B2s, B2p : s^{-1}
C2s, C2p : None
Parameters:
-----------
TCMB : float
Current CMB temperature, in eV.
Tm : float
Current matter temperature, in eV.
H : float
Current Hubble parameter, in s^{-1}.
nH : float
Current proton (free+bound) number density, in cm^{-3}.
x1s : float
Fraction of proton in the 1s bound state.
Returns:
--------
coefficients : tuple (float, float, float, float, float, float, float, float)
A tuple of the six relevant effective coefficients for the four-level-atom.
"""
# """ Step 0: Set up interpolation environment. """
# Determine which columns of alpha_tab to use based on TCMB, Tm
# Always take the smaller of the two in numerator, such that ratio is < 1.
# The if statements are only necessary if Tm > TCMB ever, but that seems not realistic?
#T_RATIO_now = jnp.where(Tm < TCMB, Tm/TCMB, TCMB/Tm) # The input point to interpolate at.
#A2s_column = jnp.where(Tm < TCMB, 0, 2)
#A2p_column = A2s_column + 1
T_RATIO_now = Tm/TCMB
A2s_column = 0
A2p_column = 1
# Determine the array position of the requested temperatures relative to the tabulated axis.
# For instance, if axis is [0, 2, 4, 6, 8], requesting 3 should be position 1.5
# This is needed to use ndimage.map_coordinates
TCMB_index = jnp.interp(jnp.log(TCMB), TR_axis, jnp.arange(TR_axis.size))
T_RATIO_index = jnp.interp(T_RATIO_now, T_RATIO_axis, jnp.arange(T_RATIO_axis.size))
# """ Step 1: Obtain recombination coefficients A2s, A2p by interpolating tabulated alpha. """
A2s = jnp.exp(map_coordinates(jnp.log(A2s_table[:, :, A2s_column]), [TCMB_index, T_RATIO_index], order=1))
A2p = jnp.exp(map_coordinates(jnp.log(A2p_table[:, :, A2p_column]), [TCMB_index, T_RATIO_index], order=1))
# """ Step 2: Obtain photoionization coefficients B2s, B2p from A2s, A2p via detailed balance. """
# See equation (22) of arXiv:1006.1355 for detail.
# nl state energy -EI/n^2, where EI is the hydrogen ionization energy, for n=2.
E2 = -rydberg / 4.
# equation (C7) of arXiv:1006.1355, times c^3, dimensions of cm^3
q = (2*jnp.pi*mu_e*TCMB / (2*jnp.pi*hbar)**2)**(3/2) / c**3
B2s = jnp.exp(E2/TCMB) * q * jnp.exp(map_coordinates(jnp.log(A2s_table[:, :, A2s_column]), [TCMB_index, NTM-1], order=1))
B2p = (1./3.)*jnp.exp(E2/TCMB) * q * jnp.exp(map_coordinates(jnp.log(A2p_table[:, :, A2p_column]), [TCMB_index, NTM-1], order=1))
# """ Step 3: Obtain 2p->2s transition rate by interpolating tabulated R, and 2s->2p rate via detailed balance. """
R2p2s = jnp.exp(jnp.interp(jnp.log(TCMB), TR_axis, jnp.log(R_tab)))
R2s2p = 3.*R2p2s # See equation (21) of arXiv:1006.1355
# """ Step 4: Obtain generalized Peebles C factors for 2s, 2p states """
# Rate of escape of lyman-alpha photons, multiplied by the bound proton fraction to regulate possible divergence.
R_Lya_times_x1s = 8.*jnp.pi*H / (3.*nH*(c/lya_freq)**3)
Gamma2s = B2s + R2s2p + R2s1s # Inverse lifetime of 2s state
Gamma2p_times_x1s = (B2p + R2p2s)*x1s + R_Lya_times_x1s # Inverse lifetime of 2p state, with x1s factor.
C2s = (R2s1s + R2s2p*R_Lya_times_x1s/Gamma2p_times_x1s) \
/ (Gamma2s - x1s*R2s2p*R2p2s/Gamma2p_times_x1s)
C2p = (R_Lya_times_x1s + x1s*R2p2s*R2s1s/Gamma2s) \
/ (Gamma2p_times_x1s - x1s*R2p2s*R2s2p/Gamma2s)
return jnp.array([A2s, A2p, B2s, B2p, C2s, C2p, R2p2s, R2s2p])