FIX: correct the Markov-jump Riccati recursion in LQMarkov (#982)

quantecon/_matrix_eqn.py

@@ -227,11 +227,25 @@ def solve_discrete_riccati(A, B, Q, R, N=None, tolerance=1e-10, max_iter=500,
 
 def solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta,
                                   tolerance=1e-10, max_iter=1000):
-    """
+    r"""
     Solves the stacked system of algebraic matrix Riccati equations
     in the Markov Jump linear quadratic control problems, by iterating
     Ps matrices until convergence.
 
+    The system consists of the m equations
+
+    .. math::
+
+        P_i = R_i + \beta A_i' \bar{P}_i A_i
+              - (\beta B_i' \bar{P}_i A_i + N_i)'
+                (Q_i + \beta B_i' \bar{P}_i B_i)^{-1}
+                (\beta B_i' \bar{P}_i A_i + N_i)
+
+    for :math:`i = 1, \ldots, m`, where
+    :math:`\bar{P}_i = \sum_{j=1}^m \Pi_{ij} P_j`. The expectation over
+    next period's Markov state enters inside the inverse because the
+    control is chosen before that state is realized.
+
     Parameters
     ----------
     Π : array_like(float, ndim=2)
@@ -265,7 +279,7 @@ def solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta,
 
     Returns
     -------
-    Ps : array_like(float, ndim=2)
+    Ps : array_like(float, ndim=3)
         The fixed point of the stacked system of algebraic matrix
         Riccati equations, consists of m n x n P(s) matrices
 
@@ -287,9 +301,6 @@ def solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta,
     else:
         fail_msg = "Convergence failed after {} iterations."
 
-    # == Prepare array for iteration == #
-    sum1, sum2 = np.empty((n, n)), np.empty((n, n))
-
     # == Main loop == #
     iteration = 0
     while error > tolerance:
@@ -299,18 +310,17 @@ def solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta,
 
         else:
             error = 0
+            # P_bar[i] = sum_j Π[i, j] Ps[j]
+            P_bar = np.einsum('ij,jkl->ikl', Π, Ps)
             for i in range(m):
-                # Initialize arrays
-                sum1[:, :] = 0.
-                sum2[:, :] = 0.
-                for j in range(m):
-                    sum1 += beta * Π[i, j] * As[i].T @ Ps[j] @ As[i]
-                    sum2 += Π[i, j] * \
-                            (beta * As[i].T @ Ps[j] @ Bs[i] + Ns[i].T) @ \
-                            solve(Qs[i] + beta * Bs[i].T @ Ps[j] @ Bs[i],
-                                  beta * Bs[i].T @ Ps[j] @ As[i] + Ns[i])
-
-                Ps1[i][:, :] = Rs[i] + sum1 - sum2
+                BPA_N = beta * Bs[i].T @ P_bar[i] @ As[i] + Ns[i]
+                # Formed directly: as BPA_N.T, the left factor would let
+                # rounding asymmetry in Ps grow from one sweep to the next
+                APB_N = beta * As[i].T @ P_bar[i] @ Bs[i] + Ns[i].T
+                Ps1[i][:, :] = \
+                    Rs[i] + beta * As[i].T @ P_bar[i] @ As[i] - \
+                    APB_N @ solve(Qs[i] + beta * Bs[i].T @ P_bar[i] @ Bs[i],
+                                  BPA_N)
                 error += np.max(np.abs(Ps1[i] - Ps[i]))
 
             Ps[:, :, :] = Ps1[:, :, :]

quantecon/tests/test_lqcontrol.py

@@ -95,6 +95,35 @@ def test_stationary_mat(self):
             assert_allclose(val_func_lq, val_func_answer, atol=1e-3)
 
 
+def _lq_markov_policy_cost(lq, Fs):
+    """
+    Return the matrices V(s) such that, when it is finite, the expected
+    discounted cost of following u = -F(s) x from (x, s) in `lq` is
+    x' V(s) x plus a term that does not depend on x. They solve the stacked
+    linear system
+
+        V(i) = R_F(i) + beta A_F(i)' [sum_j Π[i, j] V(j)] A_F(i)
+
+    where A_F(i) = A(i) - B(i) F(i) and
+    R_F(i) = R(i) + F(i)' Q(i) F(i) - F(i)' N(i) - N(i)' F(i).
+
+    """
+    m, n = lq.m, lq.n
+    M = np.identity(m * n**2)
+    b = np.empty(m * n**2)
+    for i in range(m):
+        A_F = lq.As[i] - lq.Bs[i] @ Fs[i]
+        rows = slice(i * n**2, (i+1) * n**2)
+        for j in range(m):
+            # vec(A_F' V A_F) = kron(A_F', A_F') vec(V)
+            M[rows, j * n**2:(j+1) * n**2] -= \
+                lq.beta * lq.Π[i, j] * np.kron(A_F.T, A_F.T)
+        R_F = lq.Rs[i] + Fs[i].T @ lq.Qs[i] @ Fs[i] - \
+            Fs[i].T @ lq.Ns[i] - lq.Ns[i].T @ Fs[i]
+        b[rows] = R_F.ravel()
+    return np.linalg.solve(M, b).reshape(m, n, n)
+
+
 class TestLQMarkov:
 
     def setup_method(self):
@@ -154,8 +183,8 @@ def test_scalar_sequences_with_seed(self):
         lq_markov_scalar = self.lq_markov_scalar
         x0 = 2
 
-        expected_x_seq = np.array([[2., 1.15977567, 0.6725398]])
-        expected_u_seq = np.array([[1.28044866, 0.7425166]])
+        expected_x_seq = np.array([[2., 1.15843266, 0.67098311]])
+        expected_u_seq = np.array([[1.28313468, 0.74321256]])
         expected_w_seq = np.array([[1.3486939, 0.55721062, 0.53423587]])
         expected_state = np.array([1, 1, 1])
 
@@ -172,11 +201,11 @@ def test_stationary_scalar(self):
 
         lq_markov_scalar = self.lq_markov_scalar
 
-        P_answer = np.array([[[1.51741465]],
-                             [[1.07334181]]])
+        P_answer = np.array([[[1.52007965]],
+                             [[1.07741061]]])
         d_answer = np.array([0., 0.])
-        F_answer = np.array([[[-0.54697435]],
-                             [[-0.64022433]]])
+        F_answer = np.array([[[-0.54745226]],
+                             [[-0.64156734]]])
 
         Ps, ds, Fs = lq_markov_scalar.stationary_values()
 
@@ -189,10 +218,10 @@ def test_mat_sequences(self):
         lq_markov_mat = self.lq_markov_mat1
         x0 = np.array([[1000, 1, 25]])
 
-        expected_x_seq = np.array([[1.00000000e+03, 1.01372101e+03],
+        expected_x_seq = np.array([[1.00000000e+03, 1.01342472e+03],
                                    [1.00000000e+00, 1.00000000e+00],
                                    [2.50000000e+01, 2.61845443e+01]])
-        expected_u_seq = np.array([[1013.72101253]])
+        expected_u_seq = np.array([[1013.42471946]])
         expected_w_seq = np.array([[0.41782708, 1.18454431]])
         expected_state = np.array([1, 1])
 
@@ -207,22 +236,58 @@ def test_mat_sequences(self):
     def test_stationary_mat(self):
         lq_markov_mat = self.lq_markov_mat1
 
-        d_answer = np.array([16.2474886, 16.31935939])
-        P_answer = np.array([[[4.5144056e-02, 1.8627227e+01, 1.9348906e-01],
-                              [1.8627227e+01, 7.9343733e+03, 8.0055130e+01],
-                              [1.9348906e-01, 8.0055130e+01, 8.2991316e-01]],
-                             [[5.3606323e-02, 2.0136657e+01, 2.1763323e-01],
-                              [2.0136657e+01, 7.8100167e+03, 8.1960509e+01],
-                              [2.1763323e-01, 8.1960509e+01, 8.8413147e-01]]])
-        F_answer = np.array([[[-0.98437714, 19.2051657, -0.83142157]],
-                             [[-1.01434303, 21.58480004, -0.83851124]]])
+        d_answer = np.array([16.31235346, 16.38442572])
+        P_answer = np.array([[[4.5345640e-02, 1.8660614e+01, 1.9430712e-01],
+                              [1.8660614e+01, 7.9402383e+03, 8.0190872e+01],
+                              [1.9430712e-01, 8.0190872e+01, 8.3325640e-01]],
+                             [[5.3840146e-02, 2.0171229e+01, 2.1853447e-01],
+                              [2.0171229e+01, 7.8154682e+03, 8.2093814e+01],
+                              [2.1853447e-01, 8.2093814e+01, 8.8762671e-01]]])
+        F_answer = np.array([[[-0.98417975, 19.23774624, -0.83061122]],
+                             [[-1.01410488, 21.6197524, -0.83758363]]])
 
         Ps, ds, Fs = lq_markov_mat.stationary_values()
 
         assert_allclose(F_answer, Fs, atol=1e-6)
         assert_allclose(P_answer, Ps, atol=1e-6)
         assert_allclose(d_answer, ds, atol=1e-6)
 
+    def test_stationary_values_are_policy_costs(self):
+        # x' P(s) x must be the x-dependent part of the discounted cost of
+        # following u = -F(s) x from (x, s). This fails if the expectation
+        # over next period's Markov state is taken outside the inverse in
+        # the Riccati recursion (#982).
+        Π = np.array([[0.5, 0.5],
+                      [0.3, 0.7]])
+        As = [[[1., .1], [0., .9]], [[.9, .3], [.2, 1.]]]
+        Bs = [[[1.], [0.]], [[.2], [1.]]]
+        Rs = [np.eye(2), 2 * np.eye(2)]
+        lq_markov_mat2 = LQMarkov(Π, [.1, 5.], Rs, As, Bs, beta=.95)
+
+        for lq in (self.lq_markov_scalar, self.lq_markov_mat1,
+                   lq_markov_mat2):
+            Ps, ds, Fs = lq.stationary_values()
+
+            assert_allclose(Ps, _lq_markov_policy_cost(lq, Fs), rtol=1e-6)
+
+    def test_identical_regimes_match_lq(self):
+        # With one regime, or identical regimes, LQMarkov solves a standard
+        # LQ problem. A is unstable, so this also fails if the Riccati
+        # iteration lets rounding asymmetry in Ps grow.
+        A = np.array([[1.5, .3], [.1, 1.5]])
+        B = np.array([[1.], [1.]])
+        P, F, d = LQ(1., np.eye(2), A, B, beta=.99).stationary_values()
+
+        for Π in ([[1.]], [[.9, .1], [.2, .8]]):
+            m = len(Π)
+            lq = LQMarkov(Π, [1.] * m, [np.eye(2)] * m, [A] * m, [B] * m,
+                          beta=.99)
+            Ps, ds, Fs = lq.stationary_values()
+
+            for i in range(m):
+                assert_allclose(Ps[i], P, rtol=1e-8)
+                assert_allclose(Fs[i], F, rtol=1e-8)
+
     def test_beta_one_without_shocks_one_state(self):
         lq = LQMarkov([[1.]], [1.], [1.], [.5], [0.], beta=1)