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)