guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-02-22 17:47:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #23 from hsimonfroy/main

Browse source 
2lpt
This commit is contained in:
Wassim KABALAN 2024-08-05 19:38:35 +02:00 • committed by GitHub
commit ef7a7ef5c9
2 changed files with 163 additions and 95 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

143
jaxpm/kernels.py
View file

@ -3,8 +3,9 @@ import numpy as np
def fftk(shape, symmetric=True, finite=False, dtype=np.float32): def fftk(shape, symmetric=True, finite=False, dtype=np.float32):
""" Return k_vector given a shape (nc, nc, nc) and box_size """
""" Return wave-vectors for a given shape
"""
k = [] k = []
for d in range(len(shape)): for d in range(len(shape)):
kd = np.fft.fftfreq(shape[d]) kd = np.fft.fftfreq(shape[d])
@ -22,18 +23,20 @@ def fftk(shape, symmetric=True, finite=False, dtype=np.float32):
def gradient_kernel(kvec, direction, order=1): def gradient_kernel(kvec, direction, order=1):
""" """
Computes the gradient kernel in the requested direction Computes the gradient kernel in the requested direction
Parameters:
----------- Parameters
kvec: array -----------
Array of k values in Fourier space kvec: list
direction: int List of wave-vectors in Fourier space
Index of the direction in which to take the gradient direction: int
Returns: Index of the direction in which to take the gradient
--------
wts: array Returns
Complex kernel --------
""" wts: array
Complex kernel values
"""
if order == 0: if order == 0:
wts = 1j * kvec[direction] wts = 1j * kvec[direction]
wts = jnp.squeeze(wts) wts = jnp.squeeze(wts)
@ -47,41 +50,43 @@ def gradient_kernel(kvec, direction, order=1):
return wts return wts
def laplace_kernel(kvec): def invlaplace_kernel(kvec):
"""
Compute the inverse Laplace kernel
Parameters
-----------
kvec: list
List of wave-vectors
Returns
--------
wts: array
Complex kernel values
""" """
Compute the Laplace kernel from a given K vector
Parameters:
-----------
kvec: array
Array of k values in Fourier space
Returns:
--------
wts: array
Complex kernel
"""
kk = sum(ki**2 for ki in kvec) kk = sum(ki**2 for ki in kvec)
mask = (kk == 0).nonzero() kk_nozeros = jnp.where(kk==0, 1, kk)
kk[mask] = 1 return - jnp.where(kk==0, 0, 1 / kk_nozeros)
wts = 1. / kk
imask = (~(kk == 0)).astype(int)
wts *= imask
return wts
def longrange_kernel(kvec, r_split): def longrange_kernel(kvec, r_split):
""" """
Computes a long range kernel Computes a long range kernel
Parameters:
----------- Parameters
kvec: array -----------
Array of k values in Fourier space kvec: list
r_split: float List of wave-vectors
r_split: float
Splitting radius
Returns
--------
wts: array
Complex kernel values
TODO: @modichirag add documentation TODO: @modichirag add documentation
Returns: """
--------
wts: array
kernel
"""
if r_split != 0: if r_split != 0:
kk = sum(ki**2 for ki in kvec) kk = sum(ki**2 for ki in kvec)
return np.exp(-kk * r_split**2) return np.exp(-kk * r_split**2)
@ -91,15 +96,21 @@ def longrange_kernel(kvec, r_split):
def cic_compensation(kvec): def cic_compensation(kvec):
""" """
Computes cic compensation kernel. Computes cic compensation kernel.
Adapted from https://github.com/bccp/nbodykit/blob/a387cf429d8cb4a07bb19e3b4325ffdf279a131e/nbodykit/source/mesh/catalog.py#L499 Adapted from https://github.com/bccp/nbodykit/blob/a387cf429d8cb4a07bb19e3b4325ffdf279a131e/nbodykit/source/mesh/catalog.py#L499
Itself based on equation 18 (with p=2) of Itself based on equation 18 (with p=2) of
`Jing et al 2005 <https://arxiv.org/abs/astro-ph/0409240>`_ [Jing et al 2005](https://arxiv.org/abs/astro-ph/0409240)
Args:
kvec: array of k values in Fourier space Parameters:
Returns: -----------
v: array of kernel 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)] kwts = [np.sinc(kvec[i] / (2 * np.pi)) for i in range(3)]
wts = (kwts[0] * kwts[1] * kwts[2])**(-2) wts = (kwts[0] * kwts[1] * kwts[2])**(-2)
return wts return wts
@ -107,20 +118,22 @@ def cic_compensation(kvec):
def PGD_kernel(kvec, kl, ks): def PGD_kernel(kvec, kl, ks):
""" """
Computes the PGD kernel Computes the PGD kernel
Parameters:
----------- Parameters:
kvec: array -----------
Array of k values in Fourier space kvec: list
kl: float List of wave-vectors
initial long range scale parameter kl: float
ks: float Initial long range scale parameter
initial dhort range scale parameter ks: float
Returns: Initial dhort range scale parameter
--------
v: array Returns:
kernel --------
""" v: array
Complex kernel values
"""
kk = sum(ki**2 for ki in kvec) kk = sum(ki**2 for ki in kvec)
kl2 = kl**2 kl2 = kl**2
ks4 = ks**4 ks4 = ks**4

115
jaxpm/pm.py
View file

