Chapter 9: Ordinary differential equations

Robert Johansson

Source code listings for Numerical Python - Scientific Computing and Data Science Applications with Numpy, SciPy and Matplotlib (ISBN 978-1-484242-45-2).

In [1]:
import numpy as np
In [2]:
%matplotlib inline
%config InlineBackend.figure_format='retina'
import matplotlib.pyplot as plt
import matplotlib as mpl
# mpl.rcParams['text.usetex'] = True
mpl.rcParams['mathtext.fontset'] = 'stix'
mpl.rcParams['font.family'] = 'serif'
mpl.rcParams['font.sans-serif'] = 'stix'
In [3]:
import sympy
sympy.init_printing()
In [4]:
from scipy import integrate

Symbolic ODE solving with SymPy

Newton's law of cooling

In [5]:
t, k, T0, Ta = sympy.symbols("t, k, T_0, T_a")
In [6]:
T = sympy.Function("T")
In [7]:
ode = T(t).diff(t) + k*(T(t) - Ta)
In [8]:
ode
Out[8]:
$\displaystyle k \left(- T_{a} + T{\left(t \right)}\right) + \frac{d}{d t} T{\left(t \right)}$
In [9]:
ode_sol = sympy.dsolve(ode)
In [10]:
ode_sol
Out[10]:
$\displaystyle T{\left(t \right)} = C_{1} e^{- k t} + T_{a}$
In [11]:
ode_sol.lhs
Out[11]:
$\displaystyle T{\left(t \right)}$
In [12]:
ode_sol.rhs
Out[12]:
$\displaystyle C_{1} e^{- k t} + T_{a}$
In [13]:
ics = {T(0): T0}
In [14]:
ics
Out[14]:
$\displaystyle \left\{ T{\left(0 \right)} : T_{0}\right\}$
In [15]:
C_eq = ode_sol.subs(t, 0).subs(ics)
In [16]:
C_eq
Out[16]:
$\displaystyle T_{0} = C_{1} + T_{a}$
In [17]:
C_sol = sympy.solve(C_eq)
In [18]:
C_sol
Out[18]:
$\displaystyle \left[ \left\{ C_{1} : T_{0} - T_{a}\right\}\right]$
In [19]:
ode_sol.subs(C_sol[0])
Out[19]:
$\displaystyle T{\left(t \right)} = T_{a} + \left(T_{0} - T_{a}\right) e^{- k t}$

Function for applying initial conditions

In [20]:
def apply_ics(sol, ics, x, known_params):
    """
    Apply the initial conditions (ics), given as a dictionary on
    the form ics = {y(0): y0: y(x).diff(x).subs(x, 0): yp0, ...}
    to the solution of the ODE with indepdendent variable x.
    The undetermined integration constants C1, C2, ... are extracted
    from the free symbols of the ODE solution, excluding symbols in
    the known_params list.
    """
    free_params = sol.free_symbols - set(known_params)
    eqs = [(sol.lhs.diff(x, n) - sol.rhs.diff(x, n)).subs(x, 0).subs(ics)
           for n in range(len(ics))]
    sol_params = sympy.solve(eqs, free_params)
    return sol.subs(sol_params)
In [21]:
ode_sol
Out[21]:
$\displaystyle T{\left(t \right)} = C_{1} e^{- k t} + T_{a}$
In [22]:
apply_ics(ode_sol, ics, t, [k, Ta])
Out[22]:
$\displaystyle T{\left(t \right)} = T_{a} + \left(T_{0} - T_{a}\right) e^{- k t}$
In [23]:
ode_sol = apply_ics(ode_sol, ics, t, [k, Ta]).simplify()
In [24]:
ode_sol
Out[24]:
$\displaystyle T{\left(t \right)} = \left(T_{0} + T_{a} e^{k t} - T_{a}\right) e^{- k t}$
In [25]:
y_x = sympy.lambdify((t, k), ode_sol.rhs.subs({T0: 5, Ta: 1}), 'numpy')
In [26]:
fig, ax = plt.subplots(figsize=(8, 4))

x = np.linspace(0, 4, 100)

for k in [1, 2, 3]:
    ax.plot(x, y_x(x, k), label=r"$k=%d$" % k)

