⚡️ Speed up function describe by 230%

[codeflash-ai]

📄 230% (2.30x) speedup for describe in src/statistics/descriptive.py

⏱️ Runtime : 3.16 milliseconds → 958 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 230% speedup by replacing inefficient pandas operations with vectorized NumPy operations. The key optimizations are:

What was optimized:

1. NaN filtering: Replaced the slow list comprehension [v for v in series if not pd.isna(v)] with vectorized operations: arr = series.to_numpy(), mask = ~pd.isna(arr), and values = arr[mask]
2. Sorting: Changed from Python's sorted(values) to NumPy's np.sort(values)
3. Statistical calculations: Replaced manual calculations with NumPy methods - values.mean() instead of sum(values) / n, and ((values - mean) ** 2).mean() for variance

Why it's faster:

• Vectorization: NumPy operations are implemented in C and operate on entire arrays at once, avoiding Python's interpreter overhead for each element
• Memory efficiency: NumPy arrays have better memory layout and avoid the overhead of Python objects
• Optimized algorithms: NumPy's sorting and mathematical operations use highly optimized implementations

Performance breakdown from profiling:

• Original code spent 78.4% of time on the list comprehension (20.3ms out of 25.9ms total)
• Optimized version reduces this to just 49.9% across all NumPy operations (1.99ms out of 3.99ms total)
• The variance calculation improved from 17.6% to 15.4% of runtime while being more readable

Test case performance:
The optimization particularly benefits larger datasets - the large-scale test cases with 1000+ elements will see the most dramatic improvements due to the vectorized operations scaling much better than the original element-by-element processing.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import math

# function to test
# src/statistics/descriptive.py
import numpy as np
import pandas as pd

# imports
import pytest  # used for our unit tests
from src.statistics.descriptive import describe

# unit tests


# Helper function for comparing floats, including NaN
def assert_float_equal(a, b, tol=1e-8):
    if math.isnan(a) and math.isnan(b):
        return


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


def test_describe_basic_integer_series():
    # Test with a simple integer series
    s = pd.Series([1, 2, 3, 4, 5])
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], 3.0)
    assert_float_equal(result["std"], math.sqrt(2.0))


def test_describe_basic_float_series():
    # Test with a simple float series
    s = pd.Series([1.0, 2.0, 3.0, 4.0, 5.0])
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], 3.0)
    assert_float_equal(result["std"], math.sqrt(2.0))
    assert_float_equal(result["min"], 1.0)
    assert_float_equal(result["max"], 5.0)
    assert_float_equal(result["25%"], 2.0)
    assert_float_equal(result["50%"], 3.0)
    assert_float_equal(result["75%"], 4.0)


def test_describe_basic_negative_numbers():
    # Test with negative numbers
    s = pd.Series([-5, -4, -3, -2, -1])
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], -3.0)
    assert_float_equal(result["std"], math.sqrt(2.0))


def test_describe_basic_single_value():
    # Test with a single value
    s = pd.Series([42])
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], 42)
    assert_float_equal(result["std"], 0.0)


def test_describe_basic_two_values():
    # Test with two values
    s = pd.Series([10, 20])
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], 15.0)
    assert_float_equal(result["std"], 5.0)


def test_describe_basic_repeated_values():
    # Test with repeated values
    s = pd.Series([7, 7, 7, 7])
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], 7.0)
    assert_float_equal(result["std"], 0.0)


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


def test_describe_empty_series():
    # Test with an empty series
    s = pd.Series([], dtype=float)
    codeflash_output = describe(s)
    result = codeflash_output
    # All other values should be NaN
    for key in ["mean", "std", "min", "25%", "50%", "75%", "max"]:
        pass


def test_describe_all_nan():
    # Test with all NaN values
    s = pd.Series([float("nan"), np.nan, None])
    codeflash_output = describe(s)
    result = codeflash_output
    for key in ["mean", "std", "min", "25%", "50%", "75%", "max"]:
        pass


def test_describe_some_nan():
    # Test with some NaN values mixed in
    s = pd.Series([1, np.nan, 2, None, 3])
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], 2.0)
    assert_float_equal(result["std"], math.sqrt(2 / 3))


def test_describe_mixed_types():
    # Test with mixed int and float types
    s = pd.Series([1, 2.5, 3, 4.5])
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], (1 + 2.5 + 3 + 4.5) / 4)
    # std = sqrt(mean((x-mean)^2)), check manually
    mean = (1 + 2.5 + 3 + 4.5) / 4
    variance = sum((x - mean) ** 2 for x in [1, 2.5, 3, 4.5]) / 4
    std = variance**0.5
    assert_float_equal(result["std"], std)
    assert_float_equal(result["min"], 1.0)
    assert_float_equal(result["max"], 4.5)
    sorted_vals = sorted([1, 2.5, 3, 4.5])
    # n=4, idx for 25% = int(0.25*4)=1, 50% = 2, 75% = 3
    assert_float_equal(result["25%"], sorted_vals[1])
    assert_float_equal(result["50%"], sorted_vals[2])
    assert_float_equal(result["75%"], sorted_vals[3])


