Refined manipulation of sympy expressions

Question

I am trying to work with sympy and work with manipulation of expressions.

import sympy as sym
from sympy.abc import  t

x0,v0 = sym.symbols("x0 v0 ",  real=True)
wn = sym.symbols("omega_n",  positive = True, real=True)
z = sym.symbols("zeta",  positive = True, real=True)

x = sym.Function('x')
Dx = sym.Derivative(x(t), t)
Dx2= sym.Derivative(x(t), t,2)

res = sym.dsolve(Dx2 +2*z*wn*Dx+ wn**2*x(t), x(t), 
   ics = {  x(0): x0, 
            Dx.subs(t,0):v0})

The above yields the following expression

The part circled in red can be further simplified to , however I can't figure out how is it possible to simplify selected portions of the expression.

If I take a simpler exampler wn*x0*z**2/(2*wn*z**2-2*wn) then its possible to cancel out the terms with simplify(), but I could find anywhere some good documentation on how to work and substitute parts of the equation.

---

Another similar issue is with the term , which I would like to transform to

Answer 1

First here is the rhs of your equation:

In [58]: res.rhs
Out[58]: 
                                                                 ⎛      _______   _______⎞   ⎛          2                _______   _______                         _______   _______⎞       ⎛       _______   _______⎞
⎛           x₀⋅ζ           x₀              v₀           ⎞  -ωₙ⋅t⋅⎝ζ + ╲╱ ζ - 1 ⋅╲╱ ζ + 1 ⎠   ⎜   ωₙ⋅x₀⋅ζ       ωₙ⋅x₀⋅ζ⋅╲╱ ζ - 1 ⋅╲╱ ζ + 1        ωₙ⋅x₀        v₀⋅╲╱ ζ - 1 ⋅╲╱ ζ + 1 ⎟  ωₙ⋅t⋅⎝-ζ + ╲╱ ζ - 1 ⋅╲╱ ζ + 1 ⎠
⎜- ───────────────────── + ── - ────────────────────────⎟⋅ℯ                                + ⎜────────────── + ─────────────────────────── - ────────────── + ──────────────────────⎟⋅ℯ                               
⎜      _______   _______   2           _______   _______⎟                                    ⎜      2                       2                      2                    2           ⎟                                 
⎝  2⋅╲╱ ζ - 1 ⋅╲╱ ζ + 1         2⋅ωₙ⋅╲╱ ζ - 1 ⋅╲╱ ζ + 1 ⎠                                    ⎝2⋅ωₙ⋅ζ  - 2⋅ωₙ          2⋅ωₙ⋅ζ  - 2⋅ωₙ         2⋅ωₙ⋅ζ  - 2⋅ωₙ       2⋅ωₙ⋅ζ  - 2⋅ωₙ    ⎠       

For the first part of your question what we can do is collect on wn for the rhs of the equation. This has the effect of cancelling out the wn in the terms where they can cancel:

In [45]: res.rhs.collect(wn)
Out[45]: 
                                                                 ⎛      _______   _______⎞   ⎛     2            _______   _______                   _______   _______⎞       ⎛       _______   _______⎞
⎛           x₀⋅ζ           x₀              v₀           ⎞  -ωₙ⋅t⋅⎝ζ + ╲╱ ζ - 1 ⋅╲╱ ζ + 1 ⎠   ⎜ x₀⋅ζ      x₀⋅ζ⋅╲╱ ζ - 1 ⋅╲╱ ζ + 1       x₀      v₀⋅╲╱ ζ - 1 ⋅╲╱ ζ + 1 ⎟  ωₙ⋅t⋅⎝-ζ + ╲╱ ζ - 1 ⋅╲╱ ζ + 1 ⎠
⎜- ───────────────────── + ── - ────────────────────────⎟⋅ℯ                                + ⎜──────── + ──────────────────────── - ──────── + ──────────────────────⎟⋅ℯ                               
⎜      _______   _______   2           _______   _______⎟                                    ⎜   2                  2                  2              ⎛   2    ⎞     ⎟                                 
⎝  2⋅╲╱ ζ - 1 ⋅╲╱ ζ + 1         2⋅ωₙ⋅╲╱ ζ - 1 ⋅╲╱ ζ + 1 ⎠                                    ⎝2⋅ζ  - 2           2⋅ζ  - 2           2⋅ζ  - 2       ωₙ⋅⎝2⋅ζ  - 2⎠     ⎠                                 

Now we see there are various powers that could be simplified by combining them. In general sqrt(z-1)*sqrt(z+1) is not necessarily equal to sqrt(z**2 - 1) but the assumption that z is positive means that they are equal. In that case powsimp should combine those powers:

In [46]: res.rhs.collect(wn).powsimp()
Out[46]: 
                                                                             ⎛                   ________                    ________⎞                                 
                                                 ⎛      _______   _______⎞   ⎜     2            ╱  2                        ╱  2     ⎟       ⎛       _______   _______⎞
⎛       x₀⋅ζ       x₀          v₀       ⎞  -ωₙ⋅t⋅⎝ζ + ╲╱ ζ - 1 ⋅╲╱ ζ + 1 ⎠   ⎜ x₀⋅ζ      x₀⋅ζ⋅╲╱  ζ  - 1       x₀      v₀⋅╲╱  ζ  - 1 ⎟  ωₙ⋅t⋅⎝-ζ + ╲╱ ζ - 1 ⋅╲╱ ζ + 1 ⎠
⎜- ───────────── + ── - ────────────────⎟⋅ℯ                                + ⎜──────── + ──────────────── - ──────── + ──────────────⎟⋅ℯ                               
⎜       ________   2            ________⎟                                    ⎜   2              2              2          ⎛   2    ⎞ ⎟                                 
⎜      ╱  2                    ╱  2     ⎟                                    ⎝2⋅ζ  - 2       2⋅ζ  - 2       2⋅ζ  - 2   ωₙ⋅⎝2⋅ζ  - 2⎠ ⎠                                 
⎝  2⋅╲╱  ζ  - 1         2⋅ωₙ⋅╲╱  ζ  - 1 ⎠                                                                                                                              

Here powsimp didn't apply to the exponents so we need to use deep=True to tell it to look deeper for things to combine:

In [47]: res.rhs.collect(wn).powsimp(deep=True)
Out[47]: 
                                                 ⎛       ________⎞   ⎛                   ________                    ________⎞       ⎛        ________⎞
                                                 ⎜      ╱  2     ⎟   ⎜     2            ╱  2                        ╱  2     ⎟       ⎜       ╱  2     ⎟
⎛       x₀⋅ζ       x₀          v₀       ⎞  -ωₙ⋅t⋅⎝ζ + ╲╱  ζ  - 1 ⎠   ⎜ x₀⋅ζ      x₀⋅ζ⋅╲╱  ζ  - 1       x₀      v₀⋅╲╱  ζ  - 1 ⎟  ωₙ⋅t⋅⎝-ζ + ╲╱  ζ  - 1 ⎠
⎜- ───────────── + ── - ────────────────⎟⋅ℯ                        + ⎜──────── + ──────────────── - ──────── + ──────────────⎟⋅ℯ                       
⎜       ________   2            ________⎟                            ⎜   2              2              2          ⎛   2    ⎞ ⎟                         
⎜      ╱  2                    ╱  2     ⎟                            ⎝2⋅ζ  - 2       2⋅ζ  - 2       2⋅ζ  - 2   ωₙ⋅⎝2⋅ζ  - 2⎠ ⎠                         
⎝  2⋅╲╱  ζ  - 1         2⋅ωₙ⋅╲╱  ζ  - 1 ⎠                                                                                                              

There are common factors like 2 that we can extract with factor_terms:

In [48]: factor_terms(res.rhs.collect(wn).powsimp(deep=True))
Out[48]: 
                                             ⎛       ________⎞                                                          ⎛        ________⎞
                                             ⎜      ╱  2     ⎟   ⎛    2                                         ⎞       ⎜       ╱  2     ⎟
⎛      x₀⋅ζ                 v₀      ⎞  -ωₙ⋅t⋅⎝ζ + ╲╱  ζ  - 1 ⎠   ⎜x₀⋅ζ         x₀⋅ζ        x₀           v₀      ⎟  ωₙ⋅t⋅⎝-ζ + ╲╱  ζ  - 1 ⎠
⎜- ─────────── + x₀ - ──────────────⎟⋅ℯ                        + ⎜────── + ─────────── - ────── + ──────────────⎟⋅ℯ                       
⎜     ________              ________⎟                            ⎜ 2          ________    2             ________⎟                         
⎜    ╱  2                  ╱  2     ⎟                            ⎜ζ  - 1     ╱  2        ζ  - 1        ╱  2     ⎟                         
⎝  ╲╱  ζ  - 1         ωₙ⋅╲╱  ζ  - 1 ⎠                            ⎝         ╲╱  ζ  - 1             ωₙ⋅╲╱  ζ  - 1 ⎠                         
──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
                                                                    2                                                                     

Now we can collect on powers of z**2 - 1 to reduce the repeating subexpressions:

In [55]: factor_terms(res.rhs.collect(wn).powsimp(deep=True)).collect(z**2 - 1)
Out[55]: 
⎛              v₀⎞        ⎛       ________⎞   ⎛                     v₀ ⎞       ⎛        ________⎞
⎜      -x₀⋅ζ - ──⎟        ⎜      ╱  2     ⎟   ⎜    2         x₀⋅ζ + ── ⎟       ⎜       ╱  2     ⎟
⎜              ωₙ⎟  -ωₙ⋅t⋅⎝ζ + ╲╱  ζ  - 1 ⎠   ⎜x₀⋅ζ  - x₀           ωₙ ⎟  ωₙ⋅t⋅⎝-ζ + ╲╱  ζ  - 1 ⎠
⎜x₀ + ───────────⎟⋅ℯ                          ⎜────────── + ───────────⎟⋅ℯ                       
⎜        ________⎟                            ⎜   2            ________⎟                         
⎜       ╱  2     ⎟                            ⎜  ζ  - 1       ╱  2     ⎟                         
⎝     ╲╱  ζ  - 1 ⎠                            ⎝             ╲╱  ζ  - 1 ⎠                         
─────────────────────────────────────────── + ───────────────────────────────────────────────────
                     2                                                 2 

Finally applying factor_terms to the coefficients of that last collect allows us to cancel the one remaining awkward term and normalise the minus signs:

In [56]: factor_terms(res.rhs.collect(wn).powsimp(deep=True)).collect(z**2 - 1, factor_terms)
Out[56]: 
⎛             v₀ ⎞        ⎛       ________⎞   ⎛             v₀ ⎞       ⎛        ________⎞
⎜      x₀⋅ζ + ── ⎟        ⎜      ╱  2     ⎟   ⎜      x₀⋅ζ + ── ⎟       ⎜       ╱  2     ⎟
⎜             ωₙ ⎟  -ωₙ⋅t⋅⎝ζ + ╲╱  ζ  - 1 ⎠   ⎜             ωₙ ⎟  ωₙ⋅t⋅⎝-ζ + ╲╱  ζ  - 1 ⎠
⎜x₀ - ───────────⎟⋅ℯ                          ⎜x₀ + ───────────⎟⋅ℯ                       
⎜        ________⎟                            ⎜        ________⎟                         
⎜       ╱  2     ⎟                            ⎜       ╱  2     ⎟                         
⎝     ╲╱  ζ  - 1 ⎠                            ⎝     ╲╱  ζ  - 1 ⎠                         
─────────────────────────────────────────── + ───────────────────────────────────────────
                     2                                             2

Answer 2

One way is to use pattern matching. Note that your circled term is a multiplication, containing z**2 (a power operation).

# search the expression tree and select all multiplications
# containing a power with exponent 2
w = sym.Wild("w", properties=[
    lambda e: e.is_Mul and any(t.is_Pow and t.exp == 2 for t in e.args)
])
t = list(res.find(w))[0]
print(t)
# out: omega_n*x0*zeta**2/(2*omega_n*zeta**2 - 2*omega_n)

# Perform the simplification and substitution
res = res.subs(t, t.simplify())

Now let's look for the second term in your question:

# loop over each exponential term and apply a
# powsimp to its argument.
for t in res.find(sym.exp):
    res = res.subs(t, sym.exp(t.args[0].powsimp()))
print(res)

Edit: similarly, you can apply the same technique to simplify the other square roots on your expression.