In [1]:
from sympy import *
var('x y z')

Out[1]:
(x, y, z)

Context: Expressions in Sympy are stored as a 'tree' of functions applied to tuples of more expressions (an expression without a function would be a symbol, like x). For example, the expression

In [2]:
example = 1 + 2*x**3 - x/y


is stored as

In [3]:
srepr(example)

Out[3]:
"Add(Mul(Integer(2), Pow(Symbol('x'), Integer(3))), Mul(Integer(-1), Symbol('x'), Pow(Symbol('y'), Integer(-1))), Integer(1))"

or, as a tree,

    +----------- Sum ------------+
|             |              |
Mul            Mul             1
/   \         /  |  \
2    Pow     -1   x  Pow
/ \             /  \
x   3           y   -1

Warning: the expression is not printed in the same order as it is stored in the tree.

This notebook implements and demonstrates how to do advanced manual manipulation of expressions in Sympy.

# Modify expression tree in Sympy¶

part() and inpart() work in the same way as 'part' and 'inpart' of Maxima. They represent a way to navigate the expression tree and to replace part of it with a new expression.

In [4]:
def part(expr,address):
r"""
Returns part of an expression

Arguments
---------
expr : sympy expression
address : (list of integers) indexes of the part
of the expression tree to be recovered
Returns
-------
requested part of the expression tree
"""
for num in address:
expr = expr.args[num]
return expr

In [5]:
def inpart(expr,repl,address):
r"""
Replaces a part of the tree of an expression (and returns
the copy)

Arguments
---------
expr: (sympy expression) expression to be intervened
repl: (sympy expression) modified part of the expression
address: (list of integers) indexes of the part of the
expression tree to be replaced (see 'part()')
Returns
-------
new expression with the replacement done
"""
if len(address) == 1:
largs = list(expr.args)
return expr.func(*largs)
else:
largs = list(expr.args)
new = expr.func(*largs)
return new

In [17]:
def cpart(expr,address):
r"""
makes easier to visualize walking the tree. It returns a set of two expressions:
the original expression with the part located by 'address' substituted
by the symbol 'PIECE' and the part requested.
"""
PART = Symbol(r'{\color{red}{PART}}')


# Example: forced factorization and partial series expansion¶

Consider the following expression:

In [7]:
expr = 1 + x**2 + 1/(x+y)

In [8]:
expr

Out[8]:
$\displaystyle x^{2} + 1 + \frac{1}{x + y}$

Supose we want to force the denominator to have a factor $x$ out. None of the simplification or factorization functions will do the job. So, we intervene the expression manually.

First, we explore the expression tree until we get the desired part (in this case, the denominator of the fraction):

In [9]:
piece = part(expr,[2,0])
piece

Out[9]:
$\displaystyle x + y$

Now, we do the manual transformation of the term, and substitute back into the expression (a new expression is returned, since sympy expressions are immutable)

In [19]:
new = inpart(expr,expand(piece/x)*x,[2,0])
new

Out[19]:
$\displaystyle x^{2} + 1 + \frac{1}{x \left(1 + \frac{y}{x}\right)}$

The function cpart() helps to locate the part in an interactive notebook. It returns a set of two expressions. The original one has the part substituted by the symbol 'PART', and the actual part of the expression requested. This way, it is easier to visualize where in the expression tree are we, and trying different addresses is simpler. Caveat: the expression may be printed in a different order every time it's called.

Consider the following example. In the new expression, we set y as a small number, and we want to expand the expression in parenthesis. First, we locate the term 1/(1+y/x):

In [35]:
cpart(new,[2,1])

Out[35]:
$\displaystyle Set\left(x^{2} + 1 + \frac{{\color{red}{PART}}}{x}, \frac{1}{1 + \frac{y}{x}}\right)$

Now, we expand only that term up to second order and substitute back

In [36]:
inpart(new,series(part(new,[2,1]),y,0,2).removeO(),[2,1])

Out[36]:
$\displaystyle x^{2} + 1 + \frac{1 - \frac{y}{x}}{x}$
In [ ]: