⚡️ Speed up function relative_strength_index by 1,601%

[codeflash-ai]

📄 1,601% (16.01x) speedup for relative_strength_index in gs_quant/timeseries/technicals.py

⏱️ Runtime : 530 milliseconds → 31.1 milliseconds (best of 42 runs)

📝 Explanation and details

The optimizations achieve a 1601% speedup by targeting the most expensive computational bottlenecks in the financial technical analysis functions:

Key Performance Optimizations:

1. Eliminated Pandas overhead in loops - The original smoothed_moving_average used expensive iloc assignments in tight loops (64.6% of runtime). The optimized version uses numpy arrays for computations and only creates the final pandas Series once, reducing pandas indexing overhead dramatically.

2. Vectorized array operations - In relative_strength_index, replaced element-wise pandas operations with numpy's where() for gains/losses calculation and vectorized the RSI computation using boolean masks, eliminating the expensive per-element loop that consumed 29.4% of original runtime.

3. Reduced redundant computations - Cached date offset calculations in smoothed_moving_average for non-integer windows, avoiding repeated filtering operations. Also eliminated redundant len(x) calls and _to_offset() conversions in normalize_window.

4. Memory-efficient data structures - Pre-allocated numpy arrays instead of repeatedly modifying pandas Series, reducing memory allocation overhead and improving cache locality.

Test Case Performance:

• Large series (1000+ elements): 5000-5800% speedup - the optimizations scale exceptionally well with data size
• Small series: 100-300% speedup - still significant gains even for smaller datasets
• Edge cases: Maintained correctness while achieving 100-200% improvements

The optimizations are particularly effective for high-frequency trading scenarios and batch processing of multiple time series, where these functions are called repeatedly on large datasets. The numpy-based approach scales linearly rather than quadratically with data size.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
timeseries/test_technicals.py::test_relative_strength_index	16.1ms	8.23ms	96.0%✅

🌀 Generated Regression Tests and Runtime

from datetime import datetime, timedelta

import numpy as np
import pandas as pd
# imports
import pytest
from gs_quant.timeseries.technicals import relative_strength_index

# -------------------- TESTS --------------------

# Helper to create a date-indexed pd.Series
def make_series(values, start="2023-01-01"):
    idx = pd.date_range(start, periods=len(values), freq="D")
    return pd.Series(values, index=idx)

# 1. BASIC TEST CASES

def test_rsi_typical_uptrend():
    # Prices steadily increasing
    prices = make_series([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15])
    codeflash_output = relative_strength_index(prices, 14); rsi = codeflash_output # 1.54ms -> 698μs (120% faster)

def test_rsi_typical_downtrend():
    # Prices steadily decreasing
    prices = make_series([15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1])
    codeflash_output = relative_strength_index(prices, 14); rsi = codeflash_output # 1.43ms -> 674μs (112% faster)

def test_rsi_typical_flat():
    # Prices are flat
    prices = make_series([10]*20)
    codeflash_output = relative_strength_index(prices, 14); rsi = codeflash_output # 2.04ms -> 720μs (184% faster)

def test_rsi_typical_random():
    # Random price movements
    np.random.seed(42)
    prices = make_series(np.random.normal(100, 1, 100))
    codeflash_output = relative_strength_index(prices, 14); rsi = codeflash_output # 7.57ms -> 718μs (953% faster)

def test_rsi_window_size_1():
    # Window size 1, RSI should be either 100, 0, or NaN
    prices = make_series([1, 2, 1, 2, 1, 2])
    codeflash_output = relative_strength_index(prices, 1); rsi = codeflash_output # 2.02ms -> 719μs (181% faster)

# 2. EDGE TEST CASES

def test_rsi_empty_series():
    # Empty input
    prices = pd.Series(dtype=float)
    codeflash_output = relative_strength_index(prices, 14); rsi = codeflash_output # 1.43ms -> 474μs (202% faster)



def test_rsi_all_gains():
    # All price changes are positive
    prices = make_series([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15])
    codeflash_output = relative_strength_index(prices, 5); rsi = codeflash_output # 2.37ms -> 762μs (211% faster)

def test_rsi_all_losses():
    # All price changes are negative
    prices = make_series([15, 14, 13, 12, 11, 10, 9, 8, 7, 6])
    codeflash_output = relative_strength_index(prices, 5); rsi = codeflash_output # 1.91ms -> 696μs (174% faster)

def test_rsi_zero_window():
    # Window size zero should raise ValueError
    prices = make_series([1, 2, 3, 4])
    with pytest.raises(ValueError):
        relative_strength_index(prices, 0) # 4.69μs -> 4.05μs (15.7% faster)

def test_rsi_negative_window():
    # Negative window should raise ValueError
    prices = make_series([1, 2, 3, 4])
    with pytest.raises(ValueError):
        relative_strength_index(prices, -5) # 4.50μs -> 4.07μs (10.5% faster)



def test_rsi_nan_handling():
    # Series with NaNs
    prices = make_series([1, 2, np.nan, 4, 5, 6])
    codeflash_output = relative_strength_index(prices, 2); rsi = codeflash_output # 1.98ms -> 692μs (186% faster)

def test_rsi_all_zero_changes():
    # All price changes are zero except one
    prices = make_series([10, 10, 10, 10, 10, 20])
    codeflash_output = relative_strength_index(prices, 3); rsi = codeflash_output # 1.85ms -> 715μs (159% faster)

def test_rsi_all_nan():
    # All values are NaN
    prices = make_series([np.nan]*10)
    codeflash_output = relative_strength_index(prices, 3); rsi = codeflash_output # 2.13ms -> 668μs (218% faster)

# 3. LARGE SCALE TEST CASES

def test_rsi_large_series():
    # Large input series
    np.random.seed(0)
    prices = make_series(np.random.normal(100, 2, 1000))
    codeflash_output = relative_strength_index(prices, 14); rsi = codeflash_output # 70.0ms -> 1.18ms (5812% faster)
    # All RSI values are between 0 and 100 or NaN
    valid = rsi.dropna()

def test_rsi_large_constant_series():
    # Large constant series
    prices = make_series([50]*1000)
    codeflash_output = relative_strength_index(prices, 14); rsi = codeflash_output # 62.9ms -> 1.20ms (5141% faster)

def test_rsi_large_upward_trend():
    # Large upward trend
    prices = make_series(np.arange(1000))
    codeflash_output = relative_strength_index(prices, 14); rsi = codeflash_output # 63.1ms -> 1.20ms (5170% faster)

def test_rsi_large_downward_trend():
    # Large downward trend
    prices = make_series(np.arange(1000, 0, -1))
    codeflash_output = relative_strength_index(prices, 14); rsi = codeflash_output # 69.8ms -> 1.21ms (5682% faster)

# 4. ADDITIONAL EDGE CASES


def test_rsi_with_non_monotonic_index():
    # Non-monotonic index should not cause errors, but pandas will warn
    prices = pd.Series([1, 2, 3, 4], index=[datetime(2023,1,4), datetime(2023,1,2), datetime(2023,1,3), datetime(2023,1,1)])
    prices = prices.sort_index()  # Ensure monotonic for correct behavior
    codeflash_output = relative_strength_index(prices, 2); rsi = codeflash_output # 1.76ms -> 709μs (148% faster)

def test_rsi_with_duplicate_index():
    # Duplicate index values
    idx = [datetime(2023,1,1)]*3 + [datetime(2023,1,2)]*2
    prices = pd.Series([1,2,3,4,5], index=idx)
    # Pandas will aggregate by index in rolling, but function should not error
    codeflash_output = relative_strength_index(prices, 2); rsi = codeflash_output # 1.82ms -> 697μs (161% faster)

def test_rsi_with_non_numeric():
    # Non-numeric series should raise
    prices = pd.Series(['a', 'b', 'c'])
    with pytest.raises(Exception):
        relative_strength_index(prices, 2) # 114μs -> 108μs (4.85% faster)
# 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
import pandas as pd
# imports
import pytest
from gs_quant.timeseries.technicals import relative_strength_index

# -------------------- UNIT TESTS --------------------

# Helper: Generate a simple price series
def make_series(values, start='2020-01-01'):
    idx = pd.date_range(start, periods=len(values), freq='D')
    return pd.Series(values, index=idx)

# 1. BASIC TEST CASES

def test_rsi_typical_oscillating():
    # Typical up/down price series
    prices = make_series([1, 2, 1, 2, 3, 2, 3, 4, 3, 2, 1, 2, 3, 4, 5, 6, 7, 8])
    codeflash_output = relative_strength_index(prices, w=5); rsi = codeflash_output # 2.58ms -> 749μs (245% faster)