@ -1,48 +1,80 @@
import jax import jax
import jax.numpy as jnp import jax.numpy as jnp
import jax_cosmo as jc import jax_cosmo as jc
from jax_cosmo import Cosmology
from jaxpm.growth import dGfa, growth_factor, growth_rate from jaxpm.growth import growth_factor, growth_rate, dGfa, growth_factor_second, growth_rate_second, dGf2a
from jaxpm.kernels import (PGD_kernel, fftk, gradient_kernel, laplace_kernel, from jaxpm.kernels import PGD_kernel, fftk, gradient_kernel, invlaplace_kernel, longrange_kernel
longrange_kernel)
from jaxpm.painting import cic_paint, cic_read from jaxpm.painting import cic_paint, cic_read
def pm_forces(positions, mesh_shape=None, delta=None, r_split=0):
def pm_forces(positions, mesh_shape, delta=None, r_split=0):
""" """
Computes gravitational forces on particles using a PM scheme Computes gravitational forces on particles using a PM scheme
""" """
if mesh_shape is None:
mesh_shape = delta.shape
kvec = fftk(mesh_shape)
if delta is None: if delta is None:
delta_k = jnp.fft.rfftn(cic_paint(jnp.zeros(mesh_shape), positions)) delta_k = jnp.fft.rfftn(cic_paint(jnp.zeros(mesh_shape), positions))
else: elif jnp.isrealobj(delta):
delta_k = jnp.fft.rfftn(delta) delta_k = jnp.fft.rfftn(delta)
else:
delta_k = delta
# Computes gravitational potential # Computes gravitational potential
pot_k = delta_k * laplace_kernel(kvec) * longrange_kernel(kvec, kvec = fftk(mesh_shape)
r_split=r_split) pot_k = delta_k * invlaplace_kernel(kvec) * longrange_kernel(kvec, r_split=r_split)
# Computes gravitational forces # Computes gravitational forces
return jnp.stack([ return jnp.stack([cic_read(jnp.fft.irfftn(- gradient_kernel(kvec, i) * pot_k), positions)
cic_read(jnp.fft.irfftn(gradient_kernel(kvec, i) * pot_k), positions) for i in range(3)], axis=-1)
for i in range(3)
],
axis=-1)
def lpt(cosmo, initial_conditions, positions, a): def lpt(cosmo:Cosmology, init_mesh, positions, a, 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)
""" """
initial_force = pm_forces(positions, delta=initial_conditions)
a = jnp.atleast_1d(a) a = jnp.atleast_1d(a)
dx = growth_factor(cosmo, a) * initial_force E = jnp.sqrt(jc.background.Esqr(cosmo, a))
p = a**2 * growth_rate(cosmo, a) * jnp.sqrt(jc.background.Esqr(cosmo, delta_k = jnp.fft.rfftn(init_mesh) # TODO: pass the modes directly to save one or two fft?
a)) * dx mesh_shape = init_mesh.shape
f = a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a)) * dGfa(cosmo,
a) * initial_force init_force = pm_forces(positions, mesh_shape, delta=delta_k)
dx = growth_factor(cosmo, a) * init_force
p = a**2 * growth_rate(cosmo, a) * E * dx
f = a**2 * E * dGfa(cosmo, a) * init_force
if order == 2:
kvec = fftk(mesh_shape)
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
# 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
init_force2 = pm_forces(positions, mesh_shape, delta=jnp.fft.rfftn(delta2))
# 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
dx += dx2
p += p2
f += f2
return dx, p, f return dx, p, f
@ -82,10 +114,33 @@ def make_ode_fn(mesh_shape):
return nbody_ode return nbody_ode
def get_ode_fn(cosmo:Cosmology, mesh_shape):
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) * 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): 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: args:
pos: particle positions [npart, 3] pos: particle positions [npart, 3]
params: [alpha, kl, ks] pgd parameters params: [alpha, kl, ks] pgd parameters
@ -96,9 +151,9 @@ def pgd_correction(pos, mesh_shape, params):
delta_k = jnp.fft.rfftn(delta) delta_k = jnp.fft.rfftn(delta)
PGD_range=PGD_kernel(kvec, kl, ks) PGD_range=PGD_kernel(kvec, kl, ks)
pot_k_pgd=(delta_k * laplace_kernel(kvec))*PGD_range 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) forces_pgd= jnp.stack([cic_read(jnp.fft.irfftn(- gradient_kernel(kvec, i)*pot_k_pgd), pos)
for i in range(3)],axis=-1) for i in range(3)],axis=-1)
dpos_pgd = forces_pgd*alpha dpos_pgd = forces_pgd*alpha
@ -107,7 +162,7 @@ def pgd_correction(pos, mesh_shape, params):
def make_neural_ode_fn(model, mesh_shape): 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) state is a tuple (position, velocities)
""" """
@ -119,14 +174,14 @@ def make_neural_ode_fn(model, mesh_shape):
delta_k = jnp.fft.rfftn(delta) delta_k = jnp.fft.rfftn(delta)
# Computes gravitational potential # 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 # Apply a correction filter
kk = jnp.sqrt(sum((ki/jnp.pi)**2 for ki in kvec)) 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))) pot_k = pot_k *(1. + model.apply(params, kk, jnp.atleast_1d(a)))
# Computes gravitational forces # Computes gravitational forces
forces = jnp.stack([cic_read(jnp.fft.irfftn(gradient_kernel(kvec, i)*pot_k), pos) forces = jnp.stack([cic_read(jnp.fft.irfftn(- gradient_kernel(kvec, i)*pot_k), pos)
for i in range(3)],axis=-1) for i in range(3)],axis=-1)
forces = forces * 1.5 * cosmo.Omega_m forces = forces * 1.5 * cosmo.Omega_m