This document aims to detail some of the API, implementation choices, and internal mechanism.
Provide a user-friendly framework for distributed Particle-Mesh N-body simulations.
This project would be the latest iteration of a number of past libraries that have provided differentiable N-body models.
- FlowPM: TensorFlow
- vmad FastPM: VMAD
- Borg
In addition, a number of fast N-body simulation projets exist out there:
- FastPM
- ...
Following recent trends and JAX philosophy, the library should have a functional programming type of interface.
Here is a potential illustration of what the user interface could be for the simulation code:
import jaxpm as jpm
import jax_cosmo as jc
# Instantiate differentiable cosmology object
cosmo = jc.Planck()
# Creates initial conditions
inital_conditions = jpm.generate_ic(cosmo, boxsize, nmesh, dtype='float32')
# Create a particular solver
solver = jpm.solvers.fastpm(cosmo, B=1)
# Initialize and run the simulation
state = solver.init(initial_conditions)
state = solver.nbody(state)
# Painting the results
density = jpm.zeros(boxsize, nmesh)
density = jpm.paint(density, state.positions)