⚡️ Speed up method AddScaledTradeActionImpl._scaling_level_for_date by 326%

[codeflash-ai]

📄 326% (3.26x) speedup for AddScaledTradeActionImpl._scaling_level_for_date in gs_quant/backtests/generic_engine.py

⏱️ Runtime : 2.70 milliseconds → 632 microseconds (best of 17 runs)

📝 Explanation and details

The optimization eliminates expensive pandas Series operations in the hotpath by pre-computing index mappings and caching values arrays during initialization.

Key optimizations:

1. Pre-computed index mapping: Creates a dictionary _scaling_level_signal_index_map that maps dates to integer indices, avoiding pandas __contains__ checks that showed 23.3% of runtime in the original code.

2. Direct array access: Caches _scaling_level_signal_values as a direct reference to the pandas Series values, enabling O(1) integer indexing instead of pandas __getitem__ operations that consumed 65.2% of runtime.

3. Fast dictionary lookup: Replaces d in self._scaling_level_signal (expensive pandas operation) with self._scaling_level_signal_index_map.get(d) (fast dict lookup).

Why this is faster:

• Pandas Series operations involve significant overhead for index lookups and value retrieval
• Dictionary .get() is a highly optimized O(1) operation
• Direct array indexing with integers bypasses pandas' complex indexing machinery
• The upfront cost of building the index mapping is amortized across many calls

Performance characteristics:
The optimization shows dramatic speedups (300-1000%) when using signal dictionaries, especially beneficial for:

• Exact date lookups in sparse signals (500-900% faster)
• Large dictionaries with many queries (500-700% faster)
• Interpolated date queries (400-600% faster)

For constant scaling levels, performance is essentially unchanged (slight 1-12% variance), making this a pure win optimization.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import datetime as dt

# imports
import pytest
from gs_quant.backtests.generic_engine import AddScaledTradeActionImpl

# Function to test: _scaling_level_for_date
# We define a minimal AddScaledTradeActionImpl and dependencies to test the method in isolation.

class DummyAction:
    """A minimal dummy action class to mimic AddScaledTradeAction."""
    def __init__(self, scaling_level):
        self.scaling_level = scaling_level
from gs_quant.backtests.generic_engine import AddScaledTradeActionImpl

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

def test_constant_scaling_level_basic():
    """Test with a constant scaling_level (float)."""
    action = DummyAction(scaling_level=1.5)
    impl = AddScaledTradeActionImpl(action)
    # Should always return 1.5 for any date
    for offset in range(-3, 4):
        date = dt.date(2023, 1, 1) + dt.timedelta(days=offset)
        codeflash_output = impl._scaling_level_for_date(date) # 2.27μs -> 2.50μs (9.22% slower)

def test_signal_scaling_level_exact_dates():
    """Test with a signal dict, querying on exact dates."""
    signal = {
        dt.date(2023, 1, 1): 2.0,
        dt.date(2023, 1, 3): 3.0,
        dt.date(2023, 1, 5): 5.0,
    }
    action = DummyAction(scaling_level=signal)
    impl = AddScaledTradeActionImpl(action)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 1)) # 6.56μs -> 1.21μs (443% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 3)) # 2.90μs -> 457ns (534% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 5)) # 2.25μs -> 380ns (493% faster)

def test_signal_scaling_level_interpolated_dates():
    """Test with a signal dict, querying on dates between signal points (step interpolation)."""
    signal = {
        dt.date(2023, 1, 1): 2.0,
        dt.date(2023, 1, 3): 3.0,
        dt.date(2023, 1, 5): 5.0,
    }
    action = DummyAction(scaling_level=signal)
    impl = AddScaledTradeActionImpl(action)
    # 2023-01-2 should get 2.0 (from 2023-01-1)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 2)) # 5.96μs -> 1.14μs (421% faster)
    # 2023-01-4 should get 3.0 (from 2023-01-3)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 4)) # 2.59μs -> 372ns (595% faster)

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

def test_signal_scaling_level_before_first_date():
    """Query before the earliest signal date: should return None or 0 (by implementation, 0)."""
    signal = {
        dt.date(2023, 1, 10): 10.0,
        dt.date(2023, 1, 20): 20.0,
    }
    action = DummyAction(scaling_level=signal)
    impl = AddScaledTradeActionImpl(action)
    # Query before first date: not in _scaling_level_signal, so should return 0
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 1)) # 4.69μs -> 750ns (525% faster)

