2021-07-16 23:15:27 +02:00
|
|
|
# Design Document for JaxPM
|
|
|
|
|
|
2021-07-17 00:59:03 +02:00
|
|
|
This document aims to detail some of the API, implementation choices, and internal mechanism.
|
2021-07-16 23:15:27 +02:00
|
|
|
|
2021-07-17 00:59:03 +02:00
|
|
|
## Objective
|
|
|
|
|
|
2024-07-09 14:54:34 -04:00
|
|
|
Provide a user-friendly framework for distributed Particle-Mesh N-body simulations.
|
2021-07-17 00:59:03 +02:00
|
|
|
|
|
|
|
|
## Related Work
|
|
|
|
|
|
|
|
|
|
This project would be the latest iteration of a number of past libraries that have provided differentiable N-body models.
|
|
|
|
|
|
|
|
|
|
- [FlowPM](https://github.com/DifferentiableUniverseInitiative/flowpm): TensorFlow
|
2024-07-09 14:54:34 -04:00
|
|
|
- [vmad FastPM](https://github.com/rainwoodman/vmad/blob/master/vmad/lib/fastpm.py): VMAD
|
2021-07-17 00:59:03 +02:00
|
|
|
- Borg
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
In addition, a number of fast N-body simulation projets exist out there:
|
|
|
|
|
- [FastPM](https://github.com/fastpm/fastpm)
|
|
|
|
|
- ...
|
|
|
|
|
|
|
|
|
|
## Design Overview
|
|
|
|
|
|
|
|
|
|
### Coding principles
|
|
|
|
|
|
|
|
|
|
Following recent trends and JAX philosophy, the library should have a functional programming type of interface.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
### Illustration of API
|
|
|
|
|
|
|
|
|
|
Here is a potential illustration of what the user interface could be for the simulation code:
|
|
|
|
|
```python
|
|
|
|
|
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)
|
|
|
|
|
```
|
