⚡️ Speed up function solve_discrete_riccati_system by 30%

[codeflash-ai]

📄 30% (0.30x) speedup for solve_discrete_riccati_system in quantecon/_matrix_eqn.py

⏱️ Runtime : 4.19 seconds → 3.22 seconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves a 30% speedup through strategic precomputation and memory optimization in the computationally intensive nested loops.

Key Optimizations:

1. Precompute matrix transposes: As_T, Bs_T, and Ns_T are computed once upfront instead of repeatedly calling .T in the hot loops. The profiler shows the original code spent significant time on transpose operations within the inner loops.

2. Replace array slicing with .fill(): Changed sum1[:, :] = 0. to sum1.fill(0.), which is more efficient for zeroing arrays and avoids creating temporary objects.

3. Hoist loop-invariant calculations: Variables like A_T, B, B_T, R, Q, N, N_T are extracted once per outer loop iteration rather than being repeatedly accessed from arrays.

4. Precompute reusable matrix products: In the inner loop, intermediate results like A_T_Pj, B_T_Pj, A_T_Pj_B, etc. are computed once and reused multiple times, eliminating redundant matrix multiplications.

5. Memory allocation optimization: Use np.empty_like(Ps) instead of np.copy(Ps) for Ps1, and preallocate sum1/sum2 arrays outside the main loop.

Performance Impact by Test Case:

• Large Markov chains (50+ states, small matrices): 30-33% faster - benefits most from precomputation since the inner j-loop dominates
• High-dimensional systems (10x10+ matrices): 10-20% faster - gains from reduced matrix operations
• Small systems (scalar/2x2): 20-40% slower - optimization overhead outweighs benefits for tiny problems

The optimizations are most effective for systems with many Markov states where the nested loops execute frequently, making the precomputation overhead worthwhile.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 24 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  # used for our unit tests
# function to test
from numpy.linalg import solve
from quantecon._matrix_eqn import solve_discrete_riccati_system

# =========================
# Unit Tests Start Here
# =========================

# ----------- BASIC TEST CASES ------------

def test_single_state_scalar_case():
    # Single Markov state, scalar system
    Π = np.array([[1.0]])  # Only one state
    As = np.array([[[1.0]]])  # 1x1 matrix
    Bs = np.array([[[1.0]]])  # 1x1 matrix
    Qs = np.array([[[1.0]]])  # 1x1 matrix
    Rs = np.array([[[1.0]]])  # 1x1 matrix
    Ns = np.array([[[0.0]]])  # 1x1 matrix
    beta = 0.9
    # Known solution for scalar Riccati: P = R + beta * A^2 * P - (beta * A * P * B)^2 / (Q + beta * B^2 * P)
    # For this simple case, should converge to a positive scalar
    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 393μs -> 610μs (35.6% slower)

def test_two_state_2x2_case():
    # Two Markov states, 2x2 system
    Π = np.array([[0.8, 0.2],
                  [0.1, 0.9]])
    As = np.array([np.eye(2), np.eye(2)])
    Bs = np.array([np.eye(2), np.eye(2)])
    Qs = np.array([np.eye(2), np.eye(2)])
    Rs = np.array([np.eye(2), np.eye(2)])
    Ns = np.array([np.zeros((2,2)), np.zeros((2,2))])
    beta = 0.95
    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 1.18ms -> 1.51ms (21.7% slower)
    # Check symmetry and positive definiteness
    for i in range(2):
        pass

def test_identity_transition_matrix():
    # Markov chain never switches, should behave as independent Riccati for each state
    Π = np.eye(3)
    As = np.array([np.eye(2)*2, np.eye(2)*3, np.eye(2)*4])
    Bs = np.array([np.eye(2), np.eye(2), np.eye(2)])
    Qs = np.array([np.eye(2), np.eye(2), np.eye(2)])
    Rs = np.array([np.eye(2), np.eye(2), np.eye(2)])
    Ns = np.array([np.zeros((2,2)), np.zeros((2,2)), np.zeros((2,2))])
    beta = 0.9
    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 2.62ms -> 2.93ms (10.5% slower)
    # Each Ps[i] should be symmetric and positive definite
    for i in range(3):
        pass

# ----------- EDGE TEST CASES ------------

