Cloud in cell interpolation 
\(\begin{array}{c}{{\nabla^{2}\phi=-4\pi G\rho}}\\\\ {{f(\vec{k})=i\vec{k}k^{-2}\rho(\vec{k})}}\end{array}\)
+
\(\begin{array}{c}{{\nabla^{2}\phi=-4\pi G\rho}}\\\\ {{f(\vec{k})=i\vec{k}k^{-2}\rho(\vec{k})}}\end{array}\)
\(\begin{array}{c}{{\nabla^{2}\phi=-4\pi G\rho}}\\\\ {{f(\vec{k})=i\vec{k}k^{-2}\rho(\vec{k})}}\end{array}\)
-
example of numpy code and grad
+
import numpy as np
+
+
+def multiply_and_add(a, b, c):
+ return np.dot(a, b) + c
+
+
+a, b, c = np.random.normal(size=(3, 32, 32))
+result = multiply_and_add(a, b, c)
import jax
+import jax.numpy as jnp
+
+@jax.jit
+def multiply_and_add(a, b, c):
+ return jnp.dot(a, b) + c
+
+
+key = jax.random.PRNGKey(0)
+a, b, c = jax.random.normal(key, (3, 32, 32))
+
+result = multiply_and_add(a, b, c)
+gradient = jax.grad(multiply_and_add)(a, b, c)
JAX : Numpy + Autograd + GPU
+example of JAXPM code
+
+
DifferentiableUniverseInitiative/JaxPM
+
mesh_shape = [64, 64, 64]
+box_size = [64., 64., 64.]
+snapshots = jnp.linspace(0.1, 1., 2)
+
+@jax.jit
+def run_simulation(omega_c, sigma8):
+ # Create a small function to generate the matter power spectrum
+ k = jnp.logspace(-4, 1, 128)
+ pk = jc.power.linear_matter_power(
+ jc.Planck15(Omega_c=omega_c, sigma8=sigma8), k)
+ pk_fn = lambda x: jc.scipy.interpolate.interp(x.reshape([-1]), k, pk
+ ).reshape(x.shape)
+
+ # Create initial conditions
+ initial_conditions = linear_field(mesh_shape,
+ box_size,
+ pk_fn,
+ seed=jax.random.PRNGKey(0))
+
+ # Create particles
+ particles = jnp.stack(jnp.meshgrid(*[jnp.arange(s) for s in mesh_shape]),
+ axis=-1).reshape([-1, 3])
+
+ cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)
+
+ # Initial displacement
+ dx, p, f = lpt(cosmo, initial_conditions, particles, 0.1)
+ field = dx + particles
+
+ # Evolve the simulation forward
+ res = odeint(make_ode_fn(mesh_shape), [particles + dx, p],
+ snapshots,
+ cosmo,
+ rtol=1e-5,
+ atol=1e-5)
+
+ # Return the simulation volume at requested
+ return field, res, initial_conditions
+
+field, res, initial_conditions = run_simulation(0.25, 0.8)
Is everything solved ?
+➢ only operation that requires communication is the FFT
+
+ DifferentiableUniverseInitiative/jaxDecomp
+
add animation of the decomposition
+add plots of scallig
+jaxDecomp is also provided as a standalone library
+