Create a formal linear function in Sympy

Question

I have an expression in Sympy (like

-M - n + x(n)

) and I would like to create a formal linear function, says f, and apply it to my expression, in order to get, after simplification:

-f(M) - f(n) + f(x(n))

Is it possible to tell sympy that a property such as linearity is verified?

A very hacky way to do it would be to apply the function f to every subexpression which is in a sum. For instance when given an expressions like the first one I gave, it would be nice to simply access the terms appearing in the sum (here it would be

[-M, -n , x(n)]

Then mapping f on the list and sum it to get what is expected.

Is there an easy way to do so, or have I necessarily to go trough the syntactic tree of the expression ?

Answer 1

This works:

>>> x,f = map(Function, 'xf'); n,M = symbols('n,M'); expr = -M - n + x(n)
>>> Add(*[f(a) for a in Add.make_args(expr)])
f(-M) + f(-n) + f(x(n))

If you have an expression like f(n*(M + 1)) and you expand it you will get f(n*M + n). Can you tell SymPy to apply the function to the args of f's args? Yes:

>>> expr = f(n*(M + 1))
>>> expr.expand().replace(lambda x: x.func == f,
...         lambda x: Add(*[f(a) for a in Add.make_args(x.args[0])]))
f(n) + f(M*n)

If you call such a replacement linapp you can use it for any function that you want:

def linapp(expr, *f):
    return expr.expand().replace(
      lambda x: x.func in f,
      lambda x: Add(*[x.func(a) for a in Add.make_args(x.args[0])]))

>>> print(linapp(cos(x+y) + sin(x + y), cos, sin))
sin(x) + sin(y) + cos(x) + cos(y)

(Not saying that it's a true result, just that you can do it. And if you replace a variable with something else and you want to reapply the linearization, you can:

>>> linapp(_.subs(y, z + 1), cos)
sin(x) + sin(z + 1) + cos(x) + cos(z) + cos(1)