Speed up function shorrocks_index by 10%

[misrasaurabh1]

📄 10% (0.10x) speedup for shorrocks_index in quantecon/_inequality.py

Hello, I found a few more optimizations for this project at https://github.com/codeflash-ai/QuantEcon.py/pulls?q=is%3Apr+is%3Aopen+review%3Aapproved, would love to contribute them here.

⏱️ Runtime : 474 microseconds → 432 microseconds (best of 242 runs)

📝 Explanation and details

The optimization replaces np.diag(A).sum() with np.trace(A) to compute the sum of diagonal elements. This single change provides a 9% speedup because:

What changed: The diagonal sum calculation was changed from a two-step process (np.diag() then .sum()) to a single optimized NumPy function (np.trace()).

Why it's faster: np.trace() is specifically designed to compute the sum of diagonal elements directly, avoiding the intermediate array creation that np.diag() requires. The line profiler shows the diagonal sum computation time decreased from 680,989 ns to 600,218 ns (12% faster for that line).

Performance characteristics: The optimization is most effective for:

• Small to medium matrices (2x2 to 500x500) where it shows consistent 8-16% improvements
• Dense matrices where diagonal access patterns matter
• Identity matrices and structured matrices benefit significantly

However, for very large sparse matrices or random matrices with complex access patterns, the optimization may show smaller gains or even slight regressions (as seen in the large random test cases), likely due to cache behavior differences. Overall, this is a clean micro-optimization that improves performance across most typical use cases without changing any behavior or dependencies.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_inequality.py::test_shorrocks_index	17.1μs	15.7μs	8.98%✅

🌀 Generated Regression Tests and Runtime

import numpy as np
# imports
import pytest  # used for our unit tests
from quantecon._inequality import shorrocks_index

# unit tests

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

def test_identity_matrix_2x2():
    # 2x2 identity: no mobility, index should be 0.0
    A = np.eye(2)
    codeflash_output = shorrocks_index(A) # 12.5μs -> 10.8μs (15.3% faster)

def test_identity_matrix_3x3():
    # 3x3 identity: no mobility, index should be 0.0
    A = np.eye(3)
    codeflash_output = shorrocks_index(A) # 10.7μs -> 9.90μs (8.46% faster)

def test_full_mobility_2x2():
    # All off-diagonal, no one stays in place: index should be 1.0
    A = np.array([[0,1],[1,0]])
    codeflash_output = shorrocks_index(A) # 10.6μs -> 9.65μs (9.66% faster)

def test_full_mobility_3x3():
    # No one stays in place: index should be 1.0
    A = np.array([[0,0.5,0.5],[0.5,0,0.5],[0.5,0.5,0]])
    codeflash_output = shorrocks_index(A) # 10.3μs -> 8.85μs (16.3% faster)

def test_half_mobility_2x2():
    # Half stay, half move: diag sum = 1, m=2, index = (2-1)/(2-1) = 1.0
    A = np.array([[0.5,0.5],[0.5,0.5]])
    codeflash_output = shorrocks_index(A) # 9.88μs -> 8.71μs (13.3% faster)

def test_diagonal_sum_between_0_and_m():
    # Diagonal sum between 0 and m: index between 0 and 1
    # For m=3, diag sum = 1.5 => (3-1.5)/(3-1) = 0.75
    A = np.array([[0.5,0.25,0.25],[0.25,0.5,0.25],[0.25,0.25,0.5]])
    codeflash_output = shorrocks_index(A) # 9.61μs -> 8.29μs (15.9% faster)

def test_non_numpy_input_list():
    # Accepts lists as input
    A = [[1,0],[0,1]]
    codeflash_output = shorrocks_index(A) # 12.5μs -> 11.6μs (8.04% faster)

def test_non_numpy_input_tuple():
    # Accepts tuples as input
    A = ((0,1),(1,0))
    codeflash_output = shorrocks_index(A) # 12.0μs -> 10.9μs (10.3% faster)