def test_rsi_all_increasing():
    # All gains, RSI should be 100 after warmup
    prices = make_series([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15])
    codeflash_output = relative_strength_index(prices, w=5); rsi = codeflash_output # 2.26ms -> 711μs (217% faster)
    # After w-1 warmup, all RSI values should be 100
    rsi_valid = rsi.iloc[4:]

def test_rsi_all_decreasing():
    # All losses, RSI should be 0 after warmup
    prices = make_series([15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1])
    codeflash_output = relative_strength_index(prices, w=5); rsi = codeflash_output # 2.36ms -> 709μs (232% faster)
    rsi_valid = rsi.iloc[4:]

def test_rsi_flat_series():
    # No price change, RSI should be 100 (by convention)
    prices = make_series([5]*20)
    codeflash_output = relative_strength_index(prices, w=5); rsi = codeflash_output # 2.58ms -> 715μs (261% faster)
    # All RSI values after warmup should be 100 (since avg_loss == 0)
    rsi_valid = rsi.iloc[4:]

def test_rsi_known_values():
    # Test with a known example (from Wikipedia or a trusted source)
    # Use a small series, hand-calculate expected RSI
    prices = make_series([44, 44.15, 43.61, 44.33, 44.83, 45.10, 45.42, 45.84, 46.08, 45.89])
    codeflash_output = relative_strength_index(prices, w=5); rsi = codeflash_output # 2.00ms -> 666μs (200% faster)
    # Only check that output is between 0 and 100 and not NaN after warmup
    rsi_valid = rsi.iloc[4:]

# 2. EDGE TEST CASES

def test_rsi_empty_series():
    # Empty input should return empty output
    prices = pd.Series(dtype=float)
    codeflash_output = relative_strength_index(prices, w=5); rsi = codeflash_output # 1.47ms -> 472μs (210% faster)




def test_rsi_zero_window_raises():
    # Window of 0 should raise ValueError
    prices = make_series([1, 2, 3, 4])
    with pytest.raises(ValueError):
        relative_strength_index(prices, w=0) # 5.18μs -> 4.89μs (5.91% faster)

def test_rsi_negative_window_raises():
    # Negative window should raise ValueError
    prices = make_series([1, 2, 3, 4])
    with pytest.raises(ValueError):
        relative_strength_index(prices, w=-5) # 4.62μs -> 4.54μs (1.65% faster)



def test_rsi_nan_in_input():
    # Series with NaN: should propagate through diff, output NaN where appropriate
    prices = make_series([1, 2, np.nan, 4, 5, 6])
    codeflash_output = relative_strength_index(prices, w=2); rsi = codeflash_output # 1.92ms -> 717μs (168% faster)

def test_rsi_index_preserved():
    # Output index should match input index[1:]
    prices = make_series([1, 2, 3, 2, 1])
    codeflash_output = relative_strength_index(prices, w=2); rsi = codeflash_output # 1.87ms -> 716μs (162% faster)

# 3. LARGE SCALE TEST CASES

def test_rsi_large_random_series():
    # Large series, random walk
    np.random.seed(42)
    n = 1000
    prices = pd.Series(np.cumsum(np.random.randn(n)) + 100, index=pd.date_range('2022-01-01', periods=n))
    codeflash_output = relative_strength_index(prices, w=14); rsi = codeflash_output # 68.6ms -> 1.18ms (5725% faster)
    # All RSI values (except for first w-1) should be finite and between 0 and 100
    valid = rsi.iloc[13:]

def test_rsi_large_constant_series():
    # Large constant series: all RSI values after warmup should be 100
    n = 1000
    prices = pd.Series([42]*n, index=pd.date_range('2022-01-01', periods=n))
    codeflash_output = relative_strength_index(prices, w=20); rsi = codeflash_output # 62.0ms -> 1.21ms (5047% faster)
    valid = rsi.iloc[19:]

def test_rsi_large_alternating_series():
    # Large series with alternating up/down: RSI should hover around 50
    n = 1000
    vals = np.tile([1, 2], n//2)
    prices = pd.Series(vals, index=pd.date_range('2022-01-01', periods=n))
    codeflash_output = relative_strength_index(prices, w=10); rsi = codeflash_output # 70.2ms -> 1.21ms (5724% faster)
    valid = rsi.iloc[9:]
    mean_rsi = valid.mean()
# 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-relative_strength_index-mhb5smfw and push.