⚡️ Speed up function solve_discrete_riccati_system by 544%

[codeflash-ai]

📄 544% (5.44x) speedup for solve_discrete_riccati_system in quantecon/_matrix_eqn.py

⏱️ Runtime : 138 milliseconds → 21.5 milliseconds (best of 83 runs)

📝 Explanation and details

The optimized code achieves a 543% speedup by extracting the computationally intensive inner loop into a separate function decorated with @njit(cache=True) from Numba. This enables Just-In-Time (JIT) compilation of the core mathematical operations to native machine code.

Key Optimization Applied:

• Numba JIT Compilation: The _iterate_riccati function containing the nested loops and matrix operations is compiled to optimized machine code, eliminating Python's interpreter overhead for the most expensive computations.

Performance Impact Analysis:
The line profiler shows that in the original code, the inner loop operations (lines with matrix multiplications and solve calls) consumed ~94% of total runtime. The optimized version moves all this computation into the JIT-compiled function, reducing execution time from 271ms to 21.5ms.

Hot Path Benefits:
Based on function_references, this function is called from stationary_values() in the LQControl class, which is likely used repeatedly in economic modeling workflows. The 5.4x speedup will significantly benefit:

• Iterative economic simulations where Riccati equations are solved multiple times
• Parameter sweeps and optimization routines that call this function repeatedly
• Large-scale econometric models with many Markov states or high-dimensional state spaces

Test Case Performance:
The optimization shows consistent improvements across all test scenarios:

• Small systems (2x2 matrices): 3-6x faster
• Large Markov chains (20 states): 5.4x faster
• High-dimensional state spaces (20x20 matrices): 2.1x faster
• Complex random systems: 2.5-3.5x faster

The numba compilation overhead is amortized after the first call due to caching, making this especially valuable for repeated usage patterns common in economic modeling applications.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import numpy as np
# imports
import pytest
from quantecon._matrix_eqn import solve_discrete_riccati_system

# unit tests

# ========== BASIC TEST CASES ==========

def test_basic_identity_case():
    """
    Test with all identity matrices and scalar beta=1.
    Should converge to a symmetric matrix.
    """
    m, n, k = 1, 2, 2
    Π = np.array([[1.0]])
    As = np.array([np.eye(n)])
    Bs = np.array([np.eye(n)])
    Qs = np.array([np.eye(n)])
    Rs = np.array([np.eye(n)])
    Ns = np.array([np.zeros((n, n))])
    beta = 1.0

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 242μs -> 36.8μs (559% faster)

def test_basic_two_state_case():
    """
    Test with two Markov states, different A matrices.
    """
    m, n, k = 2, 2, 2
    Π = np.array([[0.8, 0.2],
                  [0.3, 0.7]])
    As = np.array([np.eye(n), 2*np.eye(n)])
    Bs = np.array([np.eye(n), np.eye(n)])
    Qs = np.array([np.eye(n), np.eye(n)])
    Rs = np.array([np.eye(n), 2*np.eye(n)])
    Ns = np.array([np.zeros((n, n)), np.zeros((n, n))])
    beta = 0.95

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 862μs -> 93.6μs (821% faster)

def test_basic_nontrivial_N():
    """
    Test with nonzero N matrices.
    """
    m, n, k = 1, 2, 2
    Π = np.array([[1.0]])
    As = np.array([np.eye(n)])
    Bs = np.array([np.eye(n)])
    Qs = np.array([2*np.eye(n)])
    Rs = np.array([np.eye(n)])
    Ns = np.array([np.ones((n, n))])
    beta = 0.9

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 326μs -> 47.3μs (591% faster)

# ========== EDGE TEST CASES ==========

def test_edge_beta_zero():
    """
    Test with beta=0, should return Rs.
    """
    m, n, k = 1, 2, 2
    Π = np.array([[1.0]])
    As = np.array([np.eye(n)])
    Bs = np.array([np.eye(n)])
    Qs = np.array([np.eye(n)])
    Rs = np.array([2*np.eye(n)])
    Ns = np.array([np.zeros((n, n))])
    beta = 0.0

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 44.9μs -> 10.2μs (341% faster)

