A JAX simulator for N body problems.
Project description
NBODYX: Simulating N-Body Systems with Jax
This repository derives ODEs for N-Body systems and simulates them using Jax. The main goal is to provide a simple and efficient way to simulate N-Body systems using Jax. It is heavily inspired by the amazing ThreeBodyBot and specifically the NumericsTutorial.
Installation
You can install the package using pip
:
pip install nbodyx
or locally in editable mode:
pip install -e .
For animations, we require ffmpeg
to be installed. On Ubuntu, you can install it using:
sudo apt-get install ffmpeg
or using conda
:
conda install -c conda-forge ffmpeg
Usage
The following example simulates a 3-body system:
from diffrax import diffeqsolve, Dopri5, ODETerm, SaveAt
import jax.numpy as jnp
from nbodyx.constants import G
from nbodyx.ode import make_n_body_ode
from nbodyx.rendering.opencv import animate_n_body, render_n_body
if __name__ == "__main__":
body_masses = jnp.array([1.0, 1.0, 1.0]) / G
ode_fn = make_n_body_ode(body_masses)
# initial positions
x0 = jnp.array(
[
[-0.97000436, 0.24208753],
[0.0, 0.0],
[0.97000436, -0.24208753],
]
)
# initial velocities
v0 = jnp.array(
[
[0.4662036850, 0.4323657300],
[-0.93240737, -0.86473146],
[0.4662036850, 0.4323657300],
]
)
# initial state
y0 = jnp.stack([x0, v0], axis=0).reshape(-1)
# external torques
tau = jnp.zeros((6,))
# state bounds
x_min, x_max = -1.5 * jnp.ones((1,)), 1.5 * jnp.ones((1,))
# simulation settings
duration = 6.3259 # duration of the simulation [s]
ts = jnp.linspace(0.0, duration, 1001)
dt = 2e-4 # time step [s]
# solve the ODE
ode_term = ODETerm(ode_fn)
sol = diffeqsolve(ode_term, Dopri5(), ts[0], ts[-1], dt, y0, tau, saveat=SaveAt(ts=ts), max_steps=None)
# extract the solution
y_bds_ts = sol.ys.reshape((len(ts), 2, 3, 2)) # shape: (timesteps, 2, num_bodies, num_dims)
# animate the solution
animate_n_body(
ts,
x_bds_ts=y_bds_ts[:, 0, ...],
width=500,
height=500,
video_path="examples/outputs/three_body.mp4",
speed_up=ts[-1] / 10,
x_min=x_min,
x_max=x_max,
timestamp_unit="s",
)
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
nbodyx-0.0.1rc55.post1.tar.gz
(1.2 MB
view hashes)
Built Distribution
Close
Hashes for nbodyx-0.0.1rc55.post1-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 64ee6c4fe434f696fa1fac336adda1f7c68cc071f7cc3e756029b3695ee9f32b |
|
MD5 | 78074d91d861fd0f7e063e3b30014a7a |
|
BLAKE2b-256 | d1fa63a1f73e1f18bb54385175ad73cebd42691639484ca7091d393c549d6194 |