Project 5: JAX N-Body Engine

ASTR 596 Fall 2025
Instructor: Anna Rosen
Due: Wednesday, November 26, 2025 by 11:59 PM

Overview¶

In Project 2, you built an N-body simulator in NumPy. It worked, but it was slow — simulating 500 particles probably took an hour per run. Now imagine you’re a researcher generating thousands of simulations to train a machine learning model (your final project). With Project 2’s code, that’s impossible. With JAX, it’s feasible.

Beyond speed, this project teaches you to build research-grade tools, not just scripts:

• Python package structure with modules and __init__.py

• Functional programming patterns used in modern scientific computing

• Reproducible science (explicit random number generation, version control)

The transformation: From writing scripts that work once on your laptop to building organized, importable packages that work reliably for anyone.

Goals:

• 10-100× speedup vs. pure-Python Project 2 (nested loops)

• Speedups vs. vectorized NumPy on CPU are workload-dependent; you must benchmark and report. Aim for a measurable win after JIT warmup.

• JIT compilation and vectorization for significant performance gains

• Reproducible simulations with JAX’s functional random number generation

• Python package structure with modular organization

• Foundation for your final ML project (fast ground truth generation)

Academic Integrity: Hints in this document describe algorithms and patterns; do not copy scaffolds from the internet or LLMs. Your code must be your own implementation.

Learning Objectives¶

• Transform NumPy → JAX (jit, vmap, grad)

• Use jax.lax.scan for timestepping and jax.lax.while_loop for bounded rejection sampling

• Implement pure functional integrators

• Generate and evolve virialized star clusters ($\alpha_{\rm vir} = 1$)

• Design, test, and document a scientific Python package

• Profile and optimize performance

Critical Configuration¶

Part 1: Core Implementation¶

Getting started tip: A good way to familiarize yourself with JAX is to start by repurposing your Project 2 IMF and Plummer sampling modules. Convert them to use jax.random and jax.numpy, JIT-compile them, and validate they produce the same distributions. This gives you a low-stakes way to learn JAX’s functional programming patterns before tackling the more complex integrator.

Step 1: Initial Conditions in JAX¶

Critical: Clusters must start in virial equilibrium to avoid artificial dynamics.

1.1 Kroupa IMF Sampler¶

Two-segment broken power law:

Inverse-CDF for segment $[a, b]$ with slope $\alpha \neq 1$:

Requirements:

• Use jax.random.PRNGKey and jax.random.split (no global RNG)

• JIT-compile with @jit

• Sample exactly $N$ stars

• Default limits: $m_{\min} = 0.08\,M_\odot$, $m_{\max} = 100\,M_\odot$ (configurable)

• Note: $m_{\max} = 100\,M_\odot$ is conventional for star clusters (very massive stars are rare and short-lived)

• Validate: Recover slopes $\{-1.3, -2.3\}$ from log-binned histogram

JAX control flow reminder: Inside @jit, avoid Python for/while loops on traced arrays. Use jax.lax.scan for time-stepping loops and jax.lax.while_loop for bounded rejection sampling or fixed-point iterations. These compile to efficient XLA control flow.

1.2 Plummer Sphere Positions¶

The Plummer sphere is a spherically symmetric stellar system with density profile:

where $a$ is the Plummer scale length (characteristic radius).

Radius inverse-CDF:

Angles (uniform on sphere):

Requirements:

• Re-center: $\sum_i m_i \vec{r}_i = \vec{0}$

• Validate: Density profile matches $\rho(r) \propto (1 + r^2/a^2)^{-5/2}$

1.3 Virial Equilibrium Velocities¶

Goal: Initialize with virial parameter $\alpha_{\rm vir} = 1$ (equilibrium)

The virial parameter (or virial ratio $Q$) is defined as:

where $K$ is total kinetic energy and $W$ is gravitational potential energy (negative).

Physical meaning:

• $\alpha_{\rm vir} = 1$: Virial equilibrium ($2K + W = 0$), bound and stable

• $\alpha_{\rm vir} > 1$: Super-virial (unbound, will expand)

• $\alpha_{\rm vir} < 1$: Sub-virial (will contract)