def test_edge_zero_matrices():
    """
    Test with zero matrices for Q, R, N.
    Should converge to zero matrix.
    """
    m, n, k = 1, 2, 2
    Π = np.array([[1.0]])
    As = np.array([np.eye(n)])
    Bs = np.array([np.eye(n)])
    Qs = np.array([np.eye(n)])  # Q must be invertible for solve
    Rs = np.array([np.zeros((n, n))])
    Ns = np.array([np.zeros((n, n))])
    beta = 0.5

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 587μs -> 80.7μs (629% faster)

def test_edge_non_converge_raises():
    """
    Test that the function raises ValueError if not converged.
    """
    m, n, k = 1, 2, 2
    Π = np.array([[1.0]])
    As = np.array([np.eye(n)])
    Bs = np.array([np.eye(n)])
    Qs = np.array([np.eye(n)])
    Rs = np.array([np.eye(n)])
    Ns = np.array([np.zeros((n, n))])
    beta = 1.0

    # Use a very small max_iter to force failure
    with pytest.raises(ValueError):
        solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta, tolerance=1e-16, max_iter=1) # 46.6μs -> 11.9μs (291% faster)

def test_edge_invalid_shapes():
    """
    Test that invalid input shapes raise appropriate errors.
    """
    m, n, k = 1, 2, 2
    Π = np.array([[1.0]])
    As = np.array([np.eye(n)])
    Bs = np.array([np.eye(n)])
    Qs = np.array([np.eye(n)])
    Rs = np.array([np.eye(n)])
    Ns = np.array([np.zeros((n, n))])
    beta = 1.0

    # Qs wrong shape
    with pytest.raises(Exception):
        solve_discrete_riccati_system(Π, As, Bs, None, Qs.reshape((2,1,2)), Rs, Ns, beta)

    # Π not square
    with pytest.raises(Exception):
        solve_discrete_riccati_system(np.array([[1,0],[0,1],[0,0]]), As, Bs, None, Qs, Rs, Ns, beta)

def test_large_markov_chain():
    """
    Test with a large Markov chain (m=20, n=5).
    """
    m, n, k = 20, 5, 5
    Π = np.ones((m, m)) / m  # uniform transition
    As = np.array([np.eye(n) for _ in range(m)])
    Bs = np.array([np.eye(n) for _ in range(m)])
    Qs = np.array([np.eye(n) for _ in range(m)])
    Rs = np.array([np.eye(n) for _ in range(m)])
    Ns = np.array([np.zeros((n, n)) for _ in range(m)])
    beta = 0.95

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 85.2ms -> 13.2ms (543% faster)
    # Should be symmetric for all Markov states
    for i in range(m):
        pass

def test_large_state_dimension():
    """
    Test with a large state dimension (n=20, m=2).
    """
    m, n, k = 2, 20, 20
    Π = np.array([[0.9, 0.1],
                  [0.2, 0.8]])
    As = np.array([np.eye(n), np.eye(n)])
    Bs = np.array([np.eye(n), np.eye(n)])
    Qs = np.array([np.eye(n), 2*np.eye(n)])
    Rs = np.array([np.eye(n), np.eye(n)])
    Ns = np.array([np.zeros((n, n)), np.zeros((n, n))])
    beta = 0.9

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 2.03ms -> 945μs (114% faster)
    for i in range(m):
        pass

def test_large_random_system():
    """
    Test with random system matrices of moderate size.
    """
    np.random.seed(42)
    m, n, k = 5, 10, 10
    Π = np.random.rand(m, m)
    Π = Π / Π.sum(axis=1, keepdims=True)
    As = np.array([np.eye(n) + 0.01*np.random.randn(n, n) for _ in range(m)])
    Bs = np.array([np.eye(n) + 0.01*np.random.randn(n, n) for _ in range(m)])
    Qs = np.array([np.eye(n) + 0.01*np.random.randn(n, n) for _ in range(m)])
    Rs = np.array([np.eye(n) + 0.01*np.random.randn(n, n) for _ in range(m)])
    Ns = np.array([0.01*np.random.randn(n, n) for _ in range(m)])
    beta = 0.95

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 6.74ms -> 1.93ms (250% faster)
    for i in range(m):
        pass
