⚡️ Speed up function leapfrog_integration by 32,859%

[codeflash-ai]

📄 32,859% (328.59x) speedup for leapfrog_integration in code_to_optimize/sample_code.py

⏱️ Runtime : 956 milliseconds → 2.90 milliseconds (best of 101 runs)

📝 Explanation and details

The optimized code achieves a 329x speedup (32859%) by adding a single critical change: the @numba.njit(cache=True) decorator.

What Changed

1. Added Numba JIT compilation: The decorator @numba.njit(cache=True) compiles the function to native machine code
2. Imported numba library: Added import numba at the top
3. No algorithmic changes: The logic, data structures, and computation flow remain identical

Why This Optimization Works

Numba eliminates Python interpreter overhead - The original code spends most of its time in nested loops performing simple arithmetic operations. Line profiler shows the inner loops (particle-particle interactions) account for ~85% of runtime with intensive array indexing and floating-point operations. Python's interpreter adds significant overhead for:

• Array indexing operations (pos[j, 0], acc[i, 1], etc.)
• Arithmetic operations in tight loops
• Loop iteration and control flow

JIT compilation converts Python to optimized machine code - Numba translates the function to LLVM intermediate representation, then to native code that runs at C/C++ speeds. For this computation-heavy, loop-intensive code with minimal Python object overhead, JIT compilation provides near-optimal performance.

Cache=True avoids recompilation - The compiled function is cached to disk, so subsequent runs skip compilation overhead entirely.

Test Results Patterns

The optimization shows dramatic improvements for compute-intensive scenarios:

• Large-scale tests: 38000%+ speedup (100 particles, 10+ steps) - the nested O(n²) complexity amplifies JIT benefits
• Many-step simulations: 9800-37000% speedup (50-200 steps) - loop overhead compounds over iterations
• Dense clusters: 31844% speedup (25 particles in tight space) - many force calculations benefit from elimination of interpreter overhead
• Moderate speedups for trivial cases: 30-300% for edge cases (zero steps, single particles) where setup/teardown dominates

The optimization is particularly effective for the hot path: the triply-nested loop computing pairwise gravitational forces, which dominates runtime in realistic N-body simulations.

Impact Considerations

This is a drop-in optimization with minimal risk:

• Pure performance enhancement with no behavioral changes
• All regression tests pass with identical numerical results
• Numba is a mature, widely-used library for scientific computing
• The cache=True flag ensures the first-run compilation cost is amortized across executions

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 21 Passed
🌀 Generated Regression Tests	✅ 37 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	177μs	11.7μs	1412%✅
test_numba_jit_code.py::TestLeapfrogIntegration.test_momentum_conservation	7.79ms	29.4μs	26385%✅
test_numba_jit_code.py::TestLeapfrogIntegration.test_single_moving_particle	522μs	14.0μs	3637%✅
test_numba_jit_code.py::TestLeapfrogIntegration.test_single_stationary_particle	516μs	14.2μs	3534%✅
test_numba_jit_code.py::TestLeapfrogIntegration.test_two_particles_approach	785μs	13.1μs	5909%✅

🌀 Click to see Generated Regression Tests

import math
import time

# function to test
import numpy as np  # used to construct typed numeric arrays

# imports
from code_to_optimize.sample_code import leapfrog_integration

# unit tests
#
# The tests below target:
# - Basic functionality (zero steps, zero dt, single particle)
# - Edge cases (empty inputs, negative steps behave like zero)
# - Correctness in small analytic scenarios (two-body single step expected outcome)
# - Conservation properties (momentum conservation)
# - Softening behavior (larger softening -> smaller accelerations)
# - Large-scale sanity/performance (moderate number of particles)
#
# Notes on assertions:
# - We avoid numpy's testing helpers as requested; we use Python assert and pytest.approx
# - All numeric inputs are numpy arrays with explicit dtypes (float64) to be JIT-friendly


