pythonsympy

Plotting the components of a general solution returned by sympy

I have written the following code that returns a solution.

import sympy as sym
import numpy as np

z = sym.Symbol('z')
e = sym.Symbol('e')
f = sym.Function('f')

edo = sym.diff(f(z), z , 2) + e * f(z)
soln = sym.dsolve(edo, f(z))
print(soln.rhs)

the above code returns:

C1*exp(-z*sqrt(-e)) + C2*exp(z*sqrt(-e))

I want to be able to access the elements of this 'soln.rhs' directly and plot them. I could copy and paste the results, but I want to do something like:

plot(x, soln.rhs[0])

which would plot exp(-z*sqrt(-e))

The reason for this is that I am analysing many different types of ODE, some of which return solutions which are combinations of Airy functions like

C1*airyai(-e + z) + C2*airybi(-e + z)

Does anyone know how to access the elements of the general solution? I have looked through the documentation and nothing really pops out.

Solution

  • What you are looking for is the args attribute of a Sympy's symbolic expression.

    For example:

    print(soln.rhs.args[0])
# C1*exp(-z*sqrt(-e))
print(soln.rhs.args[1])
# C2*exp(z*sqrt(-e))

    You might also want to insert appropriate values for the integrations constants by using the subs method:

    C1, C2 = symbols("C1, C2")
soln.rhs.subs({C1: 2, C2: 3}) # random numeric values to show how to do it.

    Then, you can plot it:

    # plot the solution for z in [0, 10]
plot(soln.rhs.subs({C1: 2, C2: 3}), (z, 0, 10))