def test_float_elements():
    # Accepts float elements
    A = np.array([[0.8,0.2],[0.3,0.7]])
    # diag sum = 0.8 + 0.7 = 1.5, m=2, index = (2-1.5)/(2-1) = 0.5
    codeflash_output = shorrocks_index(A) # 9.49μs -> 8.25μs (15.1% faster)

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

def test_non_square_matrix_raises():
    # Should raise ValueError if not square
    A = np.array([[0.5,0.5,0],[0.5,0.5,0]])
    with pytest.raises(ValueError):
        shorrocks_index(A) # 1.30μs -> 1.28μs (1.40% faster)


def test_empty_matrix():
    # Empty matrix: should raise ValueError due to shape unpacking
    A = np.array([[]])
    with pytest.raises(ValueError):
        shorrocks_index(A) # 1.96μs -> 1.72μs (13.7% faster)

def test_negative_entries():
    # Negative entries: function does not check, so computes as usual
    A = np.array([[0.5,0.6],[-0.1,0.5]])
    # diag sum = 0.5+0.5=1.0, m=2, index = 1.0
    codeflash_output = shorrocks_index(A) # 18.6μs -> 16.5μs (12.9% faster)

def test_sum_of_rows_not_one():
    # Not a valid transition matrix, but function does not check
    A = np.array([[0.9,0.9],[0.1,0.1]])
    # diag sum = 0.9+0.1=1.0, m=2, index = 1.0
    codeflash_output = shorrocks_index(A) # 12.0μs -> 10.9μs (9.79% faster)

def test_large_values():
    # Large values: function should not overflow
    A = np.array([[1e10,0],[0,1e10]])
    # diag sum = 2e10, m=2, index = (2-2e10)/(2-1) = -1.9999999998e10
    codeflash_output = shorrocks_index(A) # 10.5μs -> 9.41μs (11.5% faster)

def test_nan_entries():
    # NaN in matrix: result should be nan
    A = np.array([[np.nan,0],[0,1]])
    codeflash_output = shorrocks_index(A); result = codeflash_output # 10.3μs -> 9.12μs (12.4% faster)

def test_inf_entries():
    # Inf in matrix: result should be -inf if on diagonal
    A = np.array([[np.inf,0],[0,1]])
    codeflash_output = shorrocks_index(A); result = codeflash_output # 9.68μs -> 9.03μs (7.15% faster)


def test_object_dtype():
    # Object dtype: should work if convertible to float
    A = np.array([[1,0],[0,1]], dtype=object)
    codeflash_output = shorrocks_index(A) # 18.2μs -> 15.7μs (15.3% faster)

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

def test_large_identity_matrix():
    # Large identity matrix (no mobility, index should be 0)
    n = 500
    A = np.eye(n)
    codeflash_output = shorrocks_index(A) # 17.6μs -> 16.0μs (10.0% faster)

def test_large_full_mobility():
    # Large matrix, all off-diagonal equal, diag=0
    n = 500
    A = np.ones((n,n)) / (n-1)
    np.fill_diagonal(A, 0)
    # diag sum = 0, index = (n-0)/(n-1) = n/(n-1)
    expected = n / (n-1)
    codeflash_output = shorrocks_index(A) # 16.1μs -> 14.6μs (10.5% faster)

def test_large_half_mobility():
    # Large matrix, half on diagonal, half off
    n = 1000
    A = np.ones((n,n)) / n
    # diag sum = n*(1/n) = 1, index = (n-1)/(n-1) = 1.0
    codeflash_output = shorrocks_index(A) # 20.7μs -> 19.1μs (8.33% faster)

def test_large_random_stochastic_matrix():
    # Large random stochastic matrix, rows sum to 1
    np.random.seed(42)
    n = 200
    A = np.random.rand(n,n)
    A = A / A.sum(axis=1, keepdims=True)
    diag_sum = np.diag(A).sum()
    expected = (n - diag_sum) / (n - 1)
    codeflash_output = shorrocks_index(A) # 4.01μs -> 5.70μs (29.7% slower)