def test_zero_steps_returns_identical_arrays():
    # Basic: When n_steps is 0, the function should return unchanged copies of inputs.
    positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.1, 0.0, 0.0], [0.0, 0.2, 0.0]], dtype=np.float64)
    masses = np.array([1.0, 2.0], dtype=np.float64)

    # Call with zero steps
    out_pos, out_vel = leapfrog_integration(
        positions, velocities, masses, dt=0.1, n_steps=0
    )  # 19.8μs -> 10.6μs (86.8% faster)


def test_dt_zero_leaves_system_unchanged_even_with_steps():
    # Edge: When dt is zero, there should be no change regardless of n_steps.
    positions = np.array([[2.0, -1.0, 0.5]], dtype=np.float64)  # single particle
    velocities = np.array([[0.3, -0.2, 1.0]], dtype=np.float64)
    masses = np.array([5.0], dtype=np.float64)

    out_pos, out_vel = leapfrog_integration(
        positions, velocities, masses, dt=0.0, n_steps=5
    )  # 46.5μs -> 10.9μs (326% faster)


def test_single_particle_moves_linearly_and_velocity_constant():
    # Basic: A single particle feels no internal forces -> velocity stays constant,
    # and positions advance linearly with pos = pos0 + n_steps * dt * velocity.
    pos0 = np.array([[0.0, 1.0, -1.0]], dtype=np.float64)
    vel0 = np.array([[0.5, -0.25, 0.0]], dtype=np.float64)
    masses = np.array([10.0], dtype=np.float64)

    dt = 0.02
    n_steps = 10  # keep under 1000 per instructions

    expected_pos = pos0 + n_steps * dt * vel0  # analytic expectation
    expected_vel = vel0.copy()  # no forces imply unchanged velocity

    out_pos, out_vel = leapfrog_integration(
        pos0, vel0, masses, dt=dt, n_steps=n_steps
    )  # 67.3μs -> 10.9μs (515% faster)


def compute_two_body_acc_i(pos_i, pos_j, mj, softening):
    # Helper (pure numpy) to compute acceleration on particle i due to particle j (vector).
    dx = pos_j[0] - pos_i[0]
    dy = pos_j[1] - pos_i[1]
    dz = pos_j[2] - pos_i[2]
    dist_sq = dx * dx + dy * dy + dz * dz + softening * softening
    dist = math.sqrt(dist_sq)
    dist_cubed = dist_sq * dist
    # G = 1.0 per implementation
    ax = (1.0 / dist_cubed) * dx * mj
    ay = (1.0 / dist_cubed) * dy * mj
    az = (1.0 / dist_cubed) * dz * mj
    return np.array([ax, ay, az], dtype=np.float64)


def test_two_body_single_step_matches_analytic_prediction():
    # Analytic small-case test:
    # For two particles with initial v=0, after one full leapfrog step:
    # vel_final = vel0 + dt * acc  (since two half-steps sum to dt*acc)
    # pos_final = pos0 + dt * (vel0 + 0.5 * dt * acc) = pos0 + 0.5 * dt^2 * acc (if vel0 = 0)
    pos0 = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)
    vel0 = np.zeros((2, 3), dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)

    dt = 0.001  # small dt for better linearization
    n_steps = 1
    softening = 0.01

    # Compute expected accelerations analytically using the same formulas
    acc0 = compute_two_body_acc_i(pos0[0], pos0[1], masses[1], softening)
    acc1 = compute_two_body_acc_i(pos0[1], pos0[0], masses[0], softening)

    # Expected final velocities after one full step
    expected_vel0 = vel0[0] + dt * acc0
    expected_vel1 = vel0[1] + dt * acc1
    expected_vel = np.vstack([expected_vel0, expected_vel1])

    # Expected positions after one step (vel0 are zero)
    expected_pos0 = pos0[0] + 0.5 * dt * dt * acc0
    expected_pos1 = pos0[1] + 0.5 * dt * dt * acc1
    expected_pos = np.vstack([expected_pos0, expected_pos1])

    out_pos, out_vel = leapfrog_integration(
        pos0, vel0, masses, dt=dt, n_steps=n_steps, softening=softening
    )  # 33.2μs -> 10.3μs (223% faster)


def test_momentum_conservation_randomized_small_system():
    # Conservation property: total linear momentum (sum m_i * v_i) should be conserved by internal forces.
    rng = np.random.RandomState(12345)
    n_particles = 10  # small but nontrivial system
    positions = rng.randn(n_particles, 3).astype(np.float64)
    velocities = rng.randn(n_particles, 3).astype(np.float64) * 0.1
    masses = rng.rand(n_particles).astype(np.float64) + 0.1  # avoid zero mass

    dt = 0.01
    n_steps = 5

    # initial total momentum
    init_momentum = (masses.reshape(-1, 1) * velocities).sum(axis=0)

    out_pos, out_vel = leapfrog_integration(
        positions, velocities, masses, dt=dt, n_steps=n_steps
    )  # 1.48ms -> 14.8μs (9876% faster)

    final_momentum = (masses.reshape(-1, 1) * out_vel).sum(axis=0)


def test_softening_reduces_effective_acceleration_between_two_bodies():
    # Edge: Increasing softening should reduce the acceleration magnitude (force softening).
    pos0 = np.array([[0.0, 0.0, 0.0], [0.5, 0.0, 0.0]], dtype=np.float64)
    vel0 = np.zeros((2, 3), dtype=np.float64)
    masses = np.array([1.0, 1.0], dtype=np.float64)
    dt = 0.01
    n_steps = 1

    # small softening
    out_pos_s, out_vel_s = leapfrog_integration(
        pos0, vel0, masses, dt=dt, n_steps=n_steps, softening=0.0001
    )  # 39.6μs -> 10.8μs (267% faster)
    # large softening
    out_pos_L, out_vel_L = leapfrog_integration(
        pos0, vel0, masses, dt=dt, n_steps=n_steps, softening=0.5
    )  # 24.0μs -> 3.49μs (588% faster)

    # compute delta v magnitudes for one of the particles (index 0)
    delta_v_small = math.sqrt(((out_vel_s[0] - vel0[0]) ** 2).sum())
    delta_v_large = math.sqrt(((out_vel_L[0] - vel0[0]) ** 2).sum())


def test_empty_system_returns_immediate_copies():
    # Edge: zero-particle system (shapes (0,3)) should be supported and return identical empty arrays.
    positions = np.zeros((0, 3), dtype=np.float64)
    velocities = np.zeros((0, 3), dtype=np.float64)
    masses = np.zeros((0,), dtype=np.float64)

    out_pos, out_vel = leapfrog_integration(
        positions, velocities, masses, dt=0.1, n_steps=10
    )  # 21.2μs -> 11.0μs (92.5% faster)


def test_negative_n_steps_treated_as_zero_behavior():
    # Edge: negative n_steps should effectively behave like zero steps because range(-1) is empty.
    positions = np.array([[0.0, 0.0, 0.0]], dtype=np.float64)
    velocities = np.array([[0.1, 0.2, 0.3]], dtype=np.float64)
    masses = np.array([1.0], dtype=np.float64)

    out_pos, out_vel = leapfrog_integration(
        positions, velocities, masses, dt=0.1, n_steps=-5
    )  # 13.8μs -> 10.6μs (30.6% faster)


def test_large_scale_sanity_and_momentum_conservation_runtime():
    # Large Scale: moderate number of particles (kept under 1000) to check performance and conservation.
    rng = np.random.RandomState(2026)
    n_particles = 100  # square is 10k pair computations, acceptable for a unit test
    positions = rng.randn(n_particles, 3).astype(np.float64) * 10.0
    velocities = rng.randn(n_particles, 3).astype(np.float64) * 0.1
    masses = (rng.rand(n_particles).astype(np.float64) + 0.1) * 5.0

    dt = 0.005
    n_steps = 5

    # measure runtime to ensure the implementation scales reasonably (sanity check)
    start = time.time()
    out_pos, out_vel = leapfrog_integration(
        positions, velocities, masses, dt=dt, n_steps=n_steps
    )  # 137ms -> 358μs (38200% faster)
    elapsed = time.time() - start

    # Check momentum conservation for the larger system
    init_momentum = (masses.reshape(-1, 1) * velocities).sum(axis=0)
    final_momentum = (masses.reshape(-1, 1) * out_vel).sum(axis=0)