def test_describe_large_identical_values():
    # Edge case: large number of identical values
    s = pd.Series([999] * 100)
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], 999)
    assert_float_equal(result["std"], 0.0)


def test_describe_unsorted_input():
    # Test with unsorted input
    s = pd.Series([5, 1, 3, 2, 4])
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], 3.0)
    assert_float_equal(result["std"], math.sqrt(2.0))


def test_describe_series_with_inf():
    # Test with inf and -inf values
    s = pd.Series([1, float("inf"), 2, float("-inf"), 3])
    # Remove infs for expected calculation
    finite_vals = [1, 2, 3]
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], sum(finite_vals) / 3)
    mean = sum(finite_vals) / 3
    variance = sum((x - mean) ** 2 for x in finite_vals) / 3
    std = variance**0.5
    assert_float_equal(result["std"], std)
    sorted_vals = sorted(finite_vals)


def test_describe_series_with_zero():
    # Test with zeros included
    s = pd.Series([0, 0, 0, 1, 2])
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], 0.6)
    mean = 0.6
    variance = sum((x - mean) ** 2 for x in [0, 0, 0, 1, 2]) / 5
    std = variance**0.5
    assert_float_equal(result["std"], std)
    sorted_vals = sorted([0, 0, 0, 1, 2])


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


def test_describe_large_series_sequential():
    # Large series of sequential numbers
    s = pd.Series(range(1, 1001))  # 1 to 1000
    codeflash_output = describe(s)
    result = codeflash_output
    # mean = (1+1000)/2 = 500.5
    assert_float_equal(result["mean"], 500.5)
    # std = sqrt(mean((x-mean)^2)), for 1..n: std = sqrt((n^2-1)/12)
    expected_std = math.sqrt((1000**2 - 1) / 12)
    assert_float_equal(result["std"], expected_std)


def test_describe_large_series_reverse():
    # Large series in reverse order
    s = pd.Series(list(range(1000, 0, -1)))  # 1000 to 1
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], 500.5)
    expected_std = math.sqrt((1000**2 - 1) / 12)
    assert_float_equal(result["std"], expected_std)


def test_describe_large_series_with_nan():
    # Large series with some NaN values
    vals = list(range(1, 501)) + [np.nan] * 10 + list(range(501, 1001))
    s = pd.Series(vals)
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], 500.5)
    expected_std = math.sqrt((1000**2 - 1) / 12)
    assert_float_equal(result["std"], expected_std)


def test_describe_large_series_all_identical():
    # Large series with all identical values
    s = pd.Series([123.456] * 1000)
    codeflash_output = describe(s)
    result = codeflash_output
    assert_float_equal(result["mean"], 123.456)
    assert_float_equal(result["std"], 0.0)
    assert_float_equal(result["min"], 123.456)
    assert_float_equal(result["max"], 123.456)
    assert_float_equal(result["25%"], 123.456)
    assert_float_equal(result["50%"], 123.456)
    assert_float_equal(result["75%"], 123.456)


def test_describe_large_series_random():
    # Large series with random values
    import random

    random.seed(42)
    vals = [random.uniform(-1000, 1000) for _ in range(1000)]
    s = pd.Series(vals)
    codeflash_output = describe(s)
    result = codeflash_output
    # mean and std should match manual calculation
    mean = sum(vals) / 1000
    variance = sum((x - mean) ** 2 for x in vals) / 1000
    std = variance**0.5
    assert_float_equal(result["mean"], mean)
    assert_float_equal(result["std"], std)
    sorted_vals = sorted(vals)
    # idx for 25% = int(0.25*1000)=250, 50% = 500, 75% = 750
    assert_float_equal(result["min"], sorted_vals[0])
    assert_float_equal(result["max"], sorted_vals[-1])
    assert_float_equal(result["25%"], sorted_vals[250])
    assert_float_equal(result["50%"], sorted_vals[500])
    assert_float_equal(result["75%"], sorted_vals[750])


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

import math

# function to test
import numpy as np
import pandas as pd

# imports
import pytest  # used for our unit tests
from src.statistics.descriptive import describe

# unit tests

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


def test_describe_basic_integers():
    # Test with a simple integer series
    s = pd.Series([1, 2, 3, 4, 5])
    codeflash_output = describe(s)
    result = codeflash_output


def test_describe_basic_floats():
    # Test with floats
    s = pd.Series([1.0, 2.0, 3.0, 4.0, 5.0])
    codeflash_output = describe(s)
    result = codeflash_output


