⚡️ Speed up function trim_mean by 20%

[codeflash-ai]

📄 20% (0.20x) speedup for trim_mean in framework/py/flwr/serverapp/strategy/fedtrimmedavg.py

⏱️ Runtime : 1.50 milliseconds → 1.25 milliseconds (best of 130 runs)

📝 Explanation and details

The optimized code introduces two key performance improvements:

1. Early return optimization for no-trim cases: Added a if lowercut == 0: check that directly computes np.mean(array, axis=axis) when there's no trimming needed. This avoids the expensive np.partition operation entirely when cut_fraction is 0 or rounds down to 0.

2. Simplified array slicing: Replaced the complex slice list creation and tuple conversion (slice_list = [slice(None)] * atmp.ndim; slice_list[axis] = slice(lowercut, uppercut); atmp[tuple(slice_list)]) with direct slicing (atmp[lowercut:uppercut]). This eliminates the overhead of creating intermediate objects and tuple conversion.

Why this leads to speedup:

• The early return path completely bypasses the costly np.partition operation (which took ~30% of original runtime) for no-trim cases
• Direct slicing is more efficient than constructing slice objects and tuples, reducing memory allocations and function call overhead
• From the line profiler, 19 out of 59 test calls hit the fast path, avoiding partition entirely

Test case performance:
This optimization is particularly effective for test cases with cut_fraction=0 (like test_trim_mean_basic_1d_no_trim, test_trim_mean_large_array_no_trim) and very small fractions that round to 0 (like test_trim_mean_large_array_cut_fraction_near_zero). The remaining cases still benefit from the simplified slicing operation.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import numpy as np
# imports
import pytest  # used for our unit tests
from serverapp.strategy.fedtrimmedavg import trim_mean

# unit tests

# --- Basic Test Cases ---

def test_trim_mean_basic_1d_no_trim():
    # No trimming, should return the mean
    arr = np.array([1, 2, 3, 4, 5])
    cut_fraction = 0.0
    expected = np.mean(arr)
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

def test_trim_mean_basic_1d_trim_20_percent():
    # Trim 20% (1 element from each end for 5 elements)
    arr = np.array([1, 2, 3, 4, 100])
    cut_fraction = 0.2
    # After trimming: [2, 3, 4]
    expected = np.mean([2, 3, 4])
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

def test_trim_mean_basic_2d_no_trim():
    # 2D array, no trimming, mean along axis=0
    arr = np.array([[1, 2], [3, 4], [5, 6]])
    cut_fraction = 0.0
    expected = np.mean(arr, axis=0)
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

def test_trim_mean_basic_2d_trim_33_percent():
    # 2D array, trim 33% (1 row from each end for 3 rows)
    arr = np.array([[1, 10], [2, 20], [100, 200]])
    cut_fraction = 1/3
    # After trimming: only the middle row remains
    expected = arr[1]
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

def test_trim_mean_basic_1d_trim_half():
    # Trim 50% (2 elements from each end for 6 elements)
    arr = np.array([1, 2, 3, 100, 101, 102])
    cut_fraction = 1/3
    # After trimming: [3, 100]
    expected = np.mean([3, 100])
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

# --- Edge Test Cases ---




def test_trim_mean_single_element():
    # Single element array, no trimming
    arr = np.array([42])
    cut_fraction = 0.0
    expected = 42
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output


def test_trim_mean_non_integer_cut_fraction():
    # Non-integer cut_fraction
    arr = np.array([1, 2, 3, 4, 5])
    cut_fraction = 0.25
    # lowercut = int(0.25*5) = 1, uppercut = 4
    # After trimming: [2,3,4]
    expected = np.mean([2,3,4])
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output


def test_trim_mean_cut_fraction_zero():
    # cut_fraction = 0, should return mean
    arr = np.array([1, 2, 3, 4, 5])
    cut_fraction = 0.0
    expected = np.mean(arr)
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

def test_trim_mean_all_elements_same():
    # All elements are the same, trimming should not affect mean
    arr = np.array([7,7,7,7,7])
    cut_fraction = 0.2
    expected = 7
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

def test_trim_mean_2d_edge_case():
    # 2D array, trimming all but one row
    arr = np.array([[1,2,3], [4,5,6], [7,8,9]])
    cut_fraction = 1/3
    expected = arr[1]
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

# --- Large Scale Test Cases ---

def test_trim_mean_large_array_no_trim():
    # Large array, no trimming
    arr = np.arange(1000)
    cut_fraction = 0.0
    expected = np.mean(arr)
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

def test_trim_mean_large_array_with_trim():
    # Large array, trim 10% from each end
    arr = np.arange(1000)
    cut_fraction = 0.1
    lowercut = int(0.1 * 1000)
    uppercut = 1000 - lowercut
    trimmed = arr[lowercut:uppercut]
    expected = np.mean(trimmed)
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