# 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

from code_to_optimize.sample_code import leapfrog_integration

# ============================================================================
# BASIC TEST CASES
# ============================================================================


class TestBasicFunctionality:
    """Test the fundamental functionality of leapfrog_integration."""

    def test_single_particle_no_movement(self):
        """Single particle with zero velocity should remain at origin with no forces."""
        # Setup: 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)

        # Execute integration for 1 step
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.1, n_steps=1
        )  # 24.0μs -> 11.8μs (104% faster)

    def test_two_particles_identical_properties(self):
        """Two identical particles should have symmetric acceleration and motion."""
        # Setup: Two equal-mass particles positioned 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)

        # Execute integration for 1 step
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.01, n_steps=1
        )  # 40.2μs -> 11.4μs (254% faster)

    def test_output_shape_preservation(self):
        """Output arrays should preserve input shape regardless of step count."""
        # Setup: 3 particles in 3D space
        n_particles = 3
        positions = np.random.rand(n_particles, 3).astype(np.float64)
        velocities = np.random.rand(n_particles, 3).astype(np.float64)
        masses = np.ones(n_particles, dtype=np.float64)

        # Execute with different step counts
        for n_steps in [1, 5, 10]:
            result_pos, result_vel = leapfrog_integration(
                positions, velocities, masses, dt=0.01, n_steps=n_steps
            )  # 523μs -> 19.5μs (2582% faster)

    def test_input_arrays_not_modified(self):
        """Input arrays should not be modified during integration."""
        # Setup: Store original copies for verification
        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([2.0, 3.0], dtype=np.float64)

        pos_copy = positions.copy()
        vel_copy = velocities.copy()
        masses_copy = masses.copy()

        # Execute integration
        leapfrog_integration(positions, velocities, masses, dt=0.1, n_steps=5)  # 100μs -> 11.6μs (769% faster)

    def test_zero_steps_returns_original_state(self):
        """Zero integration steps should return original positions and velocities."""
        # Setup: Random initial 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)

        # Execute with zero steps
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.1, n_steps=0
        )  # 13.9μs -> 10.9μs (27.8% faster)

    def test_constant_velocity_motion(self):
        """Single particle with constant velocity and no forces should move linearly."""
        # Setup: Particle with initial velocity, away from other particles
        positions = np.array([[0.0, 0.0, 0.0]], dtype=np.float64)
        velocities = np.array([[1.0, 0.0, 0.0]], dtype=np.float64)
        masses = np.array([1.0], dtype=np.float64)
        dt = 0.1
        n_steps = 5

        # Execute integration
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=dt, n_steps=n_steps
        )  # 44.4μs -> 11.3μs (293% faster)

        # Verify: Velocity unchanged (no forces), position moved by velocity*time
        expected_displacement = velocities[0] * dt * n_steps


class TestOutputDataTypes:
    """Test that output arrays have correct data types."""

    def test_float64_precision_maintained(self):
        """Results should maintain float64 precision from inputs."""
        # Setup: float64 input arrays
        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)

        # Execute integration
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.01, n_steps=1
        )  # 40.0μs -> 11.2μs (256% faster)

    def test_return_type_is_tuple(self):
        """Function should return a tuple of two arrays."""
        # Setup: Simple system
        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)

        # Execute integration
        codeflash_output = leapfrog_integration(positions, velocities, masses, dt=0.1, n_steps=1)
        result = codeflash_output  # 23.9μs -> 11.0μs (117% faster)


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


