forked from Aquila-Consortium/JaxPM_highres
271 lines
9 KiB
Python
271 lines
9 KiB
Python
import jax.numpy as jnp
|
|
import jax_cosmo as jc
|
|
|
|
from jaxpm.distributed import fft3d, 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,
|
|
invlaplace_kernel, longrange_kernel)
|
|
from jaxpm.painting import cic_paint, cic_paint_dx, cic_read, cic_read_dx
|
|
|
|
|
|
def pm_forces(positions,
|
|
mesh_shape=None,
|
|
delta=None,
|
|
r_split=0,
|
|
paint_absolute_pos=True,
|
|
halo_size=0,
|
|
fd=False,
|
|
sharding=None):
|
|
"""
|
|
Computes gravitational forces on particles using a PM scheme
|
|
"""
|
|
if mesh_shape is None:
|
|
assert (delta is not None),\
|
|
"If mesh_shape is not provided, delta should be provided"
|
|
mesh_shape = delta.shape
|
|
|
|
if paint_absolute_pos:
|
|
paint_fn = lambda pos: cic_paint(jnp.zeros(shape=mesh_shape,
|
|
device=sharding),
|
|
pos,
|
|
halo_size=halo_size,
|
|
sharding=sharding)
|
|
read_fn = lambda grid_mesh, pos: cic_read(
|
|
grid_mesh, pos, halo_size=halo_size, sharding=sharding)
|
|
else:
|
|
paint_fn = lambda disp: cic_paint_dx(
|
|
disp, halo_size=halo_size, sharding=sharding)
|
|
read_fn = lambda grid_mesh, disp: cic_read_dx(
|
|
grid_mesh, disp, halo_size=halo_size, sharding=sharding)
|
|
|
|
if delta is None:
|
|
field = paint_fn(positions)
|
|
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 * invlaplace_kernel(kvec, fd=fd) * longrange_kernel(
|
|
kvec, r_split=r_split)
|
|
# Computes gravitational forces
|
|
forces = jnp.stack([
|
|
read_fn(ifft3d(-gradient_kernel(kvec, i, fd=fd) * pot_k),positions
|
|
) for i in range(3)], axis=-1) # yapf: disable
|
|
|
|
return forces
|
|
|
|
|
|
def lpt(cosmo,
|
|
initial_conditions,
|
|
particles=None,
|
|
a=0.1,
|
|
halo_size=0,
|
|
sharding=None,
|
|
order=1):
|
|
"""
|
|
Computes first and second order LPT displacement and momentum,
|
|
e.g. Eq. 2 and 3 [Jenkins2010](https://arxiv.org/pdf/0910.0258)
|
|
"""
|
|
paint_absolute_pos = particles is not None
|
|
if particles is None:
|
|
particles = jnp.zeros_like(initial_conditions,
|
|
shape=(*initial_conditions.shape, 3))
|
|
|
|
a = jnp.atleast_1d(a)
|
|
E = jnp.sqrt(jc.background.Esqr(cosmo, a))
|
|
delta_k = fft3d(initial_conditions)
|
|
initial_force = pm_forces(particles,
|
|
delta=delta_k,
|
|
paint_absolute_pos=paint_absolute_pos,
|
|
halo_size=halo_size,
|
|
fd=False,
|
|
sharding=sharding)
|
|
dx = growth_factor(cosmo, a) * initial_force
|
|
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, fd=False)**2
|
|
shear_ii = ifft3d(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, fd=False) * gradient_kernel(
|
|
kvec, j, fd=False)
|
|
delta2 -= ifft3d(nabla_i_nabla_j * pot_k)**2
|
|
|
|
delta_k2 = fft3d(delta2)
|
|
init_force2 = pm_forces(particles,
|
|
delta=delta_k2,
|
|
paint_absolute_pos=paint_absolute_pos,
|
|
halo_size=halo_size,
|
|
fd=False,
|
|
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
|
|
|
|
dx += dx2
|
|
p += p2
|
|
f += f2
|
|
|
|
return dx, p, f
|
|
|
|
|
|
def linear_field(mesh_shape, box_size, pk, seed, sharding=None):
|
|
"""
|
|
Generate initial conditions.
|
|
"""
|
|
# Initialize a random field with one slice on each gpu
|
|
field = normal_field(mesh_shape, seed=seed, sharding=sharding)
|
|
field = fft3d(field)
|
|
kvec = fftk(field)
|
|
kmesh = sum((kk / box_size[i] * mesh_shape[i])**2
|
|
for i, kk in enumerate(kvec))**0.5
|
|
pkmesh = pk(kmesh) * (mesh_shape[0] * mesh_shape[1] * mesh_shape[2]) / (
|
|
box_size[0] * box_size[1] * box_size[2])
|
|
|
|
field = field * (pkmesh)**0.5
|
|
field = ifft3d(field)
|
|
return field
|
|
|
|
|
|
def make_ode_fn(mesh_shape,
|
|
paint_absolute_pos=True,
|
|
halo_size=0,
|
|
sharding=None):
|
|
|
|
def nbody_ode(state, a, cosmo):
|
|
"""
|
|
state is a tuple (position, velocities)
|
|
"""
|
|
pos, vel = state
|
|
|
|
forces = pm_forces(pos,
|
|
mesh_shape=mesh_shape,
|
|
paint_absolute_pos=paint_absolute_pos,
|
|
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 dpos, dvel
|
|
|
|
return nbody_ode
|
|
|
|
|
|
def make_diffrax_ode(cosmo,
|
|
mesh_shape,
|
|
paint_absolute_pos=True,
|
|
halo_size=0,
|
|
sharding=None):
|
|
|
|
def nbody_ode(a, state, args):
|
|
"""
|
|
state is a tuple (position, velocities)
|
|
"""
|
|
pos, vel = state
|
|
|
|
forces = pm_forces(pos,
|
|
mesh_shape=mesh_shape,
|
|
paint_absolute_pos=paint_absolute_pos,
|
|
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
|
|
|
|
args:
|
|
pos: particle positions [npart, 3]
|
|
params: [alpha, kl, ks] pgd parameters
|
|
"""
|
|
delta = cic_paint(jnp.zeros(mesh_shape), pos)
|
|
delta_k = fft3d(delta)
|
|
kvec = fftk(delta_k)
|
|
alpha, kl, ks = params
|
|
PGD_range = PGD_kernel(kvec, kl, ks)
|
|
|
|
pot_k_pgd = (delta_k * invlaplace_kernel(kvec)) * PGD_range
|
|
|
|
forces_pgd = jnp.stack([
|
|
cic_read(fft3d(-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: Cosmology, params):
|
|
"""
|
|
state is a tuple (position, velocities)
|
|
"""
|
|
pos, vel = state
|
|
delta = cic_paint(jnp.zeros(mesh_shape), pos)
|
|
delta_k = fft3d(delta)
|
|
kvec = fftk(delta_k)
|
|
|
|
# Computes gravitational potential
|
|
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)))
|
|
|
|
# Computes gravitational forces
|
|
forces = jnp.stack([
|
|
cic_read(fft3d(-gradient_kernel(kvec, i) * pot_k), pos)
|
|
for i in range(3)
|
|
],
|
|
axis=-1)
|
|
|
|
forces = forces * 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 dpos, dvel
|
|
|
|
return neural_nbody_ode
|