def test_zero_matrices():
    # All matrices are zero except Rs, should converge to Rs
    Π = np.array([[1.0]])
    As = np.array([[[0.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 # 93.7μs -> 140μs (33.3% slower)

def test_high_discount_factor():
    # beta very close to 1
    Π = 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.99999
    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 345μs -> 529μs (34.7% slower)

def test_low_discount_factor():
    # beta very close to 0
    Π = np.array([[1.0]])
    As = np.array([[[1.0]]])
    Bs = np.array([[[1.0]]])
    Qs = np.array([[[1.0]]])
    Rs = np.array([[[2.0]]])
    Ns = np.array([[[0.0]]])
    beta = 1e-6
    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 108μs -> 168μs (35.4% slower)


def test_convergence_failure():
    # Force convergence failure by setting max_iter very low
    Π = 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-15, max_iter=1) # 124μs -> 178μs (30.1% slower)

def test_negative_definite_Q_R():
    # Negative definite Q/R, should still run but result may not be positive definite
    Π = np.array([[1.0]])
    As = np.array([[[1.0]]])
    Bs = np.array([[[1.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 # 324μs -> 505μs (35.8% slower)

def test_non_square_matrices():
    # Bs is not square, n=2, k=1
    Π = np.array([[1.0]])
    As = np.array([[[1.0, 0.0], [0.0, 1.0]]])
    Bs = np.array([[[1.0], [0.0]]])
    Qs = np.array([[[1.0]]])
    Rs = np.array([[[1.0, 0.0], [0.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 # 5.08ms -> 8.17ms (37.7% slower)

# ----------- LARGE SCALE TEST CASES ------------

def test_large_markov_chain():
    # Large m, small n
    m = 50
    n = 2
    Π = np.eye(m) * 0.8 + 0.2 / m  # mostly self-transition, small prob to switch
    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 # 648ms -> 496ms (30.6% faster)
    # Check symmetry for a few random states
    for i in [0, m//2, m-1]:
        pass

def test_large_state_dimension():
    # Large n, small m
    m = 2
    n = 20
    Π = np.array([[0.9, 0.1], [0.2, 0.8]])
    As = np.array([np.eye(n)*1.01, np.eye(n)*0.99])
    Bs = np.array([np.eye(n), np.eye(n)])
    Qs = np.array([np.eye(n), 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 # 1.86ms -> 2.32ms (20.1% slower)
    for i in range(m):
        pass

def test_large_nonsquare_B():
    # Large n, k < n
    m = 2
    n = 10
    k = 5
    Π = np.array([[0.7, 0.3], [0.4, 0.6]])
    As = np.array([np.eye(n), np.eye(n)])
    Bs = np.array([np.eye(n, k), np.eye(n, k)])
    Qs = np.array([np.eye(k), np.eye(k)])
    Rs = np.array([np.eye(n), np.eye(n)])
    Ns = np.array([np.zeros((k,n)), np.zeros((k,n))])
    beta = 0.95
    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 40.4ms -> 50.7ms (20.3% slower)
    for i in range(m):
        pass

def test_large_random_system():
    # Random system, moderate size
    np.random.seed(42)
    m = 5
    n = 5
    k = 5
    Π = np.abs(np.random.rand(m, m))
    Π /= Π.sum(axis=1, keepdims=True)  # Normalize rows
    As = np.random.randn(m, n, n)
    Bs = np.random.randn(m, n, k)
    Qs = np.array([np.eye(k) for _ in range(m)])
    Rs = np.array([np.eye(n) for _ in range(m)])
    Ns = np.zeros((m, k, n))
    beta = 0.9
    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, None, Qs, Rs, Ns, beta); Ps = codeflash_output # 7.25ms -> 7.09ms (2.28% faster)
    # Symmetry check
    for i in range(m):
        pass

def test_performance_large_system():
    # Large system for performance (under 1000 elements)
    m = 10
    n = 10
    k = 10
    Π = np.eye(m) * 0.95 + 0.05 / m
    As = np.array([np.eye(n) for _ in range(m)])
    Bs = np.array([np.eye(n, k) 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 # 30.6ms -> 27.6ms (10.8% 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  # used for our unit tests
# function to test
from numpy.linalg import solve
from quantecon._matrix_eqn import solve_discrete_riccati_system

# ===========================
# Unit Tests for the function
# ===========================

# ----------- BASIC TEST CASES ------------

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

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta); Ps = codeflash_output # 341μs -> 553μs (38.3% slower)

def test_basic_two_states_simple():
    # Test with 2 Markov states, simple diagonal matrices
    m, n, k = 2, 2, 2
    Π = np.array([[0.7, 0.3], [0.4, 0.6]])
    As = np.array([np.eye(n), np.eye(n)])
    Bs = np.array([np.eye(n), np.eye(n)])
    Cs = None
    Qs = np.array([np.eye(k), np.eye(k)])
    Rs = np.array([np.eye(n), np.eye(n)])
    Ns = np.array([np.zeros((k, n)), np.zeros((k, n))])
    beta = 0.95

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta); Ps = codeflash_output # 1.17ms -> 1.50ms (22.2% slower)

def test_basic_nontrivial_A_B():
    # Test with non-identity A and B matrices
    m, n, k = 1, 2, 2
    Π = np.array([[1.0]])
    As = np.array([np.array([[1.1, 0.], [0., 0.9]])])
    Bs = np.array([np.array([[0.5, 0.], [0., 0.5]])])
    Cs = None
    Qs = np.array([np.eye(k)])
    Rs = np.array([np.eye(n)])
    Ns = np.array([np.zeros((k, n))])
    beta = 0.9

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta); Ps = codeflash_output # 661μs -> 1.04ms (36.7% slower)

# ----------- EDGE TEST CASES ------------

def test_edge_zero_beta():
    # Test with beta = 0 (no discounting, 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)])
    Cs = None
    Qs = np.array([np.eye(k)])
    Rs = np.array([np.eye(n) * 2])
    Ns = np.array([np.zeros((k, n))])
    beta = 0.0

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta); Ps = codeflash_output # 85.2μs -> 126μs (32.7% slower)

def test_edge_zero_transition_prob():
    # Test with a Markov chain that never transitions (Π is identity)
    m, n, k = 3, 2, 2
    Π = np.eye(m)
    As = np.array([np.eye(n) for _ in range(m)])
    Bs = np.array([np.eye(n) for _ in range(m)])
    Cs = None
    Qs = np.array([np.eye(k) for _ in range(m)])
    Rs = np.array([np.eye(n) * (i+1) for i in range(m)])
    Ns = np.array([np.zeros((k, n)) for _ in range(m)])
    beta = 0.8

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta); Ps = codeflash_output # 2.45ms -> 2.69ms (9.04% slower)
    # Each state should be independent, so Ps[i] should be symmetric and positive definite
    for i in range(m):
        pass

def test_edge_non_symmetric_Q_R():
    # Test with non-symmetric Q and R matrices (should still converge to symmetric Ps)
    m, n, k = 1, 2, 2
    Π = np.array([[1.0]])
    As = np.array([np.eye(n)])
    Bs = np.array([np.eye(n)])
    Cs = None
    Qs = np.array([np.array([[1., 2.], [0., 1.]])])
    Rs = np.array([np.array([[2., 1.], [1., 2.]])])
    Ns = np.array([np.zeros((k, n))])
    beta = 0.9

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta); Ps = codeflash_output # 320μs -> 524μs (39.0% slower)


def test_edge_max_iter_exceeded():
    # Test that function raises if max_iter is exceeded (by making tolerance impossible)
    m, n, k = 1, 2, 2
    Π = np.array([[1.0]])
    As = np.array([np.eye(n)])
    Bs = np.array([np.eye(n)])
    Cs = None
    Qs = np.array([np.eye(k)])
    Rs = np.array([np.eye(n)])
    Ns = np.array([np.zeros((k, n))])
    beta = 1.0

    # Set tolerance to zero so it never converges
    with pytest.raises(ValueError):
        solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta, tolerance=0, max_iter=5) # 213μs -> 327μs (34.8% slower)


def test_large_scale_many_states():
    # Test with 50 Markov states, 3x3 matrices
    m, n, k = 50, 3, 3
    np.random.seed(123)
    Π = np.abs(np.random.rand(m, m))
    Π = Π / Π.sum(axis=1, keepdims=True)  # Normalize rows
    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)])
    Cs = None
    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, Cs, Qs, Rs, Ns, beta); Ps = codeflash_output # 664ms -> 512ms (29.7% faster)
    # Check symmetry for a few random states
    for idx in [0, 10, 25, 49]:
        pass

def test_large_scale_high_dimensional():
    # Test with 10 Markov states, 10x10 matrices
    m, n, k = 10, 10, 10
    np.random.seed(456)
    Π = np.abs(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)])
    Cs = None
    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.99

    codeflash_output = solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta); Ps = codeflash_output # 30.8ms -> 27.8ms (10.5% faster)
    # Check symmetry for a few random states
    for idx in [0, 5, 9]:
        pass

def test_large_scale_randomized():
    # Test with 100 states, 2x2 matrices, random Q and R
    m, n, k = 100, 2, 2
    np.random.seed(789)
    Π = np.abs(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)])
    Cs = None
    Qs = np.array([np.eye(k) + 0.01*np.random.randn(k, k) for _ in range(m)])
    Rs = np.array([np.eye(n) + 0.01*np.random.randn(n, 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, Cs, Qs, Rs, Ns, beta); Ps = codeflash_output # 2.75s -> 2.08s (32.7% faster)
    # Check symmetry for a few random states
    for idx in [0, 50, 99]:
        pass
# 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-mgfzy5nb and push.