⚡️ Speed up method DiscreteDP.bellman_operator by 58%

[codeflash-ai]

📄 58% (0.58x) speedup for DiscreteDP.bellman_operator in quantecon/markov/ddp.py

⏱️ Runtime : 308 microseconds → 196 microseconds (best of 77 runs)

📝 Explanation and details

The optimization achieves a 57% speedup by replacing Python's dynamic dispatch with specialized Numba-compiled functions for different data structure types. Here's why this works:

Key Optimization: Specialized Numba Functions

The original code used a dynamically-defined closure s_wise_max that dispatched to Numba-compiled helpers. This introduced Python function call overhead and prevented full optimization. The optimized version creates three specialized @njit functions:

1. _bellman_operator_sa_pair - For sparse state-action pair formulation
2. _bellman_operator_sa_pair_dense - For dense state-action pair formulation
3. _bellman_operator_product - For product formulation (2D R, 3D Q)

Why It's Faster

Static Dispatch: The bellman_operator method now uses simple if-else logic to select the appropriate Numba function based on self._sa_pair and self._sparse flags. This eliminates the overhead of calling through the closure.

Inlined Operations: Numba can now optimize the entire computation chain—from vals = R + beta * (Q @ v) to the state-wise maximization—as a single compiled unit, enabling better register allocation and loop fusion.

Reduced Python Interpreter Interaction: The hot loop that previously spent 96.2% of time in s_wise_max (calling through Python) is now entirely within Numba-compiled code.

Test Performance Patterns

• Small problems (2-100 states): Show 65-283% speedups due to elimination of Python call overhead
• Large problems (100 states): One test shows 38% slowdown, likely because Numba's JIT compilation cost amortizes better with the original approach for very large matrix operations where NumPy's optimized BLAS already dominates

Impact Assessment

Since function_references is unavailable, we cannot definitively assess hot path usage. However, the Bellman operator is fundamental to dynamic programming solvers and typically called thousands of times in iterative algorithms (value iteration, policy iteration), making even small per-call speedups significant for cumulative performance.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import numpy as np
# imports
import pytest
from quantecon.markov.ddp import DiscreteDP

# Function to test: DiscreteDP.bellman_operator
# (Assume DiscreteDP and its dependencies are imported and available as per the user's code block.)

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

def test_bellman_operator_dense_basic_2_states_2_actions():
    """
    Basic test: 2 states, 2 actions, dense representation.
    Verifies correct Bellman update and greedy policy.
    """
    R = np.array([[5, 10], [-1, -np.inf]])  # Rewards
    Q = np.array([
        [[0.5, 0.5], [0.0, 1.0]],
        [[0.0, 1.0], [0.5, 0.5]]
    ])  # Transition probabilities
    beta = 0.95
    ddp = DiscreteDP(R, Q, beta)
    v = np.array([0.0, 0.0])
    codeflash_output = ddp.bellman_operator(v); Tv = codeflash_output # 17.7μs -> 8.21μs (115% faster)
    # Compute expected manually
    # For state 0:
    # action 0: 5 + 0.95 * (0.5*0 + 0.5*0) = 5
    # action 1: 10 + 0.95 * (0*0 + 1*0) = 10
    # max = 10
    # For state 1:
    # action 0: -1 + 0.95 * (0*0 + 1*0) = -1
    # action 1: -inf
    # max = -1
    expected = np.array([10.0, -1.0])

def test_bellman_operator_sa_pair_basic_2_states_2_actions():
    """
    Basic test: 2 states, 2 actions, state-action pair representation.
    """
    s_indices = np.array([0, 0, 1])
    a_indices = np.array([0, 1, 0])
    R = np.array([5, 10, -1])
    Q = np.array([
        [0.5, 0.5],
        [0.0, 1.0],
        [0.0, 1.0]
    ])
    beta = 0.95
    ddp = DiscreteDP(R, Q, beta, s_indices, a_indices)
    v = np.array([0.0, 0.0])
    codeflash_output = ddp.bellman_operator(v); Tv = codeflash_output # 16.3μs -> 9.84μs (65.2% faster)
    expected = np.array([10.0, -1.0])