def test_large_sparse_matrix():
    # Large sparse matrix, only diagonal and one off-diagonal element per row
    n = 500
    A = np.zeros((n,n))
    for i in range(n):
        A[i,i] = 0.9
        A[i,(i+1)%n] = 0.1
    diag_sum = np.diag(A).sum()
    expected = (n - diag_sum) / (n - 1)
    codeflash_output = shorrocks_index(A) # 4.66μs -> 6.71μs (30.5% slower)

def test_large_matrix_all_zeros():
    # All zeros: diag sum = 0, index = n/(n-1)
    n = 300
    A = np.zeros((n,n))
    expected = n / (n-1)
    codeflash_output = shorrocks_index(A) # 12.3μs -> 11.2μs (9.90% faster)
# 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
from quantecon._inequality import shorrocks_index

# unit tests

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

def test_identity_matrix():
    # Identity matrix: all probability mass stays in place, so index = 0
    A = np.eye(3)
    codeflash_output = shorrocks_index(A) # 9.21μs -> 8.21μs (12.2% faster)

def test_uniform_transition_matrix():
    # Every state is equally likely to go to any other state
    # For 3x3, each entry is 1/3, so diag_sum = 1, m=3
    # s(A) = (3 - 1) / (3 - 1) = 1.0
    A = np.ones((3, 3)) / 3
    codeflash_output = shorrocks_index(A) # 8.44μs -> 7.61μs (11.0% faster)

def test_diagonal_half_matrix():
    # Diagonal is 0.5, off-diagonal is 0.25, rows sum to 1
    # diag_sum = 0.5 * 2 = 1.0, m = 2, s(A) = (2-1)/(2-1) = 1.0
    A = np.array([[0.5, 0.5], [0.5, 0.5]])
    codeflash_output = shorrocks_index(A) # 8.88μs -> 7.94μs (11.9% faster)

def test_partial_mobility():
    # Diagonal is 0.7, off-diagonal is 0.3, 2x2
    # diag_sum = 0.7 + 0.7 = 1.4, m = 2
    # s(A) = (2-1.4)/(2-1) = 0.6
    A = np.array([[0.7, 0.3], [0.3, 0.7]])
    codeflash_output = pytest.approx(shorrocks_index(A), 0.0001) # 9.04μs -> 7.56μs (19.6% faster)

def test_non_numpy_input():
    # Input as nested list, should work
    A = [[0.8, 0.2], [0.1, 0.9]]
    # diag_sum = 0.8+0.9=1.7, m=2, s(A) = (2-1.7)/(2-1) = 0.3
    codeflash_output = pytest.approx(shorrocks_index(A), 0.0001) # 11.3μs -> 10.2μs (11.1% faster)

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


def test_non_square_matrix():
    # Should raise ValueError if not square
    A = np.ones((2, 3))
    with pytest.raises(ValueError):
        shorrocks_index(A) # 1.89μs -> 1.76μs (7.50% faster)

def test_empty_matrix():
    # Should raise error (cannot compute shape)
    A = np.array([[]])
    with pytest.raises(ValueError):
        shorrocks_index(A) # 1.34μs -> 1.30μs (3.07% faster)

def test_negative_entries():
    # Negative probabilities do not make sense, but function does not check this
    # Let's check that the function computes the formula regardless
    A = np.array([[0.5, 0.5], [-0.2, 1.2]])
    # diag_sum = 0.5+1.2=1.7, m=2, s(A) = (2-1.7)/(2-1) = 0.3
    codeflash_output = pytest.approx(shorrocks_index(A), 0.0001) # 17.9μs -> 15.9μs (12.8% faster)

def test_large_values():
    # Large values on diagonal, but still a square matrix
    A = np.array([[1000, 0], [0, 1000]])
    # diag_sum = 2000, m=2, s(A) = (2-2000)/(2-1) = -1998
    codeflash_output = pytest.approx(shorrocks_index(A), 0.0001) # 12.1μs -> 10.8μs (11.3% faster)