Your clusters must start at $\alpha_{\rm vir} = 1 \pm 0.05$ to avoid artificial dynamics.

Plummer Equilibrium Velocities (Exact Distribution Function) — REQUIRED:

At radius $r$ with escape speed $v_{\rm esc}(r) = \sqrt{2\psi(r)}$, draw a dimensionless speed $q = v/v_{\rm esc} \in [0,1]$ from:

This comes from Plummer’s exact distribution function (Dejonghe 1987). The peak is at $q^\star = \sqrt{2}/3 \approx 0.471$.

Acceptance-Rejection Implementation:

1. For star at radius $r$, compute $\psi(r) = \frac{GM_{\rm total}}{\sqrt{r^2 + a^2}}$ and $v_{\rm esc} = \sqrt{2\psi}$

2. Use jax.lax.while_loop with bounded max_trials (e.g., 256) to implement rejection sampling

3. Draw $q \sim \mathcal{U}(0,1)$ and accept/reject against $f(q) \propto q^2 (1-q^2)^{7/2}$

4. Use a constant envelope $M = f(q^\star)$ with $q^\star = \sqrt{2}/3$

5. On failure after max_trials: Fallback deterministically to $q = q^\star$ (the mode) and increment a counter

6. Record the fraction of particles that used the fallback — it should be $\ll 1\%$. If not, tighten your envelope

7. Direction: Sample uniform on sphere

8. Compute $\vec{v} = q v_{\rm esc} \hat{n}$

9. Vectorize with vmap over all particles

Important JAX note: Rejection sampling with while_loop requires careful state management. Your loop must carry the accepted value, current attempt count, and a boolean indicating whether to continue.

Critical RNG requirement: Carry a PRNG key in your state and split inside while_loop/scan. No global RNG or hidden state—JAX’s functional random number generation requires explicit key threading through all operations.

Verification:

• Clusters should achieve $\alpha_{\rm vir} = 1 \pm 0.05$ after this initialization

• Re-center velocities: $\sum_i m_i \vec{v}_i = \vec{0}$

Optional: Maxwellian Approximation (for comparison)

If you want to compare methods, you can also implement:

Draw components from $\mathcal{N}(0, v_{\rm disp}^2)$. This is an approximation to the exact distribution but is much simpler and still produces roughly virial equilibrium.

Step 2: Gravitational Force Calculation¶

Physics: The gravitational acceleration on particle $i$ is:

where $\epsilon$ is a softening length that prevents singularities in close encounters.

Why softening? Real stellar systems have finite-size stars. In simulations, point masses can get arbitrarily close, causing:

1. Numerical instability ($r \to 0 \Rightarrow a \to \infty$)

2. Unrealistically large energy changes from two-body encounters

Softening models the smoothing effect of finite stellar radii. For star clusters, typical values: $\epsilon \sim 0.1 - 0.5$ pc, or adaptively: $\epsilon \sim 0.3 r_h / \sqrt{N}$ where $r_h$ is the half-mass radius.

Deliverable for report: State your $\epsilon$ policy (constant vs. $\epsilon \propto r_h/\sqrt{N}$), the chosen value/factor $\kappa$, and provide a 2-line justification tied to energy drift and timestep collapse rates observed in your simulations.

Two Implementation Strategies¶

Strategy A: Direct Vectorization (Pairwise Distance Matrix).

Implement a function accelerations_matrix(x, m, eps2, G) that:

1. Builds all pairwise separations $\Delta \vec{r}_{ij} = \vec{r}_i - \vec{r}_j$

2. Computes softened distances $r_{ij}^2 + \epsilon^2$ for all pairs

3. Zeroes the self-interaction terms (diagonal)

4. Returns accelerations $\vec{a}_i = -G \sum_{j \neq i} m_j \frac{\Delta \vec{r}_{ij}}{(r_{ij}^2 + \epsilon^2)^{3/2}}$

   • Critical: Weight each term by the mass $m_j$ of the source particle

Requirements:

• Fully vectorized (no Python loops)

• JIT-compiled with @jax.jit

• Float64 precision

• Memory: $O(N^2)$ temporaries

