Merge pull request #489 from shizejin/add_Markov_LQ

Add `LQMarkov`.

Author: mmcky

quantecon/__init__.py

@@ -29,7 +29,7 @@
 from .kalman import Kalman
 from .lae import LAE
 from .arma import ARMA
-from .lqcontrol import LQ
+from .lqcontrol import LQ, LQMarkov
 from .filter import hamilton_filter
 from .lqnash import nnash
 from .lss import LinearStateSpace

quantecon/lqcontrol.py

@@ -1,14 +1,16 @@
 """
 Provides a class called LQ for solving linear quadratic control
-problems.
+problems, and a class called LQMarkov for solving Markov jump
+linear quadratic control problems.
 
 """
 from textwrap import dedent
 import numpy as np
 from numpy import dot
 from scipy.linalg import solve
-from .matrix_eqn import solve_discrete_riccati
+from .matrix_eqn import solve_discrete_riccati, solve_discrete_riccati_system
 from .util import check_random_state
+from .markov import MarkovChain
 
 
 class LQ:
@@ -327,3 +329,282 @@ def compute_sequence(self, x0, ts_length=None, method='doubling',
         x_path[:, T] = Ax + Bu + Cw_path[:, T]
 
         return x_path, u_path, w_path
+
+
+class LQMarkov:
+    r"""
+    This class is for analyzing Markov jump linear quadratic optimal
+    control problems of the infinite horizon form
+
+    .. math::
+
+        \min \mathbb{E}
+        \Big[ \sum_{t=0}^{\infty} \beta^t r(x_t, s_t, u_t) \Big]
+
+    with
+
+    .. math::
+
+         r(x_t, s_t, u_t) :=
+            (x_t' R(s_t) x_t + u_t' Q(s_t) u_t + 2 u_t' N(s_t) x_t)
+
+    subject to the law of motion
+
+    .. math::
+
+         x_{t+1} = A(s_t) x_t + B(s_t) u_t + C(s_t) w_{t+1}
+
+    Here :math:`x` is n x 1, :math:`u` is k x 1, :math:`w` is j x 1 and the
+    matrices are conformable for these dimensions.  The sequence :math:`{w_t}`
+    is assumed to be white noise, with zero mean and
+    :math:`\mathbb{E} [ w_t' w_t ] = I`, the j x j identity.
+
+    If :math:`C` is not supplied as a parameter, the model is assumed to be
+    deterministic (and :math:`C` is set to a zero matrix of appropriate
+    dimension).
+
+    The optimal value function :math:`V(x_t, s_t)` takes the form
+
+    .. math::
+
+         x_t' P(s_t) x_t + d(s_t)
+
+    and the optimal policy is of the form :math:`u_t = -F(s_t) x_t`.
+
+    Parameters
+    ----------
+    Π : array_like(float, ndim=2)
+        The Markov chain transition matrix with dimension m x m.
+    Qs : array_like(float)
+        Consists of m symmetric and non-negative definite payoff
+        matrices Q(s) with dimension k x k that corresponds with
+        the control variable u for each Markov state s
+    Rs : array_like(float)
+        Consists of m symmetric and non-negative definite payoff
+        matrices R(s) with dimension n x n that corresponds with
+        the state variable x for each Markov state s
+    As : array_like(float)
+        Consists of m state transition matrices A(s) with dimension
+        n x n for each Markov state s
+    Bs : array_like(float)
+        Consists of m state transition matrices B(s) with dimension
+        n x k for each Markov state s
+    Cs : array_like(float), optional(default=None)
+        Consists of m state transition matrices C(s) with dimension
+        n x j for each Markov state s. If the model is deterministic
+        then Cs should take default value of None
+    Ns : array_like(float), optional(default=None)
+        Consists of m cross product term matrices N(s) with dimension
+        k x n for each Markov state,
+    beta : scalar(float), optional(default=1)
+        beta is the discount parameter
+
+    Attributes
+    ----------
+    Π, Qs, Rs, Ns, As, Bs, Cs, beta : see Parameters
+    Ps : array_like(float)
+        Ps is part of the value function representation of
+        :math:`V(x, s) = x' P(s) x + d(s)`
+    ds : array_like(float)
+        ds is part of the value function representation of
+        :math:`V(x, s) = x' P(s) x + d(s)`
+    Fs : array_like(float)
+        Fs is the policy rule that determines the choice of control in
+        each period at each Markov state
+    m : scalar(int)
+        The number of Markov states
+    k, n, j : scalar(int)
+        The dimensions of the matrices as presented above
+
+    """
+
+    def __init__(self, Π, Qs, Rs, As, Bs, Cs=None, Ns=None, beta=1):
+
+        # == Make sure all matrices for each state are 2D arrays == #
+        def converter(Xs):
+            return np.array([np.atleast_2d(np.asarray(X, dtype='float'))
+                             for X in Xs])
+        self.As, self.Bs, self.Qs, self.Rs = list(map(converter,
+                                                      (As, Bs, Qs, Rs)))
+
+        # == Record number of states == #
+        self.m = self.Qs.shape[0]
+        # == Record dimensions == #
+        self.k, self.n = self.Qs.shape[1], self.Rs.shape[1]
+
+        if Ns is None:
+            # == No cross product term in payoff. Set N=0. == #
+            Ns = [np.zeros((self.k, self.n)) for i in range(self.m)]
+
+        self.Ns = converter(Ns)
+
+        if Cs is None:
+            # == If C not given, then model is deterministic. Set C=0. == #
+            self.j = 1
+            Cs = [np.zeros((self.n, self.j)) for i in range(self.m)]
+
+        self.Cs = converter(Cs)
+        self.j = self.Cs.shape[2]
+
+        self.beta = beta
+
+        self.Π = np.asarray(Π, dtype='float')
+
+        self.Ps = None
+        self.ds = None
+        self.Fs = None
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __str__(self):
+        m = """\
+        Markov Jump Linear Quadratic control system
+          - beta (discount parameter)       : {b}
+          - T (time horizon)                : {t}
+          - m (number of Markov states)     : {m}
+          - n (number of state variables)   : {n}
+          - k (number of control variables) : {k}
+          - j (number of shocks)            : {j}
+        """
+        t = "infinite"
+        return dedent(m.format(b=self.beta, m=self.m, n=self.n, k=self.k,
+                               j=self.j, t=t))
+
+    def stationary_values(self):
+        """
+        Computes the matrix :math:`P(s)` and scalar :math:`d(s)` that
+        represent the value function
+
+        .. math::
+
+             V(x, s) = x' P(s) x + d(s)
+
+        in the infinite horizon case.  Also computes the control matrix
+        :math:`F` from :math:`u = - F(s) x`.
+
+        Returns
+        -------
+        Ps : array_like(float)
+            Ps is part of the value function representation of
+            :math:`V(x, s) = x' P(s) x + d(s)`
+        ds : array_like(float)
+            ds is part of the value function representation of
+            :math:`V(x, s) = x' P(s) x + d(s)`
+        Fs : array_like(float)
+            Fs is the policy rule that determines the choice of control in
+            each period at each Markov state
+
+        """
+
+        # == Simplify notations == #
+        beta, Π = self.beta, self.Π
+        m, n, k = self.m, self.n, self.k
+        As, Bs, Cs = self.As, self.Bs, self.Cs
+        Qs, Rs, Ns = self.Qs, self.Rs, self.Ns
+
+        # == Solve for P(s) by iterating discrete riccati system== #
+        Ps = solve_discrete_riccati_system(Π, As, Bs, Cs, Qs, Rs, Ns, beta)
+
+        # == calculate F and d == #
+        Fs = np.array([np.empty((k, n)) for i in range(m)])
+        X = np.empty((m, m))
+        sum1, sum2 = np.empty((k, k)), np.empty((k, n))
+        for i in range(m):
+            # CCi = C_i C_i'
+            CCi = Cs[i] @ Cs[i].T
+            sum1[:, :] = 0.
+            sum2[:, :] = 0.
+            for j in range(m):
+                # for F
+                sum1 += beta * Π[i, j] * Bs[i].T @ Ps[j] @ Bs[i]
+                sum2 += beta * Π[i, j] * Bs[i].T @ Ps[j] @ As[i]
+
+                # for d
+                X[j, i] = np.trace(Ps[j] @ CCi)
+
+            Fs[i][:, :] = solve(Qs[i] + sum1, sum2 + Ns[i])
+
+        ds = solve(np.eye(m) - beta * Π,
+                   np.diag(beta * Π @ X).reshape((m, 1))).flatten()
+
+        self.Ps, self.ds, self.Fs = Ps, ds, Fs
+
+        return Ps, ds, Fs
+
+    def compute_sequence(self, x0, ts_length=None, random_state=None):
+        """
+        Compute and return the optimal state and control sequences
+        :math:`x_0, ..., x_T` and :math:`u_0,..., u_T`  under the
+        assumption that :math:`{w_t}` is iid and :math:`N(0, 1)`,
+        with Markov states sequence :math:`s_0, ..., s_T`
+
+        Parameters
+        ----------
+        x0 : array_like(float)
+            The initial state, a vector of length n
+
+        ts_length : scalar(int), optional(default=None)
+            Length of the simulation. If None, T is set to be 100
+
+        random_state : int or np.random.RandomState, optional
+            Random seed (integer) or np.random.RandomState instance to set
+            the initial state of the random number generator for
+            reproducibility. If None, a randomly initialized RandomState is
+            used.
+
+        Returns
+        -------
+        x_path : array_like(float)
+            An n x T+1 matrix, where the t-th column represents :math:`x_t`
+
+        u_path : array_like(float)
+            A k x T matrix, where the t-th column represents :math:`u_t`
+
+        w_path : array_like(float)
+            A j x T+1 matrix, where the t-th column represent :math:`w_t`
+
+        state : array_like(int)
+            Array containing the state values :math:`s_t` of the sample path
+
+        """
+
+        # === solve for optimal policies === #
+        self.stationary_values()
+
+        # === Simplify notation === #
+        As, Bs, Cs = self.As, self.Bs, self.Cs
+        Fs = self.Fs
+
+        random_state = check_random_state(random_state)
+        x0 = np.asarray(x0)
+        x0 = x0.reshape(self.n, 1)
+
+        T = ts_length if ts_length else 100
+
+        # == Simulate Markov states == #
+        chain = MarkovChain(self.Π)
+        state = chain.simulate_indices(ts_length=T+1,
+                                       random_state=random_state)
+
+        # == Prepare storage arrays == #
+        x_path = np.empty((self.n, T+1))
+        u_path = np.empty((self.k, T))
+        w_path = random_state.randn(self.j, T+1)
+        Cw_path = np.empty((self.n, T+1))
+        for i in range(T+1):
+            Cw_path[:, i] = Cs[state[i]] @ w_path[:, i]
+
+        # == Use policy sequence to generate states and controls == #
+        x_path[:, 0] = x0.flatten()
+        u_path[:, 0] = - (Fs[state[0]] @ x0).flatten()
+        for t in range(1, T):
+            Ax = As[state[t]] @ x_path[:, t-1]
+            Bu = Bs[state[t]] @ u_path[:, t-1]
+            x_path[:, t] = Ax + Bu + Cw_path[:, t]
+            u_path[:, t] = - (Fs[state[t]] @ x_path[:, t])
+        Ax = As[state[T]] @ x_path[:, T-1]
+        Bu = Bs[state[T]] @ u_path[:, T-1]
+        x_path[:, T] = Ax + Bu + w_path[:, T]
+
+        return x_path, u_path, w_path, state