This repo contains a demo of how to use JAX to create a simple gravity simulation. The script main.py uses JAX's experimental ode package to solve the differential equation.
One cool thing about this demo is that it creates a function called gravity that only calculates the force between two bodies, and then uses jax.vmap twice to transform it into a function that calculates the force between all pairs of bodies.
Additionally, the script main_toroidal.py implements a gravitational simulation in an approximation to a toroidal space.
A toroidal topology can be represented by tessellating a rectangle into an infinite plane and then using periodic boundary conditions. Here we approximate the space by using a finite number of replicas.
In this simulation we only use a grid of 3x3 replicas as its faster to compute, and the results are more satisfying to watch. To perform the simulation we use diffrax.
We can also visualize the gravitational field in this space. The following image shows the gravitational field of a single body in the bottom left of the grid, (L/4, L/4).
We see that the gravitational field in this space has 4 fixed points, 1 stable fixed on the body, and 3 unstable fixed points at (L/4, 3L/4), (3L/4, L/4), and (3L/4, 3L/4). These unstable points are equidistant from the body in this space. To get this approximation we used a grid of 401x401 replicas.
pip install -r requirements.txtpoetry installUpon running running the following command, you should the animation of the simulation.
python main.pyChange the parameters for fun :)