Repository: sympy/sympy

Issue:
How to generate results for queries like `P(Eq(X(j), 2), Eq(X(1), i))`?
In the query, `P(Eq(X(j), 2), Eq(X(1), i))`, X represents a Markov chain(discrete/continuous). Here, i and j are symbols.
Currently,
```python
>>> from sympy import symbols
>>> i, j = symbols('i, j', integer=True)
>>> P(Eq(X[j], 2), Eq(X[1], i))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/czgdp1807ssd/sympy_project/sympy/sympy/stats/stochastic_process_types.py", line 196, in __getitem__
    if time not in self.index_set:
  File "/home/czgdp1807ssd/sympy_project/sympy/sympy/sets/sets.py", line 672, in __contains__
    raise TypeError('did not evaluate to a bool: %r' % c)
TypeError: did not evaluate to a bool: None
```
This issue should discuss what all needs to be done to support such queries.

Discussion:

Patch:
diff --git a/sympy/stats/stochastic_process_types.py b/sympy/stats/stochastic_process_types.py
index e40774553d..de912f9763 100644
--- a/sympy/stats/stochastic_process_types.py
+++ b/sympy/stats/stochastic_process_types.py
@@ -10,7 +10,7 @@
                    Union, Expr, Function, exp, cacheit, sqrt, pi, gamma,
                    Ge, Piecewise, Symbol, NonSquareMatrixError, EmptySet,
                    ceiling, MatrixBase, ConditionSet, ones, zeros, Identity,
-                   Rational, Lt, Gt, Ne, BlockMatrix)
+                   Rational, Lt, Gt, Le, Ne, BlockMatrix, Sum)
 from sympy.core.relational import Relational
 from sympy.logic.boolalg import Boolean
 from sympy.utilities.exceptions import SymPyDeprecationWarning
@@ -250,7 +250,8 @@ def __getitem__(self, time):
 
         RandomIndexedSymbol
         """
-        if time not in self.index_set:
+        time = sympify(time)
+        if not time.is_symbol and time not in self.index_set:
             raise IndexError("%s is not in the index set of %s"%(time, self.symbol))
         idx_obj = Indexed(self.symbol, time)
         pspace_obj = StochasticPSpace(self.symbol, self, self.distribution)
@@ -269,7 +270,8 @@ def __call__(self, time):
 
         RandomIndexedSymbol
         """
-        if time not in self.index_set:
+        time = sympify(time)
+        if not time.is_symbol and time not in self.index_set:
             raise IndexError("%s is not in the index set of %s"%(time, self.symbol))
         func_obj = Function(self.symbol)(time)
         pspace_obj = StochasticPSpace(self.symbol, self, self.distribution)
@@ -509,6 +511,13 @@ def probability(self, condition, given_condition=None, evaluate=True, **kwargs):
         check, mat, state_index, new_given_condition = \
             self._preprocess(given_condition, evaluate)
 
+        rv = list(condition.atoms(RandomIndexedSymbol))
+        symbolic = False
+        for sym in rv:
+            if sym.key.is_symbol:
+                symbolic = True
+                break
+
         if check:
             return Probability(condition, new_given_condition)
 
@@ -526,21 +535,34 @@ def probability(self, condition, given_condition=None, evaluate=True, **kwargs):
             else:
                 gcs = (new_given_condition, )
             min_key_rv = list(new_given_condition.atoms(RandomIndexedSymbol))
-            rv = list(condition.atoms(RandomIndexedSymbol))
 
             if len(min_key_rv):
                 min_key_rv = min_key_rv[0]
                 for r in rv:
+                    if min_key_rv.key.is_symbol or r.key.is_symbol:
+                        continue
                     if min_key_rv.key > r.key:
                         return Probability(condition)
             else:
                 min_key_rv = None
                 return Probability(condition)
 
+            if symbolic:
+                return self._symbolic_probability(condition, new_given_condition, rv, min_key_rv)
+
             if len(rv) > 1:
