sympy/concrete/summations.py

   1
   2 from sympy.core import Basic, S, C, Add, Mul, Symbol, Equality, Interval, sympify
   3 from sympy.core.methods import NoRelMeths, ArithMeths
   4
   5 class Sum(Basic, NoRelMeths, ArithMeths):
   6     """Represents unevaluated summation."""
   7
   8     precedence = Basic.Apply_precedence
   9
  10     def __new__(cls, f, *symbols, **assumptions):
  11         f = sympify(f)
  12
  13         if isinstance(f, C.Number):
  14             if f is S.NaN:
  15                 return S.NaN
  16             elif f is S.Zero:
  17                 return S.Zero
  18
  19         if not symbols:
  20             limits = f.atoms(Symbol)
  21
  22             if not limits:
  23                 return f
  24         else:
  25             limits = []
  26
  27             for V in symbols:
  28                 if isinstance(V, Symbol):
  29                     limits.append(V)
  30                     continue
  31                 elif isinstance(V, Equality):
  32                     if isinstance(V.lhs, Symbol):
  33                         if isinstance(V.rhs, Interval):
  34                             limits.append((V.lhs, V.rhs.start, V.rhs.end))
  35                         else:
  36                             limits.append((V.lhs, V.rhs))
  37
  38                         continue
  39                 elif isinstance(V, (tuple, list)):
  40                     if len(V) == 1:
  41                         if isinstance(V[0], Symbol):
  42                             limits.append(V[0])
  43                             continue
  44                     elif len(V) in (2, 3):
  45                         if isinstance(V[0], Symbol):
  46                             limits.append(tuple(map(sympify, V)))
  47                             continue
  48
  49                 raise ValueError("Invalid summation variable or limits")
  50
  51         obj = Basic.__new__(cls, **assumptions)
  52         obj._args = (f, tuple(limits))
  53
  54         return obj
  55
  56     @property
  57     def function(self):
  58         return self._args[0]
  59
  60     @property
  61     def limits(self):
  62         return self._args[1]
  63
  64     def tostr(self, level=0):
  65         L = ', '.join([ str(L) for L in self.limits ])
  66         return 'Sum(%s, %s)' % (self.function.tostr(), L)
  67
  68     def doit(self, **hints):
  69         #if not hints.get('sums', True):
  70         #    return self
  71         f = self.function
  72         for i, a, b in self.limits:
  73             f = eval_sum(f, (i, a, b))
  74             if f is None:
  75                 return self
  76         return f
  77
  78     def _eval_summation(self, f, x):
  79         return
  80
  81 def sum(*args, **kwargs):
  82     summation = Sum(*args, **kwargs)
  83
  84     if isinstance(summation, Sum):
  85         return summation.doit()
  86     else:
  87         return summation
  88
  89
  90 def getab(expr):
  91     cls = expr.__class__
  92     return cls(expr.args[0]), cls(*expr.args[1:])
  93
  94 def eval_sum(f, (i, a, b)):
  95     if not f.has(i):
  96         return f*(b-a+1)
  97     definite = isinstance(a, C.Integer) and isinstance(b, C.Integer)
  98     # Doing it directly may be faster if there are very few terms.
  99     if definite and (b-a < 100):
 100         return eval_sum_direct(f, (i, a, b))
 101     # Try to do it symbolically. Even when the number of terms is known,
 102     # this can save time when b-a is big.
 103     value = eval_sum_symbolic(f.expand(), (i, a, b))
 104     if value is not None:
 105         return value
 106     # Do it directly
 107     if definite:
 108         return eval_sum_direct(f, (i, a, b))
 109
 110 def eval_sum_symbolic(f, (i, a, b)):
 111     if not f.has(i):
 112         return f*(b-a+1)
 113     # Linearity
 114     if isinstance(f, C.Mul):
 115         L, R = getab(f)
 116         if not L.has(i): return L*eval_sum_symbolic(R, (i, a, b))
 117         if not R.has(i): return R*eval_sum_symbolic(L, (i, a, b))
 118     if isinstance(f, C.Add):
 119         L, R = getab(f)
 120         lsum = eval_sum_symbolic(L, (i, a, b))
 121         rsum = eval_sum_symbolic(R, (i, a, b))
 122         if None not in (lsum, rsum):
 123             return lsum + rsum
 124     # Polynomial terms with Faulhaber's formula
 125     p = C.Wild('p')
 126     e = f.match(i**p)
 127     if e != None:
 128         c = p.subs_dict(e)
 129         B = C.bernoulli
 130         if c.is_integer and c >= 0:
 131             s = (B(c+1, b+1) - B(c+1, a))/(c+1)
 132             return s.expand()
 133     # Geometric terms
 134     if isinstance(f, C.Pow):
 135         r, k = f.args[:]
 136         if not r.has(i) and k == i:
 137             # TODO: Pow should be able to simplify x**oo depending
 138             # on whether |x| < 1 or |x| > 1 for non-rational x
 139             if (b is S.Infinity) and abs(r.evalf()) < 1:
 140                 return r**a / (1-r)
 141             else:
 142                 return (r**a - r**(b+1)) / (1-r)
 143     return None
 144
 145 def eval_sum_direct(expr, (i, a, b)):
 146     s = S.Zero
 147     if expr.has(i):
 148         for j in xrange(a, b+1):
 149             s += expr.subs(i, j)
 150     else:
 151         for j in xrange(a, b+1):
 152             s += expr
 153     return s