def test_non_float_entries():
    # Integer entries, should work as formula is agnostic to type
    A = np.array([[1, 0], [0, 1]])
    codeflash_output = shorrocks_index(A) # 10.4μs -> 9.50μs (9.11% faster)

def test_non_2d_input():
    # 1D input should raise error when trying to unpack shape
    A = np.array([1, 0, 0])
    with pytest.raises(ValueError):
        shorrocks_index(A) # 2.45μs -> 2.45μs (0.041% slower)

def test_nan_entries():
    # Matrix with NaN on diagonal should result in NaN
    A = np.array([[np.nan, 0], [0, 1]])
    codeflash_output = shorrocks_index(A); result = codeflash_output # 12.7μs -> 11.6μs (9.32% faster)

def test_inf_entries():
    # Matrix with inf on diagonal should result in -inf
    A = np.array([[np.inf, 0], [0, 1]])
    codeflash_output = shorrocks_index(A); result = codeflash_output # 10.6μs -> 9.57μs (10.4% faster)

def test_rows_do_not_sum_to_one():
    # Not a stochastic matrix, but function does not check this
    # Should compute formula regardless
    A = np.array([[0.9, 0.1], [0.4, 0.4]])
    # diag_sum = 0.9+0.4=1.3, m=2, s(A) = (2-1.3)/(2-1) = 0.7
    codeflash_output = pytest.approx(shorrocks_index(A), 0.0001) # 9.89μs -> 9.09μs (8.74% faster)

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

def test_large_identity_matrix():
    # Large identity matrix, should return 0.0
    n = 500
    A = np.eye(n)
    codeflash_output = shorrocks_index(A) # 15.8μs -> 13.8μs (15.1% faster)

def test_large_uniform_matrix():
    # Large uniform matrix, should return 1.0
    n = 500
    A = np.ones((n, n)) / n
    codeflash_output = pytest.approx(shorrocks_index(A), 0.0001) # 16.6μs -> 14.7μs (13.3% faster)

def test_large_partial_mobility_matrix():
    # Diagonal 0.6, off-diagonal equal, rows sum to 1
    n = 500
    diag = 0.6
    off_diag = (1.0 - diag) / (n - 1)
    A = np.full((n, n), off_diag)
    np.fill_diagonal(A, diag)
    diag_sum = diag * n
    expected = (n - diag_sum) / (n - 1)
    codeflash_output = pytest.approx(shorrocks_index(A), 0.0001) # 14.8μs -> 12.8μs (15.8% faster)

def test_large_random_stochastic_matrix():
    # Generate a random row-stochastic matrix and test formula
    np.random.seed(42)
    n = 300
    A = np.random.rand(n, n)
    A = A / A.sum(axis=1, keepdims=True)
    diag_sum = np.diag(A).sum()
    expected = (n - diag_sum) / (n - 1)
    codeflash_output = pytest.approx(shorrocks_index(A), 0.0001) # 4.86μs -> 6.64μs (26.8% slower)

def test_large_non_square_matrix_raises():
    # Should raise error if not square
    A = np.ones((100, 200))
    with pytest.raises(ValueError):
        shorrocks_index(A) # 1.55μs -> 1.45μs (6.98% faster)

def test_large_matrix_with_nan():
    # Should return nan if any diagonal entry is nan
    n = 100
    A = np.ones((n, n)) / n
    A[0, 0] = np.nan
    codeflash_output = shorrocks_index(A); result = codeflash_output # 11.7μs -> 10.9μs (7.42% faster)

def test_large_matrix_with_inf():
    # Should return -inf if any diagonal entry is inf
    n = 100
    A = np.ones((n, n)) / n
    A[0, 0] = np.inf
    codeflash_output = shorrocks_index(A); result = codeflash_output # 9.56μs -> 8.78μs (8.85% 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-shorrocks_index-mgg0hw20 and push.

[coveralls]

coverage: 92.579%. remained the same
when pulling 0ee281e on codeflash-ai:codeflash/optimize-shorrocks_index-mgg0hw20
into 29cda3a on QuantEcon:main.

[oyamad]

@misrasaurabh1 Thank you for your contribution!