⚡️ Speed up function leapfrog_integration by 41%

[codeflash-ai]

📄 41% (0.41x) speedup for leapfrog_integration in code_to_optimize/sample_jit_code.py

⏱️ Runtime : 1.07 seconds → 759 milliseconds (best of 6 runs)

📝 Explanation and details

The optimized code achieves a 40% speedup by leveraging Numba's JIT compilation to convert the computationally intensive N-body simulation loops from Python bytecode to native machine code. Here's why this yields substantial performance gains:

Key Optimizations

1. JIT Compilation with Numba (@njit): The core computation is extracted into _leapfrog_core() and decorated with @njit(cache=True). Numba compiles this function to optimized machine code on first execution, eliminating Python interpreter overhead for:

   • Nested loops iterating over particles (O(N²) pairwise interactions)
   • Array indexing operations
   • Arithmetic operations on distances, forces, and accelerations

   The line profiler shows the original code spends ~90% of time in the innermost pairwise interaction loop—exactly where JIT compilation provides maximum benefit.

2. Micro-optimizations within JIT context:

   • Precomputed constants: soft2 = softening * softening and half_dt = 0.5 * dt are computed once instead of repeatedly in loops
   • Local variable hoisting: Position components (pos_i0, pos_i1, pos_i2) and masses are loaded into locals in the outer loop, reducing array accesses in the hot inner loop
   • math.sqrt() vs np.sqrt(): Uses scalar math.sqrt() which is more efficient than NumPy's vectorized version for single values in JIT-compiled code
   • Explicit loop unrolling: Acceleration reset uses explicit assignments (acc[ii, 0] = 0.0) instead of fill(), which can be better optimized by Numba

3. Graceful Fallback: The try/except wrapper ensures the code works identically when Numba is unavailable, maintaining portability while gaining performance where possible.

Performance Impact Analysis

Based on annotated tests:

• Small particle counts (2-10 particles, few steps): 0-15% speedup or slight slowdown due to JIT compilation overhead
• Medium workloads (50-100 particles): 28-43% speedup as JIT benefits outweigh compilation cost
• Large workloads (200 particles): 39-42% speedup, demonstrating scalability

The optimization particularly benefits the test_hundred_particles_cluster (43% faster) and test_two_hundred_particles_large_scale (42% faster) cases, indicating strong performance for realistic N-body simulations. Edge cases with zero particles or zero steps show minor slowdowns due to unchanged overhead, but these are trivial execution paths.

Why This Works

Python's interpreted nature makes tight numerical loops extremely slow. The original code's line profiler shows array indexing operations (pos[j, 0] - pos[i, 0]) each taking 6-7% of total runtime. Numba compiles these to direct memory accesses with zero interpreter overhead, transforming the bottleneck. The cache=True flag ensures compilation happens only once per function version, amortizing startup cost across multiple calls.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 21 Passed
🌀 Generated Regression Tests	✅ 39 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

⚙️ Click to see Existing Unit Tests

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_numba_jit_code.py::TestLeapfrogIntegration.test_does_not_modify_input	68.7μs	65.6μs	4.70%✅
test_numba_jit_code.py::TestLeapfrogIntegration.test_momentum_conservation	2.82ms	2.40ms	17.3%✅
test_numba_jit_code.py::TestLeapfrogIntegration.test_single_moving_particle	210μs	231μs	-9.25%⚠️
test_numba_jit_code.py::TestLeapfrogIntegration.test_single_stationary_particle	211μs	231μs	-8.73%⚠️
test_numba_jit_code.py::TestLeapfrogIntegration.test_two_particles_approach	305μs	301μs	1.34%✅

🌀 Click to see Generated Regression Tests

import numpy as np

# imports
from code_to_optimize.sample_jit_code import leapfrog_integration

# unit tests