-                rv = rv[:2]
+                rv[0] = condition.lhs
+                rv[1] = condition.rhs
                 if rv[0].key < rv[1].key:
                         rv[0], rv[1] = rv[1], rv[0]
+                        if isinstance(condition, Gt):
+                            condition = Lt(condition.lhs, condition.rhs)
+                        elif isinstance(condition, Lt):
+                            condition = Gt(condition.lhs, condition.rhs)
+                        elif isinstance(condition, Ge):
+                            condition = Le(condition.lhs, condition.rhs)
+                        elif isinstance(condition, Le):
+                            condition = Ge(condition.lhs, condition.rhs)
                 s = Rational(0, 1)
                 n = len(self.state_space)
 
@@ -658,6 +680,37 @@ def probability(self, condition, given_condition=None, evaluate=True, **kwargs):
         raise NotImplementedError("Mechanism for handling (%s, %s) queries hasn't been "
                                 "implemented yet."%(condition, given_condition))
 
+    def _symbolic_probability(self, condition, new_given_condition, rv, min_key_rv):
+        #Function to calculate probability for queries with symbols
+        if isinstance(condition, Relational):
+            curr_state = new_given_condition.rhs if isinstance(new_given_condition.lhs, RandomIndexedSymbol) \
+                    else new_given_condition.lhs
+            next_state = condition.rhs if isinstance(condition.lhs, RandomIndexedSymbol) \
+                else condition.lhs
+
+            if isinstance(condition, Eq) or isinstance(condition, Ne):
+                if isinstance(self, DiscreteMarkovChain):
+                    P = self.transition_probabilities**(rv[0].key - min_key_rv.key)
+                else:
+                    P = exp(self.generator_matrix*(rv[0].key - min_key_rv.key))
+                prob = P[curr_state, next_state] if isinstance(condition, Eq) else 1 - P[curr_state, next_state]
+                return Piecewise((prob, rv[0].key > min_key_rv.key), (Probability(condition), True))
+            else:
+                upper = 1
+                greater = False
+                if isinstance(condition, Ge) or isinstance(condition, Lt):
+                    upper = 0
+                if isinstance(condition, Gt) or isinstance(condition, Ge):
+                    greater = True
+                k = Dummy('k')
+                condition = Eq(condition.lhs, k) if isinstance(condition.lhs, RandomIndexedSymbol)\
+                    else Eq(condition.rhs, k)
+                total = Sum(self.probability(condition, new_given_condition), (k, next_state + upper, self.state_space._sup))
+                return Piecewise((total, rv[0].key > min_key_rv.key), (Probability(condition), True)) if greater\
+                    else Piecewise((1 - total, rv[0].key > min_key_rv.key), (Probability(condition), True))
+        else:
+            return Probability(condition, new_given_condition)
+
     def expectation(self, expr, condition=None, evaluate=True, **kwargs):
         """
         Handles expectation queries for markov process.
@@ -806,6 +859,24 @@ class DiscreteMarkovChain(DiscreteTimeStochasticProcess, MarkovProcess):
     >>> P(Le(Y[15], Y[10]), Eq(Y[8], 2)).round(7)
     0.6963328
 
+    Symbolic probability queries are also supported
+
+    >>> from sympy import symbols, Matrix, Rational, Eq, Gt
+    >>> from sympy.stats import P, DiscreteMarkovChain
+    >>> a, b, c, d = symbols('a b c d')
+    >>> T = Matrix([[Rational(1, 10), Rational(4, 10), Rational(5, 10)], [Rational(3, 10), Rational(4, 10), Rational(3, 10)], [Rational(7, 10), Rational(2, 10), Rational(1, 10)]])
+    >>> Y = DiscreteMarkovChain("Y", [0, 1, 2], T)
+    >>> query = P(Eq(Y[a], b), Eq(Y[c], d))
+    >>> query.subs({a:10 ,b:2, c:5, d:1}).round(4)
+    0.3096
+    >>> P(Eq(Y[10], 2), Eq(Y[5], 1)).evalf().round(4)
+    0.3096
+    >>> query_gt = P(Gt(Y[a], b), Eq(Y[c], d))
+    >>> query_gt.subs({a:21, b:0, c:5, d:0}).evalf().round(5)
+    0.64705
+    >>> P(Gt(Y[21], 0), Eq(Y[5], 0)).round(5)
+    0.64705
+
     There is limited support for arbitrarily sized states:
 
     >>> n = symbols('n', nonnegative=True, integer=True)
@@ -813,6 +884,9 @@ class DiscreteMarkovChain(DiscreteTimeStochasticProcess, MarkovProcess):
     >>> Y = DiscreteMarkovChain("Y", trans_probs=T)
     >>> Y.state_space
     Range(0, n, 1)
+    >>> query = P(Eq(Y[a], b), Eq(Y[c], d))
+    >>> query.subs({a:10, b:2, c:5, d:1})
+    (T**5)[1, 2]
 
     References
     ==========
@@ -1418,12 +1492,59 @@ class ContinuousMarkovChain(ContinuousTimeStochasticProcess, MarkovProcess):
     Examples
     ========
 
-    >>> from sympy.stats import ContinuousMarkovChain
-    >>> from sympy import Matrix, S
+    >>> from sympy.stats import ContinuousMarkovChain, P
+    >>> from sympy import Matrix, S, Eq, Gt
     >>> G = Matrix([[-S(1), S(1)], [S(1), -S(1)]])
     >>> C = ContinuousMarkovChain('C', state_space=[0, 1], gen_mat=G)
     >>> C.limiting_distribution()
     Matrix([[1/2, 1/2]])
+    >>> C.state_space
+    FiniteSet(0, 1)
+    >>> C.generator_matrix
+    Matrix([
+    [-1,  1],
+    [ 1, -1]])
+
+    Probability queries are supported
+
+    >>> P(Eq(C(1.96), 0), Eq(C(0.78), 1)).round(5)
+    0.45279
+    >>> P(Gt(C(1.7), 0), Eq(C(0.82), 1)).round(5)
+    0.58602
+
+    Probability of expressions with multiple RandomIndexedSymbols
+    can also be calculated provided there is only 1 RandomIndexedSymbol
+    in the given condition. It is always better to use Rational instead
+    of floating point numbers for the probabilities in the
+    generator matrix to avoid errors.
+
+    >>> from sympy import Gt, Le, Rational
+    >>> G = Matrix([[-S(1), Rational(1, 10), Rational(9, 10)], [Rational(2, 5), -S(1), Rational(3, 5)], [Rational(1, 2), Rational(1, 2), -S(1)]])
+    >>> C = ContinuousMarkovChain('C', state_space=[0, 1, 2], gen_mat=G)
+    >>> P(Eq(C(3.92), C(1.75)), Eq(C(0.46), 0)).round(5)
+    0.37933
+    >>> P(Gt(C(3.92), C(1.75)), Eq(C(0.46), 0)).round(5)
+    0.34211
+    >>> P(Le(C(1.57), C(3.14)), Eq(C(1.22), 1)).round(4)
+    0.7143
+
+    Symbolic probability queries are also supported
+
+    >>> from sympy import S, symbols, Matrix, Rational, Eq, Gt
+    >>> from sympy.stats import P, ContinuousMarkovChain
+    >>> a,b,c,d = symbols('a b c d')
+    >>> G = Matrix([[-S(1), Rational(1, 10), Rational(9, 10)], [Rational(2, 5), -S(1), Rational(3, 5)], [Rational(1, 2), Rational(1, 2), -S(1)]])
+    >>> C = ContinuousMarkovChain('C', state_space=[0, 1, 2], gen_mat=G)
+    >>> query = P(Eq(C(a), b), Eq(C(c), d))
+    >>> query.subs({a:3.65 ,b:2, c:1.78, d:1}).evalf().round(10)
+    0.4002723175
+    >>> P(Eq(C(3.65), 2), Eq(C(1.78), 1)).round(10)
+    0.4002723175
+    >>> query_gt = P(Gt(C(a), b), Eq(C(c), d))
+    >>> query_gt.subs({a:43.2 ,b:0, c:3.29, d:2}).evalf().round(10)
+    0.6832579186
+    >>> P(Gt(C(43.2), 0), Eq(C(3.29), 2)).round(10)
+    0.6832579186
 
     References
     ==========


Test patch:
diff --git a/sympy/stats/tests/test_stochastic_process.py b/sympy/stats/tests/test_stochastic_process.py
index 7b3af9d9ab..444210c7e3 100644
--- a/sympy/stats/tests/test_stochastic_process.py
+++ b/sympy/stats/tests/test_stochastic_process.py
@@ -315,10 +315,27 @@ def test_DiscreteMarkovChain():
     Y = DiscreteMarkovChain("Y", [0, 1, 2], T)
     assert P(Eq(Y[7], Y[5]), Eq(Y[2], 0)).round(5) == Float(0.44428, 5)
     assert P(Gt(Y[3], Y[1]), Eq(Y[0], 0)).round(2) == Float(0.36, 2)
-    assert P(Le(Y[5], Y[10]), Eq(Y[4], 2)).round(6) == Float(0.739072, 6)
-    assert Float(P(Eq(Y[500], Y[240]), Eq(Y[120], 1)), 14) == Float(1 - P(Ne(Y[500], Y[240]), Eq(Y[120], 1)), 14)
-    assert Float(P(Gt(Y[350], Y[100]), Eq(Y[75], 2)), 14) == Float(1 - P(Le(Y[350], Y[100]), Eq(Y[75], 2)), 14)
-    assert Float(P(Lt(Y[400], Y[210]), Eq(Y[161], 0)), 14) == Float(1 - P(Ge(Y[400], Y[210]), Eq(Y[161], 0)), 14)
+    assert P(Le(Y[5], Y[10]), Eq(Y[4], 2)).round(6) == Float(0.583120, 6)
+    assert Float(P(Eq(Y[10], Y[5]), Eq(Y[4], 1)), 14) == Float(1 - P(Ne(Y[10], Y[5]), Eq(Y[4], 1)), 14)
+    assert Float(P(Gt(Y[8], Y[9]), Eq(Y[3], 2)), 14) == Float(1 - P(Le(Y[8], Y[9]), Eq(Y[3], 2)), 14)
+    assert Float(P(Lt(Y[1], Y[4]), Eq(Y[0], 0)), 14) == Float(1 - P(Ge(Y[1], Y[4]), Eq(Y[0], 0)), 14)
+    assert P(Eq(Y[5], Y[10]), Eq(Y[2], 1)) == P(Eq(Y[10], Y[5]), Eq(Y[2], 1))
+    assert P(Gt(Y[1], Y[2]), Eq(Y[0], 1)) == P(Lt(Y[2], Y[1]), Eq(Y[0], 1))
+    assert P(Ge(Y[7], Y[6]), Eq(Y[4], 1)) == P(Le(Y[6], Y[7]), Eq(Y[4], 1))
+
+    #test symbolic queries
+    a, b, c, d = symbols('a b c d')
+    T = Matrix([[Rational(1, 10), Rational(4, 10), Rational(5, 10)], [Rational(3, 10), Rational(4, 10), Rational(3, 10)], [Rational(7, 10), Rational(2, 10), Rational(1, 10)]])
+    Y = DiscreteMarkovChain("Y", [0, 1, 2], T)
+    query = P(Eq(Y[a], b), Eq(Y[c], d))
+    assert query.subs({a:10, b:2, c:5, d:1}).evalf().round(4) == P(Eq(Y[10], 2), Eq(Y[5], 1)).round(4)
+    assert query.subs({a:15, b:0, c:10, d:1}).evalf().round(4) == P(Eq(Y[15], 0), Eq(Y[10], 1)).round(4)
+    query_gt = P(Gt(Y[a], b), Eq(Y[c], d))
+    query_le = P(Le(Y[a], b), Eq(Y[c], d))
+    assert query_gt.subs({a:5, b:2, c:1, d:0}).evalf() + query_le.subs({a:5, b:2, c:1, d:0}).evalf() == 1
+    query_ge = P(Ge(Y[a], b), Eq(Y[c], d))
+    query_lt = P(Lt(Y[a], b), Eq(Y[c], d))
+    assert query_ge.subs({a:4, b:1, c:0, d:2}).evalf() + query_lt.subs({a:4, b:1, c:0, d:2}).evalf() == 1
 
 def test_sample_stochastic_process():
     if not import_module('scipy'):
@@ -382,6 +399,30 @@ def test_ContinuousMarkovChain():
     assert P(Eq(C3(1), 1), Eq(C3(0), 1)) == exp(-2)/2 + S.Half
     assert P(Eq(C3(1), Symbol('1')), Eq(C3(0), Symbol('1'))) == exp(-2)/2 + S.Half
 
+    #test probability queries
+    G = Matrix([[-S(1), Rational(1, 10), Rational(9, 10)], [Rational(2, 5), -S(1), Rational(3, 5)], [Rational(1, 2), Rational(1, 2), -S(1)]])
+    C = ContinuousMarkovChain('C', state_space=[0, 1, 2], gen_mat=G)
+    assert P(Eq(C(7.385), C(3.19)), Eq(C(0.862), 0)).round(5) == Float(0.35469, 5)
+    assert P(Gt(C(98.715), C(19.807)), Eq(C(11.314), 2)).round(5) == Float(0.32452, 5)
+    assert P(Le(C(5.9), C(10.112)), Eq(C(4), 1)).round(6) == Float(0.675214, 6)
+    assert Float(P(Eq(C(7.32), C(2.91)), Eq(C(2.63), 1)), 14) == Float(1 - P(Ne(C(7.32), C(2.91)), Eq(C(2.63), 1)), 14)
+    assert Float(P(Gt(C(3.36), C(1.101)), Eq(C(0.8), 2)), 14) == Float(1 - P(Le(C(3.36), C(1.101)), Eq(C(0.8), 2)), 14)
+    assert Float(P(Lt(C(4.9), C(2.79)), Eq(C(1.61), 0)), 14) == Float(1 - P(Ge(C(4.9), C(2.79)), Eq(C(1.61), 0)), 14)
+    assert P(Eq(C(5.243), C(10.912)), Eq(C(2.174), 1)) == P(Eq(C(10.912), C(5.243)), Eq(C(2.174), 1))
+    assert P(Gt(C(2.344), C(9.9)), Eq(C(1.102), 1)) == P(Lt(C(9.9), C(2.344)), Eq(C(1.102), 1))
+    assert P(Ge(C(7.87), C(1.008)), Eq(C(0.153), 1)) == P(Le(C(1.008), C(7.87)), Eq(C(0.153), 1))
+
+    #test symbolic queries
+    a, b, c, d = symbols('a b c d')
+    query = P(Eq(C(a), b), Eq(C(c), d))
+    assert query.subs({a:3.65, b:2, c:1.78, d:1}).evalf().round(10) == P(Eq(C(3.65), 2), Eq(C(1.78), 1)).round(10)
+    query_gt = P(Gt(C(a), b), Eq(C(c), d))
+    query_le = P(Le(C(a), b), Eq(C(c), d))
+    assert query_gt.subs({a:13.2, b:0, c:3.29, d:2}).evalf() + query_le.subs({a:13.2, b:0, c:3.29, d:2}).evalf() == 1
+    query_ge = P(Ge(C(a), b), Eq(C(c), d))
+    query_lt = P(Lt(C(a), b), Eq(C(c), d))
+    assert query_ge.subs({a:7.43, b:1, c:1.45, d:0}).evalf() + query_lt.subs({a:7.43, b:1, c:1.45, d:0}).evalf() == 1
+
 def test_BernoulliProcess():
 
     B = BernoulliProcess("B", p=0.6, success=1, failure=0)