Hint: Your validation tests will check force symmetry and agreement with Strategy B to tight tolerances — this will force you to get the broadcasting and mass weighting exactly right.

• Pros: Straightforward vectorization, easier to reason about

• Cons: Memory: $O(N^2)$ temporaries; not scalable to very large $N$

• When to use: Small-to-medium $N$ (< 1000), learning JAX fundamentals

Strategy B: vmap Over Particles

Implement a function accelerations_vmap(x, m, eps2, G) by mapping a single-particle force function over all particles.

Design considerations:

• Write a helper that computes force on one particle due to all others

• Exclude self-interaction by index, not by float equality — your tests will catch float comparison errors

• Use jax.vmap to vectorize over all particles

• Must match Strategy A within rtol=1e-12, atol=1e-14 in float64

Deliverable for report: Compare accelerations from both strategies on the same positions/masses and assert max_rel_err < 1e-12. Submit a short table showing measured error and wall time for N=32, 64, 128.

Requirements:

• JIT-compiled

• Float64 precision

• Agreement with Strategy A verified by tests

Pros: Lower peak memory ($O(N)$ per particle), composable with other transformations like grad

Cons: Slightly more complex logic, requires understanding vmap semantics

When to use: Want to learn vmap, planning to scale to larger $N$, or composing with grad for differentiable physics

Implementation Strategy: Implement both strategies (matrix-based and vmap-based)

• Compare memory usage and performance

• Ensure they produce identical results (see Validation)

• Wrap in @jit for production use

Softening guidance:

• Use $\epsilon = 0.01$ pc for validation tests

• For realistic clusters, estimate $r_h$ (half-mass radius) and use $\epsilon \sim 0.3 r_h / \sqrt{N}$

• Document your choice

Step 3: Time Integration¶

Leapfrog (DKD) Integration — REQUIRED

The leapfrog integrator is second-order symplectic, meaning it exactly conserves a “shadow Hamiltonian” close to the true Hamiltonian. This gives excellent long-term energy conservation.

DKD Algorithm (Drift-Kick-Drift):

1. Drift: $\vec{x}_{1/2} = \vec{x}_n + \vec{v}_n \cdot \frac{\Delta t}{2}$

2. Kick: $\vec{v}_{n+1} = \vec{v}_n + \vec{a}(\vec{x}_{1/2}) \cdot \Delta t$

3. Drift: $\vec{x}_{n+1} = \vec{x}_{1/2} + \vec{v}_{n+1} \cdot \frac{\Delta t}{2}$

Why leapfrog?

• Symplectic: Preserves phase-space volume (Liouville’s theorem)

• Time-reversible: Run backward to recover initial conditions (within numerical precision)

• Excellent energy conservation: Typical $|\Delta E / E| \sim 10^{-6}$ or better for sufficiently small $\Delta t$

Implementation Requirements:

Implement two functions following the DKD equations above:

1. leapfrog_step(state, dt, params) → new_state

   • Pure function (no mutation)

   • state = (x, v), params contains masses, softening, etc.

   • Returns updated (x_new, v_new)

2. integrate(state0, dt, n_steps, params) → (state_final, trajectory)

   • Use jax.lax.scan for the timestepping loop (not Python for)

   • Returns final state and trajectory of states at each timestep

Requirements:

• JIT-compile both functions

• Use jax.lax.scan (not Python loops) for timestepping

• Float64 precision

• Deliverable: Show that $|\Delta E/E| \propto \Delta t^2$ by plotting error vs. timestep

Step 4: Adaptive Timestepping¶

Fixed timesteps are inefficient: During expansion, you can take large steps. During close encounters, you need small steps.

Criterion: Choose $\Delta t$ such that:

where $C \sim 0.01 - 0.1$ is the Courant factor.

Physical meaning: Limit timestep so no particle moves more than $\sim \epsilon$ under its current acceleration.

Implementation Requirements:

• Compute $\Delta t$ at each integration step based on maximum acceleration

• Clamp: $\Delta t \in [\Delta t_{\min}, \Delta t_{\max}]$ to avoid runaway

• Integrate adaptive $\Delta t$ into your lax.scan loop (carry updated dt in state)

Validation: Plot $\Delta t(t)$ over a simulation. It should:

