Merge pull request sympy#20196 from Maelstrom6/mc_stationary_distribution

czgdp1807 · web-flow · commit 5d5b7d7be41b · 2020-10-09T19:49:50.000+05:30
Improved fixed_row_vector for Discrete Markov Chains
diff --git a/sympy/stats/stochastic_process_types.py b/sympy/stats/stochastic_process_types.py
@@ -9,7 +9,7 @@
                    linsolve, eye, Or, Not, Intersection, factorial, Contains,
                    Union, Expr, Function, exp, cacheit, sqrt, pi, gamma,
                    Ge, Piecewise, Symbol, NonSquareMatrixError, EmptySet,
-                   ceiling, MatrixBase)
+                   ceiling, MatrixBase, ConditionSet, ones, zeros, Identity)
 from sympy.core.relational import Relational
 from sympy.logic.boolalg import Boolean
 from sympy.utilities.exceptions import SymPyDeprecationWarning
@@ -1012,21 +1012,98 @@ def is_absorbing_chain(self):
         return any(self.is_absorbing_state(state) == True
                     for state in range(trans_probs.shape[0]))
 
-    def fixed_row_vector(self):
+    def stationary_distribution(self, condition_set=False) -> tUnion[ImmutableMatrix, ConditionSet, Lambda]:
+        """
+        The stationary distribution is any row vector, p, that solves p = pP,
+        is row stochastic and each element in p must be nonnegative.
+        That means in matrix form: :math:`(P-I)^T p^T = 0` and
+        :math:`(1, ..., 1) p = 1`
+        where ``P`` is the one-step transition matrix.
+
+        All time-homogeneous Markov Chains with a finite state space
+        have at least one stationary distribution. In addition, if
+        a finite time-homogeneous Markov Chain is irreducible, the
+        stationary distribution is unique.
+
+        Parameters
+        ==========
+
+        condition_set : bool
+            If the chain has a symbolic size or transition matrix,
+            it will return a ``Lambda`` if ``False`` and return a
+            ``ConditionSet`` if ``True``.
+
+        Examples
+        ========
+
+        >>> from sympy.stats import DiscreteMarkovChain
+        >>> from sympy import Matrix, S
+
+        An irreducible Markov Chain
+
+        >>> T = Matrix([[S(1)/2, S(1)/2, 0],
+        ...             [S(4)/5, S(1)/5, 0],
+        ...             [1, 0, 0]])
+        >>> X = DiscreteMarkovChain('X', trans_probs=T)
+        >>> X.stationary_distribution()
+        Matrix([[8/13, 5/13, 0]])
+
+        A reducible Markov Chain
+
+        >>> T = Matrix([[S(1)/2, S(1)/2, 0],
+        ...             [S(4)/5, S(1)/5, 0],
+        ...             [0, 0, 1]])
+        >>> X = DiscreteMarkovChain('X', trans_probs=T)
+        >>> X.stationary_distribution()
+        Matrix([[8/13 - 8*tau0/13, 5/13 - 5*tau0/13, tau0]])
+
+        >>> Y = DiscreteMarkovChain('Y')
+        >>> Y.stationary_distribution()
+        Lambda((wm, _T), Eq(wm*_T, wm))
+
+        >>> Y.stationary_distribution(condition_set=True)
+        ConditionSet(wm, Eq(wm*_T, wm))
+
+        References
+        ==========
+
+        .. [1] https://www.probabilitycourse.com/chapter11/11_2_6_stationary_and_limiting_distributions.php
+        .. [2] https://galton.uchicago.edu/~yibi/teaching/stat317/2014/Lectures/Lecture4_6up.pdf
+
+        See Also
+        ========
+
+        sympy.stats.stochastic_process_types.DiscreteMarkovChain.limiting_distribution
+        """
         trans_probs = self.transition_probabilities
-        if trans_probs == None:
-            return None
+        n = self.number_of_states
+
+        if n == 0:
+            return ImmutableMatrix(Matrix([[]]))
+
+        # symbolic matrix version
         if isinstance(trans_probs, MatrixSymbol):
-            wm = MatrixSymbol('wm', 1, trans_probs.shape[0])
-            return Lambda((wm, trans_probs), Eq(wm*trans_probs, wm))
-        w = IndexedBase('w')
-        wi = [w[i] for i in range(trans_probs.shape[0])]
-        wm = Matrix([wi])
-        eqs = (wm*trans_probs - wm).tolist()[0]
-        eqs.append(sum(wi) - 1)
-        soln = list(linsolve(eqs, wi))[0]
+            wm = MatrixSymbol('wm', 1, n)
+            if condition_set:
+                return ConditionSet(wm, Eq(wm * trans_probs, wm))
+            else:
+                return Lambda((wm, trans_probs), Eq(wm * trans_probs, wm))
+
+        # numeric matrix version
+        a = Matrix(trans_probs - Identity(n)).T
+        a[0, 0:n] = ones(1, n)
+        b = zeros(n, 1)
+        b[0, 0] = 1
+
+        soln = list(linsolve((a, b)))[0]
         return ImmutableMatrix([[sol for sol in soln]])
 
+    def fixed_row_vector(self):
+        """
+        A wrapper for ``stationary_distribution()``.
+        """
+        return self.stationary_distribution()
+
     @property
     def limiting_distribution(self):
         """
diff --git a/sympy/stats/tests/test_stochastic_process.py b/sympy/stats/tests/test_stochastic_process.py
@@ -131,13 +131,18 @@ def test_DiscreteMarkovChain():
     TS1 = MatrixSymbol('T', 3, 3)
     Y5 = DiscreteMarkovChain('Y', trans_probs=TS1)
     assert Y5.limiting_distribution(w, TO4).doit() == True
+    assert Y5.stationary_distribution(condition_set=True).subs(TS1, TO4).contains(w).doit() == S.true
     TO6 = Matrix([[S.One, 0, 0, 0, 0],[S.Half, 0, S.Half, 0, 0],[0, S.Half, 0, S.Half, 0], [0, 0, S.Half, 0, S.Half], [0, 0, 0, 0, 1]])
     Y6 = DiscreteMarkovChain('Y', trans_probs=TO6)
     assert Y6._transient2absorbing() == ImmutableMatrix([[S.Half, 0], [0, 0], [0, S.Half]])
     assert Y6._transient2transient() == ImmutableMatrix([[0, S.Half, 0], [S.Half, 0, S.Half], [0, S.Half, 0]])
     assert Y6.fundamental_matrix() == ImmutableMatrix([[Rational(3, 2), S.One, S.Half], [S.One, S(2), S.One], [S.Half, S.One, Rational(3, 2)]])
     assert Y6.absorbing_probabilities() == ImmutableMatrix([[Rational(3, 4), Rational(1, 4)], [S.Half, S.Half], [Rational(1, 4), Rational(3, 4)]])
 
+    # test for zero-sized matrix functionality
+    X = DiscreteMarkovChain('X', trans_probs=Matrix([[]]))
+    assert X.number_of_states == 0
+    assert X.stationary_distribution() == Matrix([[]])
     # test communication_class
     # see https://drive.google.com/drive/folders/1HbxLlwwn2b3U8Lj7eb_ASIUb5vYaNIjg?usp=sharing
     # tutorial 2.pdf
