JAX TPU Benchmark & Physics Simulation 🚀

1. Molecular Dynamics (molecular_dynamics_jax_single-host_workload.py)

This script implements a 2D Molecular Dynamics (MD) simulation of a Lennard-Jones fluid, written purely in JAX.

How to Run:

# (Basic) Run with default settings (N=400, 10k steps)
python3 molecular_dynamics_jax_single-host_workload.py

# (Advanced) Run a longer simulation with more particles and lower density
python3 molecular_dynamics_jax_single-host_workload.py --N 800 --rho 0.7 --prod_steps 50000 --eq_steps 20000

All Arguments:

• --N (int, default: 400): Number of particles.
• --rho (float, default: 0.8): Density.
• --kT (float, default: 1.0): Temperature (kT).
• --dt (float, default: 1e-3): Time step size.
• --eq_steps (int, default: 10000): Number of equilibration steps.
• --prod_steps (int, default: 10000): Number of production (simulation) steps.
• --sample_every (int, default: 100): Sample the state every N steps.
• --seed (int, default: 42): PRNG seed.
• --output (str, default: "g_r_plot.png"): Output filename for the g(r) plot.

Physics Concepts:

• Lennard-Jones Potential: Simulates the interaction force between two neutral particles ($U(r)$), featuring a strong short-range repulsion (Pauli exclusion) and a weaker long-range attraction (van der Waals force).
• Statistical Mechanics: The system is initialized with random positions and velocities (based on temperature $kT$). It undergoes an "Equilibration" phase to reach a stable state before the "Production" (data collection) phase begins.
• Periodic Boundary Conditions: Simulates an infinite system by using a finite number of particles in a "box" that wraps around on itself.
• Radial Distribution Function $g(r)$: The final output, $g(r)$, describes the probability of finding another particle at a distance $r$ from a reference particle. Its shape reveals the state of matter (e.g., solid, liquid, gas).

Mathematical & JAX Implementation:

• Lennard-Jones Potential $U(r)$:
  $$U(r) = 4\epsilon \left[ \left( \frac{\sigma}{r} \right)^{12} - \left( \frac{\sigma}{r} \right)^6 \right]$$
  Here, $\epsilon$ represents the depth of the potential well (energy scale of attraction), and $\sigma$ is the finite distance at which the inter-particle potential is zero (size parameter). The $r^{-12}$ term models the steep repulsive force due to overlapping electron clouds, while the $r^{-6}$ term captures the attractive dispersion forces. The total potential energy for the system is the sum over all unique particle pairs: $U_{\text{total}} = \sum_{i < j} U(r_{ij})$. This is implemented in the `total_energy_fn` function, which efficiently computes this sum using vectorized JAX operations to avoid explicit loops.
• Force Calculation (via `jax.grad`): The force on each particle is the negative gradient of the potential energy with respect to its position: $\vec{F}_i = -\nabla_i U_{\text{total}}$. Instead of manually deriving the complex analytical expression for the force from the Lennard-Jones potential (which involves differentiating the pairwise terms), the script leverages JAX's automatic differentiation:

  force_fn = jit(grad(lambda R: -total_energy_fn(R)))

  Here, `grad` computes the vector-Jacobian product automatically, treating positions $\vec{R}$ as the input. The negative sign ensures we get the force from the potential. This approach is highly advantageous in JAX for physics simulations, as it allows easy modification of potentials without re-deriving forces, and it compiles efficiently on TPUs.
• Integration (Velocity Verlet): To propagate the positions $\vec{R}$ and velocities $\vec{V}$ over time, the `verlet_step` function employs the Velocity Verlet algorithm, a second-order symplectic integrator that conserves energy better than simpler Euler methods. It alternates velocity and position updates in a leapfrog manner:
  $$\begin{align*} \vec{V}\left(t + \frac{1}{2} \Delta t \right) &= \vec{V}(t) + \frac{1}{2} \vec{A}(t) \Delta t, \\ \vec{R}(t + \Delta t) &= \vec{R}(t) + \vec{V}\left(t + \frac{1}{2} \Delta t \right) \Delta t, \\ \vec{A}(t + \Delta t) &= \frac{\vec{F}(\vec{R}(t + \Delta t))}{m}, \\ \vec{V}(t + \Delta t) &= \vec{V}\left(t + \frac{1}{2} \Delta t \right) + \frac{1}{2} \vec{A}(t + \Delta t) \Delta t. \end{align*}$$
  This method is time-reversible and preserves the symplectic structure of Hamiltonian systems, reducing long-term energy drift in simulations. The acceleration $\vec{A} = \vec{F}/m$ is computed at each half-step using the force function.
• Performance: `jax.jit` is applied to `verlet_step`, `equilibrate_fn`, `production_fn`, and `calculate_g_r`. `jax.lax.fori_loop` is used to compile the entire simulation loop onto the TPU for maximum speed.

Example Output: Radial Distribution Function

This plot shows the $g(r)$. The example image provided shows a result where $g(r) \approx 0$, which can occur if the simulation parameters (like density `rho` or temperature `kT`) are not set to form a stable liquid structure, or if the production run is too short. A typical liquid would show distinct peaks.

2. N-Body Black Hole Merger (nbody_bh_merger_sim_single-host_workload.py)

This script simulates the dynamics of N bodies (e.g., 3 black holes) under their mutual gravity and computes the resulting gravitational waves (GW).

How to Run (Interactive):

# (Basic) Run the script and press Enter to accept all default values
python3 nbody_bh_merger_sim_single-host_workload.py
# -> Number of black holes (default=3): [Enter]
# -> Mass of BH1 (default=30.0): [Enter]
# -> ... (etc.)

# (Advanced) Run the script and input custom values
python3 nbody_bh_merger_sim_single-host_workload.py
# -> Number of black holes: 5
# -> Mass of BH1: 50
# -> ... (etc.)
# -> Compute Lyapunov exponent? (y/n): n

This script is **interactive** and will prompt you for parameters in the console (Number of black holes, Mass, Separation, Velocity, etc.) instead of using command-line arguments.

Physics Concepts:

• Newtonian Gravity (N-Body): Solves the classic N-body problem by calculating the gravitational force between all pairs of objects.
• Gravitational Waves (Approximation): The script computes the GW strain ($h_+$) using the Quadrupole approximation, summing the contributions from each orbiting pair.
• Chaos Theory: For N > 2, the system is often chaotic. The script can compute the Lyapunov Exponent ($\lambda$), which measures the exponential rate at which nearby trajectories diverge, quantifying the system's chaos.

Mathematical & JAX Implementation:

• ODE System: The N-body problem is formulated as a first-order system of ordinary differential equations (ODEs). The state vector $Y$ concatenates all positions and velocities: $Y = [\vec{r}_1, \dots, \vec{r}_N, \vec{v}_1, \dots, \vec{v}_N]^T \in \mathbb{R}^{6N}$. The time derivative is $\frac{dY}{dt} = [\vec{v}_1, \dots, \vec{v}_N, \vec{a}_1, \dots, \vec{a}_N]^T$, where $\vec{a}_i$ are the accelerations due to gravity. This structure allows efficient vectorized computation in JAX.
• Gravitational Acceleration $\vec{a}_{g,i}$:
  $$\vec{a}_{g,i} = \sum_{j \neq i} G m_j \frac{\vec{r}_j - \vec{r}_i}{|\vec{r}_j - \vec{r}_i|^3}$$
  This arises from Newton's law of universal gravitation, $\vec{F}_{ij} = -G m_i m_j \frac{\vec{r}_i - \vec{r}_j}{|\vec{r}_i - \vec{r}_j|^3}$, divided by $m_i$ for acceleration. The $1/r^2$ force law leads to the $1/r^3$ form after normalization. Regularization (e.g., softening) may be added to avoid singularities at close encounters, implemented in `pairwise_acc` using JAX's vectorized broadcasting for all pairs.