def test_signal_scaling_level_after_last_date():
    """Query after the last signal date: should return the last value (step interpolation)."""
    signal = {
        dt.date(2023, 1, 10): 10.0,
        dt.date(2023, 1, 20): 20.0,
    }
    action = DummyAction(scaling_level=signal)
    impl = AddScaledTradeActionImpl(action)
    # After last date: should return last value (20.0)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 25)) # 4.50μs -> 842ns (435% faster)

def test_signal_scaling_level_on_gap():
    """Query on a date not in the original signal but within the range."""
    signal = {
        dt.date(2023, 1, 1): 1.0,
        dt.date(2023, 1, 10): 10.0,
    }
    action = DummyAction(scaling_level=signal)
    impl = AddScaledTradeActionImpl(action)
    # 2023-01-5 should get 1.0 (from 2023-01-1)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 5)) # 6.14μs -> 1.25μs (390% faster)
    # 2023-01-10 should get 10.0
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 10)) # 2.89μs -> 442ns (554% faster)

def test_signal_scaling_level_single_point():
    """Signal dict with only one date."""
    signal = {dt.date(2023, 7, 4): 7.4}
    action = DummyAction(scaling_level=signal)
    impl = AddScaledTradeActionImpl(action)
    # Should return 7.4 for the only date
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 7, 4)) # 11.7μs -> 1.20μs (879% faster)
    # Should return 7.4 for any date after (step interpolation)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 7, 5)) # 3.09μs -> 367ns (743% faster)
    # Should return 0 for any date before
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 7, 3)) # 1.11μs -> 250ns (344% faster)

def test_constant_scaling_level_zero_and_negative():
    """Test with constant scaling_level of 0 and negative values."""
    action_zero = DummyAction(scaling_level=0)
    impl_zero = AddScaledTradeActionImpl(action_zero)
    codeflash_output = impl_zero._scaling_level_for_date(dt.date(2022, 12, 31)) # 626ns -> 584ns (7.19% faster)

    action_neg = DummyAction(scaling_level=-3.5)
    impl_neg = AddScaledTradeActionImpl(action_neg)
    codeflash_output = impl_neg._scaling_level_for_date(dt.date(2022, 12, 31)) # 307ns -> 341ns (9.97% slower)

def test_signal_scaling_level_zero_and_negative():
    """Test signal dict with 0 and negative values."""
    signal = {
        dt.date(2023, 1, 1): 0,
        dt.date(2023, 1, 2): -1.5,
        dt.date(2023, 1, 3): 2.5,
    }
    action = DummyAction(scaling_level=signal)
    impl = AddScaledTradeActionImpl(action)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 1)) # 7.88μs -> 1.21μs (552% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 2)) # 2.96μs -> 405ns (631% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 3)) # 2.13μs -> 378ns (463% faster)

def test_signal_scaling_level_non_consecutive_dates():
    """Signal dict with non-consecutive dates, check interpolation."""
    signal = {
        dt.date(2023, 1, 1): 1.0,
        dt.date(2023, 1, 10): 10.0,
        dt.date(2023, 1, 20): 20.0,
    }
    action = DummyAction(scaling_level=signal)
    impl = AddScaledTradeActionImpl(action)
    # 2023-01-5 should get 1.0, 2023-01-15 should get 10.0
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 5)) # 6.01μs -> 1.18μs (409% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 15)) # 2.93μs -> 378ns (675% faster)

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

def test_large_constant_scaling_level():
    """Test constant scaling_level with many queries."""
    action = DummyAction(scaling_level=7.77)
    impl = AddScaledTradeActionImpl(action)
    base_date = dt.date(2020, 1, 1)
    for i in range(0, 1000, 100):  # Test 10 dates
        date = base_date + dt.timedelta(days=i)
        codeflash_output = impl._scaling_level_for_date(date) # 2.98μs -> 3.00μs (0.600% slower)



def test_large_signal_scaling_level_random_gaps():
    """Test signal dict with large gaps between dates, and queries in between."""
    # Signal every 100 days
    signal = {dt.date(2020, 1, 1) + dt.timedelta(days=i*100): i for i in range(10)}
    action = DummyAction(scaling_level=signal)
    impl = AddScaledTradeActionImpl(action)
    # Query exactly on signal dates
    for i in range(10):
        date = dt.date(2020, 1, 1) + dt.timedelta(days=i*100)
        codeflash_output = impl._scaling_level_for_date(date) # 25.9μs -> 4.40μs (488% faster)
    # Query in between: should get last value
    codeflash_output = impl._scaling_level_for_date(dt.date(2020, 1, 1) + dt.timedelta(days=150)) # 2.00μs -> 290ns (590% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2020, 1, 1) + dt.timedelta(days=250)) # 2.01μs -> 300ns (570% faster)
    # Before first date: 0
    codeflash_output = impl._scaling_level_for_date(dt.date(2019, 12, 31)) # 2.68μs -> 288ns (829% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import datetime as dt

