This is one of the 100 recipes of the IPython Cookbook, the definitive guide to high-performance scientific computing and data science in Python.

12.3. Simulating an Ordinary Differential Equation with SciPy

  1. Let's import NumPy, SciPy (integrate package), and matplotlib.
In [ ]:
import numpy as np
import scipy.integrate as spi
import matplotlib.pyplot as plt
%matplotlib inline
  1. We define a few parameters appearing in our model.
In [ ]:
m = 1.  # particle's mass
k = 1.  # drag coefficient
g = 9.81  # gravity acceleration
  1. We have two variables: x and y (two dimensions). We note $\mathbf{u}=(x,y)$. The ODE we are going to simulate is:
$$\ddot{\mathbf{u}} = -\frac{k}{m} \dot{\mathbf{u}} + \mathbf{g}$$

where $\mathbf{g}$ is the gravity acceleration vector. In order to simulate this second-order ODE with SciPy, we can convert it to a first-order ODE (another option would be to solve $\dot{\mathbf{u}}$ first before integrating the solution). To do this, we consider two 2D variables: $\mathbf{u}$ and $\dot{\mathbf{u}}$. We note $\mathbf{v} = (\mathbf{u}, \dot{\mathbf{u}})$. We can express $\dot{\mathbf{v}}$ as a function of $\mathbf{v}$. Now, we create the initial vector $\mathbf{v}_0$ at time $t=0$: it has four components.

In [ ]:
# The initial position is (0, 0).
v0 = np.zeros(4)
# The initial speed vector is oriented
# to the top right.
v0[2] = 4.
v0[3] = 10.
  1. We need to create a Python function $f$ that takes the current vector $\mathbf{v}(t_0)$ and a time $t_0$ as argument (with optional parameters), and that returns the derivative $\dot{\mathbf{v}}(t_0)$.
In [ ]:
def f(v, t0, k):
    # v has four components: v=[u, u'].
    u, udot = v[:2], v[2:]
    # We compute the second derivative u'' of u.
    udotdot = -k/m * udot
    udotdot[1] -= g
    # We return v'=[u', u''].
    return np.r_[udot, udotdot]
  1. Now, we simulate the system for different values of $k$. We use the SciPy function odeint, defined in the scipy.integrate package.
In [ ]:
# We want to evaluate the system on 30 linearly
# spaced times between t=0 and t=3.
t = np.linspace(0., 3., 30)
# We simulate the system for different values of k.
for k in np.linspace(0., 1., 5):
    # We simulate the system and evaluate $v$ on the 
    # given times.
    v = spi.odeint(f, v0, t, args=(k,))
    # We plot the particle's trajectory.
    plt.plot(v[:,0], v[:,1], 'o-', mew=1, ms=8, mec='w',
plt.xlim(0, 12);
plt.ylim(0, 6);

The most outward trajectory (blue) corresponds to drag-free motion (without air resistance). It is a parabola. In the other trajectories, we can observe the increasing effect of air resistance, parameterized with $k$.

You'll find all the explanations, figures, references, and much more in the book (to be released later this summer).

IPython Cookbook, by Cyrille Rossant, Packt Publishing, 2014 (500 pages).