class TestEdgeCases:
    """Test edge cases and boundary conditions."""

    def test_very_small_timestep(self):
        """Very small timestep should produce minimal displacement."""
        # Setup: System with gravitational interaction
        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)

        # Execute with very small timestep
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=1e-6, n_steps=1
        )  # 39.8μs -> 11.2μs (256% faster)

        # Verify: Small timestep produces small changes
        displacement = np.linalg.norm(result_pos - positions)

    def test_very_large_timestep(self):
        """Large timestep should produce larger displacements."""
        # Setup: Simple system
        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)

        # Execute with large timestep
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=1.0, n_steps=1
        )  # 39.6μs -> 11.1μs (258% faster)

    def test_very_small_softening_parameter(self):
        """Small softening should not cause numerical errors for nearby particles."""
        # Setup: Particles at moderate distance
        positions = np.array([[0.0, 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)

        # Execute with very small softening
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.01, n_steps=1, softening=1e-6
        )  # 40.1μs -> 10.6μs (278% faster)

    def test_large_softening_parameter(self):
        """Large softening should prevent singularities at small distances."""
        # Setup: Particles very close together
        positions = np.array([[0.0, 0.0, 0.0], [1e-3, 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)

        # Execute with large softening
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.01, n_steps=1, softening=1.0
        )  # 39.9μs -> 10.2μs (291% faster)

    def test_very_large_mass_difference(self):
        """Large mass differences should not cause numerical instability."""
        # Setup: One very massive particle and one light particle
        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([1e6, 1e-6], dtype=np.float64)

        # Execute integration
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.01, n_steps=1
        )  # 40.2μs -> 11.3μs (256% faster)

    def test_very_small_mass_values(self):
        """Small mass values should be handled correctly."""
        # Setup: System with very small masses
        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)

        # Execute integration
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.01, n_steps=1
        )  # 39.9μs -> 11.0μs (262% faster)

    def test_particles_at_same_position(self):
        """Particles at identical positions with softening should not cause singularities."""
        # Setup: Two particles at same location
        positions = np.array([[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]], dtype=np.float64)
        masses = np.array([1.0, 1.0], dtype=np.float64)

        # Execute with non-zero softening
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.01, n_steps=1, softening=0.1
        )  # 39.4μs -> 9.77μs (303% faster)

    def test_zero_mass_particle(self):
        """Zero-mass particles should not affect system dynamics but should move."""
        # Setup: Massive particle and massless test particle
        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.1, 0.0, 0.0]], dtype=np.float64)
        masses = np.array([1.0, 0.0], dtype=np.float64)

        # Execute integration
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.1, n_steps=1
        )  # 39.1μs -> 11.2μs (250% faster)

    def test_many_steps_accumulation(self):
        """Multiple steps should accumulate correctly."""
        # Setup: Simple 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)

        # Execute with many steps
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.01, n_steps=100
        )  # 1.58ms -> 16.0μs (9805% faster)
        # Verify particles have moved due to gravitational attraction
        distance_changed = not np.allclose(
            np.linalg.norm(result_pos[1] - result_pos[0]), np.linalg.norm(positions[1] - positions[0])
        )

    def test_high_velocity_particles(self):
        """Particles with very high velocities should be handled correctly."""
        # Setup: Particles with high velocities
        positions = np.array([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]], dtype=np.float64)
        velocities = np.array([[100.0, 100.0, 100.0], [-100.0, -100.0, -100.0]], dtype=np.float64)
        masses = np.array([1.0, 1.0], dtype=np.float64)

        # Execute integration
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.01, n_steps=1
        )  # 39.7μs -> 11.2μs (256% faster)


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


