⚡️ Speed up function solve_discrete_lyapunov by 10%

[codeflash-ai]

📄 10% (0.10x) speedup for solve_discrete_lyapunov in quantecon/_matrix_eqn.py

⏱️ Runtime : 4.13 milliseconds → 3.75 milliseconds (best of 39 runs)

📝 Explanation and details

The optimization achieves a 9% speedup by reordering NumPy operations in the convergence check within the iterative doubling algorithm loop.

Key Change

Line 63 was transformed from:

diff = np.max(np.abs(gamma1 - gamma0))

to:

diff = np.abs(gamma1 - gamma0).max()

Why This is Faster

In NumPy, np.max() is a function call that:

1. Accepts the array as an argument
2. Performs additional argument parsing and validation
3. Then finds the maximum value

In contrast, .max() is a method call directly on the ndarray object that:

1. Operates directly on the already-validated array
2. Bypasses the function call overhead
3. Uses optimized C code paths with less Python interpreter involvement

The line profiler results confirm this: the optimized version takes 2.74 seconds (24.2% of total time) versus 6.04 seconds (40% of total time) in the original—a 54% reduction in time spent on this single line. This line is executed 446 times per run in the hot loop, so the per-call overhead compounds significantly.

Performance Context

Based on the function_references, this function is called from DLE.compute_sequence() to solve for asset pricing matrices (self.Q = solve_discrete_lyapunov(...)), making it part of economic simulation workflows. Since compute_sequence may be called repeatedly with different initial states or time series lengths, and the Lyapunov solver iterates hundreds of times internally (446 iterations in profiled runs), this optimization meaningfully reduces computational cost in typical workloads.

Test Case Impact

The annotated tests show consistent improvements across all test cases (1-14% faster), with the most significant gains in:

• test_basic_diagonal_solution: 14.1% faster—benefits most as it runs the full doubling iteration
• test_exceeded_max_iterations_raises: 11.1% faster—even short iteration counts benefit from reduced per-iteration overhead
• test_both_methods_return_similar_for_random_matrix: 10.5% faster for doubling method—demonstrates real-world use case improvement

The optimization is universally beneficial with no downsides: it maintains identical numerical results, doesn't change the API, and works well for both small matrices (2×2) and larger matrices (8×8) tested.

✅ Correctness verification report:

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

⚙️ Click to see Existing Unit Tests

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_lyapunov.py::test_dlyap_scalar	75.3μs	65.9μs	14.4%✅
test_lyapunov.py::test_dlyap_simple_ones	32.3μs	29.5μs	9.80%✅
test_matrix_eqn.py::test_solve_discrete_lyapunov_B	30.0μs	27.4μs	9.49%✅
test_matrix_eqn.py::test_solve_discrete_lyapunov_complex	121μs	105μs	14.4%✅
test_matrix_eqn.py::test_solve_discrete_lyapunov_zero	30.9μs	27.6μs	12.2%✅

🌀 Click to see Generated Regression Tests

import numba
import numpy as np  # used to build test matrices and compare results
# imports
import pytest  # used for our unit tests
from quantecon._matrix_eqn import solve_discrete_lyapunov

# unit tests

def _make_stable_matrix(n, seed=0, scale=0.5):
    """
    Helper to create a random stable (spectral radius < 1) matrix.
    Uses a fixed seed for determinism.
    """
    rng = np.random.RandomState(seed)
    M = rng.randn(n, n).astype(np.float64)  # random Gaussian matrix
    # scale so that spectral radius < scale (<1)
    eigvals = np.linalg.eigvals(M)
    max_abs = np.max(np.abs(eigvals))
    if max_abs == 0:
        return M * (scale)
    return M * (scale / max_abs)

def _max_abs_diff(X, Y):
    """Return maximum absolute elementwise difference as a Python float."""
    # convert to ndarray and compute difference
    X = np.asarray(X)
    Y = np.asarray(Y)
    return float(np.max(np.abs(X - Y)))

def test_basic_diagonal_solution():
    # Basic sanity check with diagonal A where analytic solution is known.
    # A = 0.5 * I, B = I, so X = sum_{k>=0} 0.5^{2k} I = 1/(1 - 0.25) I = 4/3 I
    A = np.array([[0.5, 0.0], [0.0, 0.5]], dtype=np.float64)
    B = np.eye(2, dtype=np.float64)
    expected = (4.0 / 3.0) * np.eye(2, dtype=np.float64)
    codeflash_output = solve_discrete_lyapunov(A, B, method="doubling"); X = codeflash_output # 84.3μs -> 73.8μs (14.1% faster)
    # Use a strict tolerance for small matrices
    diff = _max_abs_diff(X, expected)

def test_invalid_non_square_A_raises():
    # A is not square -> should raise ValueError with expected message
    A = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])  # 2x3
    B = np.eye(2)
    with pytest.raises(ValueError) as exc:
        solve_discrete_lyapunov(A, B, method="doubling") # 12.5μs -> 12.3μs (1.80% faster)

def test_invalid_non_square_B_raises():
    # B is not square -> should raise ValueError
    A = np.eye(2)
    B = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])  # 2x3
    with pytest.raises(ValueError) as exc:
        solve_discrete_lyapunov(A, B, method="doubling") # 21.2μs -> 21.4μs (0.982% slower)

def test_mismatched_shapes_raises():
    # A and B are square but have different shapes -> should raise ValueError
    A = np.eye(2)
    B = np.eye(3)
    with pytest.raises(ValueError) as exc:
        solve_discrete_lyapunov(A, B, method="doubling") # 18.9μs -> 18.8μs (0.563% faster)

def test_exceeded_max_iterations_raises():
    # Force an exceeded iterations error by setting max_it very small.
    # A simple stable matrix will not converge to the 1e-15 threshold in 1 iteration.
    A = 0.5 * np.eye(2)
    B = np.eye(2)
    with pytest.raises(ValueError) as exc:
        # Setting max_it=1 will cause the internal loop to exceed allowed iterations
        solve_discrete_lyapunov(A, B, max_it=1, method="doubling") # 33.9μs -> 30.5μs (11.1% faster)

def test_invalid_method_string_raises():
    # An invalid method string should raise ValueError with helpful message
    A = np.eye(2)
    B = np.eye(2)
    with pytest.raises(ValueError) as exc:
        solve_discrete_lyapunov(A, B, method="not-a-method") # 1.86μs -> 1.73μs (8.12% faster)

def test_both_methods_return_similar_for_random_matrix():
    # Ensure the "doubling" option and "bartels-stewart" give similar outputs
    n = 8
    A = _make_stable_matrix(n, seed=2021, scale=0.4)
    B = np.eye(n, dtype=np.float64)
    codeflash_output = solve_discrete_lyapunov(A, B, method="doubling"); X_doubling = codeflash_output # 85.9μs -> 77.8μs (10.5% faster)
    codeflash_output = solve_discrete_lyapunov(A, B, method="bartels-stewart"); X_bartels = codeflash_output # 304μs -> 288μs (5.58% faster)
    diff = _max_abs_diff(X_doubling, X_bartels)
# 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
import pytest
from quantecon._matrix_eqn import solve_discrete_lyapunov

class TestSolveDiscreteLyapunovBasic:
    """Basic test cases for solve_discrete_lyapunov function."""

    

To edit these changes git checkout codeflash/optimize-solve_discrete_lyapunov-mkpas4sv and push.