# imports
import pytest  # used for our unit tests
from gs_quant.backtests.generic_engine import AddScaledTradeActionImpl


# Minimal stub for AddScaledTradeAction to use in tests
class AddScaledTradeAction:
    def __init__(self, scaling_level):
        self.scaling_level = scaling_level

# Minimal stub for interpolate_signal to use in tests
def interpolate_signal(signal: dict, method=None):
    # For testing, we just return the dict itself (simulate step interpolation)
    # In real code, this would interpolate, but for our tests, we want deterministic lookups
    return signal
from gs_quant.backtests.generic_engine import AddScaledTradeActionImpl

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

# 1. BASIC TEST CASES

def test_constant_scaling_level_basic():
    """Test with a constant scaling level (float)"""
    action = AddScaledTradeAction(1.5)
    impl = AddScaledTradeActionImpl(action)
    # Should always return the constant
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 6, 1)) # 641ns -> 732ns (12.4% slower)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 1, 1)) # 315ns -> 340ns (7.35% slower)
    codeflash_output = impl._scaling_level_for_date(dt.date(2050, 12, 31)) # 268ns -> 289ns (7.27% slower)

def test_dict_scaling_level_basic():
    """Test with a dictionary scaling level for specific dates"""
    signal = {
        dt.date(2024, 6, 1): 2.0,
        dt.date(2024, 6, 2): 3.0,
        dt.date(2024, 6, 3): 4.0
    }
    action = AddScaledTradeAction(signal)
    impl = AddScaledTradeActionImpl(action)
    # Should return the value for exact dates
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 6, 1)) # 7.70μs -> 1.24μs (520% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 6, 2)) # 2.99μs -> 441ns (578% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 6, 3)) # 2.34μs -> 355ns (559% faster)
    # Should return 0 for dates not in the dict
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 6, 4)) # 2.68μs -> 289ns (826% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 5, 31)) # 924ns -> 269ns (243% faster)

def test_dict_scaling_level_single_date():
    """Test with a dictionary containing a single date"""
    signal = {dt.date(2024, 1, 1): 10.0}
    action = AddScaledTradeAction(signal)
    impl = AddScaledTradeActionImpl(action)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 1, 1)) # 11.8μs -> 1.08μs (991% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 1, 2)) # 3.01μs -> 344ns (774% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2023, 12, 31)) # 1.16μs -> 235ns (392% faster)

# 2. EDGE TEST CASES

def test_scaling_level_zero_and_negative():
    """Test with zero and negative scaling levels"""
    # Constant zero
    action = AddScaledTradeAction(0)
    impl = AddScaledTradeActionImpl(action)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 6, 1)) # 575ns -> 578ns (0.519% slower)

    # Constant negative
    action = AddScaledTradeAction(-1.25)
    impl = AddScaledTradeActionImpl(action)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 6, 1)) # 337ns -> 328ns (2.74% faster)

    # Dict with zero and negative
    signal = {
        dt.date(2024, 6, 1): 0,
        dt.date(2024, 6, 2): -2.5
    }
    action = AddScaledTradeAction(signal)
    impl = AddScaledTradeActionImpl(action)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 6, 1)) # 7.28μs -> 1.05μs (595% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 6, 2)) # 2.92μs -> 474ns (516% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 6, 3)) # 2.62μs -> 319ns (722% faster)


def test_scaling_level_dict_with_non_consecutive_dates():
    """Test with a dictionary with non-consecutive dates"""
    signal = {
        dt.date(2024, 1, 1): 5,
        dt.date(2024, 1, 10): 10,
        dt.date(2024, 2, 1): 20
    }
    action = AddScaledTradeAction(signal)
    impl = AddScaledTradeActionImpl(action)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 1, 1)) # 6.14μs -> 1.33μs (362% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 1, 10)) # 2.90μs -> 458ns (533% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 2, 1)) # 2.37μs -> 351ns (576% faster)
    # Dates not in dict
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 1, 2)) # 2.09μs -> 296ns (605% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 1, 9)) # 2.06μs -> 280ns (634% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 2, 2)) # 2.64μs -> 259ns (920% faster)