• Decrease during collapse/close encounters (high $|\vec{a}|$)

• Increase during expansion (low $|\vec{a}|$)

• Never reach $\Delta t_{\min}$ or $\Delta t_{\max}$ for reasonable parameters

Discussion point for report: Compare energy conservation and runtime for fixed vs. adaptive timesteps.

Step 5: Diagnostic Utilities¶

Implement helper functions for energy analysis and validation:

Required Functions:

1. kinetic_energy(v, m) → K

   • Total kinetic energy of the system

2. potential_energy(x, m, eps2, G) → W

   • Total gravitational potential energy (negative)

   • Formula: $W = -G \sum_{i<j} \frac{m_i m_j}{\sqrt{r_{ij}^2 + \epsilon^2}}$

   • Critical: Avoid double-counting pairs (each pair $i,j$ counted once)

   • Must handle softening correctly and include gravitational constant $G$ explicitly

3. total_energy(x, v, m, eps2, G) → (K, W, E)

   • Returns kinetic, potential, and total energy

   • $E = K + W$

4. virial_ratio(K, W) → α_vir

   • Computes virial parameter $\alpha_{\rm vir} = 2K / |W|$

Requirements:

• All functions JIT-compiled

• Float64 precision

• Numerically stable

Test Case (validates your implementation):

• Two-body system with specific setup should give: $K = 0.25$, $W = -0.5$ (in your units), $\alpha_{\rm vir} = 1.0$

• Use this to verify you’re not double-counting potential energy

Usage: These standardize your energy diagnostics and make plotting/validation consistent. Use them in your validation tests and include virial ratio plots in your technical report.

Part 2: Production Pipeline¶

Goal: Generate an ensemble of 10 simulations. (This production pipeline will generate the training data for your final ML project for a much larger ensemble with specific parameter ranges.)

Requirements¶

1. Parameter sweep:

   • Vary $N \in [200, 500]$

   • Vary $a$ (Plummer scale length) over $\sim 10\times$ range

   • All clusters start at $\alpha_{\rm vir} = 1 \pm 0.05$

2. Automated pipeline:

   • A simple Python script (e.g., run_ensemble.py) that generates all simulations

   • Explicit random seed management (log seeds for reproducibility)

   • Can be sequential (no need for complex parallelization)

3. Data storage:

   • Organized directory structure (e.g., sim_0000/, sim_0001/, ...)

   • Each simulation saves:

     • Initial conditions: $(x_0, v_0, m)$

     • Trajectories: $(x(t), v(t))$ at $\sim 50-100$ snapshots

     • Metadata: $N$, $a$, $M_{\rm total}$, $\alpha_{\rm vir,0}$, seed

     • Diagnostics: $E(t)$, $\alpha_{\rm vir}(t)$, $\Delta t(t)$

   • Format: NumPy .npz files (simple and efficient) or HDF5 if you prefer. Document your choice in the report.

   • Make sure you include a .gitignore to avoid committing large data files to your repo

4. Statistical validation:

   • Create a single summary plot showing your ensemble is physically consistent

   • Example: Energy conservation histogram across all runs, or $\alpha_{\rm vir}$ distribution

   • Quantify: “95% of simulations achieve $|\Delta E/E| < 10^{-5}$”

Suggested Workflow¶

Your pipeline should:

1. Define parameter grid: Vary $N \in [200, 500]$, Plummer scale $a$ over $\sim 10\times$ range, total mass via IMF sampling

2. Manage random seeds: Use a master seed, split for each simulation, log all seeds for reproducibility

3. Generate initial conditions: For each parameter combination, create virialized cluster

4. Integrate: Run simulation, save trajectories and diagnostics

5. Organize output: Save to structured directories with metadata

6. Statistical validation: Aggregate results and verify ensemble consistency

Implementation tips:

• Use itertools.product or nested loops to generate parameter combinations

• Save incrementally (don’t wait until all runs complete)

• Include progress logging so you know how many simulations are done

• Keep it simple—a straightforward loop over parameters is perfectly fine

Part 3: Performance Analysis¶

Benchmarking Requirements¶

