⚡️ Speed up method DiscreteDP.compute_greedy by 17%

[codeflash-ai]

📄 17% (0.17x) speedup for DiscreteDP.compute_greedy in quantecon/markov/ddp.py

⏱️ Runtime : 513 microseconds → 437 microseconds (best of 135 runs)

📝 Explanation and details

The optimized code achieves a 17% speedup by replacing pure Python loops in the state-wise maximization operations with Numba JIT-compiled functions (@njit(cache=True)).

Key Optimizations:

1. Numba JIT Compilation for State-wise Max Operations: The original code called _s_wise_max and _s_wise_max_argmax directly (which are already Numba-compiled in utilities.py), but added wrapper functions _njit_s_wise_max_1d and _njit_s_wise_max_argmax_1d that are explicitly JIT-compiled. More importantly, for the 2D case (product formulation), it replaced NumPy's vals.max(axis=1) and vals.argmax(axis=1) with a custom _njit_s_wise_max_2d function that uses explicit loops compiled by Numba.

2. Why This Is Faster:

   • Avoiding NumPy overhead: For the 2D case, np.max(axis=1) and np.argmax(axis=1) incur Python interpreter overhead and temporary array allocations. The Numba-compiled loop directly iterates over the data with no intermediate allocations.
   • Cache directive: @njit(cache=True) means the compiled machine code is cached to disk, eliminating compilation overhead on subsequent runs.
   • Tight loops: Numba generates LLVM machine code that's optimized for the CPU, with efficient memory access patterns and no GIL overhead.

Performance Impact by Test Case:

• Product formulation (2D) tests show the most significant gains (27-40% faster), since they benefit directly from the custom _njit_s_wise_max_2d replacing NumPy operations.
• State-action pair formulation (1D) tests show smaller gains (2-4% faster), as they already used Numba-compiled utilities but now have an additional wrapper layer that might have minimal caching benefits.
• Large-scale tests (n=100-200) benefit moderately (8-28% faster), demonstrating that the optimization scales well with problem size.

Workload Considerations:
The optimization is particularly beneficial when:

• bellman_operator or compute_greedy are called repeatedly (e.g., in value iteration, policy iteration loops)
• The state-action space is moderately sized (tests show gains even at n=100-200)
• Product formulation is used (2D arrays), where NumPy's axis operations are replaced with tight Numba loops

The line profiler shows that 95-99% of time is spent in s_wise_max, making it the critical hot path. By optimizing this bottleneck with JIT compilation, the overall runtime improves significantly.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import numba
import numpy as np
# imports
import pytest
import scipy.sparse as sp
from quantecon.markov.ddp import DiscreteDP

# --- Unit Tests ---

# 1. Basic Test Cases

def test_basic_product_formulation():
    # Simple 2-state, 2-action product formulation
    R = np.array([[5, 10], [-1, -np.inf]])
    Q = np.array([
        [[0.5, 0.5], [0.0, 1.0]],
        [[0.0, 1.0], [0.5, 0.5]]
    ])
    beta = 0.95
    ddp = DiscreteDP(R, Q, beta)
    v = np.array([0.0, 0.0])  # Initial value function
    codeflash_output = ddp.compute_greedy(v); sigma = codeflash_output # 26.3μs -> 18.8μs (40.0% faster)

def test_basic_sa_pair_formulation():
    # Simple 2-state, 2-action, 3 feasible pairs
    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.compute_greedy(v); sigma = codeflash_output # 19.2μs -> 18.6μs (3.23% faster)

def test_basic_return_type_and_shape():
    # Check that output is correct shape and dtype
    R = np.array([[1, 2], [3, 4]])
    Q = np.zeros((2, 2, 2))
    beta = 0.9
    ddp = DiscreteDP(R, Q, beta)
    v = np.array([0.0, 0.0])
    codeflash_output = ddp.compute_greedy(v); sigma = codeflash_output # 28.9μs -> 21.1μs (36.8% faster)

def test_basic_custom_sigma_argument():
    # Check that passing a pre-allocated sigma array works
    R = np.array([[1, 2], [3, 4]])
    Q = np.zeros((2, 2, 2))
    beta = 0.9
    ddp = DiscreteDP(R, Q, beta)
    v = np.array([0.0, 0.0])
    sigma = np.full(2, -1, dtype=int)
    codeflash_output = ddp.compute_greedy(v, sigma=sigma); result = codeflash_output # 27.9μs -> 20.7μs (35.0% faster)

# 2. Edge Test Cases

def test_edge_all_actions_infeasible():
    # All actions infeasible for one state (should raise error at init)
    R = np.array([[1, -np.inf], [-np.inf, -np.inf]])
    Q = np.zeros((2, 2, 2))
    beta = 0.9
    with pytest.raises(Exception):
        DiscreteDP(R, Q, beta)

def test_edge_single_state_single_action():
    # Only one state and one action
    R = np.array([[42]])
    Q = np.array([[[1.0]]])
    beta = 0.5
    ddp = DiscreteDP(R, Q, beta)
    v = np.array([0.0])
    codeflash_output = ddp.compute_greedy(v); sigma = codeflash_output # 28.2μs -> 20.5μs (37.2% faster)