def test_scaling_level_dict_with_duplicate_values():
    """Test with a dictionary where multiple dates have the same value"""
    signal = {
        dt.date(2024, 1, 1): 7,
        dt.date(2024, 1, 2): 7,
        dt.date(2024, 1, 3): 7
    }
    action = AddScaledTradeAction(signal)
    impl = AddScaledTradeActionImpl(action)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 1, 1)) # 7.25μs -> 1.13μs (541% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 1, 2)) # 3.04μs -> 391ns (677% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 1, 3)) # 2.24μs -> 320ns (600% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2024, 1, 4)) # 2.59μs -> 364ns (612% faster)

def test_scaling_level_dict_with_far_future_and_past_dates():
    """Test with dates far in the future and past"""
    signal = {
        dt.date(1900, 1, 1): 1,
        dt.date(2100, 12, 31): 2
    }
    action = AddScaledTradeAction(signal)
    impl = AddScaledTradeActionImpl(action)
    codeflash_output = impl._scaling_level_for_date(dt.date(1900, 1, 1)) # 21.2μs -> 4.72μs (349% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2100, 12, 31)) # 3.29μs -> 485ns (578% faster)
    codeflash_output = impl._scaling_level_for_date(dt.date(2000, 1, 1)) # 2.66μs -> 557ns (378% faster)


def test_large_dict_scaling_level():
    """Test with a large dictionary of scaling levels"""
    base_date = dt.date(2024, 1, 1)
    signal = {base_date + dt.timedelta(days=i): float(i) for i in range(1000)}  # 1000 dates
    action = AddScaledTradeAction(signal)
    impl = AddScaledTradeActionImpl(action)
    # Test first, middle, last, and out-of-range dates
    codeflash_output = impl._scaling_level_for_date(base_date) # 9.96μs -> 1.70μs (485% faster)
    codeflash_output = impl._scaling_level_for_date(base_date + dt.timedelta(days=499)) # 3.36μs -> 412ns (715% faster)
    codeflash_output = impl._scaling_level_for_date(base_date + dt.timedelta(days=999)) # 2.66μs -> 305ns (772% faster)
    codeflash_output = impl._scaling_level_for_date(base_date + dt.timedelta(days=1000)) # 3.91μs -> 329ns (1087% faster)
    codeflash_output = impl._scaling_level_for_date(base_date - dt.timedelta(days=1)) # 1.06μs -> 263ns (302% faster)

def test_large_constant_scaling_level():
    """Test with a constant scaling level and many queries"""
    action = AddScaledTradeAction(42.0)
    impl = AddScaledTradeActionImpl(action)
    for i in range(1000):
        # Should always return the constant value
        codeflash_output = impl._scaling_level_for_date(dt.date(2024, 1, 1) + dt.timedelta(days=i)) # 243μs -> 246μs (1.11% slower)

def test_large_dict_with_sparse_values():
    """Test with a large dict but only a few dates have values"""
    base_date = dt.date(2024, 1, 1)
    signal = {base_date + dt.timedelta(days=i*100): float(i) for i in range(10)}  # 10 sparse dates
    action = AddScaledTradeAction(signal)
    impl = AddScaledTradeActionImpl(action)
    for i in range(10):
        d = base_date + dt.timedelta(days=i*100)
        codeflash_output = impl._scaling_level_for_date(d) # 27.1μs -> 4.54μs (496% faster)
    # Test some dates not in dict
    for i in range(99):
        d = base_date + dt.timedelta(days=i)
        if d not in signal:
            codeflash_output = impl._scaling_level_for_date(d)

def test_large_dict_all_zero_values():
    """Test with a large dict where all values are zero"""
    base_date = dt.date(2024, 1, 1)
    signal = {base_date + dt.timedelta(days=i): 0 for i in range(1000)}
    action = AddScaledTradeAction(signal)
    impl = AddScaledTradeActionImpl(action)
    for i in range(1000):
        d = base_date + dt.timedelta(days=i)
        codeflash_output = impl._scaling_level_for_date(d) # 1.98ms -> 304μs (548% faster)
    # Test a date not in dict
    codeflash_output = impl._scaling_level_for_date(base_date + dt.timedelta(days=1001)) # 3.50μs -> 332ns (955% 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-AddScaledTradeActionImpl._scaling_level_for_date-mhavruix and push.