def test_bellman_operator_greedy_policy_output():
    """
    Basic test: Check that sigma returns greedy actions.
    """
    R = np.array([[1, 2], [3, 4]])
    Q = np.array([
        [[1, 0], [0, 1]],
        [[0, 1], [1, 0]]
    ])
    beta = 0.5
    ddp = DiscreteDP(R, Q, beta)
    v = np.array([10, 20])
    sigma = np.empty(ddp.num_states, dtype=int)
    codeflash_output = ddp.bellman_operator(v, sigma=sigma); Tv = codeflash_output # 28.3μs -> 8.11μs (249% faster)

def test_bellman_operator_tv_argument_mutation():
    """
    Basic test: Check that Tv argument is mutated and returned.
    """
    R = np.array([[0, 1], [2, 3]])
    Q = np.zeros((2,2,2))
    beta = 0.0
    ddp = DiscreteDP(R, Q, beta)
    v = np.array([0,0])
    Tv = np.zeros(ddp.num_states)
    codeflash_output = ddp.bellman_operator(v, Tv=Tv); ret = codeflash_output # 21.3μs -> 7.17μs (197% faster)

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

def test_bellman_operator_negative_rewards_and_zero_discount():
    """
    Edge case: negative rewards, zero discount factor.
    """
    R = np.array([[-5, -10], [-1, -2]])
    Q = np.ones((2,2,2)) / 2
    beta = 0.0
    ddp = DiscreteDP(R, Q, beta)
    v = np.array([100, 200])
    codeflash_output = ddp.bellman_operator(v); Tv = codeflash_output # 20.9μs -> 7.64μs (174% faster)
    # Should ignore v, just take max reward per state
    expected = np.array([-5, -1])

def test_bellman_operator_all_negative_inf_rewards():
    """
    Edge case: all rewards -inf (no feasible actions).
    Should raise error at construction.
    """
    R = np.full((3, 2), -np.inf)
    Q = np.ones((3,2,3)) / 3
    beta = 0.9
    # Should raise ValueError in _check_action_feasibility
    with pytest.raises(ValueError):
        DiscreteDP(R, Q, beta)

def test_bellman_operator_single_state_single_action():
    """
    Edge case: single state, single action.
    """
    R = np.array([[42]])
    Q = np.array([[[1.0]]])
    beta = 0.5
    ddp = DiscreteDP(R, Q, beta)
    v = np.array([7])
    codeflash_output = ddp.bellman_operator(v); Tv = codeflash_output # 20.5μs -> 7.03μs (192% faster)

def test_bellman_operator_zero_transition_probabilities():
    """
    Edge case: zero transition probabilities (absorbing state).
    """
    R = np.array([[1, 2], [3, 4]])
    Q = np.zeros((2,2,2))
    beta = 0.9
    ddp = DiscreteDP(R, Q, beta)
    v = np.array([10, 20])
    codeflash_output = ddp.bellman_operator(v); Tv = codeflash_output # 21.8μs -> 7.57μs (188% faster)
    # Should just be rewards
    expected = np.array([2, 4])

def test_bellman_operator_sparse_q_matrix():
    """
    Edge case: Q as scipy.sparse matrix.
    """
    import scipy.sparse as sp
    s_indices = np.array([0, 0, 1])
    a_indices = np.array([0, 1, 0])
    R = np.array([5, 10, -1])
    Q_dense = np.array([
        [0.5, 0.5],
        [0.0, 1.0],
        [0.0, 1.0]
    ])
    Q_sparse = sp.csr_matrix(Q_dense)
    beta = 0.95
    ddp = DiscreteDP(R, Q_sparse, beta, s_indices, a_indices)
    v = np.array([0.0, 0.0])
    codeflash_output = ddp.bellman_operator(v); Tv = codeflash_output # 34.5μs -> 9.03μs (283% faster)
    expected = np.array([10.0, -1.0])

