⚡️ Speed up function leapfrog_integration by 2,611%

[codeflash-ai]

📄 2,611% (26.11x) speedup for leapfrog_integration in code_to_optimize/sample_jit_code.py

⏱️ Runtime : 506 milliseconds → 18.7 milliseconds (best of 249 runs)

📝 Explanation and details

The optimized code achieves a 26x speedup by replacing nested Python loops with vectorized NumPy operations.

Key optimizations:

1. Vectorized pairwise distance calculations: Instead of computing distances between particles using nested for loops over i and j, the optimized version creates difference matrices for all particle pairs simultaneously using NumPy broadcasting (dx = x[None, :] - x[:, None]). This creates an (N×N) matrix where element [i,j] contains the x-component difference between particles j and i.

2. Bulk acceleration computation: Rather than accumulating forces particle-by-particle in Python loops, the optimized code computes all accelerations at once with np.sum(factor * dx, axis=1). This leverages NumPy's highly optimized C implementations for array operations.

3. Consolidated velocity/position updates: The three separate loops that updated velocity and position components are replaced with single vectorized operations (vel += 0.5 * dt * acc), eliminating loop overhead.

Why this is faster:

• Python loop elimination: The original code's innermost loop runs ~260k times per execution (visible in line profiler), with each iteration incurring Python interpreter overhead. Vectorization moves this work into compiled NumPy routines.

• Cache efficiency: Vectorized operations process contiguous memory blocks, improving CPU cache utilization compared to scattered array indexing in nested loops.

• SIMD utilization: NumPy can leverage CPU vector instructions (SIMD) to process multiple elements simultaneously, which explicit Python loops cannot.

Performance characteristics from tests:

• Massive gains for larger systems: Tests with 50-200 particles show 60-150x speedups (e.g., test_large_scale_system_50_particles_100_steps: 235ms → 3ms), as vectorization benefits scale with problem size.

• Slower for tiny systems: Small cases (2-3 particles, few steps) run 16-70% slower due to NumPy's fixed overhead of allocating temporary (N×N) matrices, which dominates when N is small.

• Break-even around N=10-20: The optimization starts showing gains when the nested loop cost exceeds vectorization overhead.

Impact consideration:

Since N-body simulations typically involve many particles over many timesteps (the intended use case for leapfrog integration), this optimization would significantly benefit production workloads, despite the slight regression on toy examples.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 21 Passed
🌀 Generated Regression Tests	✅ 30 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	66.7μs	149μs	-55.2%⚠️
test_numba_jit_code.py::TestLeapfrogIntegration.test_momentum_conservation	2.79ms	1.44ms	93.5%✅
test_numba_jit_code.py::TestLeapfrogIntegration.test_single_moving_particle	204μs	1.24ms	-83.5%⚠️
test_numba_jit_code.py::TestLeapfrogIntegration.test_single_stationary_particle	208μs	1.24ms	-83.2%⚠️
test_numba_jit_code.py::TestLeapfrogIntegration.test_two_particles_approach	298μs	711μs	-58.1%⚠️

🌀 Click to see Generated Regression Tests

import numpy as np

# imports
from code_to_optimize.sample_jit_code import leapfrog_integration

# unit tests


