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)
        largs[address[0]] = repl
        return expr.func(*largs)
    else:
        largs = list(expr.args)
        largs[address[0]] = inpart(expr.args[address[0]],repl,address[1:])
        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}}')
    return Set(inpart(expr,PART,address),part(expr,address))

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 [ ]: