⚡️ Speed up function rolling_mean by 163%

[codeflash-ai]

📄 163% (1.63x) speedup for rolling_mean in src/data_processing/series.py

⏱️ Runtime : 4.01 milliseconds → 1.53 milliseconds (best of 234 runs)

📝 Explanation and details

The optimized code achieves a 162% speedup by replacing an inefficient nested loop with vectorized NumPy operations while preserving exact behavioral compatibility with the original implementation.

Key Performance Optimizations:

1. Cumulative Sum Algorithm: The core optimization replaces the O(n×w) nested loop with O(n) cumulative sum operations. Instead of recalculating window sums from scratch, it uses cumsum[i] - cumsum[i-window] to compute rolling sums in constant time per window.

2. Vectorized NumPy Operations: Pre-allocates result arrays with np.full() and leverages NumPy's optimized C implementations for cumulative sum calculations, eliminating Python loop overhead.

Behavioral Preservation:
The optimization carefully maintains the original's edge case handling through fallback logic:

• Window = 0: Preserves the original's ZeroDivisionError behavior
• Non-numeric data: Falls back to original logic to maintain TypeError exceptions
• Negative windows: Uses original slow path for exact compatibility
• Large windows: Optimizes the common case where window > series_length

Performance Impact Analysis:
From the line profiler results, the optimization eliminates the expensive nested loop (lines accounting for ~90% of original runtime) and replaces it with efficient NumPy operations. The test results show significant gains for larger datasets:

• Large series (1000 elements): 387% faster
• Large windows: 111% faster
• Performance-critical scenarios benefit most from the O(n×w) → O(n) algorithmic improvement

When This Optimization Matters:
This optimization is particularly valuable for time-series analysis, financial data processing, or any scenario requiring rolling statistics on large datasets, where the quadratic time complexity of the original implementation becomes a bottleneck.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import numpy as np
import pandas as pd

# imports
import pytest  # used for our unit tests
from src.data_processing.series import rolling_mean

# unit tests

# ----------- BASIC TEST CASES -----------


def test_rolling_mean_basic_small_window():
    # Test with a simple series and window=2
    s = pd.Series([1, 2, 3, 4, 5])
    expected = [np.nan, 1.5, 2.5, 3.5, 4.5]
    codeflash_output = rolling_mean(s, 2)
    result = codeflash_output  # 2.62μs -> 8.88μs (70.4% slower)
    # Compare each value, allowing for nan
    for r, e in zip(result, expected):
        if np.isnan(e):
            pass
        else:
            pass


def test_rolling_mean_basic_window_1():
    # Window size 1 should return the original values
    s = pd.Series([10, 20, 30])
    expected = [10, 20, 30]
    codeflash_output = rolling_mean(s, 1)
    result = codeflash_output  # 2.25μs -> 8.54μs (73.7% slower)
    for r, e in zip(result, expected):
        pass


def test_rolling_mean_basic_full_window():
    # Window size equal to the length of the series
    s = pd.Series([2, 4, 6, 8])
    expected = [np.nan, np.nan, np.nan, 5.0]
    codeflash_output = rolling_mean(s, 4)
    result = codeflash_output  # 2.33μs -> 7.71μs (69.7% slower)
    for r, e in zip(result, expected):
        if np.isnan(e):
            pass
        else:
            pass


def test_rolling_mean_basic_floats():
    # Test with float values
    s = pd.Series([1.5, 2.5, 3.5, 4.5])
    expected = [np.nan, 2.0, 3.0, 4.0]
    codeflash_output = rolling_mean(s, 2)
    result = codeflash_output  # 3.00μs -> 8.42μs (64.4% slower)
    for r, e in zip(result, expected):
        if np.isnan(e):
            pass
        else:
            pass


# ----------- EDGE TEST CASES -----------


def test_rolling_mean_empty_series():
    # Empty series should return empty list
    s = pd.Series([], dtype=float)
    codeflash_output = rolling_mean(s, 3)
    result = codeflash_output  # 1.38μs -> 1.75μs (21.4% slower)


def test_rolling_mean_window_zero():
    # Window size zero is invalid, should raise error
    s = pd.Series([1, 2, 3])
    with pytest.raises(ZeroDivisionError):
        rolling_mean(s, 0)  # 2.04μs -> 2.29μs (10.9% slower)


def test_rolling_mean_window_larger_than_series():
    # Window larger than series: all elements should be nan
    s = pd.Series([1, 2, 3])
    expected = [np.nan, np.nan, np.nan]
    codeflash_output = rolling_mean(s, 5)
    result = codeflash_output  # 2.17μs -> 1.88μs (15.6% faster)
    for r in result:
        pass


def test_rolling_mean_with_nans_in_series():
    # Series contains np.nan values
    s = pd.Series([1, np.nan, 3, 4])
    expected = [np.nan, np.nan, np.nan, np.nan]  # sum with nan always gives nan
    codeflash_output = rolling_mean(s, 4)
    result = codeflash_output  # 2.75μs -> 8.62μs (68.1% slower)
    for r in result:
        pass


def test_rolling_mean_with_inf_in_series():
    # Series contains inf values
    s = pd.Series([1, 2, np.inf, 4])
    expected = [np.nan, np.nan, np.nan, np.inf]
    codeflash_output = rolling_mean(s, 4)
    result = codeflash_output  # 2.58μs -> 7.79μs (66.9% slower)
    for r, e in zip(result, expected):
        if np.isnan(e):
            pass
        elif np.isinf(e):
            pass
        else:
            pass


def test_rolling_mean_series_with_negatives():
    # Series contains negative numbers
    s = pd.Series([-1, -2, -3, -4])
    expected = [np.nan, -1.5, -2.5, -3.5]
    codeflash_output = rolling_mean(s, 2)
    result = codeflash_output  # 2.71μs -> 8.88μs (69.5% slower)
    for r, e in zip(result, expected):
        if np.isnan(e):
            pass
        else:
            pass


def test_rolling_mean_series_with_repeated_values():
    # All values are the same
    s = pd.Series([7, 7, 7, 7])
    expected = [np.nan, 7.0, 7.0, 7.0]
    codeflash_output = rolling_mean(s, 2)
    result = codeflash_output  # 2.50μs -> 8.62μs (71.0% slower)
    for r, e in zip(result, expected):
        if np.isnan(e):
            pass
        else:
            pass


# ----------- LARGE SCALE TEST CASES -----------


def test_rolling_mean_large_series():
    # Large series, window=10
    s = pd.Series(range(1000))
    window = 10
    codeflash_output = rolling_mean(s, window)
    result = codeflash_output  # 405μs -> 307μs (31.9% faster)
    # First window-1 should be nan
    for i in range(window - 1):
        pass
    # The rest should be correct
    for i in range(window - 1, 1000):
        expected = sum(range(i - window + 1, i + 1)) / window


def test_rolling_mean_large_window():
    # Window almost as large as the series
    s = pd.Series(range(1, 1001))  # 1 to 1000
    window = 999
    codeflash_output = rolling_mean(s, window)
    result = codeflash_output  # 132μs -> 66.1μs (101% faster)
    # First 998 should be nan
    for i in range(window - 1):
        pass
    # Last two values should be the mean of 1..999 and 2..1000
    expected1 = sum(range(1, 1000)) / window
    expected2 = sum(range(2, 1001)) / window


def test_rolling_mean_performance():
    # Test that function completes in reasonable time for 1000 elements
    import time

    s = pd.Series([i for i in range(1000)])
    start = time.time()
    codeflash_output = rolling_mean(s, 50)
    result = codeflash_output  # 1.43ms -> 294μs (387% faster)
    end = time.time()


# ----------- ADDITIONAL EDGE CASES -----------