def test_bellman_operator_beta_equals_one_warning():
    """
    Edge case: beta=1 triggers warning and disables infinite horizon methods.
    Bellman operator should still work.
    """
    R = np.array([[1, 2], [3, 4]])
    Q = np.ones((2,2,2)) / 2
    beta = 1.0
    with pytest.warns(UserWarning):
        ddp = DiscreteDP(R, Q, beta)
    v = np.array([0.0, 0.0])
    codeflash_output = ddp.bellman_operator(v); Tv = codeflash_output # 21.2μs -> 7.97μs (166% faster)
    expected = np.array([2, 4])

def test_bellman_operator_non_sorted_sa_indices():
    """
    Edge case: s_indices and a_indices not sorted.
    Should still produce correct results.
    """
    s_indices = np.array([1, 0, 0])
    a_indices = np.array([0, 1, 0])
    R = np.array([-1, 10, 5])
    Q = np.array([
        [0.0, 1.0],
        [0.0, 1.0],
        [0.5, 0.5]
    ])
    beta = 0.95
    ddp = DiscreteDP(R, Q, beta, s_indices, a_indices)
    v = np.array([0.0, 0.0])
    codeflash_output = ddp.bellman_operator(v); Tv = codeflash_output # 16.3μs -> 9.95μs (63.8% faster)
    expected = np.array([10.0, -1.0])

def test_bellman_operator_invalid_beta_value():
    """
    Edge case: beta out of [0,1] should raise error.
    """
    R = np.array([[1, 2], [3, 4]])
    Q = np.ones((2,2,2)) / 2
    beta = 1.5
    with pytest.raises(ValueError):
        DiscreteDP(R, Q, beta)

def test_bellman_operator_invalid_q_dimensions():
    """
    Edge case: Q with invalid shape should raise error.
    """
    R = np.array([[1, 2], [3, 4]])
    Q = np.ones((2,2))  # Should be (2,2,2)
    beta = 0.9
    with pytest.raises(ValueError):
        DiscreteDP(R, Q, beta)

def test_bellman_operator_invalid_r_dimensions():
    """
    Edge case: R with invalid shape should raise error.
    """
    R = np.ones((2,2,2))
    Q = np.ones((2,2,2))
    beta = 0.9
    with pytest.raises(ValueError):
        DiscreteDP(R, Q, beta)

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

def test_bellman_operator_large_dense():
    """
    Large scale: 100 states, 5 actions, dense representation.
    """
    n, m = 100, 5
    np.random.seed(42)
    R = np.random.uniform(-10, 10, (n, m))
    Q = np.random.dirichlet(np.ones(n), (n, m))
    beta = 0.99
    ddp = DiscreteDP(R, Q, beta)
    v = np.random.uniform(-100, 100, n)
    codeflash_output = ddp.bellman_operator(v); Tv = codeflash_output # 42.7μs -> 69.0μs (38.1% slower)

def test_bellman_operator_large_greedy_policy():
    """
    Large scale: Check sigma output for large problem.
    """
    n, m = 100, 3
    np.random.seed(789)
    R = np.random.uniform(-5, 5, (n, m))
    Q = np.random.dirichlet(np.ones(n), (n, m))
    beta = 0.95
    ddp = DiscreteDP(R, Q, beta)
    v = np.random.uniform(-10, 10, n)
    sigma = np.empty(ddp.num_states, dtype=int)
    codeflash_output = ddp.bellman_operator(v, sigma=sigma); Tv = codeflash_output # 46.7μs -> 44.0μs (6.31% 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-DiscreteDP.bellman_operator-mjw1i8pk and push.