# 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
from quantecon._matrix_eqn import solve_discrete_riccati_system

# unit tests

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

def test_basic_scalar_case():
    # Test with m=1, n=1, k=1 (all scalars)
    Π = np.array([[1.0]])
    As = np.array([[[1.0]]])
    Bs = np.array([[[1.0]]])
    Qs = np.array([[[1.0]]])
    Rs = np.array([[[1.0]]])
    Ns = np.array([[[0.0]]])
    beta = 0.9

    # Analytical solution for scalar Riccati
    # P = R + beta*A^2*P - (beta*A*B*P + N)^2 / (Q + beta*B^2*P)
    # We'll check if the result is close to the fixed point
    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 252μs -> 34.7μs (628% faster)
    P = Ps[0,0,0]
    # The fixed point should satisfy the equation (approximately)
    lhs = P
    rhs = 1 + 0.9*1*P - ((0.9*1*1*P + 0)**2) / (1 + 0.9*1*P)

def test_basic_two_state_case():
    # Test with m=2, n=1, k=1 (two Markov states, scalar system)
    Π = np.array([[0.7, 0.3],
                  [0.4, 0.6]])
    As = np.array([[[1.0]], [[0.5]]])
    Bs = np.array([[[1.0]], [[2.0]]])
    Qs = np.array([[[2.0]], [[1.0]]])
    Rs = np.array([[[1.0]], [[3.0]]])
    Ns = np.array([[[0.0]], [[0.0]]])
    beta = 0.8

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 885μs -> 80.1μs (1005% faster)

def test_basic_matrix_case():
    # Test with m=1, n=2, k=2 (single Markov state, 2x2 matrices)
    Π = np.array([[1.0]])
    As = np.array([np.eye(2)])
    Bs = np.array([np.eye(2)])
    Qs = np.array([np.eye(2)])
    Rs = np.array([np.eye(2)])
    Ns = np.array([np.zeros((2,2))])
    beta = 0.95

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 240μs -> 36.8μs (552% faster)

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

def test_edge_zero_discount():
    # beta = 0, solution should be Ps = Rs
    Π = np.array([[1.0]])
    As = np.array([[[1.0]]])
    Bs = np.array([[[1.0]]])
    Qs = np.array([[[2.0]]])
    Rs = np.array([[[5.0]]])
    Ns = np.array([[[0.0]]])
    beta = 0.0

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 45.0μs -> 10.2μs (341% faster)

def test_edge_zero_B():
    # B=0, so control has no effect, solution should be Lyapunov equation
    Π = np.array([[1.0]])
    As = np.array([[[1.0]]])
    Bs = np.array([[[0.0]]])
    Qs = np.array([[[1.0]]])
    Rs = np.array([[[2.0]]])
    Ns = np.array([[[0.0]]])
    beta = 0.9

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 3.81ms -> 431μs (784% faster)

def test_edge_non_convergence():
    # Should raise ValueError if max_iter is too small to converge
    Π = np.array([[1.0]])
    As = np.array([[[1.0]]])
    Bs = np.array([[[1.0]]])
    Qs = np.array([[[1.0]]])
    Rs = np.array([[[1.0]]])
    Ns = np.array([[[0.0]]])
    beta = 0.9

    with pytest.raises(ValueError):
        solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta, tolerance=1e-16, max_iter=1) # 46.3μs -> 11.2μs (313% faster)

def test_edge_zero_N():
    # N=0 should be handled correctly
    Π = np.array([[1.0]])
    As = np.array([[[1.0]]])
    Bs = np.array([[[1.0]]])
    Qs = np.array([[[2.0]]])
    Rs = np.array([[[3.0]]])
    Ns = np.array([[[0.0]]])
    beta = 0.95

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 246μs -> 52.8μs (367% faster)

def test_edge_highly_coupled_markov_chain():
    # Markov chain with strong off-diagonal transitions
    Π = np.array([[0.01, 0.99],
                  [0.99, 0.01]])
    As = np.array([[[1.0]], [[2.0]]])
    Bs = np.array([[[1.0]], [[1.0]]])
    Qs = np.array([[[1.0]], [[1.0]]])
    Rs = np.array([[[1.0]], [[1.0]]])
    Ns = np.array([[[0.0]], [[0.0]]])
    beta = 0.9

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 943μs -> 85.9μs (998% faster)

def test_edge_non_square_A_B():
    # A: 2x2, B: 2x1, Q: 1x1, R: 2x2, N: 1x2
    Π = np.array([[1.0]])
    As = np.array([np.eye(2)])
    Bs = np.array([[[1.0],[0.0]]])
    Qs = np.array([[[2.0]]])
    Rs = np.array([np.eye(2)])
    Ns = np.array([[[0.0, 0.0]]])
    beta = 0.95

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 7.56ms -> 968μs (681% faster)

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

def test_large_markov_chain():
    # m=10, n=2, k=2, random but stable system
    np.random.seed(42)
    m = 10
    n = 2
    k = 2
    Π = np.ones((m, m)) / m  # uniform Markov chain
    As = np.array([np.eye(n)*0.9 for _ in range(m)])
    Bs = np.array([np.eye(n) for _ in range(m)])
    Qs = np.array([np.eye(k) for _ in range(m)])
    Rs = np.array([np.eye(n) for _ in range(m)])
    Ns = np.array([np.zeros((k, n)) for _ in range(m)])
    beta = 0.9

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 18.6ms -> 1.74ms (968% faster)
    for i in range(m):
        pass

def test_large_state_space():
    # m=2, n=10, k=10
    np.random.seed(123)
    m = 2
    n = 10
    k = 10
    Π = np.array([[0.7, 0.3], [0.3, 0.7]])
    As = np.array([np.eye(n)*0.95 for _ in range(m)])
    Bs = np.array([np.eye(n) for _ in range(m)])
    Qs = np.array([np.eye(k) for _ in range(m)])
    Rs = np.array([np.eye(n) for _ in range(m)])
    Ns = np.array([np.zeros((k, n)) for _ in range(m)])
    beta = 0.9

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 1.09ms -> 322μs (237% faster)
    for i in range(m):
        pass

def test_large_coupled_chain_and_state():
    # m=5, n=5, k=5, non-diagonal A and B
    np.random.seed(321)
    m = 5
    n = 5
    k = 5
    Π = np.random.dirichlet(np.ones(m), m)
    As = np.array([np.eye(n)*0.8 + 0.05*np.random.randn(n, n) for _ in range(m)])
    Bs = np.array([np.eye(n)*0.5 + 0.05*np.random.randn(n, n) for _ in range(m)])
    Qs = np.array([np.eye(k) for _ in range(m)])
    Rs = np.array([np.eye(n) for _ in range(m)])
    Ns = np.array([np.zeros((k, n)) for _ in range(m)])
    beta = 0.95

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 8.18ms -> 1.27ms (546% faster)
    for i in range(m):
        pass

# ----------------------------- Miscellaneous -------------------------------

def test_output_is_numpy_array():
    # Output should be a numpy ndarray
    Π = np.array([[1.0]])
    As = np.array([[[1.0]]])
    Bs = np.array([[[1.0]]])
    Qs = np.array([[[1.0]]])
    Rs = np.array([[[1.0]]])
    Ns = np.array([[[0.0]]])
    beta = 0.9

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 250μs -> 34.0μs (639% faster)

def test_invalid_shapes_raise():
    # Should raise ValueError or broadcast error for mismatched shapes
    Π = np.array([[1.0]])
    As = np.array([[[1.0]]])
    Bs = np.array([[[1.0]]])
    Qs = np.array([[[1.0]]])
    Rs = np.array([[[1.0]]])
    Ns = np.array([[[0.0]]])
    beta = 0.9

    # Qs wrong shape
    with pytest.raises(Exception):
        solve_discrete_riccati_system(Π, As, Bs, None, np.array([1.0]), Rs, Ns, beta) # 709ns -> 666ns (6.46% 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-solve_discrete_riccati_system-mj9nm2yu and push.