@pytest.mark.parametrize(
    "series,window,expected",
    [
        (pd.Series([1]), 1, [1]),  # Single element, window 1
        (pd.Series([1]), 2, [np.nan]),  # Single element, window 2
        (pd.Series([0, 0, 0]), 2, [np.nan, 0.0, 0.0]),  # All zeros
    ],
)
def test_rolling_mean_parametrized(series, window, expected):
    codeflash_output = rolling_mean(series, window)
    result = codeflash_output  # 6.71μs -> 19.7μs (66.0% slower)
    for r, e in zip(result, expected):
        if np.isnan(e):
            pass
        else:
            pass


def test_rolling_mean_series_with_object_dtype():
    # Series with object dtype (should still work for numbers)
    s = pd.Series([1, 2, 3, 4], dtype=object)
    expected = [np.nan, 1.5, 2.5, 3.5]
    codeflash_output = rolling_mean(s, 2)
    result = codeflash_output  # 2.62μs -> 8.96μs (70.7% slower)
    for r, e in zip(result, expected):
        if np.isnan(e):
            pass
        else:
            pass


def test_rolling_mean_series_with_non_numeric():
    # Series with non-numeric values should raise TypeError
    s = pd.Series(["a", "b", "c"])
    with pytest.raises(TypeError):
        rolling_mean(s, 2)  # 2.71μs -> 6.25μs (56.7% slower)


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

import math
from typing import List

# function to test
# src/data_processing/series.py
import numpy as np
import pandas as pd

# imports
import pytest  # used for our unit tests
from src.data_processing.series import rolling_mean

# unit tests

# ------------------- Basic Test Cases -------------------


def test_rolling_mean_basic_integers():
    # Test with a simple integer series and window size 3
    s = pd.Series([1, 2, 3, 4, 5])
    expected = [math.nan, math.nan, 2.0, 3.0, 4.0]
    codeflash_output = rolling_mean(s, 3)
    result = codeflash_output  # 2.83μs -> 9.17μs (69.1% slower)
    for r, e in zip(result, expected):
        if math.isnan(e):
            pass
        else:
            pass


def test_rolling_mean_basic_floats():
    # Test with a float series and window size 2
    s = pd.Series([1.0, 2.0, 3.0, 4.0])
    expected = [math.nan, 1.5, 2.5, 3.5]
    codeflash_output = rolling_mean(s, 2)
    result = codeflash_output  # 3.04μs -> 8.62μs (64.7% slower)
    for r, e in zip(result, expected):
        if math.isnan(e):
            pass
        else:
            pass


def test_rolling_mean_window_1():
    # Window size 1 should return the original values
    s = pd.Series([10, 20, 30])
    expected = [10.0, 20.0, 30.0]
    codeflash_output = rolling_mean(s, 1)
    result = codeflash_output  # 2.42μs -> 8.58μs (71.9% slower)
    for r, e in zip(result, expected):
        pass


def test_rolling_mean_window_equals_length():
    # Window size equal to series length: all but last value are nan, last is mean
    s = pd.Series([2, 4, 6, 8])
    expected = [math.nan, math.nan, math.nan, 5.0]
    codeflash_output = rolling_mean(s, 4)
    result = codeflash_output  # 2.33μs -> 7.71μs (69.7% slower)
    for r, e in zip(result, expected):
        if math.isnan(e):
            pass
        else:
            pass


def test_rolling_mean_negative_numbers():
    # Test with negative numbers
    s = pd.Series([-1, -2, -3, -4, -5])
    expected = [math.nan, math.nan, -2.0, -3.0, -4.0]
    codeflash_output = rolling_mean(s, 3)
    result = codeflash_output  # 2.92μs -> 8.58μs (66.0% slower)
    for r, e in zip(result, expected):
        if math.isnan(e):
            pass
        else:
            pass


# ------------------- Edge Test Cases -------------------


def test_rolling_mean_empty_series():
    # Empty series should return empty list
    s = pd.Series([])
    codeflash_output = rolling_mean(s, 3)
    result = codeflash_output  # 1.54μs -> 1.79μs (14.0% slower)