1. Force calculation comparison:

   • Matrix-based vs. vmap: Time both implementations for $N \in \{50, 100, 200, 300, 400, 500\}$

   • Plot: runtime vs. $N$ on log-log scale

   • Quantify: Which is faster? When does the crossover occur?

2. JAX vs. NumPy speedup:

   • Compare your JAX implementation to your Project 2 vectorized NumPy code

   • Fair comparison: Same force calculation strategy, same timestep, measure after JIT warmup

   • Quantify: “JAX is Xfactor faster than NumPy for N=500”

   • Note: Speedups depend strongly on your Project 2 implementation quality. The order-of-magnitude of the runtime depends on your machine—measure and report on your hardware.

3. JIT warmup:

   • Measure first call (includes compilation) vs. subsequent calls

   • Demonstrate the cost of compilation is amortized over many runs

4. Scaling analysis:

   • Plot: total integration time vs. $N$ for fixed $t_{\rm final}$ and adaptive $\Delta t$

   • Quantify: How does your integrator scale? $O(N^2)$? Better?

Profiling¶

Use jax.profiler.trace to identify bottlenecks. This generates a Chrome trace file that visualizes where time is spent (compilation, execution, data transfer).

Steps:

1. Wrap your integration call with jax.profiler.trace(output_path)

2. Load the trace in Chrome at chrome://tracing

3. Analyze: Is force calculation JIT-compiled? Does it dominate runtime? Any unexpected Python overhead?

Report: Include one profiling screenshot showing your force calculation is JIT-compiled and dominates runtime (not Python overhead).

Validation Requirements¶

Required tests (implement at least these 6):

1. Two-body orbit:

• Set up circular or elliptical orbit with known analytic solution

• Integrate for 1000 steps

• Verify: trajectories match theory, energy conserved to $\lesssim 10^{-6}$ over 10³ steps at sufficiently small $\Delta t$. Show the integrator exhibits $\propto \Delta t^2$ error scaling.

2. Force method agreement:

• Generate a random cluster (N=100)

• Compute forces with both matrix and vmap strategies

• Verify: jnp.allclose(a1, a2, rtol=1e-12, atol=1e-14) in float64

3. Force symmetry:

• Set up an equilateral triangle of equal masses

• Verify zero net force at the centroid, or verify pairwise force antisymmetry: $\vec{F}_{ij} = -\vec{F}_{ji}$

4. Units sanity check:

• With $G=1$, place two unit masses at positions $\pm 1$ (separation $d=2$)

• For this check, set softening to $\epsilon = 0$ (i.e., eps2 = 0.0)—with any $\epsilon > 0$ the analytic value changes

• By symmetry, each particle experiences $|\vec{a}| = G m / d^2 = 1/4$

• Verify your force calculation returns exactly this value (tests unit handling and self-exclusion)

5. Sampler validation:

• IMF: Sample 10⁴ masses, verify power-law slopes in log-binned histogram

• Plummer: Sample 10³ positions, compare density profile to analytical form

• Velocities: Both methods (exact and approximate) should produce $\alpha_{\rm vir} \approx 1 \pm 0.05$

6. Gradient safety (mini-check):

• Define a scalar objective $J = \|x_T\|^2$ for a 2-body run of 10 steps

• Verify jax.grad(lambda x0: loss(run(x0))) returns finite values in float64

• Purpose: Ensures your integrator is differentiable (critical for final project physics-informed learning)

N-body cluster validation (included in your 50-100 simulation ensemble):

• Generate realistic clusters (N=200-500, Kroupa IMF, Plummer sphere, $\alpha_{\rm vir} = 1$)

• Integrate for several dynamical times

• Check that $\alpha_{\rm vir}$ remains near 1 (some evolution is physical)

• Verify momentum conservation

Deliverables¶

1. Code Repository¶

Structure:

Recommended directory layout (your first Python package!):

project5-jax-nbody/
├── jax_nbody/              # Your package directory
│   ├── __init__.py         # Makes this a package (can be empty)
│   ├── samplers.py         # IMF, Plummer sampling
│   ├── forces.py           # Force calculations
│   ├── integrators.py      # Leapfrog, adaptive dt
│   └── diagnostics.py      # Energy, virial ratio
├── tests/
│   ├── test_forces.py      # At least 6 tests total
│   ├── test_samplers.py
│   └── ...
├── quickstart.py           # Demo script for grading
├── run_ensemble.py         # Generate 50-100 simulations
├── data/                   # Simulation outputs
│   ├── sim_0000/
│   └── ...
└── README.md