ax.set_title(r"$%s$" % sympy.latex(ode_sol), fontsize=18)
ax.set_xlabel(r"$x$", fontsize=18)
ax.set_ylabel(r"$y$", fontsize=18)
ax.legend()

fig.tight_layout()

Damped harmonic oscillator

In [27]:
t, omega0 = sympy.symbols("t, omega_0", positive=True)
gamma = sympy.symbols("gamma", complex=True)
In [28]:
x = sympy.Function("x")
In [29]:
ode = x(t).diff(t, 2) + 2 * gamma * omega0 * x(t).diff(t) + omega0**2 * x(t)
In [30]:
ode
Out[30]:
$\displaystyle 2 \gamma \omega_{0} \frac{d}{d t} x{\left(t \right)} + \omega_{0}^{2} x{\left(t \right)} + \frac{d^{2}}{d t^{2}} x{\left(t \right)}$
In [31]:
ode_sol = sympy.dsolve(ode)
In [32]:
ode_sol
Out[32]:
$\displaystyle x{\left(t \right)} = C_{1} e^{\omega_{0} t \left(- \gamma - \sqrt{\gamma^{2} - 1}\right)} + C_{2} e^{\omega_{0} t \left(- \gamma + \sqrt{\gamma^{2} - 1}\right)}$
In [33]:
ics = {x(0): 1, x(t).diff(t).subs(t, 0): 0}
In [34]:
ics
Out[34]:
$\displaystyle \left\{ x{\left(0 \right)} : 1, \ \left. \frac{d}{d t} x{\left(t \right)} \right|_{\substack{ t=0 }} : 0\right\}$
In [35]:
x_t_sol = apply_ics(ode_sol, ics, t, [omega0, gamma])
In [36]:
x_t_sol
Out[36]:
$\displaystyle x{\left(t \right)} = \left(- \frac{\gamma}{2 \sqrt{\gamma^{2} - 1}} + \frac{1}{2}\right) e^{\omega_{0} t \left(- \gamma - \sqrt{\gamma^{2} - 1}\right)} + \left(\frac{\gamma}{2 \sqrt{\gamma^{2} - 1}} + \frac{1}{2}\right) e^{\omega_{0} t \left(- \gamma + \sqrt{\gamma^{2} - 1}\right)}$
In [37]:
x_t_critical = sympy.limit(x_t_sol.rhs, gamma, 1)
In [38]:
x_t_critical
Out[38]:
$\displaystyle \left(\omega_{0} t + 1\right) e^{- \omega_{0} t}$
In [39]:
fig, ax = plt.subplots(figsize=(8, 4))

tt = np.linspace(0, 3, 250)
for g in [0.1, 0.5, 1, 2.0, 5.0]:
    if g == 1:
        x_t = sympy.lambdify(t, x_t_critical.subs({omega0: 2.0 * sympy.pi}), 'numpy')
    else:
        x_t = sympy.lambdify(t, x_t_sol.rhs.subs({omega0: 2.0 * sympy.pi, gamma: g}), 'numpy')
    ax.plot(tt, x_t(tt).real, label=r"$\gamma = %.1f$" % g)

ax.set_xlabel(r"$t$", fontsize=18)
ax.set_ylabel(r"$x(t)$", fontsize=18)
ax.set_xlim(0, 3)
ax.legend()

fig.tight_layout()
fig.savefig('ch9-harmonic-oscillator.pdf')
In [40]:
# example of equation that sympy cannot solve
In [41]:
x = sympy.symbols("x")
In [42]:
y = sympy.Function("y")
In [43]:
f = y(x)**2 + x
In [44]:
sympy.Eq(y(x).diff(x, x), f)
Out[44]:
$\displaystyle \frac{d^{2}}{d x^{2}} y{\left(x \right)} = x + y^{2}{\left(x \right)}$
In [45]:
# sympy is unable to solve this differential equation
sympy.dsolve(y(x).diff(x, x) - f)
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-45-811796dbc135> in <module>
      1 # sympy is unable to solve this differential equation
----> 2 sympy.dsolve(y(x).diff(x, x) - f)