def test_zero_steps_returns_initial():
    """Basic case:
    If n_steps is zero, the integration loop should not run and the
    function should return copies of the original positions and velocities.
    """
    # initial conditions for two particles
    pos = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    vel = np.array([[0.1, 0.0, 0.0], [-0.1, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    # zero steps should result in identical arrays (but not necessarily the same object)
    res_pos, res_vel = leapfrog_integration(pos, vel, masses, dt=0.01, n_steps=0)  # 2.96μs -> 2.67μs (10.9% faster)


def test_two_body_center_of_mass_conservation():
    """Basic physical property:
    For an isolated two-body system with equal masses initially at rest,
    the center of mass should remain (approximately) constant after integration.
    """
    # two equal masses at symmetric positions, zero initial velocities
    pos = np.array([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]], dtype=np.float64)
    vel = np.zeros_like(pos)
    masses = np.array([2.0, 2.0], dtype=np.float64)

    # small time step, few steps to keep numerical error small
    dt = 0.001
    n_steps = 5

    # compute initial center of mass
    total_mass = masses.sum()
    com_initial = (masses.reshape(-1, 1) * pos).sum(axis=0) / total_mass

    res_pos, res_vel = leapfrog_integration(pos, vel, masses, dt=dt, n_steps=n_steps)  # 36.8μs -> 77.9μs (52.8% slower)

    # compute final center of mass
    com_final = (masses.reshape(-1, 1) * res_pos).sum(axis=0) / total_mass


def test_zero_masses_no_forces_particles_move_with_constant_velocity():
    """Edge case:
    If all particle masses are zero, gravitational acceleration terms should be zero,
    so velocities remain constant and positions update linearly with time.
    """
    pos = np.array([[0.0, 0.0, 0.0], [1.0, 2.0, -1.0], [-0.5, 0.2, 0.1]], dtype=np.float64)
    vel = np.array([[1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.5, 0.5, 0.5]], dtype=np.float64)
    # all masses zero -> no gravitational influence
    masses = np.zeros(3, dtype=np.float64)

    dt = 1.0  # choose dt = 1 so arithmetic is exact in many cases
    n_steps = 3

    expected_pos = pos + vel * (dt * n_steps)  # straight-line motion
    expected_vel = vel.copy()  # velocities unchanged

    res_pos, res_vel = leapfrog_integration(pos, vel, masses, dt=dt, n_steps=n_steps)  # 40.5μs -> 50.5μs (19.8% slower)


def test_overlapping_particles_softening_avoids_singularities():
    """Edge case:
    Two particles initialized at the exact same position could cause division by zero
    without softening. Ensure that with the default softening the function
    returns finite numbers (no inf or NaN).
    """
    pos = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=np.float64)  # identical positions
    vel = np.zeros_like(pos)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    res_pos, res_vel = leapfrog_integration(
        pos, vel, masses, dt=0.1, n_steps=1, softening=0.01
    )  # 9.75μs -> 19.8μs (50.8% slower)


def test_negative_dt_steps_backwards_in_time_for_free_particle():
    """Edge case:
    With a single particle and zero accelerations (or zero masses),
    a negative time step should move the particle backward along its velocity vector.
    """
    pos = np.array([[1.0, 0.0, 0.0]], dtype=np.float64)
    vel = np.array([[2.0, 0.0, 0.0]], dtype=np.float64)
    masses = np.array([0.0], dtype=np.float64)  # zero mass => no acceleration

    dt = -0.5  # negative time step
    n_steps = 1

    expected_pos = pos + vel * (dt * n_steps)  # moves backward
    expected_vel = vel.copy()  # unchanged because no forces

    res_pos, res_vel = leapfrog_integration(pos, vel, masses, dt=dt, n_steps=n_steps)  # 5.54μs -> 18.4μs (69.8% slower)


def test_negative_n_steps_treated_as_zero_iteration():
    """Edge case:
    If n_steps is negative, the loop range(n_steps) yields no iterations.
    The function should behave as if n_steps == 0 and return initial conditions.
    """
    pos = np.array([[0.0, 1.0, 2.0]], dtype=np.float64)
    vel = np.array([[0.1, 0.2, 0.3]], dtype=np.float64)
    masses = np.array([1.0], dtype=np.float64)

    res_pos, res_vel = leapfrog_integration(pos, vel, masses, dt=0.01, n_steps=-5)  # 2.50μs -> 2.54μs (1.65% slower)


def test_large_scale_sanity_center_of_mass_conserved_many_particles():
    """Large scale scenario (but kept below the 1000-element cap):
    Create many particles with carefully zeroed total momentum, and verify that
    after several steps the center of mass remains approximately constant and
    no NaNs are introduced. This checks scalability (within reasonable limits)
    and numerical stability.
    """
    rng = np.random.RandomState(12345)  # deterministic RNG for reproducibility
    n_particles = 200  # well under the 1000 limit
    # random positions in a cube, dtype float64
    pos = rng.normal(scale=1.0, size=(n_particles, 3)).astype(np.float64)
    # random velocities
    vel = rng.normal(scale=0.1, size=(n_particles, 3)).astype(np.float64)
    # random positive masses
    masses = rng.uniform(low=0.1, high=5.0, size=(n_particles,)).astype(np.float64)

    # Adjust velocities so total momentum is exactly zero (or extremely close)
    total_momentum = (masses.reshape(-1, 1) * vel).sum(axis=0)
    vel -= total_momentum.reshape(1, 3) / masses.sum()  # now net momentum ~ 0

    # compute initial center of mass
    com_initial = (masses.reshape(-1, 1) * pos).sum(axis=0) / masses.sum()

    dt = 0.001
    n_steps = 5  # keep number of steps modest to limit test time (and stay under loop limit guidance)

    res_pos, res_vel = leapfrog_integration(
        pos, vel, masses, dt=dt, n_steps=n_steps, softening=0.01
    )  # 190ms -> 1.25ms (15145% faster)

    # center of mass should remain approximately constant (tolerance accounts for numerical error)
    com_final = (masses.reshape(-1, 1) * res_pos).sum(axis=0) / masses.sum()


# 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
# ============================================================================


def test_single_particle_no_motion():
    """Test that a single particle with zero velocity doesn't move."""
    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 = 1

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 5.75μs -> 19.5μs (70.5% slower)


def test_output_shape_two_particles():
    """Test that output shapes match input shapes for 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.001
    n_steps = 1

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 10.8μs -> 20.5μs (47.7% slower)


def test_two_particle_symmetry():
    """Test that two particles of equal mass interact symmetrically."""
    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 = 1

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 10.3μs -> 20.1μs (48.6% slower)


def test_zero_timestep_no_change():
    """Test that zero timestep produces no change in state."""
    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.0
    n_steps = 10

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 63.5μs -> 149μs (57.4% slower)


def test_constant_velocity_motion():
    """Test that a particle with constant velocity moves correctly without external forces."""
    # Place particles far apart so forces are negligible
    positions = np.array([[0.0, 0.0, 0.0], [100.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[1.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 = 1
    softening = 10.0  # Large softening to minimize forces

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening=softening
    )  # 10.6μs -> 20.3μs (47.6% slower)


def test_different_softening_values():
    """Test that softening parameter affects acceleration magnitude."""
    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.001
    n_steps = 1

    # Small softening leads to larger forces
    final_pos_small, final_vel_small = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening=0.001
    )  # 10.5μs -> 20.2μs (47.9% slower)

    # Large softening leads to smaller forces
    final_pos_large, final_vel_large = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening=0.1
    )  # 8.12μs -> 17.6μs (53.9% slower)

    # Velocity changes should be larger with smaller softening
    vel_change_small = abs(final_vel_small[0, 0])
    vel_change_large = abs(final_vel_large[0, 0])


# ============================================================================
# EDGE CASE TEST CASES
# ============================================================================


def test_single_particle_with_initial_velocity():
    """Test that a single particle moves with constant velocity (no forces)."""
    positions = np.array([[0.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[1.0, 2.0, 3.0]], dtype=np.float64)
    masses = np.array([1.0], dtype=np.float64)
    dt = 0.01
    n_steps = 1

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 5.58μs -> 18.6μs (70.0% slower)


def test_very_small_timestep():
    """Test that very small timesteps produce minimal changes."""
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.5, 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-8
    n_steps = 1

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 10.4μs -> 19.8μs (47.7% slower)


def test_very_large_timestep():
    """Test that very large timesteps still produce valid results."""
    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 = 1

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 10.2μs -> 20.0μs (49.2% slower)


def test_many_particles_same_position():
    """Test behavior when multiple particles are at the same position (edge case)."""
    positions = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.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.01
    n_steps = 1
    softening = 0.01  # Softening prevents singularity

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening=softening
    )  # 17.7μs -> 21.1μs (16.2% slower)


def test_highly_disparate_masses():
    """Test interaction between particles with very different masses."""
    # Heavy particle at origin, light particle far away
    positions = np.array([[0.0, 0.0, 0.0], [10.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([1000.0, 0.001], dtype=np.float64)  # Huge mass difference
    dt = 0.001
    n_steps = 1

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 10.0μs -> 19.8μs (49.5% slower)


def test_negative_coordinates():
    """Test that particles can have negative coordinates."""
    positions = np.array([[-5.0, -3.0, -1.0], [2.0, 4.0, 6.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 = 1

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 9.96μs -> 19.9μs (50.0% slower)


def test_zero_mass_particle():
    """Test behavior with a particle of very small (but non-zero) mass."""
    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, 1e-10], dtype=np.float64)  # Very small but non-zero
    dt = 0.001
    n_steps = 1

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 9.96μs -> 20.0μs (50.2% slower)


def test_multiple_steps_convergence():
    """Test that energy remains stable over multiple steps (qualitative check)."""
    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.001
    n_steps = 10

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 64.1μs -> 148μs (56.9% slower)


def test_large_softening_dominates():
    """Test that very large softening parameter effectively removes gravitational forces."""
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[1.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 = 1
    softening = 1e6  # Extremely large softening

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening=softening
    )  # 10.0μs -> 20.1μs (50.2% slower)

    # With huge softening, gravitational forces should be negligible
    # Particle 1 should move roughly along its initial velocity
    expected_x = 0.0 + 1.0 * 0.01


# ============================================================================
# LARGE SCALE TEST CASES
# ============================================================================


def test_large_number_of_particles():
    """Test system with many particles (100 particles)."""
    n_particles = 100
    # Create random initial conditions within a bounded region
    np.random.seed(42)
    positions = np.random.uniform(-10.0, 10.0, size=(n_particles, 3)).astype(np.float64)
    velocities = np.random.uniform(-0.1, 0.1, size=(n_particles, 3)).astype(np.float64)
    masses = np.ones(n_particles, dtype=np.float64)
    dt = 0.001
    n_steps = 2  # Keep steps low for performance

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 18.4ms -> 194μs (9377% faster)


def test_many_integration_steps():
    """Test system with many integration steps (500 steps)."""
    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.0001
    n_steps = 500

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 2.88ms -> 7.05ms (59.2% slower)


def test_large_scale_system_50_particles_100_steps():
    """Test scalability with 50 particles and 100 integration steps."""
    n_particles = 50
    np.random.seed(123)
    positions = np.random.normal(0.0, 5.0, size=(n_particles, 3)).astype(np.float64)
    velocities = np.random.normal(0.0, 0.05, size=(n_particles, 3)).astype(np.float64)
    masses = np.random.uniform(0.5, 2.0, size=n_particles).astype(np.float64)
    dt = 0.001
    n_steps = 100

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 235ms -> 3.02ms (7693% faster)


def test_large_scale_with_high_precision():
    """Test numerical precision is maintained over large scale computations."""
    n_particles = 20
    np.random.seed(999)
    positions = np.random.uniform(-1.0, 1.0, size=(n_particles, 3)).astype(np.float64)
    velocities = np.zeros((n_particles, 3), dtype=np.float64)  # Start at rest
    masses = np.ones(n_particles, dtype=np.float64)
    dt = 0.001
    n_steps = 50

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 19.7ms -> 863μs (2185% faster)
    # Particles should have moved from initial positions (gravity should act)
    pos_change = np.linalg.norm(final_pos - positions)


def test_large_coordinates_scale():
    """Test with particles at large spatial coordinates."""
    positions = np.array([[0.0, 0.0, 0.0], [1000.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 = 10
    softening = 50.0  # Scaled softening for large distances

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening=softening
    )  # 70.2μs -> 149μs (53.0% slower)


def test_dense_particle_system():
    """Test system where particles are close together (tests softening effectiveness)."""
    n_particles = 30
    np.random.seed(456)
    # Create particles in small region to test softening
    positions = np.random.normal(0.0, 0.1, size=(n_particles, 3)).astype(np.float64)
    velocities = np.random.normal(0.0, 0.001, size=(n_particles, 3)).astype(np.float64)
    masses = np.ones(n_particles, dtype=np.float64)
    dt = 0.0001
    n_steps = 5
    softening = 0.05  # Softening prevents singularities in dense system

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps, softening=softening
    )  # 4.25ms -> 117μs (3520% faster)
    # Particles should interact significantly
    vel_change = np.linalg.norm(final_vel - velocities)


def test_mixed_mass_large_system():
    """Test system with varied masses at scale."""
    n_particles = 40
    np.random.seed(789)
    positions = np.random.uniform(-5.0, 5.0, size=(n_particles, 3)).astype(np.float64)
    velocities = np.random.uniform(-0.1, 0.1, size=(n_particles, 3)).astype(np.float64)
    # Create varied mass distribution
    masses = np.concatenate(
        [
            np.full(10, 0.1, dtype=np.float64),  # Light particles
            np.full(20, 1.0, dtype=np.float64),  # Medium particles
            np.full(10, 10.0, dtype=np.float64),  # Heavy particles
        ]
    )
    dt = 0.001
    n_steps = 20

    final_pos, final_vel = leapfrog_integration(
        positions, velocities, masses, dt, n_steps
    )  # 31.1ms -> 504μs (6069% faster)


# 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-mkg7vhgc and push.