def imwofz_naive(x, y):
"""Imaginary part of wofz function based on Algorithm 916 (naive
implementation)
We apply a=0.5 for Algorithm 916. We apply a=0.5 for Algorithm 916. This function is a slower version of faddeeva.rewofz (about 2 times slower). See PRs #117 and #118.
Args:
x: x < ncut/2
y:
Returns:
jnp.array: Imag(wofz(x+iy))
"""
ncut = 27
an = 0.5 * jnp.arange(1, ncut + 1)
a2n2 = an * an
xy = x * y
xyp = 2.0 * xy / jnp.pi
exx = jnp.exp(-x * x)
f = -exx * erfcx(y) * jnp.sin(2.0 * xy) + x / jnp.pi * exx * jnp.sinc(xyp)
vec0 = 1.0 / (a2n2 + y * y)
vec1 = jnp.exp(-(a2n2 + x * x))
Sigma1 = jnp.dot(vec0, vec1)
vecm = an * vec0
vec4 = jnp.exp(-(an + x) * (an + x))
vec5 = jnp.exp(-(an - x) * (an - x))
Sigma4 = jnp.dot(vecm, vec4)
Sigma5 = jnp.dot(vecm, vec5)
f = f + 1.0 / jnp.pi * (y * jnp.sin(2.0 * xy) * Sigma1 + 0.5 * Sigma5 -
0.5 * Sigma4)
return f
def rewofz_nonvector(x, y):
"""Real part of wofz function based on Algorithm 916.
We apply a=0.5 for Algorithm 916.
Args:
x: x < ncut/2
y:
Returns:
f: Real(wofz(x+iy))
"""
ncut = 27
xy = x * y
xyp = xy / jnp.pi
exx = jnp.exp(-x * x)
f = exx * erfcx(y) * jnp.cos(
2.0 * xy) + x * jnp.sin(xy) / jnp.pi * exx * jnp.sinc(xyp)
n = jnp.arange(1, ncut + 1)
n2 = n * n
vec0 = 1.0 / (0.25 * n2 + y * y)
vec1 = jnp.exp(-(0.25 * n2 + x * x))
vec2 = jnp.exp(-(0.5 * n + x) * (0.5 * n + x))
vec3 = jnp.exp(-(0.5 * n - x) * (0.5 * n - x))
Sigma1 = jnp.dot(vec0, vec1)
Sigma2 = jnp.dot(vec0, vec2)
Sigma3 = jnp.dot(vec0, vec3)
f = f + 1.0 / jnp.pi * (-y * jnp.cos(2.0 * xy) * Sigma1 +
0.5 * y * Sigma2 + 0.5 * y * Sigma3)
return f
def rewofz_naive(x, y):
"""Real part of wofz function based on Algorithm 916 (naive implementation)
We apply a=0.5 for Algorithm 916. This function is a slower version of faddeeva.rewofz (about 2 times slower). See PRs #117 and #118.
Args:
x: x < ncut/2
y:
Returns:
jnp.array: Real(wofz(x+iy))
"""
ncut = 27
an = 0.5 * jnp.arange(1, ncut + 1)
a2n2 = an * an
xy = x * y
xyp = xy / jnp.pi
exx = jnp.exp(-x * x)
f = exx * erfcx(y) * jnp.cos(
2.0 * xy) + x * jnp.sin(xy) / jnp.pi * exx * jnp.sinc(xyp)
vec0 = 1.0 / (a2n2 + y * y)
vec1 = jnp.exp(-(a2n2 + x * x))
vec2 = jnp.exp(-(0.5 * n + x) * (0.5 * n + x))
vec3 = jnp.exp(-(0.5 * n - x) * (0.5 * n - x))
Sigma1 = jnp.dot(vec0, vec1)
Sigma2 = jnp.dot(vec0, vec2)
Sigma3 = jnp.dot(vec0, vec3)
f = f + 1.0 / jnp.pi * (-y * jnp.cos(2.0 * xy) * Sigma1 +
0.5 * y * Sigma2 + 0.5 * y * Sigma3)
return f
def imwofz_nonvector(x, y):
"""Imaginary part of wofz function based on Algorithm 916.
We apply a=0.5 for Algorithm 916.
Args:
x: x < ncut/2
y:
Returns:
f: Imag(wofz(x+iy))
"""
ncut = 27
xy = x * y
xyp = 2.0 * xy / jnp.pi
exx = jnp.exp(-x * x)
f = -exx * erfcx(y) * jnp.sin(2.0 * xy) + x / jnp.pi * exx * jnp.sinc(xyp)
n = jnp.arange(1, ncut + 1)
n2 = n * n
vec0 = 0.5 * n / (0.25 * n2 + y * y)
vec1 = jnp.exp(-(0.25 * n2 + x * x))
vec4 = jnp.exp(-(0.5 * n + x) * (0.5 * n + x))
vec5 = jnp.exp(-(0.5 * n - x) * (0.5 * n - x))
Sigma1 = jnp.dot(vec0, vec1)
Sigma4 = jnp.dot(vec0, vec4)
Sigma5 = jnp.dot(vec0, vec5)
f = f + 1.0 / jnp.pi * (y * jnp.sin(2.0 * xy) * Sigma1 + 0.5 *
(Sigma5 - Sigma4))
return f
def imwofz_vectorized(x, y):
"""Imaginary part of wofz function based on Algorithm 916.
We apply a=0.5 for Algorithm 916.
Args:
x: x < ncut/2
y:
Returns:
f: Imag(wofz(x+iy))
"""
ncut = 27
xy = x * y
xyp = 2.0 * xy / jnp.pi
exx = jnp.exp(-x * x)
f = -exx * erfcx(y) * jnp.sin(2.0 * xy) + x / jnp.pi * exx * jnp.sinc(xyp)
n = jnp.arange(1, ncut + 1)
n2 = n * n
x2 = x * x
y2 = y * y
vec0 = 0.5 * n / (0.25 * n2 + y2)
vec1 = jnp.exp(-(0.25 * n2[None, :] + x2[:, None]))
vec4 = jnp.exp(-(0.5 * n[None, :] + x[:, None]) *
(0.5 * n[None, :] + x[:, None]))
vec5 = jnp.exp(-(0.5 * n[None, :] - x[:, None]) *
(0.5 * n[None, :] - x[:, None]))
Sigma1 = jnp.sum(vec0 * vec1, axis=1)
Sigma4 = jnp.sum(vec0 * vec4, axis=1)
Sigma5 = jnp.sum(vec0 * vec5, axis=1)
f = f + 1.0 / jnp.pi * (y * jnp.sin(2.0 * xy) * Sigma1 + 0.5 *
(Sigma5 - Sigma4))
return f
def rewofz_vectorized(x, y):
"""Real part of wofz function based on Algorithm 916.
We apply a=0.5 for Algorithm 916.
Args:
x: x < ncut/2
y:
Returns:
f: Real(wofz(x+iy))
"""
ncut = 27
xy = x * y
xyp = xy / jnp.pi
exx = jnp.exp(-x * x)
f = exx * erfcx(y) * jnp.cos(
2.0 * xy) + x * jnp.sin(xy) / jnp.pi * exx * jnp.sinc(xyp)
n = jnp.arange(1, ncut + 1)
n2 = n * n
x2 = x * x
y2 = y * y
vec0 = 1.0 / (0.25 * n2 + y2)
vec1 = jnp.exp(-(0.25 * n2[None, :] + x2[:, None]))
vec2 = jnp.exp(-(0.5 * n[None, :] + x[:, None]) *
(0.5 * n[None, :] + x[:, None]))
vec3 = jnp.exp(-(0.5 * n[None, :] - x[:, None]) *
(0.5 * n[None, :] - x[:, None]))
Sigma1 = jnp.sum(vec0 * vec1, axis=1)
Sigma2 = jnp.sum(vec0 * vec2, axis=1)
Sigma3 = jnp.sum(vec0 * vec3, axis=1)
f = f + 1.0 / jnp.pi * (-y * jnp.cos(2.0 * xy) * Sigma1 +
0.5 * y * Sigma2 + 0.5 * y * Sigma3)
return f
def rewofzx_naive(x, y):
"""[VJP custom defined] Real part of wofz function based on Algorithm 916
(naive implementation)
We apply a=0.5 for Algorithm 916.
Args:
x: x < ncut/2
y:
Returns:
jnp.array: Real(wofz(x+iy))
"""
xy = x * y
xyp = xy / jnp.pi
exx = jnp.exp(-x * x)
f = exx * erfcx(y) * jnp.cos(
2.0 * xy) + x * jnp.sin(xy) / jnp.pi * exx * jnp.sinc(xyp)
ncut = 27
n = jnp.arange(1, ncut + 1)
n2 = n * n
vec0 = 1.0 / (0.25 * n2 + y * y)
vec1 = jnp.exp(-(0.25 * n2 + x * x))
vec2 = jnp.exp(-(0.5 * n + x) * (0.5 * n + x))
vec3 = jnp.exp(-(0.5 * n - x) * (0.5 * n - x))
Sigma2 = jnp.dot(vec0, vec2)
Sigma3 = jnp.dot(vec0, vec3)
f = f + 1.0 / jnp.pi * (-y * jnp.cos(2.0 * xy) * Sigma1 +
0.5 * y * Sigma2 + 0.5 * y * Sigma3)
return f
def imwofz(x, y):
"""Imaginary part of wofz (Faddeeva) function based on Algorithm 916.
We apply a=0.5 for Algorithm 916.
Args:
x: x < ncut/2
y:
Returns:
jnp.array: Imag(wofz(x+iy))
"""
wxy = 2.0 * x * y
exx = jnp.exp(-x * x)
f = -exx * erfcx(y) * jnp.sin(wxy) + x / jnp.pi * exx * jnp.sinc(
wxy / jnp.pi)
y2 = y * y
Sigma1 = exx * (7.78800786e-01 / (0.25 + y2) + 3.67879450e-01 /
(1. + y2) + 1.05399221e-01 / (2.25 + y2) + 1.83156393e-02 /
(4. + y2) + 1.93045416e-03 / (6.25 + y2) + 1.23409802e-04 /
(9. + y2) + 4.78511765e-06 /
(12.25 + y2) + 1.12535176e-07 / (16. + y2))
Sigma45 = jnp.sum(an * (-jnp.exp(-(an + x)**2) + jnp.exp(-(an - x)**2)) /
(a2n2 + y2))
f = f + 1.0 / jnp.pi * (y * jnp.sin(wxy) * Sigma1 + 0.5 * Sigma45)
return f
def rewofz(x, y):
"""Real part of wofz (Faddeeva) function based on Algorithm 916.
We apply a=0.5 for Algorithm 916.
Args:
x: x < ncut/2
y:
Returns:
jnp.array: Real(wofz(x+iy))
"""
xy = x * y
exx = jnp.exp(-x * x)
f = exx * (erfcx(y) * jnp.cos(2.0 * xy) +
x * jnp.sin(xy) / jnp.pi * jnp.sinc(xy / jnp.pi))
y2 = y * y
Sigma23 = jnp.sum(
(jnp.exp(-(an + x)**2) + jnp.exp(-(an - x)**2)) / (a2n2 + y2))
Sigma1 = exx * (7.78800786e-01 / (0.25 + y2) + 3.67879450e-01 /
(1. + y2) + 1.05399221e-01 / (2.25 + y2) + 1.83156393e-02 /
(4. + y2) + 1.93045416e-03 / (6.25 + y2) + 1.23409802e-04 /
(9. + y2) + 4.78511765e-06 /
(12.25 + y2) + 1.12535176e-07 / (16. + y2))
f = f + y / jnp.pi * (-jnp.cos(2.0 * xy) * Sigma1 + 0.5 * Sigma23)
return f
def sinc2_2d(width=1.0,
height=None,
wavelength=1e-6,
shape=(512, 512),
pixelscale=0.010,
center=None):
"""
Create a 2D sinc function PSF, representing the PSF of a square or rectangular aperture
Parameters
-----------
width : float
Width in meters of the aperture.
height : float, optional
height in meters of the aperture. If not specified, the aperture is assumed
to be a square so height=width
wavelength : float
wavelength in meters
shape : tuple with 2 elements
shape of array to create
pixelscale : float
pixel scale in arcseconds per pixel
center : tuple with 2 elements, optional
Center coordinates of the PSF. Defaults to center of array.
"""
if height is None:
height = width
halfwidth = float(width) / 2
halfheight = float(height) / 2
if center is None:
center = (np.asarray(shape) - 1.) / 2
y, x = np.indices(shape, float)
y -= center[0]
x -= center[1]
y *= pixelscale
x *= pixelscale
k = 2 * np.pi / wavelength # wavenumber
alpha = k * x * halfwidth * _ARCSECtoRAD
beta = k * y * halfheight * _ARCSECtoRAD
psf = (np.sinc(alpha))**2 * (np.sinc(beta))**2
return psf
def sincinterp(x):
N = x.shape[0]
y = jnp.zeros(2 * N -1, dtype=x.dtype)
y = index_update(y, index[:2 * N:2], x)
xint = fftconvolve(
y[:2 * N],
jnp.sinc(jnp.arange(-(2 * N - 3), (2 * N - 2)).T / 2),
)
return xint[2 * N - 3: -2 * N + 3]
def rewofz_scan(x, y):
ncut = 27
xy = x * y
xyp = xy / jnp.pi
exx = jnp.exp(-x * x)
f = exx * erfcx(y) * jnp.cos(
2.0 * xy) + x * jnp.sin(xy) / jnp.pi * exx * jnp.sinc(xyp)
narr = jnp.arange(1, ncut + 1)
def fscan(sv, n):
n2 = n * n
fac0 = 1.0 / (0.25 * n2 + y * y)
Sigma1, Sigma2, Sigma3 = sv
Sigma1 = Sigma1 + fac0 * jnp.exp(-(0.25 * n2 + x * x))
Sigma2 = Sigma2 + fac0 * jnp.exp(-(0.5 * n + x) * (0.5 * n + x))
Sigma3 = Sigma3 + fac0 * jnp.exp(-(0.5 * n - x) * (0.5 * n - x))
return (Sigma1, Sigma2, Sigma3), None
s0 = jnp.zeros(jnp.shape(x))
sv, null = scan(fscan, (s0, s0, s0), narr)
Sigma1, Sigma2, Sigma3 = sv
f = f + 1.0 / jnp.pi * (-y * jnp.cos(2.0 * xy) * Sigma1 +
0.5 * y * Sigma2 + 0.5 * y * Sigma3)
return f
def Eisenstein_Hu(cosmo, k, type='eisenhu_osc'):
""" Computes the Eisenstein & Hu matter transfer function.
Parameters
----------
cosmo: Background
Background cosmology
k: array_like
Wave number in h Mpc^{-1}
type: str, optional
Type of transfer function. Either 'eisenhu' or 'eisenhu_osc'
(def: 'eisenhu_osc')
Returns
-------
T: array_like
Value of the transfer function at the requested wave number
Notes
-----
The Eisenstein & Hu transfer functions are computed using the fitting
formulae of :cite:`1998:EisensteinHu`
"""
#############################################
# Quantities computed from 1998:EisensteinHu
# Provides : - k_eq : scale of the particle horizon at equality epoch
# - z_eq : redshift of equality epoch
# - R_eq : ratio of the baryon to photon momentum density
# at z_eq
# - z_d : redshift of drag epoch
# - R_d : ratio of the baryon to photon momentum density
# at z_d
# - sh_d : sound horizon at drag epoch
# - k_silk : Silk damping scale
T_2_7_sqr = (const.tcmb/2.7)**2
h2 = cosmo.h**2
w_m = cosmo.Omega_m*h2
w_b = cosmo.Omega_b*h2
fb = cosmo.Omega_b / cosmo.Omega_m
fc = (cosmo.Omega_m - cosmo.Omega_b) / cosmo.Omega_m
k_eq = 7.46e-2*w_m/T_2_7_sqr / cosmo.h # Eq. (3) [h/Mpc]
z_eq = 2.50e4*w_m/(T_2_7_sqr)**2 # Eq. (2)
# z drag from Eq. (4)
b1 = 0.313*np.power(w_m, -0.419)*(1.0+0.607*np.power(w_m, 0.674))
b2 = 0.238*np.power(w_m, 0.223)
z_d = 1291.0*np.power(w_m, 0.251)/(1.0+0.659*np.power(w_m, 0.828)) * \
(1.0 + b1*np.power(w_b, b2))
# Ratio of the baryon to photon momentum density at z_d Eq. (5)
R_d = 31.5 * w_b / (T_2_7_sqr)**2 * (1.e3/z_d)
# Ratio of the baryon to photon momentum density at z_eq Eq. (5)
R_eq = 31.5 * w_b / (T_2_7_sqr)**2 * (1.e3/z_eq)
# Sound horizon at drag epoch in h^-1 Mpc Eq. (6)
sh_d = 2.0/(3.0*k_eq) * np.sqrt(6.0/R_eq) * \
np.log((np.sqrt(1.0 + R_d) + np.sqrt(R_eq + R_d)) /
(1.0 + np.sqrt(R_eq)))
# Eq. (7) but in [hMpc^{-1}]
k_silk = 1.6 * np.power(w_b, 0.52) * np.power(w_m, 0.73) * \
(1.0 + np.power(10.4*w_m, -0.95)) / cosmo.h
#############################################
alpha_gamma = 1.-0.328*np.log(431.*w_m)*w_b/w_m + \
0.38*np.log(22.3*w_m)*(cosmo.Omega_b/cosmo.Omega_m)**2
gamma_eff = cosmo.Omega_m*cosmo.h * \
(alpha_gamma + (1.-alpha_gamma)/(1.+(0.43*k*sh_d)**4))
if(type == 'eisenhu'):
q = k * np.power(const.tcmb/2.7, 2)/gamma_eff
# EH98 (29) #
L = log(2.*exp(1.0) + 1.8*q)
C = 14.2 + 731.0/(1.0 + 62.5*q)
res = L/(L + C*q*q)
elif(type == 'eisenhu_osc'):
# Cold dark matter transfer function
# EH98 (11, 12)
a1 = np.power(46.9*w_m, 0.670) * (1.0 + np.power(32.1*w_m, -0.532))
a2 = np.power(12.0*w_m, 0.424) * (1.0 + np.power(45.0*w_m, -0.582))
alpha_c = np.power(a1, -fb) *np.power(a2, -fb**3)
b1 = 0.944 / (1.0 + np.power(458.0*w_m, -0.708))
b2 = np.power(0.395*w_m, -0.0266)
beta_c = 1.0 + b1*(np.power(fc, b2) - 1.0)
beta_c = 1.0 / beta_c
# EH98 (19). [k] = h/Mpc
def T_tilde(k1, alpha, beta):
# EH98 (10); [q] = 1 BUT [k] = h/Mpc
q = k1 / (13.41 * k_eq)
L = np.log(np.exp(1.0) + 1.8 * beta * q)
C = 14.2 / alpha + 386.0 / (1.0 + 69.9 * np.power(q, 1.08))
T0 = L/(L + C*q*q)
return T0
# EH98 (17, 18)
f = 1.0 / (1.0 + (k * sh_d / 5.4)**4)
Tc = f * T_tilde(k, 1.0, beta_c) + \
(1.0 - f) * T_tilde(k, alpha_c, beta_c)
# Baryon transfer function
# EH98 (19, 14, 21)
y = (1.0 + z_eq) / (1.0 + z_d)
x = np.sqrt(1.0 + y)
G_EH98 = y * (-6.0 * x +
(2.0 + 3.0*y) * np.log((x + 1.0) / (x - 1.0)))
alpha_b = 2.07 * k_eq * sh_d * \
np.power(1.0 + R_d, -0.75) * G_EH98
beta_node = 8.41 * np.power(w_m, 0.435)
tilde_s = sh_d / np.power(1.0 + (beta_node /
(k * sh_d))**3, 1.0/3.0)
beta_b = 0.5 + fb + (3.0 - 2.0 * fb) * np.sqrt((17.2 * w_m)**2 + 1.0)
# [tilde_s] = Mpc/h
Tb = (T_tilde(k, 1.0, 1.0) / (1.0 + (k * sh_d / 5.2)**2) +
alpha_b / (1.0 + (beta_b/(k * sh_d))**3) *
np.exp(-np.power(k / k_silk, 1.4))) * np.sinc(k*tilde_s/np.pi)
# Total transfer function
res = fb * Tb + fc * Tc
else:
raise NotImplementedError
return res
def sinc(x): if isinstance(x, JaxArray): x = x.value return JaxArray(jnp.sinc(x))
def testSinc(self): # Regression test for #6936 self.assertEqual(jnp.sinc(0.0), 1.0)