CODEWITHSUNDEEP
python-3.xsympy
Uwe.Schneider
Uwe.Schneider

Reputation: 1415

Sympy - Find all expressions of a certain level

in the documentation it is shown, how one can use srepr(expr) to create a string containing all elements of the formula tree.

I would like to have all elements of a certain level of such a tree in a list, not as a string but as sympy objects. For the example

from sympy import *
x,y,z = symbols('x y z')

expr = sin(x*y)/2 - x**2 + 1/y
s(repr) 
"""
gives out
=> "Add(Mul(Integer(-1), Pow(Symbol('x'), Integer(2))), Mul(Rational(1, 2),
sin(Mul(Symbol('x'), Symbol('y')))), Pow(Symbol('y'), Integer(-1)))"
"""

This should lead to

first_level     = [sin(x*y)/2,-x**2,1/y]
second_level[0] = [sin(x*y),2]

I assume, the fundamental question is, how to index a sympy expression according to the expression tree such that all higher levels starting from an arbitrary point are sumarized in a list element.

For the first level, it is possible to solve the task by getting the type of the first level (Add in this case and then run) 

print(list(Add.make_args(expr)))

but how does it work for expressions deeper in the tree?

Upvotes: 1

Views: 577

Answers (1)

Chris du Plessis
Chris du Plessis

Reputation: 1370

I think the thing you are looking for is .args. For example your example we have:

from sympy import *

x, y, z = symbols('x y z')

expr = sin(x * y) / 2 - x ** 2 + 1 / y

first_level = expr.args
second_level = [e.args for e in first_level]
print(first_level)
print(second_level[0])
"""
(1/y, sin(x*y)/2, -x**2)
(y, -1)
"""

These give tuples, but I trust that you can convert them to lists no problem.

If you want to get directly to the bottom of the tree, use .atoms() which produces a set:

print(expr.atoms())
"""
{x, 2, y, -1, 1/2}
"""

Another thing that seems interesting is perorder_traversal() which seems to give a tree kind of object and if you iterate through it, you get something like a power set:

print([e for e in preorder_traversal(expr)])
"""
[-x**2 + sin(x*y)/2 + 1/y, 1/y, y, -1, sin(x*y)/2, 1/2, sin(x*y), x*y, x, y, -x**2, -1, x**2, x, 2]
"""

Here is the beginning of the docstring of the function:

Do a pre-order traversal of a tree.

This iterator recursively yields nodes that it has visited in a pre-order fashion. That is, it yields the current node then descends through the tree breadth-first to yield all of a node's children's pre-order traversal.

For an expression, the order of the traversal depends on the order of .args, which in many cases can be arbitrary.

If any of these are not what you want then I probably don't understand your question correctly. Apologies for that.

Upvotes: 1

Related Questions