Merge pull request #658 from QuantEcon/markov_est

Add MLE estimation for Markov chains

docs/source/markov.rst

@@ -7,6 +7,7 @@ Markov
    markov/approximation
    markov/core
    markov/ddp
+   markov/estimate
    markov/gth_solve
    markov/random
    markov/utilities

docs/source/markov/estimate.rst

@@ -0,0 +1,7 @@
+estimate
+========
+
+.. automodule:: quantecon.markov.estimate
+    :members:
+    :undoc-members:
+    :show-inheritance:

quantecon/markov/__init__.py

@@ -11,3 +11,4 @@
 from .approximation import tauchen, rouwenhorst
 from .ddp import DiscreteDP, backward_induction
 from .utilities import sa_indices
+from .estimate import estimate_mc

quantecon/markov/estimate.py

@@ -0,0 +1,54 @@
+import numpy as np
+from numba import njit
+from .core import MarkovChain
+
+
+def estimate_mc(X):
+    r"""
+    Estimate the Markov chain associated with a time series :math:`X =
+    (X_0, \ldots, X_{T-1})` assuming that the state space is the finite
+    set :math:`\{X_0, \ldots, X_{T-1}\}` (duplicates removed). The
+    estimation is by maximum likelihood. The estimated transition
+    probabilities are given by the matrix :math:`P` such that
+    :math:`P[i, j] = N_{ij} / N_i`, where :math:`N_{ij} =
+    \sum_{t=0}^{T-1} 1_{\{X_t=s_i, X_{t+1}=s_j\}}`, the number of
+    transitions from state :math:`s_i` to state :math:`s_j`, while
+    :math:`N_i` is the total number of visits to :math:`s_i`. The result
+    is returned as a `MarkovChain` instance.
+
+    Parameters
+    ----------
+    X : array_like
+        A time series of state values, from which the transition matrix
+        will be estimated, where `X[t]` contains the t-th observation.
+
+    Returns
+    -------
+    mc : MarkovChain
+        A MarkovChain instance where `mc.P` is a stochastic matrix
+        estimated from the data `X` and `mc.state_values` is an array of
+        values that appear in `X` (sorted in ascending order).
+
+    """
+    X = np.asarray(X)
+    axis = 0 if X.ndim > 1 else None
+    state_values, indices = np.unique(X, return_inverse=True, axis=axis)
+
+    n = len(state_values)
+    P = np.zeros((n, n))  # dtype=float to modify in place upon normalization
+    P = _count_transition_frequencies(indices, P)
+    P /= P.sum(1)[:, np.newaxis]
+
+    mc = MarkovChain(P, state_values=state_values)
+    return mc
+
+
+@njit(cache=True)
+def _count_transition_frequencies(index_series, trans_counter):
+    T = len(index_series)
+    i = index_series[0]
+    for t in range(1, T):
+        j = index_series[t]
+        trans_counter[i, j] += 1
+        i = j
+    return trans_counter

quantecon/markov/tests/test_estimate.py

@@ -0,0 +1,53 @@
+"""
+Tests for markov/estimate.py
+
+"""
+import numpy as np
+from numpy.testing import assert_array_equal, assert_allclose
+from quantecon.markov.estimate import estimate_mc
+
+
+class TestEstimateMCDiscrete:
+    def setup_method(self):
+        self.test_index_series = np.array((0, 1, 1, 1, 1, 0, 2, 1))
+        self.initial_state_values = np.array(['a', 'b', 'c', 'd'])
+        self.test_value_series = \
+            self.initial_state_values[self.test_index_series]
+
+        self.P = np.array([[0.,   0.5,  0.5 ],
+                           [0.25, 0.75, 0.  ],
+                           [0.,   1.,   0.  ]])
+
+        self.final_state_indices = np.array([0, 1, 2])
+        self.final_state_values = \
+            self.initial_state_values[self.final_state_indices]
+
+    def test_integer_state(self):
+        mc = estimate_mc(self.test_index_series)
+        assert_allclose(mc.P, self.P)
+        assert_array_equal(mc.state_values, self.final_state_indices)
+
+    def test_non_integer_state(self):
+        mc = estimate_mc(self.test_value_series)
+        assert_allclose(mc.P, self.P)
+        assert_array_equal(mc.state_values, self.final_state_values)
+
+        mc = estimate_mc(self.test_index_series)
+        mc.state_values = self.initial_state_values[mc.state_values]
+        assert_allclose(mc.P, self.P)
+        assert_array_equal(mc.state_values, self.final_state_values)
+
+    def test_mult_dim_state(self):
+        initial_state_values = np.array([[0.97097089, 0.76167618],
+                                         [0.61878456, 0.41691572],
+                                         [0.42137226, 0.09307409],
+                                         [0.62696609, 0.40898893]])
+        X = initial_state_values[self.test_index_series]
+        ind = np.lexsort(
+            np.rot90(initial_state_values[self.final_state_indices])
+        )
+        final_state_values = initial_state_values[ind]
+
+        mc = estimate_mc(X)
+        assert_allclose(mc.P, self.P[np.ix_(ind, ind)])
+        assert_array_equal(mc.state_values, final_state_values)