Chapter 9: Ordinary differential equations

Robert Johansson

Source code listings for Numerical Python - A Practical Techniques Approach for Industry (ISBN 978-1-484205-54-9).

The source code listings can be downloaded from http://www.apress.com/9781484205549

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
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]:
$$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]:
$$T{\left (t \right )} = C_{1} e^{- k t} + T_{a}$$
In [11]:
ode_sol.lhs
Out[11]:
$$T{\left (t \right )}$$
In [12]:
ode_sol.rhs
Out[12]:
$$C_{1} e^{- k t} + T_{a}$$
In [13]:
ics = {T(0): T0}
In [14]:
ics
Out[14]:
$$\left \{ T{\left (0 \right )} : T_{0}\right \}$$
In [15]:
C_eq = sympy.Eq(ode_sol.lhs.subs(t, 0).subs(ics), ode_sol.rhs.subs(t, 0))
In [16]:
C_eq
Out[16]:
$$T_{0} = C_{1} + T_{a}$$
In [17]:
C_sol = sympy.solve(C_eq)
In [18]:
C_sol
Out[18]:
$$\left [ \left \{ C_{1} : T_{0} - T_{a}\right \}\right ]$$
In [19]:
ode_sol.subs(C_sol[0])
Out[19]:
$$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]:
$$T{\left (t \right )} = C_{1} e^{- k t} + T_{a}$$
In [22]:
apply_ics(ode_sol, ics, t, [k, Ta])
Out[22]:
$$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]:
$$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]:
$$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]:
$$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]:
$$\left \{ x{\left (0 \right )} : 1, \quad \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]:
$$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]:
$$\frac{\omega_{0} t + 1}{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.legend()

fig.tight_layout()
fig.savefig('ch9-harmonic-oscillator.pdf')

Direction fields

In [40]:
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 [41]:
x = sympy.symbols("x")
In [42]:
y = sympy.Function("y")
In [43]:
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 [44]:
x = sympy.symbols("x")
In [45]:
y = sympy.Function("y")
In [46]:
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 [47]:
sympy.Eq(y(x).diff(x), f)
Out[47]:
$$\frac{d}{d x} y{\left (x \right )} = x + y^{2}{\left (x \right )}$$
In [48]:
ics = {y(0): 0}
In [49]:
ode_sol = sympy.dsolve(y(x).diff(x) - f)
In [50]:
ode_sol
Out[50]:
$$y{\left (x \right )} = \frac{x^{2}}{2} \left(2 C_{1} + 1\right) + \frac{x^{5}}{60} \left(C_{1}^{2} \left(C_{1} + 45\right) + 20 C_{1} + 3\right) + C_{1} + C_{1} x + \frac{7 C_{1}}{6} x^{3} + \frac{C_{1} x^{4}}{12} \left(C_{1} + 5\right) + \mathcal{O}\left(x^{6}\right)$$
In [51]:
ode_sol = apply_ics(ode_sol, {y(0): 0}, x, [])
In [52]:
ode_sol
Out[52]:
$$y{\left (x \right )} = \frac{x^{2}}{2} + \frac{x^{5}}{20} + \mathcal{O}\left(x^{6}\right)$$
In [53]:
ode_sol = sympy.dsolve(y(x).diff(x) - f, ics=ics)
In [54]:
ode_sol
Out[54]:
$$y{\left (x \right )} = \frac{x^{2}}{2} + \frac{x^{5}}{20} + \mathcal{O}\left(x^{6}\right)$$
In [55]:
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-dx, -x0, 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 [56]:
t = sympy.symbols("t", positive=True)
In [57]:
s, Y = sympy.symbols("s, Y", real=True)
In [58]:
y = sympy.Function("y")
In [59]:
ode = y(t).diff(t, 2) + 2 * y(t).diff(t) + 10 * y(t) - 2 * sympy.sin(3*t)
In [60]:
ode
Out[60]:
$$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 [61]:
L_y = sympy.laplace_transform(y(t), t, s)
In [62]:
L_y
Out[62]:
$$\mathcal{L}_{t}\left[y{\left (t \right )}\right]\left(s\right)$$
In [63]:
L_ode = sympy.laplace_transform(ode, t, s, noconds=True)
In [64]:
L_ode
Out[64]:
$$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 [65]:
def laplace_transform_derivatives(e):
    """
    Evaluate the unevaluted laplace transforms of derivatives
    of functions
    """
    if isinstance(e, sympy.LaplaceTransform):
        if isinstance(e.args[0], sympy.Derivative):
            d, t, s = e.args
            n = len(d.args) - 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 [66]:
L_ode_2 = laplace_transform_derivatives(L_ode)
In [67]:
L_ode_2
Out[67]:
$$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 [68]:
L_ode_3 = L_ode_2.subs(L_y, Y)
In [69]:
L_ode_3
Out[69]:
$$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 [70]:
ics = {y(0): 1, y(t).diff(t).subs(t, 0): 0}
In [71]:
ics
Out[71]:
$$\left \{ y{\left (0 \right )} : 1, \quad \left. \frac{d}{d t} y{\left (t \right )} \right|_{\substack{ t=0 }} : 0\right \}$$
In [72]:
L_ode_4 = L_ode_3.subs(ics)
In [73]:
L_ode_4
Out[73]:
$$Y s^{2} + 2 Y s + 10 Y - s - 2 - \frac{6}{s^{2} + 9}$$
In [74]:
Y_sol = sympy.solve(L_ode_4, Y)
In [75]:
Y_sol
Out[75]:
$$\left [ \frac{s^{3} + 2 s^{2} + 9 s + 24}{s^{4} + 2 s^{3} + 19 s^{2} + 18 s + 90}\right ]$$
In [76]:
sympy.apart(Y_sol[0])
Out[76]:
$$- \frac{12 s - 6}{37 s^{2} + 333} + \frac{49 s + 92}{37 s^{2} + 74 s + 370}$$
In [77]:
y_sol = sympy.inverse_laplace_transform(Y_sol[0], s, t)
In [78]:
sympy.simplify(y_sol)
Out[78]:
$$\frac{1}{111 e^{t}} \left(6 \left(\sin{\left (3 t \right )} - 6 \cos{\left (3 t \right )}\right) e^{t} + 43 \sin{\left (3 t \right )} + 147 \cos{\left (3 t \right )}\right)$$
In [79]:
y_t = sympy.lambdify(t, y_sol, 'numpy')
In [80]:
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 [81]:
x = sympy.symbols("x")
In [82]:
y = sympy.Function("y")
In [83]:
f = y(x)**2 + x
In [84]:
f_np = sympy.lambdify((y(x), x), f, 'math')
In [85]:
y0 = 0
In [86]:
xp = np.linspace(0, 1.9, 100)
In [87]:
xp.shape
Out[87]:
$$\left ( 100\right )$$
In [88]:
yp = integrate.odeint(f_np, y0, xp)
In [89]:
xm = np.linspace(0, -5, 100)
In [90]:
ym = integrate.odeint(f_np, y0, xm)
In [91]:
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 [92]:
a, b, c, d = 0.4, 0.002, 0.001, 0.7
In [93]:
def f(xy, t):
    x, y = xy
    return [a * x - b * x * y,
            c * x * y - d * y]
In [94]:
xy0 = [600, 400]
In [95]:
t = np.linspace(0, 50, 250)
In [96]:
xy_t = integrate.odeint(f, xy0, t)
In [97]:
xy_t.shape
Out[97]:
$$\left ( 250, \quad 2\right )$$
In [98]:
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 [99]:
def f(xyz, t, rho, sigma, beta):
    x, y, z = xyz
    return [sigma * (y - x),
            x * (rho - z) - y,
            x * y - beta * z]
In [100]:
rho = 28
sigma = 8
beta = 8/3.0
In [101]:
t = np.linspace(0, 25, 10000)
In [102]:
xyz0 = [1.0, 1.0, 1.0]
In [103]:
xyz1 = integrate.odeint(f, xyz0, t, args=(rho, sigma, beta))
In [104]:
xyz2 = integrate.odeint(f, xyz0, t, args=(rho, sigma, 0.6*beta))
In [105]:
xyz3 = integrate.odeint(f, xyz0, t, args=(rho, 2*sigma, 0.6*beta))
In [106]:
xyz3.shape
Out[106]:
$$\left ( 10000, \quad 3\right )$$
In [107]:
from mpl_toolkits.mplot3d.axes3d import Axes3D
In [108]:
fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12, 3.5), subplot_kw={'projection': '3d'})

for ax, xyz, c in [(ax1, xyz1, 'r'), (ax2, xyz2, 'b'), (ax3, xyz3, 'g')]:
    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])

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