Simulate Gravity with JAX
Watch planets orbit each other in a simulation you built from scratch
Last reviewed: March 2026Overview
Every planet, moon, and satellite in our solar system moves according to Newton's law of gravity. In this project you'll write a small Python program using JAX — Google's numerical computing library — to simulate gravitational motion from first principles.
You'll start with two bodies (like the Earth and Moon) and watch them orbit each other, then add a third body and see how the motion becomes chaotic and unpredictable. Along the way you'll learn why JAX is special: its automatic differentiation feature lets you compute gradients without writing any calculus by hand, which is the foundation of modern machine learning.
By the end you'll have animated plots of orbital trajectories and an intuition for why three-body problems are so hard — a topic that has fascinated mathematicians for centuries and remains unsolved today.
What You'll Learn
- ✓ Implement Newton's law of universal gravitation in code
- ✓ Use numerical integration (Euler and RK4 methods) to simulate motion over time
- ✓ Understand what automatic differentiation is and why JAX provides it
- ✓ Create animated plots of orbital trajectories using matplotlib
- ✓ Explore the difference between stable two-body orbits and chaotic three-body motion
Step-by-Step Guide
Install JAX and Set Up
Install JAX with pip install jax jaxlib. Create a Jupyter notebook and verify JAX works by running import jax.numpy as jnp. JAX works like NumPy but can automatically compute derivatives of any function — you'll see why that matters later.
Code the Gravity Force
Write a function that computes the gravitational force between two masses given their positions. Use F = G * m1 * m2 / r² where r is the distance between them. Return the force as a vector pointing from one mass to the other.
Simulate Two Bodies
Use Euler integration to step the simulation forward in time: update velocity from force, then update position from velocity. Start with Earth and Moon masses and distances. Run for one full orbit and plot the trajectory of both bodies.
Upgrade to RK4 Integration
Euler integration drifts over time. Implement the 4th-order Runge-Kutta (RK4) method for much better accuracy. Compare: run both integrators for 100 orbits and see how Euler's orbit spirals while RK4 stays stable.
Add a Third Body
Add a third mass (like a small asteroid) and watch the chaos. Three-body motion is fundamentally unpredictable — tiny changes in starting position lead to completely different trajectories. Try a few different starting positions and overlay the resulting paths.
Use JAX Autodiff
Use jax.grad to automatically compute the gradient of gravitational potential energy with respect to position. Verify it matches your hand-coded force function. This is the same autodiff technology that powers all modern neural network training.
Career Connection
See how this project connects to real aerospace careers.
Aerospace Engineer →
Orbital mechanics simulation is fundamental to spacecraft mission design and trajectory planning
Space Operations →
Understanding orbital dynamics helps mission controllers predict satellite positions and plan maneuvers
Astronaut →
Astronauts study orbital mechanics to understand rendezvous, docking, and re-entry trajectories
Go Further
Extend your simulator:
- Solar system model — add all 8 planets and simulate the full solar system
- Lagrange points — find the 5 stable/unstable equilibrium points in the Earth-Sun system
- Hohmann transfers — simulate the fuel-efficient orbit transfer used by spacecraft
- GPU acceleration — use
jax.jitto compile your simulation and run it 100x faster