PERF: Compute dense Q @ v with a single gemv in bellman_operator

[oyamad]

This PR is the Python side of the DDP performance pass (the Julia side is QuantEcon/QuantEcon.jl#386, merged). It contains no behavioral changes — all solutions are bit-identical.

Note: rebased on main after the merge of #855; the diff is now the PERF commit plus the review-fix commit (per-call view instead of a cached one).

Change

For the product formulation, Q @ v with the 3-dimensional Q is a batched matmul: n separate (m, n) @ (n,) products, with per-slice dispatch overhead. Since Q is C-contiguous in the standard case, the same computation is a single gemv over a 2-dimensional (n*m, n) view. bellman_operator now takes that 2-D view per call when Q is C-contiguous (a view is free, and taking it per call means a rebound ddp.Q is picked up — no cached state); a non-C-contiguous Q (where the reshape would silently copy) falls back to the previous path, with regression tests covering both paths and Q rebinding. The np.arange(num_states) used for indexing by state in RQ_sigma and s_wise_max is also cached instead of allocated per call.

This also removes a historical asymmetry between the two dense representations: the state-action-pair form has always computed Q @ v as a single gemv, which is why it used to outperform the product form by ~1.5-2x on identical data (visible e.g. in the 2015 benchmark notebook oyamad/mdp/ddp_performance.ipynb); with this change the two dense forms are performance-equivalent.

Measured effect (identical random models across runs; values and iteration counts unchanged):

	before	after	speedup
bellman_operator, n=500, m=100	2.79 ms	1.03 ms	2.7x
solve('vi'), n=500, m=100 (206 iter)	601 ms	405 ms	1.5x
bellman_operator, n=100, m=50	66 us	22 us	3.0x
solve('vi'), n=100, m=50	14.0 ms	7.6 ms	1.8x

PFI/MPFI improve slightly (14.1 → 11.8 ms and ~unchanged at n=500); the state-action-pair formulation is unaffected.

Considered and rejected, for the record

• A fused Numba max+argmax kernel to replace the two-pass argmax + fancy-index gather in the dense s_wise_max: measured slower than NumPy (105 us vs 60 us at n=3000, m=50) — NumPy's SIMD-vectorized argmax beats the branchy scalar loop. Kept the current implementation.
• SuperLU tuning for the sparse evaluate_policy (permc_spec variants): at most 1.4x on a random-pattern test case — not worth the complexity.

A finding worth separate discussion (not included here since it changes solver semantics): for the sparse formulation, evaluate_policy's direct spsolve suffers badly from LU fill-in on scattered sparsity patterns, while Krylov methods exploit the structure of A = I - beta Q_sigma — strictly diagonally dominant, condition number bounded by (1+beta)/(1-beta) — extremely well. On a random-pattern case with n=3000 and 5 nonzeros per row: spsolve 522 ms vs bicgstab(rtol=1e-12) 0.8 ms, accurate to ~1e-12. Since an iterative solve makes policy evaluation tolerance-based rather than direct (which touches policy iteration's exactness property), I'd propose it as an opt-in (e.g. a solver option) in a separate issue rather than a silent default change. Happy to open that issue if there is interest.

🤖 Generated with Claude Code (Claude Fable 5)

[coveralls]

coverage: 93.635% (+0.03%) from 93.607% — ddp-perf-py into main

[Copilot]

Pull request overview

This PR updates DiscreteDP internals in quantecon.markov to (a) reduce overhead in dense Bellman updates by using a cached 2D view for a single GEMV, and (b) tighten correctness around several input-validation and iteration edge cases (notably for SA-pair indices and max_iter).

Changes:

• Optimize dense product-form bellman_operator by caching Q as a (n*m, n) view when C-contiguous (single GEMV) and caching np.arange(n) for state indexing.
• Strengthen SA-pair construction validation (range checks, duplicate (s, a) detection, trailing-states-with-no-actions handling).
• Add/expand regression tests covering the new validations and the non-C-contiguous Q fallback path.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
quantecon/markov/ddp.py	Adds cached dense Q 2D view for faster Bellman updates; adds/adjusts validation and iteration-edge-case behavior.
quantecon/markov/utilities.py	Fixes _generate_a_indptr to avoid out-of-bounds when trailing states have no SA pairs.
quantecon/markov/tests/test_ddp.py	Adds tests for trailing-state feasibility, duplicate SA pairs, out-of-range indices, max_iter validation, and non-C-contiguous Q fallback.

[mmcky]

Validated locally on a clean worktree of the PR head — this is a clean, correct, no-behavior-change optimization. Approving. ✅

What I checked

• Full DDP test suite passes (quantecon/markov/tests/test_ddp.py): 19 passed, including the two new regression tests (non-C-contiguous Q fallback, and Q rebinding).
• Numerically bit-identical. The gemv over the 2-D (n*m, n) view reproduces the batched Q @ v element-for-element (allclose), and solve('vi') values/policies match the fallback path exactly. Since Q is C-contiguous the reshape is a true view (row i*m+j = Q[i,j,:]), so the equivalence is exact, not approximate.
• Speedup reproduced. On bellman_operator at n=500, m=100 I measured ~2.2x here (batched matmul 4.5 ms → gemv 2.1 ms); consistent with the 2.7x reported (BLAS/machine-dependent).
• Guards are sound. if not self._sa_pair and self.Q.flags.c_contiguous short-circuits before touching .flags for the sparse/sa-pair form, and a non-C-contiguous Q (where reshape would silently copy) correctly falls back to the original path. The contiguity check is a performance guard, not a correctness one.
• Rebinding fix is real. Taking the 2-D view per call from self.Q means a rebound ddp.Q is picked up (covered by the new test); the earlier staleness class is gone.
• _s_arange caching is consistent. Defined only in the non-sa-pair branch and referenced only from non-sa-pair paths (s_wise_max, RQ_sigma), where n == num_states, so it's identical to the previous per-call np.arange(num_states).

Nice work — the two dense representations are now performance-equivalent as intended, and the separate Krylov evaluate_policy idea is rightly left for its own issue.

One tiny housekeeping note (non-blocking): the PR description still mentions a cached self._Q_2d, but the final diff uses a per-call view — worth tidying that line so the description matches the merged code.

[mmcky]

thanks @oyamad this looks great. Please merge when you're ready.

[oyamad]

@mmcky Thanks! (PR description edited (by Claude).)