def test_rolling_mean_window_larger_than_series():
    # Window larger than series: all values should be nan
    s = pd.Series([1, 2])
    expected = [math.nan, math.nan]
    codeflash_output = rolling_mean(s, 5)
    result = codeflash_output  # 1.83μs -> 1.79μs (2.29% faster)
    for r, e in zip(result, expected):
        pass


def test_rolling_mean_window_zero():
    # Window size 0 should raise an error (division by zero)
    s = pd.Series([1, 2, 3])
    with pytest.raises(ZeroDivisionError):
        rolling_mean(s, 0)  # 2.17μs -> 2.46μs (11.9% slower)


def test_rolling_mean_with_nans_in_input():
    # Test with NaN values in the input series
    s = pd.Series([1.0, math.nan, 3.0, 4.0])
    expected = [math.nan, math.nan, math.nan, math.nan]
    # The rolling window sum will fail if window includes nan, so output should be nan
    # But per the implementation, nan + number = nan, so all windows with nan will be nan
    codeflash_output = rolling_mean(s, 3)
    result = codeflash_output  # 3.25μs -> 8.96μs (63.7% slower)
    for r in result:
        pass


def test_rolling_mean_with_inf_in_input():
    # Series with inf values
    s = pd.Series([1.0, float("inf"), 3.0, 4.0])
    expected = [math.nan, math.nan, float("inf"), float("inf")]
    codeflash_output = rolling_mean(s, 3)
    result = codeflash_output  # 2.88μs -> 8.38μs (65.7% slower)
    for r, e in zip(result, expected):
        if math.isnan(e):
            pass
        elif math.isinf(e):
            pass
        else:
            pass


def test_rolling_mean_series_with_one_element():
    # Series with one element, window 1
    s = pd.Series([42])
    expected = [42.0]
    codeflash_output = rolling_mean(s, 1)
    result = codeflash_output  # 2.00μs -> 7.75μs (74.2% slower)


def test_rolling_mean_series_with_one_element_large_window():
    # Series with one element, window > 1
    s = pd.Series([42])
    expected = [math.nan]
    codeflash_output = rolling_mean(s, 2)
    result = codeflash_output  # 1.71μs -> 1.79μs (4.63% slower)


# ------------------- Large Scale Test Cases -------------------


def test_rolling_mean_large_series():
    # Test with a large series and moderate window
    data = list(range(1000))
    s = pd.Series(data)
    window = 10
    codeflash_output = rolling_mean(s, window)
    result = codeflash_output  # 404μs -> 310μs (30.3% faster)
    # The first (window-1) results should be nan
    for i in range(window - 1):
        pass


def test_rolling_mean_large_window():
    # Series length 1000, window size 999
    data = [1] * 1000
    s = pd.Series(data)
    codeflash_output = rolling_mean(s, 999)
    result = codeflash_output  # 121μs -> 57.8μs (111% faster)
    # First 998 values are nan
    for i in range(998):
        pass


def test_rolling_mean_performance():
    # This test is not a benchmark, but ensures function completes on large input
    s = pd.Series(range(1000))
    window = 50
    codeflash_output = rolling_mean(s, window)
    result = codeflash_output  # 1.44ms -> 297μs (384% faster)
    # Spot check a value
    idx = 100
    expected = sum(range(idx - 49, idx + 1)) / 50


# ------------------- Mutation Testing Guards -------------------


def test_rolling_mean_mutation_guard():
    # If the function skips the window sum or uses wrong indices, this test will fail
    s = pd.Series([1, 2, 3, 4, 100])
    # Window 3, expected: [nan, nan, 2.0, 3.0, 35.666...]
    expected = [math.nan, math.nan, 2.0, 3.0, 35.666666666666664]
    codeflash_output = rolling_mean(s, 3)
    result = codeflash_output  # 2.71μs -> 8.75μs (69.1% slower)
    for r, e in zip(result, expected):
        if math.isnan(e):
            pass
        else:
            pass


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