guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-06-29 16:41:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

merge hugos LPT2 code

Browse source
This commit is contained in:
Wassim KABALAN 2024-10-22 12:57:05 -04:00
parent 2f509932f5
commit 86233081e2
2 changed files with 149 additions and 153 deletions
Download patch file Download diff file Expand all files Collapse all files

116
jaxpm/kernels.py
View file

@ -1,5 +1,3 @@
from enum import Enum
import jax.numpy as jnp
import jax_cosmo as jc
import numpy as np
@ -45,16 +43,18 @@ def interpolate_power_spectrum(input, k, pk, sharding=None):
def gradient_kernel(kvec, direction, order=1):
"""
Computes the gradient kernel in the requested direction
Parameters:
Parameters
-----------
kvec: array
Array of k values in Fourier space
kvec: list
List of wave-vectors in Fourier space
direction: int
Index of the direction in which to take the gradient
Returns:
Index of the direction in which to take the gradient
Returns
--------
wts: array
Complex kernel
Complex kernel values
"""
if order == 0:
wts = 1j * kvec[direction]
@ -69,37 +69,43 @@ def gradient_kernel(kvec, direction, order=1):
return wts
def laplace_kernel(kvec):
def invlaplace_kernel(kvec):
"""
Compute the Laplace kernel from a given K vector
Parameters:
Compute the inverse Laplace kernel
Parameters
-----------
kvec: array
Array of k values in Fourier space
Returns:
kvec: list
List of wave-vectors
Returns
--------
wts: array
Complex kernel
Complex kernel values
"""
kk = sum(ki**2 for ki in kvec)
wts = jnp.where(kk == 0, 1., 1. / kk)
return wts
kk_nozeros = jnp.where(kk==0, 1, kk)
return - jnp.where(kk==0, 0, 1 / kk_nozeros)
def longrange_kernel(kvec, r_split):
"""
Computes a long range kernel
Parameters:
-----------
kvec: array
Array of k values in Fourier space
r_split: float
Computes a long range kernel
Parameters
-----------
kvec: list
List of wave-vectors
r_split: float
Splitting radius
Returns
--------
wts: array
Complex kernel values
TODO: @modichirag add documentation
Returns:
--------
wts: array
kernel
"""
"""
if r_split != 0:
kk = sum(ki**2 for ki in kvec)
return np.exp(-kk * r_split**2)
@ -109,15 +115,21 @@ def longrange_kernel(kvec, r_split):
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
"""
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)
Parameters:
-----------
kvec: list
List of wave-vectors
Returns:
--------
wts: array
Complex kernel values
"""
kwts = [np.sinc(kvec[i] / (2 * np.pi)) for i in range(3)]
wts = (kwts[0] * kwts[1] * kwts[2])**(-2)
return wts
@ -125,20 +137,22 @@ def cic_compensation(kvec):
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
"""
Computes the PGD kernel
Parameters:
-----------
kvec: list
List of wave-vectors
kl: float
Initial long range scale parameter
ks: float
Initial dhort range scale parameter
Returns:
--------
v: array
Complex kernel values
"""
kk = sum(ki**2 for ki in kvec)
kl2 = kl**2
ks4 = ks**4