def test_trim_mean_large_array_with_outliers():
    # Large array with outliers at both ends
    arr = np.concatenate(([10000], np.arange(1,999), [20000]))
    cut_fraction = 1/1000  # trim 1 element from each end
    trimmed = arr[1:-1]
    expected = np.mean(trimmed)
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

def test_trim_mean_large_2d_array():
    # Large 2D array, trim 10% from each end along axis=0
    arr = np.tile(np.arange(1000).reshape(-1,1), (1,5))  # shape (1000,5)
    cut_fraction = 0.1
    lowercut = int(0.1 * 1000)
    uppercut = 1000 - lowercut
    trimmed = arr[lowercut:uppercut]
    expected = np.mean(trimmed, axis=0)
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

def test_trim_mean_large_array_cut_fraction_near_zero():
    # Large array, cut_fraction just above zero
    arr = np.arange(1000)
    cut_fraction = 1e-6  # lowercut = 0
    expected = np.mean(arr)
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

def test_trim_mean_large_array_cut_fraction_near_half():
    # Large array, cut_fraction just below 0.5
    arr = np.arange(1000)
    cut_fraction = 0.499
    lowercut = int(0.499 * 1000)
    uppercut = 1000 - lowercut
    trimmed = arr[lowercut:uppercut]
    expected = np.mean(trimmed)
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

# --- Additional Edge Cases ---

def test_trim_mean_cut_fraction_rounding():
    # Test rounding of lowercut/uppercut
    arr = np.array([1,2,3,4,5,6,7])
    cut_fraction = 0.3  # int(0.3*7) = 2, uppercut = 5
    trimmed = arr[2:5]
    expected = np.mean(trimmed)
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output

def test_trim_mean_2d_large_trim():
    # 2D array, cut_fraction removes all but one row
    arr = np.array([[10,20],[30,40],[50,60],[70,80],[90,100]])
    cut_fraction = 0.4  # int(0.4*5)=2, uppercut=3
    trimmed = arr[2:3]
    expected = trimmed[0]
    codeflash_output = trim_mean(arr, cut_fraction); result = codeflash_output
# 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
import pytest  # used for our unit tests
from serverapp.strategy.fedtrimmedavg import trim_mean

# unit tests

# --- Basic Test Cases ---

def test_trim_mean_basic_1d_float():
    # Simple 1D array, cut_fraction=0, should return mean of all elements
    arr = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
    codeflash_output = trim_mean(arr, 0); result = codeflash_output
    expected = np.mean(arr)

def test_trim_mean_basic_1d_int():
    # 1D array of ints, cut_fraction=0.2, trims 1 from each end
    arr = np.array([1, 2, 3, 4, 100])
    codeflash_output = trim_mean(arr, 0.2); result = codeflash_output
    # After trimming 1 from each end: [2, 3, 4]
    expected = np.mean([2, 3, 4])

def test_trim_mean_basic_2d():
    # 2D array, cut_fraction=0.25, trims 1 from each end along axis=0
    arr = np.array([[1, 10],
                    [2, 20],
                    [3, 30],
                    [4, 40]])
    codeflash_output = trim_mean(arr, 0.25); result = codeflash_output
    # After trimming 1 from each end: rows [2, 20], [3, 30]
    expected = np.mean(arr[1:3], axis=0)

def test_trim_mean_basic_fraction_zero():
    # cut_fraction=0, no trimming
    arr = np.arange(10)
    codeflash_output = trim_mean(arr, 0); result = codeflash_output
    expected = np.mean(arr)



def test_trim_mean_single_element():
    # Single element, cut_fraction=0, should return the element itself
    arr = np.array([42])
    codeflash_output = trim_mean(arr, 0); result = codeflash_output
    expected = 42


def test_trim_mean_cut_fraction_negative():
    # Negative cut_fraction should not trim, but should not error
    arr = np.array([1, 2, 3, 4])
    codeflash_output = trim_mean(arr, 0); result = codeflash_output
    expected = np.mean(arr)


def test_trim_mean_non_integer_trim():
    # Non-integer trim, e.g. cut_fraction=0.33, nobs=6, trims int(0.33*6)=1 from each end
    arr = np.arange(6)
    codeflash_output = trim_mean(arr, 0.33); result = codeflash_output
    # After trimming 1 from each end: [1,2,3,4]
    expected = np.mean([1,2,3,4])

def test_trim_mean_all_equal():
    # All elements equal, any cut_fraction should return the same value
    arr = np.full(10, 7.5)
    for frac in [0, 0.1, 0.25, 0.49]:
        codeflash_output = trim_mean(arr, frac); result = codeflash_output

