guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-04-28 13:12:44 +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
Show stats 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

186
jaxpm/pm.py
View file

@ -9,7 +9,7 @@ from jaxpm.distributed import (autoshmap, fft3d, get_local_shape, ifft3d,
normal_field)
from jaxpm.growth import (dGf2a, dGfa, growth_factor, growth_factor_second,
growth_rate, growth_rate_second)
from jaxpm.kernels import (PGD_kernel, fftk, gradient_kernel, laplace_kernel,
from jaxpm.kernels import (PGD_kernel, fftk, gradient_kernel, invlaplace_kernel,
longrange_kernel)
from jaxpm.painting import cic_paint, cic_paint_dx, cic_read, cic_read_dx
@ -29,18 +29,20 @@ def pm_forces(positions,
mesh_shape = delta.shape
if delta is None:
delta_k = fft3d(
cic_paint_dx(positions, halo_size=halo_size, sharding=sharding))
else:
field = cic_paint_dx(positions, halo_size=halo_size, sharding=sharding)
delta_k = fft3d(field)
elif jnp.isrealobj(delta):
delta_k = fft3d(delta)
else:
delta_k = delta
kvec = fftk(delta_k)
# Computes gravitational potential
pot_k = delta_k * laplace_kernel(kvec) * longrange_kernel(kvec,
pot_k = delta_k * invlaplace_kernel(kvec) * longrange_kernel(kvec,
r_split=r_split)
# Computes gravitational forces
forces = jnp.stack([
cic_read_dx(ifft3d(gradient_kernel(kvec, i) * pot_k),
cic_read_dx(ifft3d( - gradient_kernel(kvec, i) * pot_k),
halo_size=halo_size,
sharding=sharding) for i in range(3)
],
@ -49,44 +51,10 @@ def pm_forces(positions,
return forces
def lpt2_source(mesh_size, initial_conditions):
kvec = fftk(mesh_size)
# TODO : this has already been done for LPT1, we should reuse it
delta_k = fft3d(initial_conditions)
source = jnp.zeros_like(delta_k)
D1 = [1, 2, 0]
D2 = [2, 0, 1]
# laplace_kernel should be actually inv laplace_kernel
# adding a minus sign here that will be negated when computing forces
# because F = -grad(phi)
# and phi = -laplace_kernel(delta_k)
pot_k = delta_k * laplace_kernel(delta_k)
nabla_i_nabla_i = [
ifft3d(gradient_kernel(kvec, i)**2 * pot_k) for i in range(3)
]
# for diagonal terms
source += nabla_i_nabla_i[D1[0]] * nabla_i_nabla_i[D2[0]]
source += nabla_i_nabla_i[D1[1]] * nabla_i_nabla_i[D2[1]]
source += nabla_i_nabla_i[D1[2]] * nabla_i_nabla_i[D2[2]]
# off diag terms
for i in range(3):
nabla_i_nabla_j = gradient_kernel(kvec, D1[i]) * gradient_kernel(
kvec, D2[i])
phi = ifft3d(nabla_i_nabla_j * pot_k)
source -= phi**2
return source
def lpt(cosmo, initial_conditions, a, halo_size=0, sharding=None):
def lpt(cosmo, initial_conditions, a, halo_size=0, sharding=None,order=1):
"""
Computes first order LPT displacement
Computes first and second order LPT displacement and momentum,
e.g. Eq. 2 and 3 [Jenkins2010](https://arxiv.org/pdf/0910.0258)
"""
gpu_mesh = sharding.mesh if sharding is not None else None
spec = sharding.spec if sharding is not None else P()
@ -99,48 +67,48 @@ def lpt(cosmo, initial_conditions, a, halo_size=0, sharding=None):
out_specs=spec)() # yapf: disable
a = jnp.atleast_1d(a)
E = jnp.sqrt(jc.background.Esqr(cosmo, a))
delta_k = fft3d(initial_conditions)
initial_force = pm_forces(displacement,
delta=initial_conditions,
delta=delta_k,
halo_size=halo_size,
sharding=sharding)
a = jnp.atleast_1d(a)
dx = growth_factor(cosmo, a) * initial_force
p = a**2 * growth_rate(cosmo, a) * jnp.sqrt(jc.background.Esqr(cosmo,
a)) * dx
f = a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a)) * dGfa(cosmo,
a) * initial_force
return dx, p, f
p = a**2 * growth_rate(cosmo, a) * E * dx
f = a**2 * E * dGfa(cosmo,a) * initial_force
if order == 2:
kvec = fftk(delta_k)
pot_k = delta_k * invlaplace_kernel(kvec)
delta2 = 0
shear_acc = 0
# for i, ki in enumerate(kvec):
for i in range(3):
# Add products of diagonal terms = 0 + s11*s00 + s22*(s11+s00)...
# shear_ii = jnp.fft.irfftn(- ki**2 * pot_k)
nabla_i_nabla_i = gradient_kernel(kvec, i)**2
shear_ii = jnp.fft.irfftn(nabla_i_nabla_i * pot_k)
delta2 += shear_ii * shear_acc
shear_acc += shear_ii
# @Credit Hugo Simon https://github.com/hsimonfroy/montecosmo
def lpt2(cosmo, initial_conditions, dx, p, f, a, halo_size=0):
# for kj in kvec[i+1:]:
for j in range(i+1, 3):
# Substract squared strict-up-triangle terms
# delta2 -= jnp.fft.irfftn(- ki * kj * pot_k)**2
nabla_i_nabla_j = gradient_kernel(kvec, i) * gradient_kernel(kvec, j)
delta2 -= jnp.fft.irfftn(nabla_i_nabla_j * pot_k)**2
delta_k2 = fft3d(delta2)
init_force2 = pm_forces(displacement, delta=delta_k2,halo_size=halo_size,sharding=sharding)
# NOTE: growth_factor_second is renormalized: - D2 = 3/7 * growth_factor_second
dx2 = 3/7 * growth_factor_second(cosmo, a) * init_force2
p2 = a**2 * growth_rate_second(cosmo, a) * E * dx2
f2 = a**2 * E * dGf2a(cosmo, a) * init_force2
mesh_size = initial_conditions.shape
local_mesh_shape = (*get_local_shape(initial_conditions.shape), 3)
# TODO
# Displacements have been created in the previous step
# find a way to reuse them
displacement = autoshmap(
partial(jnp.zeros, shape=(local_mesh_shape), dtype='float32'),
in_specs=(),
out_specs=P('x', 'y'))() # yapf: disable
lpt2_delta = lpt2_source(mesh_size, initial_conditions)
delta2_k = fft3d(lpt2_delta)
lpt2_forces = pm_forces(displacement,
mesh_size,
delta_k=delta2_k,
halo_size=halo_size)
dx2 = 3 / 7 * growth_factor_second(cosmo, a) * lpt2_forces
p2 = a**2 * growth_rate_second(cosmo, a) * jnp.sqrt(
jc.background.Esqr(cosmo, a)) * dx2
f2 = a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a)) * dGf2a(cosmo,
a) * lpt2_forces
dx += dx2
p += p2
f += f2
dx += dx2
p += p2
f += f2
return dx, p, f
@ -185,10 +153,33 @@ def make_ode_fn(mesh_shape, halo_size=0, sharding=None):
return nbody_ode
def get_ode_fn(cosmo, mesh_shape, halo_size=0, sharding=None):
def nbody_ode(a, state, args):
"""
State is an array [position, velocities]
Compatible with [Diffrax API](https://docs.kidger.site/diffrax/)
"""
pos, vel = state
forces = pm_forces(pos, mesh_shape, halo_size=halo_size, sharding=sharding) * 1.5 * cosmo.Omega_m
# Computes the update of position (drift)
dpos = 1. / (a**3 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * vel
# Computes the update of velocity (kick)
dvel = 1. / (a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * forces
return jnp.stack([dpos, dvel])
return nbody_ode
def pgd_correction(pos, mesh_shape, params):
"""
improve the short-range interactions of PM-Nbody simulations with potential gradient descent method, based on https://arxiv.org/abs/1804.00671
improve the short-range interactions of PM-Nbody simulations with potential gradient descent method,
based on https://arxiv.org/abs/1804.00671
args:
pos: particle positions [npart, 3]
params: [alpha, kl, ks] pgd parameters
@ -197,24 +188,20 @@ def pgd_correction(pos, mesh_shape, params):
delta = cic_paint(jnp.zeros(mesh_shape), pos)
alpha, kl, ks = params
delta_k = jnp.fft.rfftn(delta)
PGD_range = PGD_kernel(kvec, kl, ks)
pot_k_pgd = (delta_k * laplace_kernel(kvec)) * PGD_range
forces_pgd = jnp.stack([
cic_read(jnp.fft.irfftn(gradient_kernel(kvec, i) * pot_k_pgd), pos)
for i in range(3)
],
axis=-1)
dpos_pgd = forces_pgd * alpha
PGD_range=PGD_kernel(kvec, kl, ks)
pot_k_pgd=(delta_k * invlaplace_kernel(kvec))*PGD_range
forces_pgd= jnp.stack([cic_read(jnp.fft.irfftn(- gradient_kernel(kvec, i)*pot_k_pgd), pos)
for i in range(3)],axis=-1)
dpos_pgd = forces_pgd*alpha
return dpos_pgd
def make_neural_ode_fn(model, mesh_shape):
def neural_nbody_ode(state, a, cosmo, params):
def neural_nbody_ode(state, a, cosmo:Cosmology, params):
"""
state is a tuple (position, velocities)
"""
@ -226,19 +213,15 @@ def make_neural_ode_fn(model, mesh_shape):
delta_k = jnp.fft.rfftn(delta)
# Computes gravitational potential
pot_k = delta_k * laplace_kernel(kvec) * longrange_kernel(kvec,
r_split=0)
pot_k = delta_k * invlaplace_kernel(kvec) * longrange_kernel(kvec, r_split=0)
# Apply a correction filter
kk = jnp.sqrt(sum((ki / jnp.pi)**2 for ki in kvec))
pot_k = pot_k * (1. + model.apply(params, kk, jnp.atleast_1d(a)))
kk = jnp.sqrt(sum((ki/jnp.pi)**2 for ki in kvec))
pot_k = pot_k *(1. + model.apply(params, kk, jnp.atleast_1d(a)))
# Computes gravitational forces
forces = jnp.stack([
cic_read(jnp.fft.irfftn(gradient_kernel(kvec, i) * pot_k), pos)
for i in range(3)
],
axis=-1)
forces = jnp.stack([cic_read(jnp.fft.irfftn(- gradient_kernel(kvec, i)*pot_k), pos)
for i in range(3)],axis=-1)
forces = forces * 1.5 * cosmo.Omega_m
@ -249,5 +232,4 @@ def make_neural_ode_fn(model, mesh_shape):
dvel = 1. / (a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * forces
return dpos, dvel
return neural_nbody_ode