Why this structure?

• __init__.py: Makes jax_nbody/ a Python package so you can import jax_nbody.forces

• Modular files: Separate concerns (sampling vs. forces vs. integration)

• tests/ directory: Keep tests organized

• Scripts at root: quickstart.py and run_ensemble.py import from your package

Example imports in your scripts:

# In quickstart.py
from jax_nbody.samplers import sample_kroupa_imf, sample_plummer_positions
from jax_nbody.forces import compute_forces_matrix
from jax_nbody.integrators import leapfrog_integrate

# Or if you organize differently:
import jax_nbody.samplers as samplers

What you DON’T need:

• ❌ pyproject.toml or setup.py (no pip installation)

• ❌ Complex packaging machinery

• ❌ Version numbers, dependencies lists

What you DO need:

• ✅ __init__.py in your package directory (can be empty)

• ✅ Organized modules with clear responsibilities

• ✅ Tests that import and verify your functions

• ✅ Scripts that use your package

Code quality (optional): Consider using ruff check jax_nbody/ for linting, but not required

Git: 20+ meaningful commits showing development progression

RNG threading: All samplers and integrators must accept and return a JAX PRNGKey, with no hidden global RNG state. Tests will fail if randomness isn’t threaded explicitly through function calls

Quickstart for grading: quickstart.py at project root that runs N=64 for 100 steps on CPU and prints: wall-time (after JIT warmup), max|ΔE/E|, and $\alpha_{\rm vir}$ summary. I will run python quickstart.py to verify basic functionality

2. Research Memo (1-2 pages)¶

Write a concise research memo demonstrating your implementation works correctly and that you understand JAX’s performance advantages. The primary deliverable is your complete working, validated pipeline—the memo provides essential evidence and commentary.

Required Components:

1. Validation Evidence (≈1 page):

   • All 6 required test results with key figures showing correctness

   • Two-body orbit plot: Energy conservation ($\max|\Delta E/E| < 10^{-8}$)

   • Virial equilibrium plot: $\alpha_{\rm vir} = 1 \pm 0.05$ for virialized cluster

   • Multi-scale validation: Show correct physics from N=2 → N=100+ → production-scale clusters

   • Use informative captions explaining what each plot demonstrates

2. Performance Analysis (≈0.5 pages):

   • JAX vs. NumPy speedup table: Quantified timings for representative N values

   • Force strategy comparison: Matrix vs. vmap measured error and wall time (N=32, 64, 128)

   • Brief commentary: What dominates runtime? How does performance scale with N?

3. Key Design Decisions (≈0.5 pages):

   • Softening policy: State your $\epsilon$ choice (constant vs. $\epsilon \propto r_h/\sqrt{N}$), the value/factor used, and brief justification tied to energy drift and timestep stability

   • Integration approach: Note your adaptive timestepping strategy

   • Production ensemble: Brief description of automated pipeline (50-100 runs, parameter coverage, data organization)

Format: Emphasize showing results over prose. Use tables and figures effectively. Quantify performance gains. Demonstrate you understand the tradeoffs through your results, not lengthy explanations.

3. Growth Memo¶

Submit a Growth Memo following the standard template at 03-growth-memo-template.md. This is where you reflect on your learning journey, including:

• Technical skills developed: What JAX patterns can you now implement?

• Key challenges & solutions: What JAX-specific pitfalls did you encounter?

• AI usage reflection: How did you use AI tools to understand JAX? What worked? What didn’t? How did you verify AI suggestions? This detailed reflection belongs in the Growth Memo, NOT the research memo.

• Conceptual insights: How did implementing N-body dynamics in JAX deepen your understanding of functional programming and JIT compilation?

• What got you excited: Any “aha” moments with JAX transformations (jit, vmap, grad)?

The Growth Memo is graded pass/fail based on honest, thoughtful reflection. See the template for detailed prompts.

