mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-02-23 10:00:54 +00:00
94 lines
2.9 KiB
Python
94 lines
2.9 KiB
Python
import jax
|
|
from jax.lax import linear_solve_p
|
|
import jax.numpy as jnp
|
|
from jax.experimental.maps import xmap
|
|
from functools import partial
|
|
import jax_cosmo as jc
|
|
|
|
from jaxpm.kernels import fftk
|
|
import jaxpm.experimental.distributed_ops as dops
|
|
from jaxpm.growth import growth_factor, growth_rate, dGfa
|
|
|
|
|
|
def pm_forces(positions, mesh_shape=None, delta_k=None, halo_size=16):
|
|
"""
|
|
Computes gravitational forces on particles using a PM scheme
|
|
"""
|
|
if mesh_shape is None:
|
|
mesh_shape = delta_k.shape
|
|
kvec = [k.squeeze() for k in fftk(mesh_shape, symmetric=False)]
|
|
|
|
if delta_k is None:
|
|
delta = dops.cic_paint(positions, mesh_shape, halo_size)
|
|
delta_k = dops.fft3d(dops.reshape_split_to_dense(delta))
|
|
|
|
forces_k = dops.gradient_laplace_kernel(delta_k, kvec)
|
|
|
|
# Recovers forces at particle positions
|
|
forces = [dops.cic_read(dops.reshape_dense_to_split(dops.ifft3d(f)),
|
|
positions, halo_size) for f in forces_k]
|
|
|
|
return dops.stack3d(*forces)
|
|
|
|
|
|
def linear_field(cosmo, mesh_shape, box_size, seed, return_Fourier=True):
|
|
"""
|
|
Generate initial conditions.
|
|
Seed should have the dimension of the computational mesh
|
|
"""
|
|
|
|
# Sample normal field
|
|
field = dops.normal(seed, shape=mesh_shape)
|
|
|
|
# Go to Fourier space
|
|
field = dops.fft3d(dops.reshape_split_to_dense(field))
|
|
|
|
# Rescaling k to physical units
|
|
kvec = [k.squeeze() / box_size[i] * mesh_shape[i]
|
|
for i, k in enumerate(fftk(mesh_shape, symmetric=False))]
|
|
k = jnp.logspace(-4, 2, 256)
|
|
pk = jc.power.linear_matter_power(cosmo, k)
|
|
pk = pk * (mesh_shape[0] * mesh_shape[1] * mesh_shape[2]
|
|
) / (box_size[0] * box_size[1] * box_size[2])
|
|
|
|
field = dops.scale_by_power_spectrum(field, kvec, k, jnp.sqrt(pk))
|
|
|
|
if return_Fourier:
|
|
return field
|
|
else:
|
|
return dops.reshape_dense_to_split(dops.ifft3d(field))
|
|
|
|
|
|
def lpt(cosmo, initial_conditions, positions, a):
|
|
"""
|
|
Computes first order LPT displacement
|
|
"""
|
|
initial_force = pm_forces(positions, delta_k=initial_conditions)
|
|
a = jnp.atleast_1d(a)
|
|
dx = dops.scalar_multiply(initial_force, growth_factor(cosmo, a))
|
|
p = dops.scalar_multiply(dx, a**2 * growth_rate(cosmo, a) *
|
|
jnp.sqrt(jc.background.Esqr(cosmo, a)))
|
|
return dx, p
|
|
|
|
|
|
def make_ode_fn(mesh_shape):
|
|
|
|
def nbody_ode(state, a, cosmo):
|
|
"""
|
|
state is a tuple (position, velocities)
|
|
"""
|
|
pos, vel = state
|
|
|
|
forces = pm_forces(pos, mesh_shape=mesh_shape) * 1.5 * cosmo.Omega_m
|
|
|
|
# Computes the update of position (drift)
|
|
dpos = dops.scalar_multiply(
|
|
vel, 1. / (a**3 * jnp.sqrt(jc.background.Esqr(cosmo, a))))
|
|
|
|
# Computes the update of velocity (kick)
|
|
dvel = dops.scalar_multiply(
|
|
forces, 1. / (a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a))))
|
|
|
|
return dpos, dvel
|
|
|
|
return nbody_ode
|