~/miniconda3/envs/py3.6/lib/python3.6/site-packages/sympy/solvers/ode.py in dsolve(eq, func, hint, simplify, ics, xi, eta, x0, n, **kwargs)
    639         hints = _desolve(eq, func=func,
    640             hint=hint, simplify=True, xi=xi, eta=eta, type='ode', ics=ics,
--> 641             x0=x0, n=n, **kwargs)
    642 
    643         eq = hints.pop('eq', eq)

~/miniconda3/envs/py3.6/lib/python3.6/site-packages/sympy/solvers/deutils.py in _desolve(eq, func, hint, ics, simplify, **kwargs)
    233             raise ValueError(string + str(eq) + " does not match hint " + hint)
    234         else:
--> 235             raise NotImplementedError(dummy + "solve" + ": Cannot solve " + str(eq))
    236     if hint == 'default':
    237         return _desolve(eq, func, ics=ics, hint=hints['default'], simplify=simplify,

NotImplementedError: solve: Cannot solve -x - y(x)**2 + Derivative(y(x), (x, 2))

Direction fields

In [46]:
def  plot_direction_field(x, y_x, f_xy, x_lim=(-5, 5), y_lim=(-5, 5), ax=None):
    
    f_np = sympy.lambdify((x, y_x), f_xy, 'numpy')
    
    x_vec = np.linspace(x_lim[0], x_lim[1], 20)
    y_vec = np.linspace(y_lim[0], y_lim[1], 20)
    
    if ax is None:
        _, ax = plt.subplots(figsize=(4, 4))

    dx = x_vec[1] - x_vec[0]
    dy = y_vec[1] - y_vec[0]

    for m, xx in enumerate(x_vec):
        for n, yy in enumerate(y_vec):
            Dy = f_np(xx, yy) * dx
            Dx = 0.8 * dx**2 / np.sqrt(dx**2 + Dy**2)
            Dy = 0.8 * Dy*dy / np.sqrt(dx**2 + Dy**2)
            ax.plot([xx - Dx/2, xx + Dx/2],
                    [yy - Dy/2, yy + Dy/2], 'b', lw=0.5)
    ax.axis('tight')

    ax.set_title(r"$%s$" %
                 (sympy.latex(sympy.Eq(y(x).diff(x), f_xy))),
                 fontsize=18)
    
    return ax
In [47]:
x = sympy.symbols("x")
In [48]:
y = sympy.Function("y")
In [49]:
fig, axes = plt.subplots(1, 3, figsize=(12, 4))

plot_direction_field(x, y(x), y(x)**2 + x, ax=axes[0])
plot_direction_field(x, y(x), -x / y(x), ax=axes[1])
plot_direction_field(x, y(x), y(x)**2 / x, ax=axes[2])

fig.tight_layout()
fig.savefig('ch9-direction-field.pdf')

Inexact solutions to ODEs

In [50]:
x = sympy.symbols("x")
In [51]:
y = sympy.Function("y")
In [52]:
f = y(x)**2 + x
#f = sympy.cos(y(x)**2) + x
#f = sympy.sqrt(y(x)) * sympy.cos(x**2)
#f = y(x)**2 / x
In [53]:
sympy.Eq(y(x).diff(x), f)
Out[53]:
$\displaystyle \frac{d}{d x} y{\left(x \right)} = x + y^{2}{\left(x \right)}$
In [54]:
ics = {y(0): 0}
In [55]:
ode_sol = sympy.dsolve(y(x).diff(x) - f)
In [56]:
ode_sol
Out[56]:
$\displaystyle y{\left(x \right)} = \frac{x^{2} \left(2 C_{1}^{3} + 1\right)}{2} + \frac{x^{5} \left(10 C_{1}^{3} \left(6 C_{1}^{3} + 1\right) + 20 C_{1}^{3} + 3\right)}{60} + C_{1} + \frac{C_{1} x^{3} \left(3 C_{1}^{3} + 1\right)}{3} + C_{1}^{2} x + \frac{C_{1}^{2} x^{4} \left(12 C_{1}^{3} + 5\right)}{12} + O\left(x^{6}\right)$
In [57]:
ode_sol = apply_ics(ode_sol, {y(0): 0}, x, [])
In [58]:
ode_sol
Out[58]:
$\displaystyle y{\left(x \right)} = \frac{x^{2}}{2} + \frac{x^{5}}{20} + O\left(x^{6}\right)$
In [59]:
ode_sol = sympy.dsolve(y(x).diff(x) - f, ics=ics)
In [60]:
ode_sol
Out[60]:
$\displaystyle y{\left(x \right)} = \frac{x^{2}}{2} + \frac{x^{5}}{20} + O\left(x^{6}\right)$
In [61]:
fig, axes = plt.subplots(1, 2, figsize=(8, 4))

plot_direction_field(x, y(x), f, ax=axes[0])
x_vec = np.linspace(-3, 3, 100)
axes[0].plot(x_vec, sympy.lambdify(x, ode_sol.rhs.removeO())(x_vec), 'b', lw=2)
axes[0].set_ylim(-5, 5)

plot_direction_field(x, y(x), f, ax=axes[1])
x_vec = np.linspace(-1, 1, 100)
axes[1].plot(x_vec, sympy.lambdify(x, ode_sol.rhs.removeO())(x_vec), 'b', lw=2)

ode_sol_m = ode_sol_p = ode_sol
dx = 0.125
for x0 in np.arange(1, 2., dx):
    x_vec = np.linspace(x0, x0 + dx, 100)
    ics = {y(x0): ode_sol_p.rhs.removeO().subs(x, x0)}
    ode_sol_p = sympy.dsolve(y(x).diff(x) - f, ics=ics, n=6)
    axes[1].plot(x_vec, sympy.lambdify(x, ode_sol_p.rhs.removeO())(x_vec), 'r', lw=2)

for x0 in np.arange(-1, -5, -dx):
    x_vec = np.linspace(x0, x0 - dx, 100)
    ics = {y(x0): ode_sol_m.rhs.removeO().subs(x, x0)}
    ode_sol_m = sympy.dsolve(y(x).diff(x) - f, ics=ics, n=6)
    axes[1].plot(x_vec, sympy.lambdify(x, ode_sol_m.rhs.removeO())(x_vec), 'r', lw=2)
    
fig.tight_layout()
fig.savefig("ch9-direction-field-and-approx-sol.pdf")

Laplace transformation method

In [62]:
t = sympy.symbols("t", positive=True)
In [63]:
s, Y = sympy.symbols("s, Y", real=True)
In [64]:
y = sympy.Function("y")
In [65]:
ode = y(t).diff(t, 2) + 2 * y(t).diff(t) + 10 * y(t) - 2 * sympy.sin(3*t)
In [66]:
ode
Out[66]:
$\displaystyle 10 y{\left(t \right)} - 2 \sin{\left(3 t \right)} + 2 \frac{d}{d t} y{\left(t \right)} + \frac{d^{2}}{d t^{2}} y{\left(t \right)}$
In [67]:
L_y = sympy.laplace_transform(y(t), t, s)
In [68]:
L_y
Out[68]:
$\displaystyle \mathcal{L}_{t}\left[y{\left(t \right)}\right]\left(s\right)$
In [69]:
L_ode = sympy.laplace_transform(ode, t, s, noconds=True)
In [70]:
L_ode
Out[70]:
$\displaystyle 10 \mathcal{L}_{t}\left[y{\left(t \right)}\right]\left(s\right) + 2 \mathcal{L}_{t}\left[\frac{d}{d t} y{\left(t \right)}\right]\left(s\right) + \mathcal{L}_{t}\left[\frac{d^{2}}{d t^{2}} y{\left(t \right)}\right]\left(s\right) - \frac{6}{s^{2} + 9}$
In [71]:
def laplace_transform_derivatives(e):
    """
    Evaluate the laplace transforms of derivatives of functions
    """
    if isinstance(e, sympy.LaplaceTransform):
        if isinstance(e.args[0], sympy.Derivative):
            d, t, s = e.args
            n = d.args[1][1]
            return ((s**n) * sympy.LaplaceTransform(d.args[0], t, s) - 
                    sum([s**(n-i) * sympy.diff(d.args[0], t, i-1).subs(t, 0)
                         for i in range(1, n+1)]))
        
    if isinstance(e, (sympy.Add, sympy.Mul)):
        t = type(e)
        return t(*[laplace_transform_derivatives(arg) for arg in e.args])
    
    return e
In [72]:
L_ode
Out[72]:
$\displaystyle 10 \mathcal{L}_{t}\left[y{\left(t \right)}\right]\left(s\right) + 2 \mathcal{L}_{t}\left[\frac{d}{d t} y{\left(t \right)}\right]\left(s\right) + \mathcal{L}_{t}\left[\frac{d^{2}}{d t^{2}} y{\left(t \right)}\right]\left(s\right) - \frac{6}{s^{2} + 9}$
In [73]:
L_ode_2 = laplace_transform_derivatives(L_ode)
In [74]:
L_ode_2
Out[74]:
$\displaystyle s^{2} \mathcal{L}_{t}\left[y{\left(t \right)}\right]\left(s\right) + 2 s \mathcal{L}_{t}\left[y{\left(t \right)}\right]\left(s\right) - s y{\left(0 \right)} + 10 \mathcal{L}_{t}\left[y{\left(t \right)}\right]\left(s\right) - 2 y{\left(0 \right)} - \left. \frac{d}{d t} y{\left(t \right)} \right|_{\substack{ t=0 }} - \frac{6}{s^{2} + 9}$
In [75]:
L_ode_3 = L_ode_2.subs(L_y, Y)
In [76]:
L_ode_3
Out[76]:
$\displaystyle Y s^{2} + 2 Y s + 10 Y - s y{\left(0 \right)} - 2 y{\left(0 \right)} - \left. \frac{d}{d t} y{\left(t \right)} \right|_{\substack{ t=0 }} - \frac{6}{s^{2} + 9}$
In [77]:
ics = {y(0): 1, y(t).diff(t).subs(t, 0): 0}
In [78]:
ics
Out[78]:
$\displaystyle \left\{ y{\left(0 \right)} : 1, \ \left. \frac{d}{d t} y{\left(t \right)} \right|_{\substack{ t=0 }} : 0\right\}$
In [79]:
L_ode_4 = L_ode_3.subs(ics)
In [80]:
L_ode_4
Out[80]:
$\displaystyle Y s^{2} + 2 Y s + 10 Y - s - 2 - \frac{6}{s^{2} + 9}$
In [81]:
Y_sol = sympy.solve(L_ode_4, Y)
In [82]:
Y_sol
Out[82]:
$\displaystyle \left[ \frac{s^{3} + 2 s^{2} + 9 s + 24}{s^{4} + 2 s^{3} + 19 s^{2} + 18 s + 90}\right]$
In [83]:
sympy.apart(Y_sol[0])
Out[83]:
$\displaystyle - \frac{6 \left(2 s - 1\right)}{37 \left(s^{2} + 9\right)} + \frac{49 s + 92}{37 \left(s^{2} + 2 s + 10\right)}$
In [84]:
y_sol = sympy.inverse_laplace_transform(Y_sol[0], s, t)
In [85]:
sympy.simplify(y_sol)
Out[85]:
$\displaystyle \frac{\left(6 e^{t} \sin{\left(3 t \right)} - 36 e^{t} \cos{\left(3 t \right)} + 43 \sin{\left(3 t \right)} + 147 \cos{\left(3 t \right)}\right) e^{- t}}{111}$
In [86]:
sympy.simplify(y_sol).evalf()
Out[86]:
$\displaystyle 0.00900900900900901 \left(6.0 e^{t} \sin{\left(3 t \right)} - 36.0 e^{t} \cos{\left(3 t \right)} + 43.0 \sin{\left(3 t \right)} + 147.0 \cos{\left(3 t \right)}\right) e^{- t}$
In [87]:
y_t = sympy.lambdify(t, y_sol.evalf(), 'numpy')
In [88]:
fig, ax = plt.subplots(figsize=(8, 4))

tt = np.linspace(0, 10, 500)
ax.plot(tt, y_t(tt).real)
ax.set_xlabel(r"$t$", fontsize=18)
ax.set_ylabel(r"$y(t)$", fontsize=18)
fig.tight_layout()

Numerical integration of ODEs using SciPy

In [89]:
x = sympy.symbols("x")
In [90]:
y = sympy.Function("y")
In [91]:
f = y(x)**2 + x
In [92]:
f_np = sympy.lambdify((y(x), x), f, 'math')
In [93]:
y0 = 0
In [94]:
xp = np.linspace(0, 1.9, 100)
In [95]:
xp.shape
Out[95]:
$\displaystyle \left( 100\right)$
In [96]:
yp = integrate.odeint(f_np, y0, xp)
In [97]:
xm = np.linspace(0, -5, 100)
In [98]:
ym = integrate.odeint(f_np, y0, xm)
In [99]:
fig, ax = plt.subplots(1, 1, figsize=(4, 4))
plot_direction_field(x, y(x), f, ax=ax)
ax.plot(xm, ym, 'b', lw=2)
ax.plot(xp, yp, 'r', lw=2)
fig.savefig('ch9-odeint-single-eq-example.pdf')

Lotka-Volterra equations for predator/pray populations

$$ x'(t) = a x - b x y $$$$ y'(t) = c x y - d y $$
In [100]:
a, b, c, d = 0.4, 0.002, 0.001, 0.7
In [101]:
def f(xy, t):
    x, y = xy
    return [a * x - b * x * y,
            c * x * y - d * y]
In [102]:
xy0 = [600, 400]
In [103]:
t = np.linspace(0, 50, 250)
In [104]:
xy_t = integrate.odeint(f, xy0, t)
In [105]:
xy_t.shape
Out[105]:
$\displaystyle \left( 250, \ 2\right)$
In [106]:
fig, axes = plt.subplots(1, 2, figsize=(8, 4))

axes[0].plot(t, xy_t[:,0], 'r', label="Prey")
axes[0].plot(t, xy_t[:,1], 'b', label="Predator")
axes[0].set_xlabel("Time")
axes[0].set_ylabel("Number of animals")
axes[0].legend()

axes[1].plot(xy_t[:,0], xy_t[:,1], 'k')
axes[1].set_xlabel("Number of prey")
axes[1].set_ylabel("Number of predators")
fig.tight_layout()
fig.savefig('ch9-lokta-volterra.pdf')

Lorenz equations

$$ x'(t) = \sigma(y - x) $$$$ y'(t) = x(\rho - z) - y $$$$ z'(t) = x y - \beta z $$
In [107]:
def f(xyz, t, rho, sigma, beta):
    x, y, z = xyz
    return [sigma * (y - x),
            x * (rho - z) - y,
            x * y - beta * z]
In [108]:
rho = 28
sigma = 8
beta = 8/3.0
In [109]:
t = np.linspace(0, 25, 10000)
In [110]:
xyz0 = [1.0, 1.0, 1.0]
In [111]:
xyz1 = integrate.odeint(f, xyz0, t, args=(rho, sigma, beta))
In [112]:
xyz2 = integrate.odeint(f, xyz0, t, args=(rho, sigma, 0.6*beta))
In [113]:
xyz3 = integrate.odeint(f, xyz0, t, args=(rho, 2*sigma, 0.6*beta))
In [114]:
xyz3.shape
Out[114]:
$\displaystyle \left( 10000, \ 3\right)$
In [115]:
from mpl_toolkits.mplot3d.axes3d import Axes3D
In [116]:
fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12, 3.5), subplot_kw={'projection': '3d'})

for ax, xyz, c, p in [(ax1, xyz1, 'r', (rho, sigma, beta)),
                      (ax2, xyz2, 'b', (rho, sigma, 0.6*beta)),
                      (ax3, xyz3, 'g', (rho, 2*sigma, 0.6*beta))]:
    ax.plot(xyz[:,0], xyz[:,1], xyz[:,2], c, alpha=0.5)
    ax.set_xlabel('$x$', fontsize=16)
    ax.set_ylabel('$y$', fontsize=16)
    ax.set_zlabel('$z$', fontsize=16)
    ax.set_xticks([-15, 0, 15])
    ax.set_yticks([-20, 0, 20])
    ax.set_zticks([0, 20, 40])
    ax.set_title(r"$\rho=%.1f, \sigma=%.1f, \beta=%.1f$" % (p[0], p[1], p[2]))

fig.tight_layout()
fig.savefig('ch9-lorenz-equations.pdf')