4. AI as a Learning Accelerator¶

Through Projects 1-4, you’ve demonstrated solid Python expertise. Now it’s time to leverage that foundation with modern learning tools. JAX has a sharp learning curve—its functional programming paradigm and compilation constraints are fundamentally different from NumPy. But you have an advantage: you already understand the underlying physics and computational concepts.

Use AI to understand, not to generate code. You need to write your own implementation to truly learn JAX. Instead, use AI tools (e.g., ChatGPT, Claude, Gemini) to:

• Explain difficult programming concepts: “What does it mean for JAX arrays to be immutable?” or “Why can’t I use Python if statements in JIT-compiled functions?”

• Decode error messages: JAX errors can be verbose and cryptic. Paste the error and ask “What is this tracer error telling me?”

• Trace issues: “My code runs but energy conservation is terrible—what are common causes in JAX?”

• Compare paradigms: “I know how to do X in NumPy with mutable arrays, how does JAX’s functional approach differ conceptually?”

Learning from mistakes is how you master new frameworks. It’s common—expected, even—to encounter JAX-specific pitfalls. These errors (tracer errors, array immutability violations, JIT compilation failures) are part of the learning process. When you hit one, use AI to understand why it happened and what concept you’re missing, then fix it yourself.

Critical caveat: Maintain a “docs first” mindset. Check the JAX documentation and Equinox API before asking AI. Always fact-check AI responses—they’re often confident but wrong about version-specific details or edge cases. Use AI to understand concepts and documentation, not as a replacement for understanding.

AI usage reflection belongs in your Growth Memo (see Deliverable 3 below), not in the research memo. The Growth Memo has dedicated sections for reflecting on AI tool usage, what worked, what didn’t, and how your approach evolved.

Common Pitfalls & Debugging Guide¶

Debugging Strategy¶

When something goes wrong (and it will):

1. Start small: Does your code work for N=2? N=3? If not, fix that first.

2. Check conservation laws: Energy, momentum, angular momentum are your physics smoke detectors

3. Visualize: Plot trajectories, $\alpha_{\rm vir}(t)$, $\Delta t(t)$ — your eyes catch patterns tests miss

4. Compare: Does your JAX code give the same answer as Project 2 NumPy code for identical initial conditions?

5. Bisect: Comment out half your code. Which half has the bug?

6. Read error messages carefully: JAX errors are verbose but informative

Most common mistake: Assuming initial conditions are correct. Always validate IMF slopes, Plummer density, and $\alpha_{\rm vir}$ before running long integrations.

Essential JAX Functions & Tools¶

Core JAX¶

Note: jax.lax.scan and jax.lax.while_loop are critical for JIT-compiled loops—do not use Python for or while inside jitted functions. Additionally, when compiled jax.lax.scan is preferred for fixed-iteration loops (like timestepping) and is faster than jax.lax.while_loop and jax.lax.ifor_loop, while jax.lax.while_loop is for data-dependent loops (like rejection sampling).

Performance Tools¶

Optional Tools¶

Grading Approach¶

Evaluation Method: Your grade is based on your working validated pipeline (code quality), research memo, quickstart demo, and validation evidence.

Note: Grading will be based on your quickstart demo output, repository code reading, and research memo. I will not be installing your package or running your full test suite—your memo must provide convincing evidence that your implementation works correctly. The Growth Memo is graded separately as pass/fail.

Looking Ahead: Final Project¶

The automated simulation pipeline you built is the foundation for your final ML project:

• Training data generation: Your 50-100 simulations are a prototype. For the final project, you’ll scale to 1000+ simulations to train Gaussian Processes and Neural Networks

• Fast ground truth: JAX’s 10-100× speedup makes this computationally feasible

• Differentiable physics: JAX’s grad enables physics-informed learning (gradients through entire simulations)

• Research-grade infrastructure: Professional package + automated pipeline = reproducible science

What’s next in the final project:

• Gaussian Processes: Learn cluster evolution from your simulation ensemble

• Neural Networks: Emulate N-body dynamics at 1000× speedup

• Uncertainty quantification: How confident should we be in ML predictions?

The computational infrastructure you build now directly enables machine learning for astrophysics.