Sum over Piecewise, simplify

[albertz]

I have:

n = Symbol("n", integer=True, negative=False)
i = Symbol("i", integer=True, negative=False)

I expect these possible simplifications:

• Sum(Piecewise((1, Eq(i, 0)), (0, True)), (i, 0, n)) should simplify to 1.
• Sum(Piecewise((2, Eq(i, 0)), (1, True)), (i, 0, n)) should simplify to 2 + n.

SymPy does not do that.

[oscarbenjamin]

It's usually better to express assumptions as True rather than False although it is equivalent in this example (because you also have integer=True):

n = Symbol('n', integer=True, nonnegative=True)

Those sums should simplify ideally. Probably the way to fix this is to have the sum split into parts so

In [25]: Sum(Piecewise((f(i), Eq(i, 0)), (g(i), True)), (i, 0, n))
Out[25]: 
  n                   
 ____                 
 ╲                    
  ╲                   
   ╲  ⎧f(i)  for i = 0
   ╱  ⎨               
  ╱   ⎩g(i)  otherwise
 ╱                    
 ‾‾‾‾                 
i = 0                 

In [26]: Sum(Piecewise((f(i), Eq(i, 0)), (g(i), True)), (i, 0, n)).subs(n, 3).doit()
Out[26]: f(0) + g(1) + g(2) + g(3)

In [30]: s2 = Sum(f(i), (i, 0, 0)) + Sum(g(i), (i, 1, n))

In [31]: s2
Out[31]: 
  0            n       
 ___          ___      
 ╲            ╲        
  ╲            ╲       
  ╱   f(i) +   ╱   g(i)
 ╱            ╱        
 ‾‾‾          ‾‾‾      
i = 0        i = 1     

In [32]: s2.subs(n, 3).doit()
Out[32]: f(0) + g(1) + g(2) + g(3)

We also need a way to fix piecewise_fold so it handles a Sum of a Piecewise where the conditions depend on the summation variable:

In [9]: Sum(Piecewise((1, Eq(i, 0)), (0, True)), (i, 0, n))
Out[9]: 
  n                
 ____              
 ╲                 
  ╲                
   ╲  ⎧1  for i = 0
   ╱  ⎨            
  ╱   ⎩0  otherwise
 ╱                 
 ‾‾‾‾              
i = 0              

In [10]: Sum(Piecewise((1, Eq(i, 0)), (0, True)), (i, 0, n)).doit()
Out[10]: 
  n                
 ____              
 ╲                 
  ╲                
   ╲  ⎧1  for i = 0
   ╱  ⎨            
  ╱   ⎩0  otherwise
 ╱                 
 ‾‾‾‾              
i = 0              

In [11]: piecewise_fold(_)
Out[11]: 
⎧  n               
⎪ ___              
⎪ ╲                
⎪  ╲               
⎪  ╱   1  for i = 0
⎪ ╱                
⎪ ‾‾‾              
⎪i = 0             
⎨                  
⎪  n               
⎪ ___              
⎪ ╲                
⎪  ╲               
⎪  ╱   0  otherwise
⎪ ╱                
⎪ ‾‾‾              
⎩i = 0             

In [12]: _.doit()  # wrong!
Out[12]: 
⎧n + 1  for i = 0
⎨                
⎩  0    otherwise

[j3parera]

Just in case it's useful:
Sum(Piecewise((n, n > 0), (0, True)), (n, -oo, 0)).doit()
simplifies correctly, but
Sum(Piecewise((n, n >= 0), (0, True)), (n, -oo, 0)).doit()
does not.
Another example:
Sum(Heaviside(k, 1), (k, -oo, -1)).doit()
simplifies to 0, but
Sum(Heaviside(k, 1), (k, -oo, 0)).doit()
does not.