def test_edge_zero_discount():
    # beta = 0, so only immediate reward matters
    R = np.array([[1, 2], [3, 4]])
    Q = np.zeros((2, 2, 2))
    beta = 0.0
    ddp = DiscreteDP(R, Q, beta)
    v = np.array([100, 100])  # Shouldn't affect greedy choice
    codeflash_output = ddp.compute_greedy(v); sigma = codeflash_output # 29.6μs -> 22.6μs (31.1% faster)

def test_edge_negative_rewards():
    # All rewards are negative
    R = np.array([[-10, -5], [-1, -20]])
    Q = np.zeros((2, 2, 2))
    beta = 0.9
    ddp = DiscreteDP(R, Q, beta)
    v = np.array([0, 0])
    codeflash_output = ddp.compute_greedy(v); sigma = codeflash_output # 29.6μs -> 22.6μs (30.8% faster)

def test_edge_ties():
    # Tie-breaking: if two actions yield same value, pick lowest index
    R = np.array([[1, 1], [2, 2]])
    Q = np.zeros((2, 2, 2))
    beta = 0.5
    ddp = DiscreteDP(R, Q, beta)
    v = np.array([0, 0])
    codeflash_output = ddp.compute_greedy(v); sigma = codeflash_output # 29.4μs -> 22.8μs (29.4% faster)

def test_edge_sparse_Q():
    # Test with sparse Q matrix in state-action formulation
    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.compute_greedy(v); sigma = codeflash_output # 31.6μs -> 30.3μs (4.23% faster)

def test_edge_invalid_beta():
    # beta outside [0, 1] should raise
    R = np.array([[1, 2], [3, 4]])
    Q = np.zeros((2, 2, 2))
    with pytest.raises(ValueError):
        DiscreteDP(R, Q, 1.1)

def test_edge_invalid_shapes():
    # R and Q shape mismatch
    R = np.array([[1, 2], [3, 4]])
    Q = np.zeros((2, 2, 3))
    beta = 0.9
    with pytest.raises(ValueError):
        DiscreteDP(R, Q, beta)

def test_edge_unsorted_sa_indices():
    # Unsorted s_indices/a_indices should be sorted internally
    s_indices = np.array([1, 0, 0])
    a_indices = np.array([0, 1, 0])
    R = np.array([10, 5, -1])
    Q = np.array([
        [0.0, 1.0],
        [0.5, 0.5],
        [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.compute_greedy(v); sigma = codeflash_output # 18.9μs -> 18.3μs (3.22% faster)

# 3. Large Scale Test Cases

def test_large_product_formulation():
    # Large n, m; product formulation
    n = 100
    m = 5
    R = np.random.rand(n, m)
    Q = np.zeros((n, m, n))
    for s in range(n):
        for a in range(m):
            Q[s, a, :] = np.random.dirichlet(np.ones(n))
    beta = 0.95
    ddp = DiscreteDP(R, Q, beta)
    v = np.random.rand(n)
    codeflash_output = ddp.compute_greedy(v); sigma = codeflash_output # 49.8μs -> 38.9μs (27.9% faster)

def test_large_sa_pair_formulation():
    # Large L; state-action pair formulation
    n = 100
    m = 5
    s_indices = np.repeat(np.arange(n), m)
    a_indices = np.tile(np.arange(m), n)
    L = len(s_indices)
    R = np.random.rand(L)
    Q = np.zeros((L, n))
    for i in range(L):
        Q[i, :] = np.random.dirichlet(np.ones(n))
    beta = 0.95
    ddp = DiscreteDP(R, Q, beta, s_indices, a_indices)
    v = np.random.rand(n)
    codeflash_output = ddp.compute_greedy(v); sigma = codeflash_output # 35.9μs -> 35.2μs (2.01% faster)

def test_large_sparse_Q():
    # Large L, sparse Q
    n = 100
    m = 5
    s_indices = np.repeat(np.arange(n), m)
    a_indices = np.tile(np.arange(m), n)
    L = len(s_indices)
    R = np.random.rand(L)
    # Sparse Q: each row has only 10 nonzero entries
    data = []
    rows = []
    cols = []
    for i in range(L):
        idxs = np.random.choice(n, 10, replace=False)
        vals = np.random.dirichlet(np.ones(10))
        for j, v in zip(idxs, vals):
            rows.append(i)
            cols.append(j)
            data.append(v)
    Q_sparse = sp.csr_matrix((data, (rows, cols)), shape=(L, n))
    beta = 0.95
    ddp = DiscreteDP(R, Q_sparse, beta, s_indices, a_indices)
    v = np.random.rand(n)
    codeflash_output = ddp.compute_greedy(v); sigma = codeflash_output # 42.2μs -> 40.6μs (4.05% faster)

def test_large_performance():
    # Quick performance check: should finish in <1s for n=200, m=5
    import time
    n = 200
    m = 5
    R = np.random.rand(n, m)
    Q = np.zeros((n, m, n))
    for s in range(n):
        for a in range(m):
            Q[s, a, :] = np.random.dirichlet(np.ones(n))
    beta = 0.95
    ddp = DiscreteDP(R, Q, beta)
    v = np.random.rand(n)
    t0 = time.time()
    codeflash_output = ddp.compute_greedy(v); sigma = codeflash_output # 115μs -> 105μs (8.90% faster)
    elapsed = time.time() - t0
# 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.compute_greedy-mjw20ykt and push.