FIX: put the next-state expectation inside the inverse in solve_discrete_riccati_system

[jadhavgaurav]

Fixes #982.

solve_discrete_riccati_system averaged the quadratic correction over next states outside the matrix inverse. The control is chosen before the next state is known, so the expectation belongs inside: with P_bar = sum_j Π[i, j] P_j, the update is the single-state Riccati update evaluated at P_bar. This is the form LQMarkov.stationary_values already uses for F, so F is now evaluated at the correct Ps. The change is the replacement suggested in the issue; the now-unused sum1/sum2 buffers are removed.

Tests

• New TestLQMarkov.test_stationary_P_is_cost_of_own_policy, built from the issue's reproducer. It checks that Ps equals the discounted cost of following the returned Fs (policy evaluation iterated to a fixed point), and that perturbing Fs by +/-0.01 does not lower that cost in any state. It fails on main and passes with the fix.
• Reference values in test_stationary_scalar, test_stationary_mat, test_scalar_sequences_with_seed and test_mat_sequences are regenerated from the fixed code. I checked them against a separate NumPy iteration of the corrected recursion from the issue (max difference 2e-11 scalar case, 9e-10 matrix case, whose P entries reach 8e3). The w and state sequences are unchanged, as expected.
• With one Markov state, LQMarkov and LQ give the same P (difference about 1e-10).

flake8 --select=F401,F405,E231 quantecon is clean and pytest quantecon/tests/test_lqcontrol.py quantecon/tests/test_matrix_eqn.py passes (20 tests). I did not run the full suite.

Assisted-by: Claude Code (Claude Sonnet 5.5)

[coveralls]

coverage: 94.136% (+0.05%) from 94.088% — jadhavgaurav:fix-982-lqmarkov-riccati into QuantEcon:main

[mmcky]

Thanks for this, @jadhavgaurav, and for the clear write-up and regression test.

One quick request that we now make of every pull request. You'll remember an earlier version of it from #964; the format has changed slightly since then. QuantEcon has adopted a Code of AI Use (QEP-5), which welcomes AI-assisted contributions and asks each pull request to say whether AI tools were used. Could you:

1. Add one line to the description above: Assisted-by: <tool> (<model>) if AI tools meaningfully contributed (for example Assisted-by: Claude Code (Claude Fable 5.1)), or Assisted-by: none if they didn't. The description is enough, so there's no need to amend or force-push the commit.
2. Reply here to confirm that you chose the task, read the result, and will answer review on it.

We're asking everyone, so this isn't a flag on your pull request. Thanks!

[jadhavgaurav]

Thanks @mmcky, I have added Assisted-by: Claude Code (Claude Sonnet 5.5) to the description.

On the confirmation: I chose this issue (#982) from a list of candidates, the change and its tests were written by Claude Code under my direction, and I will deal with review feedback on this PR. The evidence for the change is the regression test and the independent NumPy and policy-iteration check described above, so please weigh it on those.

Thanks for #984 too. If you want the extras from it on this PR (docstring, the ndim=3 fix, the comment on forming the left factor directly, the helper and the two tests), tell me and I will port them here.

[Copilot]

🟢 Approval recommended

The corrected recursion matches the issue's derivation and is consistent with the existing F formula, and I verified the regenerated scalar and matrix reference values exactly while a new regression test guards the fix.

0 open findings

What changed in this PR

This PR fixes a correctness bug (#982) in solve_discrete_riccati_system, the solver behind LQMarkov.stationary_values. The previous code averaged the quadratic correction term over next-period Markov states outside the matrix inverse, which is incorrect because the control is chosen before the next state is realized. The expectation must be taken inside the inverse via P_bar = Σ_j Π[i,j] P_j, reducing each state's update to the standard single-state Riccati update evaluated at P_bar. This aligns the value matrices Ps with the actual cost of the returned policy Fs, and makes Fs optimal. I numerically verified the regenerated reference values (both scalar and matrix cases) match the corrected recursion exactly.

Changes:

• Replaced the inner sum1/sum2 loop over next states with a single P_bar-based Riccati update inside the inverse, and removed the now-unused buffers.
• Regenerated reference values in four existing TestLQMarkov tests to reflect the corrected recursion.
• Added test_stationary_P_is_cost_of_own_policy, a regression test verifying Ps equals the policy-evaluation fixed point of Fs and that perturbing Fs never lowers cost.

File	Description
quantecon/​_matrix_eqn.py	Moves the next-state expectation inside the inverse using P_bar; removes unused sum1/sum2 buffers.
quantecon/​tests/​test_lqcontrol.py	Updates reference P/d/F and sequence values for the corrected recursion; adds a policy-cost regression test for #982.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[mmcky]

Thanks @jadhavgaurav, both for adding the line and confirming, and for offering to port the extras. Yes please, that would be great!

Your fix and #984's are the same at the core (both build the left factor directly, and the results agree up to rounding), and the regenerated reference values are identical, so only the additions need to come over. You can see them in #984's Files changed tab, or locally with:

git fetch https://github.com/QuantEcon/QuantEcon.py.git pull/984/head:pr-984
git show pr-984

1. quantecon/_matrix_eqn.py: copy the solve_discrete_riccati_system docstring from FIX: correct the Markov-jump Riccati recursion in LQMarkov (#982) #984. It becomes a raw string (r""", which \beta and \bar need), states the equation, and corrects the return type to ndim=3. Then add a short comment above your Ps1[i][:, :] = ( assignment saying why the left factor is formed directly (FIX: correct the Markov-jump Riccati recursion in LQMarkov (#982) #984's comment, adapted to your inline expression). Your loop body itself stays as it is.
2. quantecon/tests/test_lqcontrol.py: copy the _lq_markov_policy_cost helper (module level, just above class TestLQMarkov) and the two tests test_stationary_values_are_policy_costs and test_identical_regimes_match_lq. Please keep your test_stationary_P_is_cost_of_own_policy as well; its perturbation check is the only direct optimality test in either PR.
3. Check with pytest quantecon/tests/test_lqcontrol.py quantecon/tests/test_matrix_eqn.py (22 tests) and flake8 --select=F401,F405,E231 quantecon, then push it as a new commit. There's no need to squash or force-push, since we squash on merge.

A quick way to see test_identical_regimes_match_lq earn its place: swap the issue's transposed left factor, (beta * Bs[i].T @ P_bar @ As[i] + Ns[i]).T, into your code and it fails. That's the regression it guards against.

Thanks again!

[jadhavgaurav]

Ported the items from #984 in ae54500: the solve_discrete_riccati_system docstring (raw string, equation, ndim=3 return type), a comment above the left-factor assignment, and the _lq_markov_policy_cost helper with test_stationary_values_are_policy_costs and test_identical_regimes_match_lq. test_stationary_P_is_cost_of_own_policy is kept. pytest quantecon/tests/test_lqcontrol.py quantecon/tests/test_matrix_eqn.py gives 22 passed, flake8 is clean, and test_identical_regimes_match_lq fails with the transposed left factor.

[mmcky]

Thanks @jadhavgaurav, and thanks for porting the extras from #984 so quickly. This looks good to me.

I checked the fix independently. The returned Ps satisfy the Markov-jump Bellman equation to about 1e-11 on both test fixtures and the #982 example. I checked this by minimising its right-hand side numerically rather than using the closed form. A separate policy-iteration solver also reproduces the regenerated reference values. On 785 random solvable problems, the result agrees with #984 to rounding, and with LQ on every single-regime one.

The one CI failure was an unrelated Nelder–Mead test with an unseeded starting point, and it passed on re-run.