def adf_step(state, xs, q, lbound, ubound):
mu_t, tau_t = state
Phi_t, y_t = xs
mu_t_cond = mu_t
tau_t_cond = tau_t + q
# prior predictive distribution
m_t_cond = (Phi_t * mu_t_cond).sum()
v_t_cond = (Phi_t**2 * tau_t_cond).sum()
v_t_cond_sqrt = jnp.sqrt(v_t_cond)
# Moment-matched Gaussian approximation elements
Zt = integrate.romb(lambda eta: Zt_func(eta, y_t, m_t_cond, v_t_cond_sqrt),
lbound, ubound)
mt = integrate.romb(lambda eta: mt_func(eta, y_t, m_t_cond, v_t_cond_sqrt),
lbound, ubound)
mt = mt / Zt
vt = integrate.romb(lambda eta: vt_func(eta, y_t, m_t_cond, v_t_cond_sqrt),
lbound, ubound)
vt = vt / Zt - mt**2
# Posterior estimation
delta_m = mt - m_t_cond
delta_v = vt - v_t_cond
a = Phi_t * tau_t_cond / jnp.power(Phi_t * tau_t_cond, 2).sum()
mu_t = mu_t_cond + a * delta_m
tau_t = tau_t_cond + a**2 * delta_v
return (mu_t, tau_t), (mu_t, tau_t)
def sigmasqr(cosmo,
R,
transfer_fn,
kmin=0.0001,
kmax=1000.0,
ksteps=5,
**kwargs):
""" Computes the energy of the fluctuations within a sphere of R h^{-1} Mpc
.. math::
\\sigma^2(R)= \\frac{1}{2 \\pi^2} \\int_0^\\infty \\frac{dk}{k} k^3 P(k,z) W^2(kR)
where
.. math::
W(kR) = \\frac{3j_1(kR)}{kR}
"""
def int_sigma(logk):
k = np.exp(logk)
x = k * R
w = 3.0 * (np.sin(x) - x * np.cos(x)) / (x * x * x)
pk = transfer_fn(cosmo, k, **kwargs)**2 * primordial_matter_power(
cosmo, k)
return k * (k * w)**2 * pk
y = romb(int_sigma, np.log10(kmin), np.log10(kmax), divmax=7)
return 1.0 / (2.0 * np.pi**2.0) * y