# 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


## 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')