QuantEcon / QuantEcon.py / 18963640808

93%

push

github

Optimize shorrocks_index (#810)

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.

Co-authored-by: codeflash-ai[bot] <148906541+codeflash-ai[bot]@users.noreply.github.com>

1 of 1 new or added line in 1 file covered. (100.0%)

7710 of 8328 relevant lines covered (92.58%)

2.78 hits per line