class TestLargeScale:
    """Test performance and scalability with larger systems."""

    def test_many_particles_scalability(self):
        """System with many particles should compute without errors."""
        # Setup: 100 particles in random configuration
        n_particles = 100
        np.random.seed(42)
        positions = np.random.uniform(-10, 10, 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)

        # Execute integration for multiple steps
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.01, n_steps=10
        )  # 275ms -> 715μs (38387% faster)

    def test_extended_simulation(self):
        """Extended simulation with reasonable number of particles."""
        # Setup: 50 particles with 50 integration steps
        n_particles = 50
        np.random.seed(123)
        positions = np.random.uniform(-5, 5, size=(n_particles, 3)).astype(np.float64)
        velocities = np.zeros((n_particles, 3), dtype=np.float64)
        masses = np.ones(n_particles, dtype=np.float64)

        # Execute extended integration
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.005, n_steps=50
        )  # 346ms -> 907μs (38142% faster)
        # Verify: Total momentum is conserved (should be near zero)
        total_momentum = np.sum(result_vel * masses[:, np.newaxis], axis=0)

    def test_large_coordinate_ranges(self):
        """System with particles spread over large coordinate ranges."""
        # Setup: Particles spread from -1000 to 1000
        n_particles = 20
        np.random.seed(456)
        positions = np.random.uniform(-1000, 1000, size=(n_particles, 3)).astype(np.float64)
        velocities = np.zeros((n_particles, 3), dtype=np.float64)
        masses = np.ones(n_particles, dtype=np.float64)

        # Execute integration
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.01, n_steps=5
        )  # 5.68ms -> 25.6μs (22053% faster)

    def test_mixed_mass_distribution(self):
        """System with varied mass distribution should work correctly."""
        # Setup: Particles with exponentially distributed masses
        n_particles = 30
        np.random.seed(789)
        positions = np.random.uniform(-10, 10, size=(n_particles, 3)).astype(np.float64)
        velocities = np.zeros((n_particles, 3), dtype=np.float64)
        masses = np.exp(np.random.uniform(0, 3, n_particles)).astype(np.float64)

        # Execute integration
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.01, n_steps=10
        )  # 25.2ms -> 80.4μs (31181% faster)

    def test_dense_particle_cluster(self):
        """Particles in dense cluster should not produce numerical errors."""
        # Setup: 25 particles in small volume
        n_particles = 25
        np.random.seed(321)
        positions = np.random.uniform(-0.1, 0.1, size=(n_particles, 3)).astype(np.float64)
        velocities = np.random.uniform(-0.01, 0.01, size=(n_particles, 3)).astype(np.float64)
        masses = np.ones(n_particles, dtype=np.float64)

        # Execute with appropriate softening
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.001, n_steps=10, softening=0.05
        )  # 17.4ms -> 54.6μs (31844% faster)

    def test_energy_conservation_trend(self):
        """System should show reasonable energy behavior over many steps."""
        # Setup: 10 particles with initial velocities
        n_particles = 10
        np.random.seed(654)
        positions = np.random.uniform(-5, 5, size=(n_particles, 3)).astype(np.float64)
        velocities = np.random.uniform(-0.05, 0.05, size=(n_particles, 3)).astype(np.float64)
        masses = np.ones(n_particles, dtype=np.float64)

        # Execute integration
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.01, n_steps=20
        )  # 5.84ms -> 25.6μs (22677% faster)

        # Verify: System remains stable (energy doesn't explode)
        kinetic_energy = 0.5 * np.sum(masses[:, np.newaxis] * result_vel**2)

    def test_performance_with_many_steps(self):
        """Integration with many steps should complete successfully."""
        # Setup: Moderate system with many steps
        n_particles = 15
        np.random.seed(987)
        positions = np.random.uniform(-5, 5, size=(n_particles, 3)).astype(np.float64)
        velocities = np.zeros((n_particles, 3), dtype=np.float64)
        masses = np.ones(n_particles, dtype=np.float64)

        # Execute integration with 200 steps
        result_pos, result_vel = leapfrog_integration(
            positions, velocities, masses, dt=0.005, n_steps=200
        )  # 127ms -> 345μs (36984% 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-mkgfarwh and push.