def test_describe_basic_negative_numbers():
    # Test with negative numbers
    s = pd.Series([-5, -3, -1, 0, 1, 3, 5])
    codeflash_output = describe(s)
    result = codeflash_output


def test_describe_basic_duplicates():
    # Test with duplicate values
    s = pd.Series([2, 2, 2, 2, 2])
    codeflash_output = describe(s)
    result = codeflash_output


def test_describe_basic_single_value():
    # Test with a single value
    s = pd.Series([42])
    codeflash_output = describe(s)
    result = codeflash_output


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


def test_describe_empty_series():
    # Test with empty series
    s = pd.Series([], dtype=float)
    codeflash_output = describe(s)
    result = codeflash_output
    for key in ["mean", "std", "min", "25%", "50%", "75%", "max"]:
        pass


def test_describe_all_nan():
    # Test with all NaN values
    s = pd.Series([np.nan, np.nan, np.nan])
    codeflash_output = describe(s)
    result = codeflash_output
    for key in ["mean", "std", "min", "25%", "50%", "75%", "max"]:
        pass


def test_describe_some_nan():
    # Test with some NaN values mixed in
    s = pd.Series([1, np.nan, 2, np.nan, 3])
    codeflash_output = describe(s)
    result = codeflash_output


def test_describe_integers_and_floats():
    # Test with mixed int and float types
    s = pd.Series([1, 2.5, 3, 4.5, 5])
    codeflash_output = describe(s)
    result = codeflash_output
    # Manual std calculation
    mean = 3.2
    expected_std = math.sqrt(sum((x - mean) ** 2 for x in [1, 2.5, 3, 4.5, 5]) / 5)


def test_describe_with_inf_values():
    # Test with inf and -inf values
    s = pd.Series([1, 2, np.inf, -np.inf, 3])
    codeflash_output = describe(s)
    result = codeflash_output


def test_describe_non_numeric():
    # Test with non-numeric types (should ignore them as pd.Series of object type)
    s = pd.Series(["a", "b", "c"])
    with pytest.raises(TypeError):
        describe(s)  # Should fail because math operations on strings are invalid


def test_describe_bool_series():
    # Test with boolean values
    s = pd.Series([True, False, True, False])
    # Booleans are treated as 1 and 0
    codeflash_output = describe(s)
    result = codeflash_output


def test_describe_single_nan():
    # Test with a single NaN value
    s = pd.Series([np.nan])
    codeflash_output = describe(s)
    result = codeflash_output
    for key in ["mean", "std", "min", "25%", "50%", "75%", "max"]:
        pass


def test_describe_single_inf():
    # Test with a single inf value
    s = pd.Series([np.inf])
    codeflash_output = describe(s)
    result = codeflash_output


def test_describe_single_minus_inf():
    # Test with a single -inf value
    s = pd.Series([-np.inf])
    codeflash_output = describe(s)
    result = codeflash_output


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


def test_describe_large_series_uniform():
    # Test with a large uniform series
    s = pd.Series([10] * 1000)
    codeflash_output = describe(s)
    result = codeflash_output


def test_describe_large_series_range():
    # Test with a large range of numbers
    s = pd.Series(range(1000))  # 0 to 999
    codeflash_output = describe(s)
    result = codeflash_output
    # std for 0..999 is sqrt(sum((x-499.5)^2)/1000)
    expected_std = math.sqrt(sum((x - 499.5) ** 2 for x in range(1000)) / 1000)


def test_describe_large_series_with_nans():
    # Test with a large series and some NaNs
    s = pd.Series([float(i) if i % 10 != 0 else np.nan for i in range(1000)])
    # There are 100 NaNs (i=0,10,20,...,990)
    codeflash_output = describe(s)
    result = codeflash_output
    # Check percentiles
    sorted_values = sorted([float(i) for i in range(1000) if i % 10 != 0])
    n = len(sorted_values)

    def percentile(p):
        idx = int(p * n / 100)
        if idx >= n:
            idx = n - 1
        return sorted_values[idx]


def test_describe_large_series_reverse():
    # Test with a large series in reverse order
    s = pd.Series(list(reversed(range(1000))))
    codeflash_output = describe(s)
    result = codeflash_output


def test_describe_large_series_random():
    # Test with a large random series
    import random

    random.seed(42)
    data = [random.uniform(-1000, 1000) for _ in range(1000)]
    s = pd.Series(data)
    codeflash_output = describe(s)
    result = codeflash_output
    # Check min/max match sorted values
    sorted_data = sorted(data)

    # Check percentiles
    def percentile(p):
        idx = int(p * 1000 / 100)
        if idx >= 1000:
            idx = 999
        return sorted_data[idx]

    # Mean and std
    expected_mean = sum(data) / 1000
    expected_std = math.sqrt(sum((x - expected_mean) ** 2 for x in data) / 1000)


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