You can deploy this model and start earning money today!

asgeirr89/
orbits

Views

No views yet

OrbitMLP: Neural Network Orbital Trajectory Predictor

Overview

OrbitMLP is a deep learning model that learns to predict orbital trajectories using Keplerian dynamics. Instead of numerically integrating orbits step-by-step with RK4, the neural network learns to directly predict the next state given the current state.

The model is trained using physics-informed loss functions that enforce energy conservation and angular momentum conservation, in addition to standard MSE regression on trajectory data.

Key Features:

Pure JAX/Flax implementation for automatic differentiation and GPU acceleration
Physics-informed training with energy and angular momentum constraints
Lightweight MLP architecture with residual blocks
Direct trajectory prediction without iterative solvers

Demo

The following results show OrbitMLP compared against the ground-truth RK4 integrator:

Left Panel: Trajectory comparison showing an elliptic orbit. Blue = RK4 (ground truth), Red = OrbitMLP prediction. The neural network learns to reproduce the orbital shape with high fidelity.
Center Panel: Position vs time for x and y coordinates. The NN closely follows RK4 across the entire trajectory.
Right Panel: Energy conservation comparison. Both RK4 and the neural network maintain approximately constant total energy throughout the orbit.

Installation

pip install jax jaxlib flax optax numpy matplotlib pyyaml

Usage

Training

Configure your hyperparameters in config.yaml and run:

python main.py

Inference / Prediction

python

import jax
import jax.numpy as jnp
import numpy as np
from flax import serialization
from model import OrbitMLP
from train import make_predict_trajectory

# Load the model
with open("models/orbitmlp_20260505_033302.flax", "rb") as f:
    params = serialization.from_bytes(jax.random.PRNGKey(0), f.read())

# Create model and prediction function
model = OrbitMLP()
predict_trajectory = make_predict_trajectory(model)

# Predict 500 steps from initial state [x, y, vx, vy]
init_state = jnp.array([1.5, 0.0, 0.0, 0.8], dtype=jnp.float32)
num_steps = 500
nn_traj = predict_trajectory(params, init_state, num_steps)

print(f"Trajectory shape: {nn_traj.shape}")  # (501, 4)

Compare with RK4 Ground Truth

python

from physics_engine import rk4_step, energy

dt = 0.05
gm = 1.0
num_steps = 500

# RK4 integration
rk4_traj = np.zeros((num_steps + 1, 4), dtype=np.float32)
rk4_traj[0] = np.array(init_state)
s = init_state
for i in range(num_steps):
    s, _ = rk4_step(s, dt, gm)
    rk4_traj[i + 1] = np.array(s)

# Compute energies
nn_energies = np.array([energy(nn_traj[i], gm) for i in range(num_steps + 1)])
rk4_energies = np.array([energy(rk4_traj[i], gm) for i in range(num_steps + 1)])

mse = np.mean((nn_traj - rk4_traj) ** 2)
energy_drift = nn_energies[-1] - nn_energies[0]

print(f"Position MSE vs RK4: {mse:.6e}")
print(f"Energy drift (NN): {energy_drift:.6e}")

Architecture

OrbitMLP

Input (4) → Dense(128) → ResidualBlock × 3 → Dense(4)

ResidualBlock

x → Dense → LayerNorm → GELU → Dense → LayerNorm → GELU → Add → output

The model uses He normal initialization and LayerNorm for stability.

Component	Value
Hidden dimension	128
Number of residual blocks	3
Activation	GELU
Initialization	He normal

Training Details

Hyperparameters

Parameter	Value
Epochs	5000
Batch size	64
Learning rate	1e-3
Optimizer	AdamW with cosine decay
Initial decay steps	2000
Final learning rate ratio	1e-4

Loss Function

L_total = MSE + λ_energy × L_energy + λ_angular × L_angular

Where:

MSE: Mean squared error between predicted and target states
L_energy: Mean absolute error of orbital energy (|E_pred - E_target|)
L_angular: Variance of angular momentum (encourages conservation)
λ_energy = 0.1
λ_angular = 0.1

Data Generation

Training data is generated by integrating random initial conditions using RK4:

Random radii: uniform(0.8, 2.0)
Random velocities: uniform(0.4, 1.2) with perpendicular direction
100 integration steps per trajectory at dt=0.05
64 trajectories per training run

Physics

Kepler's Equations

The model learns the two-body problem gravitational dynamics:

a = -GM/r³ × r

Where:

r = (x, y) is the position vector
GM = 1.0 (normalized units)
a = (ax, ay) is the acceleration

Energy

Total orbital energy (conserved in bound orbits):

E = 0.5 × (vx² + vy²) - GM/r

Angular Momentum

Angular momentum per unit mass (also conserved):

L = x × vy - y × vx

Model Files

File	Description
`orbitmlp_20260505_033302.flax`	Latest trained model

Project Structure

orbitas/
├── main.py              # Training pipeline
├── train.py             # Training utilities and loss functions
├── model.py             # OrbitMLP architecture
├── physics_engine.py    # Keplerian dynamics and RK4 integrator
├── predict.py           # Inference script
├── checks.py            # Pre-flight checks
├── config.yaml          # Hyperparameters
├── requirements.txt    # Dependencies
└── orbit_comparison.png # Example results

Technologies Used

Library	Purpose
JAX	Autodiff, XLA compilation, GPU acceleration
Flax	Neural network framework
Optax	Optimizers (AdamW + cosine decay)
NumPy	Numerical computation
Matplotlib	Visualization

License

MIT License - see LICENSE file for details.

Citation

If you use this model in your research, please cite:

bibtex

@software{orbitas,
  author = {asgeirr89},
  title = {OrbitMLP: Neural Network Orbital Trajectory Predictor},
  url = {https://huggingface.co/asgeirr89/orbitas},
  year = {2026},
}

asgeirr89/orbits