# 1) Basic functionality: two-body center-of-mass conservation and symmetry.
def test_two_body_center_of_mass_conservation_and_symmetry():
    # Two identical masses placed symmetrically around origin with zero initial velocity.
    pos = np.array([[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    vel = np.zeros_like(pos)
    masses = np.array([1.0, 1.0], dtype=np.float64)
    dt = 0.01
    n_steps = 10

    # Compute initial center of mass
    com_initial = (masses[:, None] * pos).sum(axis=0) / masses.sum()

    pos_final, vel_final = leapfrog_integration(pos, vel, masses, dt, n_steps)  # 67.8μs -> 65.1μs (4.03% faster)

    # Center of mass should remain at origin (within numerical tolerance)
    com_final = (masses[:, None] * pos_final).sum(axis=0) / masses.sum()


# 2) Edge case: zero particles - should return empty arrays of proper shape and not fail.
def test_zero_particles_returns_empty_arrays():
    pos = np.zeros((0, 3), dtype=np.float64)
    vel = np.zeros((0, 3), dtype=np.float64)
    masses = np.zeros((0,), dtype=np.float64)

    pos_final, vel_final = leapfrog_integration(pos, vel, masses, dt=0.1, n_steps=5)  # 4.46μs -> 5.08μs (12.3% slower)


# 3) Edge case: single particle should move linearly with constant velocity (no internal forces).
def test_single_particle_linear_motion():
    pos = np.array([[0.0, 0.0, 0.0]], dtype=np.float64)
    vel = np.array([[1.0, -2.0, 0.5]], dtype=np.float64)
    masses = np.array([1.0], dtype=np.float64)
    dt = 0.1
    n_steps = 5

    expected_pos = pos + vel * (dt * n_steps)  # since no forces act on a single particle
    pos_final, vel_final = leapfrog_integration(pos, vel, masses, dt, n_steps)  # 14.6μs -> 16.8μs (13.4% slower)


# 4) Edge case: overlapping particles — softening should prevent singularities (no NaN/Inf).
def test_softening_prevents_singularity_for_overlapping_particles():
    pos = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)  # exact overlap
    vel = np.zeros_like(pos)
    masses = np.array([1.0, 2.0], dtype=np.float64)
    dt = 0.01
    n_steps = 3
    softening = 1e-3

    pos_final, vel_final = leapfrog_integration(
        pos, vel, masses, dt, n_steps, softening=softening
    )  # 22.2μs -> 25.3μs (12.3% slower)


# 5) Edge case: dt == 0 should produce no change in positions or velocities.
def test_dt_zero_no_evolution():
    rng = np.random.RandomState(123)
    n = 4
    pos = rng.randn(n, 3).astype(np.float64)
    vel = rng.randn(n, 3).astype(np.float64)
    masses = np.abs(rng.randn(n)).astype(np.float64) + 0.1  # ensure positive masses

    pos_copy = pos.copy()
    vel_copy = vel.copy()

    pos_final, vel_final = leapfrog_integration(pos, vel, masses, dt=0.0, n_steps=10)  # 196μs -> 173μs (13.5% faster)


# 6) Reversibility: integrating forward then backward (with negative dt) should approximately recover initial state.
def test_reversibility_with_negative_dt():
    pos = np.array([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]], dtype=np.float64)
    vel = np.array([[0.0, 0.1, 0.0], [0.0, -0.1, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)
    dt = 0.01
    n_steps = 20

    # Integrate forward
    pos_mid, vel_mid = leapfrog_integration(pos, vel, masses, dt, n_steps)  # 123μs -> 123μs (0.034% slower)

    # Integrate backward starting from mid state
    pos_recovered, vel_recovered = leapfrog_integration(
        pos_mid, vel_mid, masses, -dt, n_steps
    )  # 120μs -> 120μs (0.138% slower)


# 7) Conservation property: total momentum should be conserved (within numerical precision).
def test_total_momentum_conservation():
    rng = np.random.RandomState(42)
    n = 6
    pos = rng.randn(n, 3).astype(np.float64)
    vel = rng.randn(n, 3).astype(np.float64)
    masses = np.abs(rng.randn(n)).astype(np.float64) + 0.2  # positive masses

    total_momentum_initial = (masses[:, None] * vel).sum(axis=0)

    pos_final, vel_final = leapfrog_integration(pos, vel, masses, dt=0.02, n_steps=15)  # 588μs -> 492μs (19.5% faster)

    total_momentum_final = (masses[:, None] * vel_final).sum(axis=0)


# 8) Large scale: moderate number of particles and steps to check scalability and stability.
def test_large_scale_moderate_particles_and_steps():
    # Use a deterministic RNG for reproducibility
    rng = np.random.RandomState(0)
    n_particles = 200  # well under the 1000 element cap requested
    pos = rng.randn(n_particles, 3).astype(np.float64) * 0.1
    vel = rng.randn(n_particles, 3).astype(np.float64) * 0.01
    masses = np.abs(rng.randn(n_particles)).astype(np.float64) + 0.1

    # Keep steps small to respect the instruction (avoid >1000 steps)
    dt = 0.005
    n_steps = 5

    pos_final, vel_final = leapfrog_integration(pos, vel, masses, dt, n_steps)  # 184ms -> 132ms (39.5% faster)

    # Momentum should be approximately conserved even for larger N
    total_momentum_initial = (masses[:, None] * vel).sum(axis=0)
    total_momentum_final = (masses[:, None] * vel_final).sum(axis=0)


# 9) Incorrect input shapes should raise clear exceptions (defensive programming checks above).


def test_function_does_not_mutate_inputs():
    rng = np.random.RandomState(7)
    n = 5
    pos = rng.randn(n, 3).astype(np.float64)
    vel = rng.randn(n, 3).astype(np.float64)
    masses = np.abs(rng.randn(n)).astype(np.float64) + 0.1

    pos_before = pos.copy()
    vel_before = vel.copy()
    masses_before = masses.copy()

    _pos_out, _vel_out = leapfrog_integration(pos, vel, masses, dt=0.02, n_steps=3)  # 95.0μs -> 80.6μs (17.8% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import numpy as np

# imports
from code_to_optimize.sample_jit_code import leapfrog_integration

# ============================================================================
# BASIC TEST CASES: Test fundamental functionality under normal conditions
# ============================================================================


def test_single_particle_no_forces():
    """Test that a single particle with zero initial velocity and no other particles
    remains stationary (no forces act on it).
    """
    # Create a single particle at origin with zero velocity
    positions = np.array([[0.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0], dtype=np.float64)

    dt = 0.01
    n_steps = 10

    # Run integration
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 26.2μs -> 28.7μs (8.56% slower)


def test_two_particles_stationary():
    """Test two particles at rest with equal masses. They should accelerate toward
    each other due to gravitational attraction.
    """
    # Two particles at distance 1.0 apart, with equal masses
    positions = np.array([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    dt = 0.001
    n_steps = 10

    # Run integration
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 69.2μs -> 65.4μs (5.80% faster)


def test_return_types():
    """Test that the function returns two numpy arrays with correct shapes."""
    # Simple two-particle system
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    dt = 0.01
    n_steps = 5

    codeflash_output = leapfrog_integration(positions, velocities, masses, dt, n_steps)
    result = codeflash_output  # 36.3μs -> 34.7μs (4.81% faster)

    # Check both elements are numpy arrays
    final_pos, final_vel = result


def test_original_arrays_not_modified():
    """Test that the original input arrays are not modified by the function
    (function should work on copies).
    """
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    # Create copies to compare later
    positions_copy = positions.copy()
    velocities_copy = velocities.copy()

    dt = 0.01
    n_steps = 10

    # Run integration
    leapfrog_integration(positions, velocities, masses, dt, n_steps)  # 67.3μs -> 64.9μs (3.79% faster)


def test_energy_conservation_simple():
    """Test approximate energy conservation for a simple two-body system.
    In an ideal gravitational system, total energy should be conserved.
    """
    # Two particles in opposite positions with zero velocity
    positions = np.array([[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    dt = 0.001
    n_steps = 50
    softening = 0.01

    # Calculate initial energy
    G = 1.0
    dx = positions[1, 0] - positions[0, 0]
    dy = positions[1, 1] - positions[0, 1]
    dz = positions[1, 2] - positions[0, 2]
    r = np.sqrt(dx * dx + dy * dy + dz * dz)

    initial_PE = -G * masses[0] * masses[1] / r
    initial_KE = 0.0
    initial_E = initial_PE + initial_KE

    # Run integration
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 303μs -> 296μs (2.13% faster)

    # Calculate final energy
    dx = final_pos[1, 0] - final_pos[0, 0]
    dy = final_pos[1, 1] - final_pos[0, 1]
    dz = final_pos[1, 2] - final_pos[0, 2]
    r_final = np.sqrt(dx * dx + dy * dy + dz * dz)

    final_KE = 0.5 * masses[0] * (final_vel[0, 0] ** 2 + final_vel[0, 1] ** 2 + final_vel[0, 2] ** 2)
    final_KE += 0.5 * masses[1] * (final_vel[1, 0] ** 2 + final_vel[1, 1] ** 2 + final_vel[1, 2] ** 2)
    final_PE = -G * masses[0] * masses[1] / r_final
    final_E = final_PE + final_KE


def test_zero_timestep():
    """Test that zero timestep results in no change to positions and velocities."""
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]], dtype=np.float64)
    velocities = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    dt = 0.0
    n_steps = 100

    # Run integration with dt=0
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 606μs -> 583μs (3.85% faster)


def test_zero_steps():
    """Test that running for zero steps returns original positions and velocities."""
    positions = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float64)
    velocities = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], dtype=np.float64)
    masses = np.array([1.0, 2.0], dtype=np.float64)

    dt = 0.01
    n_steps = 0

    # Run integration with n_steps=0
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 2.25μs -> 3.21μs (29.9% slower)


# ============================================================================
# EDGE TEST CASES: Test extreme or unusual conditions
# ============================================================================


def test_very_small_softening():
    """Test with very small softening parameter (closer to true gravitational singularities).
    Should still produce valid results without NaN or Inf.
    """
    positions = np.array([[0.0, 0.0, 0.0], [0.1, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    dt = 0.001
    n_steps = 5
    softening = 1e-6

    # Run integration with very small softening
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 36.5μs -> 35.0μs (4.28% faster)


def test_large_softening():
    """Test with large softening parameter (forces diminished significantly).
    Particles should move very slowly.
    """
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    dt = 0.01
    n_steps = 10
    softening = 10.0

    # Run integration with large softening
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 68.2μs -> 65.0μs (5.00% faster)

    # With large softening, force should be greatly reduced
    # Movement should be minimal
    displacement_0 = np.linalg.norm(final_pos[0] - positions[0])
    displacement_1 = np.linalg.norm(final_pos[1] - positions[1])


def test_very_large_mass():
    """Test with very large particle masses (should produce proportionally larger forces)."""
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1e10, 1e10], dtype=np.float64)

    dt = 0.001
    n_steps = 5
    softening = 0.01

    # Run integration with large masses
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 36.2μs -> 34.7μs (4.57% faster)

    # Should have larger velocities/accelerations than with unit masses
    max_vel = np.max(np.abs(final_vel))


def test_very_small_mass():
    """Test with very small particle masses (forces should be proportionally smaller)."""
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1e-10, 1e-10], dtype=np.float64)

    dt = 0.01
    n_steps = 10
    softening = 0.01

    # Run integration with small masses
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 68.1μs -> 64.8μs (5.21% faster)

    # With tiny masses, forces and accelerations should be tiny
    max_vel = np.max(np.abs(final_vel))


def test_very_small_timestep():
    """Test with very small timestep (more accurate but slower integration)."""
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    dt = 1e-6
    n_steps = 10
    softening = 0.01

    # Run integration with tiny timestep
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 67.6μs -> 64.8μs (4.37% faster)

    # With tiny timestep, movement should be very small
    displacement = np.linalg.norm(final_pos[0] - positions[0])


def test_very_large_timestep():
    """Test with large timestep (may produce instability or large displacements)."""
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    dt = 1.0
    n_steps = 5
    softening = 0.01

    # Run integration with large timestep
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 36.0μs -> 38.2μs (5.88% slower)


def test_negative_velocity():
    """Test particles with negative (opposite direction) initial velocities."""
    # Two particles moving toward each other
    positions = np.array([[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    dt = 0.01
    n_steps = 10
    softening = 0.01

    # Run integration
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 68.1μs -> 65.8μs (3.42% faster)

    # Particles should move closer due to attractive force and initial velocities
    final_distance = np.linalg.norm(final_pos[1] - final_pos[0])
    initial_distance = np.linalg.norm(positions[1] - positions[0])


def test_non_aligned_axes():
    """Test particles positioned off-axis (in 3D space with all coordinates active)."""
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    dt = 0.001
    n_steps = 10
    softening = 0.01

    # Run integration
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 65.7μs -> 66.3μs (0.943% slower)


def test_three_particles_collinear():
    """Test three particles arranged in a line."""
    # Three particles on x-axis
    positions = np.array([[-1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0, 1.0], dtype=np.float64)

    dt = 0.001
    n_steps = 10
    softening = 0.01

    # Run integration
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 123μs -> 114μs (7.91% faster)

    # Check that all particles moved
    for i in range(3):
        displacement = np.linalg.norm(final_pos[i] - positions[i])


def test_three_particles_triangular():
    """Test three particles arranged in a triangle (symmetric configuration)."""
    # Three particles forming equilateral triangle
    sqrt3_2 = np.sqrt(3) / 2
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.5, sqrt3_2, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0, 1.0], dtype=np.float64)

    dt = 0.001
    n_steps = 10
    softening = 0.01

    # Run integration
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 123μs -> 113μs (8.15% faster)


def test_equal_masses_vs_unequal():
    """Test that particles with equal masses produce symmetric results."""
    positions = np.array([[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)

    dt = 0.01
    n_steps = 10
    softening = 0.01

    # Equal masses
    masses_equal = np.array([1.0, 1.0], dtype=np.float64)
    final_pos_eq, final_vel_eq = leapfrog_integration(
        positions, velocities, masses_equal, dt, n_steps, softening
    )  # 68.6μs -> 64.5μs (6.46% faster)

    # Unequal masses (heavier on right)
    masses_unequal = np.array([1.0, 2.0], dtype=np.float64)
    final_pos_uneq, final_vel_uneq = leapfrog_integration(
        positions, velocities, masses_unequal, dt, n_steps, softening
    )  # 64.7μs -> 62.4μs (3.67% faster)

    # With equal masses, particles should move symmetrically
    # Distance from initial position should be equal
    disp_0_eq = np.linalg.norm(final_pos_eq[0] - positions[0])
    disp_1_eq = np.linalg.norm(final_pos_eq[1] - positions[1])

    # With unequal masses, heavier particle should move less
    disp_0_uneq = np.linalg.norm(final_pos_uneq[0] - positions[0])
    disp_1_uneq = np.linalg.norm(final_pos_uneq[1] - positions[1])


# ============================================================================
# LARGE SCALE TEST CASES: Assess performance and scalability
# ============================================================================


def test_ten_particles_random():
    """Test with 10 particles in random positions (small-scale ensemble).
    Verifies correctness with multiple-body interactions.
    """
    np.random.seed(42)
    n_particles = 10

    # Random positions in [-1, 1] cube
    positions = np.random.uniform(-1.0, 1.0, (n_particles, 3)).astype(np.float64)
    # Small random velocities
    velocities = np.random.uniform(-0.1, 0.1, (n_particles, 3)).astype(np.float64)
    # Unit masses
    masses = np.ones(n_particles, dtype=np.float64)

    dt = 0.001
    n_steps = 20
    softening = 0.01

    # Run integration
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 2.07ms -> 1.61ms (28.5% faster)

    # All particles should have moved from initial positions
    for i in range(n_particles):
        displacement = np.linalg.norm(final_pos[i] - positions[i])


def test_fifty_particles_sparse():
    """Test with 50 particles in a sparse distribution (medium-scale ensemble)."""
    np.random.seed(123)
    n_particles = 50

    # Positions spread over larger space [-10, 10]
    positions = np.random.uniform(-10.0, 10.0, (n_particles, 3)).astype(np.float64)
    # Zero initial velocities
    velocities = np.zeros((n_particles, 3), dtype=np.float64)
    # Random masses in [0.5, 2.0]
    masses = np.random.uniform(0.5, 2.0, n_particles).astype(np.float64)

    dt = 0.001
    n_steps = 15
    softening = 0.05

    # Run integration
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 35.8ms -> 25.7ms (39.4% faster)

    # All particles should have non-negative kinetic energy
    for i in range(n_particles):
        ke = final_vel[i, 0] ** 2 + final_vel[i, 1] ** 2 + final_vel[i, 2] ** 2


def test_hundred_particles_cluster():
    """Test with 100 particles in a clustered distribution (larger-scale ensemble)."""
    np.random.seed(456)
    n_particles = 100

    # Positions in tighter cluster around origin
    positions = np.random.normal(0.0, 1.0, (n_particles, 3)).astype(np.float64)
    # Small random velocities
    velocities = np.random.normal(0.0, 0.05, (n_particles, 3)).astype(np.float64)
    # All unit masses
    masses = np.ones(n_particles, dtype=np.float64)

    dt = 0.0005
    n_steps = 20
    softening = 0.01

    # Run integration
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 191ms -> 133ms (43.0% faster)

    # Check that center of mass makes physical sense
    center_of_mass_initial = np.sum(masses[:, np.newaxis] * positions, axis=0) / np.sum(masses)
    center_of_mass_final = np.sum(masses[:, np.newaxis] * final_pos, axis=0) / np.sum(masses)

    # Center of mass should move (no external forces, but internal dynamics)
    com_displacement = np.linalg.norm(center_of_mass_final - center_of_mass_initial)


def test_two_hundred_particles_large_scale():
    """Test with 200 particles for larger-scale performance evaluation."""
    np.random.seed(789)
    n_particles = 200

    # Spread over reasonable volume
    positions = np.random.uniform(-5.0, 5.0, (n_particles, 3)).astype(np.float64)
    velocities = np.zeros((n_particles, 3), dtype=np.float64)
    masses = np.ones(n_particles, dtype=np.float64)

    dt = 0.0005
    n_steps = 10
    softening = 0.02

    # Run integration (should complete in reasonable time)
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 377ms -> 265ms (42.4% faster)


def test_diverse_masses_large():
    """Test with 100 particles having diverse mass distributions (high variance)."""
    np.random.seed(999)
    n_particles = 100

    positions = np.random.uniform(-2.0, 2.0, (n_particles, 3)).astype(np.float64)
    velocities = np.zeros((n_particles, 3), dtype=np.float64)
    # Masses ranging over several orders of magnitude
    masses = np.power(10.0, np.random.uniform(-2.0, 2.0, n_particles)).astype(np.float64)

    dt = 0.001
    n_steps = 15
    softening = 0.01

    # Run integration
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 141ms -> 99.9ms (41.2% faster)


def test_multiple_timesteps_accumulated():
    """Test that running many steps produces accumulated displacement."""
    np.random.seed(111)
    n_particles = 20

    positions = np.random.uniform(-1.0, 1.0, (n_particles, 3)).astype(np.float64)
    velocities = np.zeros((n_particles, 3), dtype=np.float64)
    masses = np.ones(n_particles, dtype=np.float64)

    dt = 0.001

    # Run with fewer steps
    final_pos_10, final_vel_10 = leapfrog_integration(
        positions, velocities, masses, dt, 10, 0.01
    )  # 3.84ms -> 2.88ms (33.6% faster)

    # Run with more steps
    final_pos_100, final_vel_100 = leapfrog_integration(
        positions, velocities, masses, dt, 100, 0.01
    )  # 38.9ms -> 28.6ms (35.9% faster)

    # Longer simulation should generally produce larger displacements
    disp_10 = np.linalg.norm(final_pos_10 - positions)
    disp_100 = np.linalg.norm(final_pos_100 - positions)


def test_different_softening_values():
    """Test that varying softening parameter produces expected differences in behavior."""
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    dt = 0.01
    n_steps = 10

    # Small softening
    final_pos_small, final_vel_small = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening=0.001
    )  # 66.0μs -> 65.2μs (1.21% faster)

    # Large softening
    final_pos_large, final_vel_large = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening=1.0
    )  # 61.6μs -> 60.9μs (1.23% faster)

    # Smaller softening should produce larger accelerations and velocities
    vel_norm_small = np.linalg.norm(final_vel_small)
    vel_norm_large = np.linalg.norm(final_vel_large)


def test_stability_long_run():
    """Test that a longer simulation remains numerically stable without blow-up."""
    np.random.seed(222)
    n_particles = 30

    positions = np.random.uniform(-1.0, 1.0, (n_particles, 3)).astype(np.float64)
    velocities = np.random.uniform(-0.1, 0.1, (n_particles, 3)).astype(np.float64)
    masses = np.ones(n_particles, dtype=np.float64)

    dt = 0.0001
    n_steps = 100
    softening = 0.01

    # Run long integration
    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening
    )  # 86.2ms -> 62.6ms (37.8% faster)

    # Velocities should remain in reasonable bounds (not growing unboundedly)
    max_vel = np.max(np.abs(final_vel))


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-leapfrog_integration-mkg8mals and push.