• Numerical Integration (RK4): To solve the ODE system $\frac{dY}{dt} = f(Y, t)$, the script uses the classical 4th-order Runge-Kutta (RK4) method in `rk4_step`. RK4 approximates the solution at $t + \Delta t$ via four evaluations of $f$ at intermediate points: $$k_1 = f(Y, t), \quad k_2 = f(Y + \frac{\Delta t}{2} k_1, t + \frac{\Delta t}{2}),$$ $$k_3 = f(Y + \frac{\Delta t}{2} k_2, t + \frac{\Delta t}{2}), \quad k_4 = f(Y + \Delta t k_3, t + \Delta t),$$ $$Y(t + \Delta t) = Y(t) + \frac{\Delta t}{6} (k_1 + 2k_2 + 2k_3 + k_4).$$ This provides a local error of $O(\Delta t^5)$, making it accurate and stable for non-stiff systems like Newtonian gravity.
• GW Strain $h_+$: In the post-Newtonian quadrupole approximation for inspiraling binaries, the plus polarization strain is
  $$h_+ = \frac{4}{D} \left( \frac{G \mathcal{M}}{c^2} \right)^{5/3} \left( \frac{\pi f_{\text{GW}}}{c^3} \right)^{2/3} \cos \Phi(t),$$
  where $\mathcal{M}$ is the chirp mass $\mathcal{M} = \frac{(m_1 m_2)^{3/5}}{(m_1 + m_2)^{1/5}}$, $f_{\text{GW}} = 2 f_{\text{orb}}$ is the GW frequency (twice the orbital due to quadrupole), $D$ is the distance, and $\Phi(t) = 2 \int \omega(t) dt$ is the phase with instantaneous angular frequency $\omega = \sqrt{G(m_i + m_j)/r^3}$ from Kepler's law. The amplitude scales as $A \propto (\mathcal{M} \omega)^{2/3}$. The script sums contributions over pairs, assuming far-field propagation.
• Lyapunov Exponent $\lambda$:
  $$\lambda \approx \frac{1}{t} \ln \left( \frac{|\delta(t)|}{|\delta(0)|} \right)$$
  This quantifies sensitivity to initial conditions in chaotic systems. $\delta(0)$ is a small perturbation (e.g., $10^{-10}$ in position), and $|\delta(t)|$ is the Euclidean norm of the deviation after time $t$ between the nominal trajectory $Y(t)$ and perturbed $Y'(t)$. The exponent $\lambda > 0$ indicates chaos, with divergence $e^{\lambda t}$. The script runs parallel simulations for $Y$ and $Y'$ using JAX's `vmap`.
• Performance: The entire simulation loop `simulate_nbody` is JIT-compiled using `jax.jit` and `jax.lax.scan`, allowing the whole trajectory to be computed efficiently on the TPU.

Example Simulation Outputs:

3. Three Particles in Non-Uniform E/M Field (three_particles_em_nonuni_single-host_workload.py)

This is a simple simulation of 3 charged particles moving under their mutual gravity and an external, non-uniform electromagnetic field.

How to Run:

# (Basic) Run with default settings (only gravity and constant B-field)
python3 three_particles_em_nonuni_single-host_workload.py

# (Advanced) Add non-uniform B-field, E-field, and run for more steps
python3 three_particles_em_nonuni_single-host_workload.py --Bz 5.0 --Bk 0.5 --Ex 0.2 --n_steps 2000

All Arguments:

• --dt (float, default: 0.01): Time step size.
• --n_steps (int, default: 1000): Total number of steps to simulate.
• --G (float, default: 1.0): Gravitational constant.
• --Bz (float, default: 1.0): Constant component of the magnetic field (Z-axis).
• --Bk (float, default: 0.0): Gradient of the magnetic field along x (Bz = Bz + Bk*x).
• --Ex (float, default: 0.0): Electric field strength (X-axis).
• --Ey (float, default: 0.0): Electric field strength (Y-axis).

Physics Concepts:

• Superposition of Forces: The net force on each particle is the vector sum of gravity from other particles and the Lorentz force from the E/M field: $\vec{F}_i = \sum_{j \neq i} \vec{F}_{g,ij} + \vec{F}_{Lorentz,i}$
• Lorentz Force: The force from the E/M field is $\vec{F}_{Lorentz}=q(\vec{E}+\vec{v}\times\vec{B})$.
• Non-Uniform Field: The magnetic field $B_z$ is position-dependent ($B_z = B_0 + B_k \cdot x$), making the dynamics more complex than a simple circular or helical motion.

Mathematical & JAX Implementation:

• Total Acceleration $\vec{a}=\vec{F}/m$: Newton's second law gives $\vec{a}_i = \vec{F}_i / m_i$. The `acceleration` function decomposes this into three additive components, computed in parallel:
  • Gravitational ($a_g$):
    $$\vec{a}_{g,i} = \sum_{j \neq i} G m_j \frac{\vec{r}_j - \vec{r}_i}{|\vec{r}_j - \vec{r}_i|^3}$$
    Identical to the N-body case, this captures pairwise inverse-square attractions (in `pairwise_acc`). For small N=3, it's computed directly without approximation.
  • Electric ($a_E$): $$\vec{a}_E = \frac{q}{m} \vec{E}$$ A constant acceleration if $\vec{E}$ is uniform (e.g., a constant field); implemented in `elec_acc` as a simple scaling.
  • Magnetic ($a_B$): $$\vec{a}_B = \frac{q}{m} (\vec{v} \times \vec{B})$$ Velocity-dependent, causing deflection perpendicular to $\vec{v}$ and $\vec{B}$. No work is done ($\vec{a}_B \cdot \vec{v} = 0$), preserving kinetic energy in pure magnetic fields.
