mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-04-24 11:50:53 +00:00
94 lines
2.7 KiB
Python
94 lines
2.7 KiB
Python
import jax
|
|
from jax.experimental.maps import xmap
|
|
import numpy as np
|
|
import jax.numpy as jnp
|
|
from functools import partial
|
|
import jaxdecomp
|
|
|
|
def fftk(shape, symmetric=False, dtype=np.float32, sharding_info=None):
|
|
""" Return k_vector given a shape (nc, nc, nc)
|
|
"""
|
|
k = []
|
|
|
|
if sharding_info is not None:
|
|
nx = sharding_info.pdims[1]
|
|
ny = sharding_info.pdims[0]
|
|
# nx = sharding_info[0].Get_size()
|
|
# ix = sharding_info[0].Get_rank()
|
|
# ny = sharding_info[1].Get_size()
|
|
# iy = sharding_info[1].Get_rank()
|
|
ix = sharding_info.rank
|
|
iy = 0
|
|
shape = sharding_info.global_shape
|
|
|
|
for d in range(len(shape)):
|
|
kd = np.fft.fftfreq(shape[d])
|
|
kd *= 2 * np.pi
|
|
|
|
if symmetric and d == len(shape) - 1:
|
|
kd = kd[:shape[d] // 2 + 1]
|
|
|
|
if (sharding_info is not None) and d == 0:
|
|
kd = kd.reshape([nx, -1])[ix]
|
|
|
|
if (sharding_info is not None) and d == 1:
|
|
kd = kd.reshape([ny, -1])[iy]
|
|
|
|
k.append(kd.astype(dtype))
|
|
return k
|
|
|
|
|
|
@partial(xmap,
|
|
in_axes=[['x', 'y', ...],
|
|
[['x'], ['y'], [...]]],
|
|
out_axes=['x', 'y', ...])
|
|
def apply_gradient_laplace(kfield, kvec):
|
|
kx, ky, kz = kvec
|
|
kk = (kx**2 + ky**2 + kz**2)
|
|
kernel = jnp.where(kk == 0, 1., 1./kk)
|
|
return jnp.stack([kfield * kernel * 1j * 1 / 6.0 * (8 * jnp.sin(kz) - jnp.sin(2 * kz)),
|
|
kfield * kernel * 1j * 1 / 6.0 * (8 * jnp.sin(kx) - jnp.sin(2 * kx)),
|
|
kfield * kernel * 1j * 1 / 6.0 * (8 * jnp.sin(ky) - jnp.sin(2 * ky))], axis=-1)
|
|
|
|
|
|
def cic_compensation(kvec):
|
|
"""
|
|
Computes cic compensation kernel.
|
|
Adapted from https://github.com/bccp/nbodykit/blob/a387cf429d8cb4a07bb19e3b4325ffdf279a131e/nbodykit/source/mesh/catalog.py#L499
|
|
Itself based on equation 18 (with p=2) of
|
|
`Jing et al 2005 <https://arxiv.org/abs/astro-ph/0409240>`_
|
|
Args:
|
|
kvec: array of k values in Fourier space
|
|
Returns:
|
|
v: array of kernel
|
|
"""
|
|
kwts = [np.sinc(kvec[i] / (2 * np.pi)) for i in range(3)]
|
|
wts = (kwts[0] * kwts[1] * kwts[2])**(-2)
|
|
return wts
|
|
|
|
|
|
def PGD_kernel(kvec, kl, ks):
|
|
"""
|
|
Computes the PGD kernel
|
|
Parameters:
|
|
-----------
|
|
kvec: array
|
|
Array of k values in Fourier space
|
|
kl: float
|
|
initial long range scale parameter
|
|
ks: float
|
|
initial dhort range scale parameter
|
|
Returns:
|
|
--------
|
|
v: array
|
|
kernel
|
|
"""
|
|
kk = sum(ki**2 for ki in kvec)
|
|
kl2 = kl**2
|
|
ks4 = ks**4
|
|
mask = (kk == 0).nonzero()
|
|
kk[mask] = 1
|
|
v = jnp.exp(-kl2 / kk) * jnp.exp(-kk**2 / ks4)
|
|
imask = (~(kk == 0)).astype(int)
|
|
v *= imask
|
|
return v
|