def test_trim_mean_cut_fraction_just_under_half():
    # cut_fraction just under 0.5, should leave one element
    arr = np.array([1, 2, 3, 4, 5])
    # 0.49*5 = 2.45, int() = 2, lowercut=2, uppercut=3
    codeflash_output = trim_mean(arr, 0.49); result = codeflash_output
    expected = np.mean([3])

def test_trim_mean_2d_large_cut_fraction():
    # 2D array, cut_fraction that trims almost all rows
    arr = np.array([[1, 2],
                    [3, 4],
                    [5, 6],
                    [7, 8],
                    [9, 10]])
    codeflash_output = trim_mean(arr, 0.4); result = codeflash_output
    # lowercut=2, uppercut=3, only row 3 remains
    expected = arr[2]

# --- Large Scale Test Cases ---

def test_trim_mean_large_array_no_trim():
    # Large array, no trimming
    arr = np.arange(1000)
    codeflash_output = trim_mean(arr, 0); result = codeflash_output
    expected = np.mean(arr)

def test_trim_mean_large_array_trim_10_percent():
    # Large array, trim 10% from each end
    arr = np.arange(1000)
    codeflash_output = trim_mean(arr, 0.1); result = codeflash_output
    lowercut = int(0.1 * 1000)
    uppercut = 1000 - lowercut
    expected = np.mean(arr[lowercut:uppercut])

def test_trim_mean_large_2d_array():
    # Large 2D array, trim 20% from each end along axis=0
    arr = np.tile(np.arange(1000).reshape(-1,1), (1,3))  # shape (1000,3)
    codeflash_output = trim_mean(arr, 0.2); result = codeflash_output
    lowercut = int(0.2 * 1000)
    uppercut = 1000 - lowercut
    expected = np.mean(arr[lowercut:uppercut], axis=0)

def test_trim_mean_large_array_all_equal():
    # Large array, all elements equal, any cut_fraction should return the same value
    arr = np.full(1000, 123.456)
    for frac in [0, 0.1, 0.3, 0.49]:
        codeflash_output = trim_mean(arr, frac); result = codeflash_output

def test_trim_mean_large_array_cut_fraction_near_half():
    # Large array, cut_fraction just under 0.5, should leave one element
    arr = np.arange(999)
    frac = 0.499
    lowercut = int(frac * 999)
    uppercut = 999 - lowercut
    codeflash_output = trim_mean(arr, frac); result = codeflash_output
    expected = np.mean(arr[lowercut:uppercut])

# --- Additional Edge Cases ---

def test_trim_mean_dtype_preservation():
    # Result should be float even if input is int
    arr = np.array([1, 2, 3, 4, 5])
    codeflash_output = trim_mean(arr, 0.2); result = codeflash_output

def test_trim_mean_fraction_zero_with_2d():
    # 2D array, cut_fraction=0, should return mean along axis=0
    arr = np.array([[1, 2],
                    [3, 4],
                    [5, 6]])
    codeflash_output = trim_mean(arr, 0); result = codeflash_output
    expected = np.mean(arr, axis=0)

def test_trim_mean_fraction_rounding_behavior():
    # Test rounding behavior for cut_fraction that does not divide nobs evenly
    arr = np.arange(7)
    frac = 0.3  # 0.3*7=2.1, int()=2, so trims 2 from each end
    codeflash_output = trim_mean(arr, frac); result = codeflash_output
    expected = np.mean(arr[2:5])

def test_trim_mean_large_2d_all_equal():
    # Large 2D array, all elements equal, any cut_fraction should return the same value
    arr = np.full((1000, 10), 9.99)
    for frac in [0, 0.1, 0.3, 0.49]:
        codeflash_output = trim_mean(arr, frac); result = codeflash_output

def test_trim_mean_returns_scalar_for_1d():
    # 1D input should return a scalar
    arr = np.array([1, 2, 3, 4, 5])
    codeflash_output = trim_mean(arr, 0.2); result = codeflash_output

def test_trim_mean_returns_1d_for_2d():
    # 2D input should return 1D array
    arr = np.arange(12).reshape(6,2)
    codeflash_output = trim_mean(arr, 0.2); result = codeflash_output


def test_trim_mean_cut_fraction_just_enough_rows():
    # cut_fraction such that exactly one row remains
    arr = np.arange(5).reshape(5,1)
    frac = 0.4  # 0.4*5=2, lowercut=2, uppercut=3, only row 2 remains
    codeflash_output = trim_mean(arr, frac); result = codeflash_output
    expected = arr[2]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from serverapp.strategy.fedtrimmedavg import trim_mean

To edit these changes git checkout codeflash/optimize-trim_mean-mhcw4ywy and push.