Sympy : simplification with expression substitution

Question

I have several expressions involving the norm or norm squared of a vector u. I'd like to simplify these expressions by substituting a known value for the norm of u. However, it seems that obvious expressions involving even simple multiples of the norm are not simplified.

As an example, this code does what I would expect :

import sympy as sp

u1,u2,u3 = sp.symbols('u_1, u_2, u_3',real=True,positive=True)

utu = u1**2 + u2**2 + u3**2

print("Ex. 1")
print(utu.subs(utu,1))

This produces the expected output

Ex. 1
1

However, 2*utu does not simplify in the way I would expect :

print("Ex 2")
print((2*utu).subs(utu,1))

Ex 2
2*u_1**2 + 2*u_2**2 + 2*u_3**2

I can explicitly force the substitution with this :

print("Ex 3")
print((2*utu).subs(2*utu,2))

which produces the expected output :

Ex 3
2

Ideally, I'd like to substitute under a norm function, but the run into the same issue.

u = sp.Matrix(3, 1, [u1,u2,u3])

print("Ex 4")
print(u.norm().subs(utu,1))

print("Ex 5")
print((2*u).norm().subs(utu,1))

print("Ex 6")
print((2*u).norm().subs(4*utu,4))

which produces

Ex 4
1
Ex 5
sqrt(4*u_1**2 + 4*u_2**2 + 4*u_3**2)
Ex 6
2

Are there tricks I am missing that will catch these obvious (to me at least - maybe not to Sympy?) simplifications? I've tried factor and expand, without much luck.

Answer 1

Let's analyze this expression:

expr = 2*utu
# out: 2*u_1**2 + 2*u_2**2 + 2*u_3**2

The multiplication has been evaluated. This is SymPy's default behavior: it evaluates things. We can work with the expression manipulation functions to achieve our goal.

For example:

expr = collect_const(expr)
# out: 2*(u_1**2 + u_2**2 + u_3**2)
expr.subs(utu, 1)
# out: 2

Another example:

expr = (2 * u).norm()
# out: sqrt(4*u_1**2 + 4*u_2**2 + 4*u_3**2)
expr = expr.simplify() # Note that expr.factor() produces the same result with this expression
# out: 2*sqrt(u_1**2 + u_2**2 + u_3**2)
expr.subs(utu, 1)
# out: 2

If you play (and modify) with these examples, you will realize that the same result can be achieved with different functions (factor, simplify, collect, collect_const, ...), but even one little change in the expression might prevent one function from "doing its work", while others might be able to. Expression manipulation is kind of an art that one should practice (a lot).

For completeness, I'm going to show you UnevaluatedExpr, which allows a particular expression to remain unevaluated during expression manipulation, though it might not always be the best choice:

n = UnevaluatedExpr(utu)
# out: u_1**2 + u_2**2 + u_3**2
expr = 4 * n
# out: 4*(u_1**2 + u_2**2 + u_3**2)

Note that SymPy didn't proceed with the full evaluation. Now:

expr.subs(utu, 1)
# out: 4*1

Why is there a 4*1 instead of 4? The 1 refers to the UnevaluateExpr object that we created earlier: to evaluate it we can use the doit() method:

expr.subs(utu, 1).doit()
# 4

Keep in mind that while using UnevaluateExpr, the expression becomes non-commutative (I think it's a bug with SymPy), which will prevent other functions to produce the expected results.

Answer 2

Substituting compound expressions is problematic. For the most part you should only expect subs to work if the expression to be replaced is known to always appear literally as part of the expression tree that you are substituting into. When possible then it is better to rearrange for a single symbol like:

In [10]: utu
Out[10]: 
  2     2     2
u?  + u?  + u? 

In [11]: (2*utu).subs(u1**2, 1 - u2**2 - u3**2)
Out[11]: 2

Even here we are substituting for a power of a symbol (u1**2) which is potentially fragile if we can't be sure that exactly that power will always appear in the expression. More generally there are functions that can simplify expressions based on knowing some polynomial relation e.g. ratsimpmodprime:

In [16]: e = (1 - u1**2) / (u1**2 + u2**2 + u3**2)

In [17]: e
Out[17]: 
          2    
    1 - u?     
???????????????
  2     2     2
u?  + u?  + u? 

In [18]: ratsimpmodprime(e, [u1**2 + u2**2 + u3**2 - 1])
Out[18]: 
  2     2
u?  + u? 

Other possibilities could be using resultants or Groebner bases to do similar things. Note that u.norm() has a square root which is symbolically awkward so it is better to work with the square of the norm (same as when working on pen and paper):

In [20]: ratsimpmodprime((2*u).norm()**2, [u1**2 + u2**2 + u3**2 - 1])
Out[20]: 4

Also if you just want a more powerful version of subs then you can use replace but with patterns:

In [21]: a = Wild('a')

In [22]: p = a*u1**2 + a*u2**2 + a*u3**2

In [23]: (2*utu).replace(p, a)
Out[23]: 2

In [24]: (2*u).norm().replace(p, a)
Out[24]: 2

Answer 3

Both solid answers already. If you have an arbitrary expression that you expect to be a factor in another, factor_terms is what I try first to make that factor appear. It will collect common factors without doing factoring. But if this doesn't work and you know you have a factor, div is a nice way to check and see the expression with the factor removed:

>>> expr = 2*(x + y)
>>> factor_terms(expr)
2*(x + y)
>>> e2 = expand(expr*(x -y))  # 2*x**2 - y**2
>>> factor_terms(e2)
2*(x**2 - y**2)
>>> div(_,-x-y)
(-2*x + 2*y, 0)
>>> _[0]*z  # if you wanted to replace factor -x-y with z
z*(-2*x + 2*y)