• Cross Product in 2D: The simulation is in the xy-plane with $\vec{B} = [0, 0, B_z(x)]$ (out-of-plane, position-dependent) and $\vec{v} = [v_x, v_y, 0]$. The vector cross product simplifies via the right-hand rule:
  $$\vec{v} \times \vec{B} = \det \begin{vmatrix} \hat{i} & \hat{j} & \hat{k} \\ v_x & v_y & 0 \\ 0 & 0 & B_z \end{vmatrix} = \hat{i}(v_y B_z) - \hat{j}(v_x B_z) + \hat{k}(0) = [v_y B_z, -v_x B_z, 0]$$
  Thus, $\vec{a}_B = \frac{q}{m} [v_y B_z, -v_x B_z]$, inducing cyclotron motion with radius $r_c = \frac{m v_\perp}{q B_z}$. The non-uniformity $B_z(x) = B_0 + B_k x$ introduces gradients, leading to drifts like $\vec{E} \times \vec{B}$ or gradient drifts. This matches the JAX code in `mag_acc`: `qm * jnp.array([vy * bz, -vx * bz])`, where `qm = q/m`.
• Integration (Leapfrog/Velocity Verlet): Similar to the MD simulation, the `step` function applies Velocity Verlet for the combined gravitational + Lorentz forces. This integrator handles the velocity-dependent magnetic term accurately, maintaining second-order accuracy and symplectic properties for the conservative gravitational part (though magnetism adds dissipation in some senses).
• Performance: `jax.vmap` is used to calculate the acceleration for all particles in parallel, and `jax.jit` compiles the entire `step` function.

Example Initial State:

Initial positions of the three particles in the 2D plane.

4. Variational & Diffusion Monte Carlo (vmc_dmc_jax_quantum_harmonic_oscillator.py)

This script implements Variational Monte Carlo (VMC) followed by Diffusion Monte Carlo (DMC) to approximate the ground state energy and wavefunction of a D-dimensional isotropic quantum harmonic oscillator using JAX for high-performance computation. It optimizes a variational parameter α via stochastic gradient descent in VMC, then refines the distribution via branching diffusion in DMC. Outputs include energy convergence plots, marginal probability density histograms vs. exact ground state, and optional GIF animations of walker distributions.

How to Run:

# (Basic) Run with default settings (N=10k walkers, 3k VMC epochs, 500 DMC steps, 3D)
python3 vmc_dmc_jax_quantum_harmonic_oscillator.py

# (Advanced) Run a longer simulation in 1D with more epochs and steps
python3 vmc_dmc_jax_quantum_harmonic_oscillator.py --n_epochs 5000 --n_dmc 1000 --dim 1

All Arguments:

• --n_walkers (int, default: 10000): Number of Monte Carlo walkers.
• --n_epochs (int, default: 3000): Number of VMC optimization epochs.
• --n_equil (int, default: 100): Number of equilibration steps per VMC epoch.
• --step_size (float, default: 2.0): Proposal step size for Metropolis-Hastings sampling.
• --lr (float, default: 0.02): Learning rate for Adam optimizer in VMC.
• --n_dmc (int, default: 500): Number of DMC propagation steps.
• --dmc_dt (float, default: 0.01): Time step size for DMC diffusion/branching.
• --dim (int, default: 3): Spatial dimensionality of the harmonic oscillator.
• --no-gif (flag): Disable generation of VMC and DMC animation GIFs (vmc_animation.gif and dmc_animation.gif).
• --no-plot (flag): Disable display of final matplotlib plots (energy convergence, histograms).

Physics Concepts:

• Quantum Harmonic Oscillator (QHO): The D-dimensional isotropic QHO has a Hamiltonian $H = -\frac{\hbar^2}{2m} \nabla^2 + \frac{1}{2} m \omega^2 \mathbf{r}^2$ (set $\hbar = m = \omega = 1$), with exact ground state energy $E_0 = D/2$ and Gaussian wavefunction $\Psi_0(\mathbf{r}) \propto e^{-\mathbf{r}^2/2}$. The script approximates this via stochastic sampling of the wavefunction.
• Variational Monte Carlo (VMC): Uses a trial wavefunction $\Psi_T(\mathbf{r}; \alpha) = e^{-\alpha \mathbf{r}^2}$ to minimize the expectation value $\langle E \rangle = \frac{\langle \Psi_T | H | \Psi_T \rangle}{\langle \Psi_T | \Psi_T \rangle}$ via Monte Carlo integration over samples from $|\Psi_T|^2$. Optimizes $\alpha$ using stochastic gradients.
• Diffusion Monte Carlo (DMC): Projects onto the ground state via imaginary-time evolution of the Schrödinger equation, interpreted as a branching random walk with drift (from the quantum force) and diffusion (from the kinetic term), biased by local energy branching.
• Walker Sampling: "Walkers" are Monte Carlo samples representing the probability distribution. Metropolis-Hastings ensures ergodic sampling in VMC, while DMC uses resampling based on weights to focus on low-energy configurations.

Mathematical & JAX Implementation:

• Local Energy $E_L(\mathbf{r})$: For a trial wavefunction, the local energy is $E_L = \frac{H \Psi_T}{\Psi_T} = -\frac{1}{2} \nabla^2 \ln \Psi_T - \frac{1}{2} |\nabla \ln \Psi_T|^2 + V(\mathbf{r})$, where $V(\mathbf{r}) = \frac{1}{2} \mathbf{r}^2$ is the potential. For the Gaussian trial, $\ln \Psi_T = -\alpha \mathbf{r}^2$, so $\nabla \ln \Psi_T = -2\alpha \mathbf{r}$ and $\nabla^2 \ln \Psi_T = -2\alpha D$. Thus, $E_L = \alpha D - 2\alpha^2 \mathbf{r}^2 + \frac{1}{2} \mathbf{r}^2$. The VMC energy is the Monte Carlo average $\langle E_L \rangle \approx E_0$. For $\alpha=0.5$, $\langle E_L \rangle = 0.5 D$ exactly. Implemented in `local_energy` using JAX's `grad` for the Laplacian and gradient terms.
• Metropolis-Hastings Sampling (VMC): Samples from $|\Psi_T|^2$ using a symmetric proposal $\mathbf{r}' = \mathbf{r} + \delta \mathbf{\epsilon}$ ($\delta =$ `step_size`), with acceptance $A = \min(1, |\Psi_T(\mathbf{r}')|^2 / |\Psi_T(\mathbf{r})|^2) = \min(1, e^{2(\ln \Psi_T(\mathbf{r}') - \ln \Psi_T(\mathbf{r}))})$. Equilibration discards initial steps to ensure sampling from the target distribution. Vectorized over walkers via `vmap(metropolis_step)`, and looped with `fori_loop` for efficiency. The variational gradient is $\frac{\partial \langle E \rangle}{\partial \alpha} \approx 2 \langle (E_L - \langle E_L \rangle) \frac{\partial \ln \Psi_T}{\partial \alpha} \rangle$, optimized with Adam (`optax.adam`).
• DMC Propagation: Solves the imaginary-time Schrödinger equation $\frac{\partial \Psi}{\partial \tau} = -\left( H - E_T \right) \Psi$ (with reference energy $E_T \approx E_0$) via a Green's function approximation: drift-diffusion with velocity $\mathbf{F} = \nabla \ln \Psi_T = -2\alpha \mathbf{r}$ and branching weight $w = e^{-(E_L - E_T) \Delta \tau}$. Walkers are resampled via multinomial choice based on normalized weights, then updated as $\mathbf{r}' = \mathbf{r} + \mathbf{F} \Delta \tau + \sqrt{\Delta \tau} \mathbf{\eta}$ ($\mathbf{\eta} \sim \mathcal{N}(0,1)$). The mixed estimator for energy is the average $E_T$ over steps, with statistical error from variance. Implemented with `jax.lax.scan` for the full trajectory, vectorized via `vmap` for local energy and drift.
• Performance: All core functions (`local_energy`, `metropolis_step`, `dmc_step_body`) are JIT-compiled. `vmap` parallelizes over walkers (up to 10k), and `fori_loop`/`scan` compile loops for TPU acceleration. Progress tracking uses `rich` for console output, with optional GIFs via `imageio` for visualizing walker histograms evolving toward the exact Gaussian marginal.