JAX Gravity

This repo contains a demo of how to use JAX to create a simple gravity simulation. It 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.

Instalation

pip

pip install -r requirements.txt

poetry

poetry install

Usage

Upon running running the following command, you should the animation of the simulation.

python main.py

Change the parameters for fun :)