I have a collection of functions defined using Piecewise from the sympy library in python. They have various different expressions and conditions under which these expressions hold. How can I, for example, add them all together into one piecewise function, which will likely be more finely divided into intervals? With the code below, I have defined 3 piecewise functions:
a = Piecewise((2, x<2), (x,True))
b = Piecewise((x**2, x<0), (x,True))
c = Piecewise((x**2-x, x<3), (6, x<4), (x+2, True))
The sum of these should be as follows:
Piecewise((2x**2 - x + 2, x<0), (x**2 + 2, x<2), (x**2 + x, x<3), (2*x + 6, x<4), (3*x + 2, True))
I have tried the following (below) to reduce to the arguments, but it gets very messy when trying to identify the conditions where there will be a list of singular expressions to add up. Is there a simple way to do this?
functions = []
for f in (a,b,c):
functions.append(list(f.args))
You should not have to worry about whether an expression is defined in a piecewise fashion or not. Just add the expressions to get a sum of expressions (in this case a sum of Piecewise) and then, since there are Piecewise, use piecewise_fold to get the expression as a single Piecewise:
from sympy import *
a = Piecewise((2, x<2), (x,True))
b = Piecewise((x**2, x<0), (x,True))
c = Piecewise((x**2-x, x<3), (6, x<4), (x+2, True))
>>> a+b+c
Piecewise((2, x < 2), (x, True)) + Piecewise((x**2, x < 0), (x, True)) + Piecewise((x**2 - x, x < 3), (6, x < 4), (x + 2, True))
>>> piecewise_fold(_)
Piecewise((2*x**2 - x + 2, (x < 0) & (x < 2) & (x < 3)), (x**2 + 8, (x < 0) & (x < 2) & (x < 4)), (x**2 + x + 4, (x < 0) & (x < 2)), (2*x**2, (x < 0) & (x < 3)), (x**2 + x + 6, (x < 0) & (x < 4)), (x**2 + 2*x + 2, x < 0), (x**2 + 2, (x < 2) & (x < 3)), (x + 8, (x < 2) & (x < 4)), (2*x + 4, x < 2), (x**2 + x, x < 3), (2*x + 6, x < 4), (3*x + 2, True))
>>> simplify(_)
Piecewise(
(2*x**2 - x + 2, x < 0),
(x**2 + 2, x < 2),
(x*(x + 1), x < 3),
(2*x + 6, x < 4),
(3*x + 2, True))