This is one of the 100 recipes of the IPython Cookbook, the definitive guide to high-performance scientific computing and data science in Python.
integrate
package), and matplotlib.import numpy as np
import scipy.integrate as spi
import matplotlib.pyplot as plt
%matplotlib inline
m = 1. # particle's mass
k = 1. # drag coefficient
g = 9.81 # gravity acceleration
x
and y
(two dimensions). We note $\mathbf{u}=(x,y)$. The ODE we are going to simulate is: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.
# 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.
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]
odeint
, defined in the scipy.integrate
package.plt.figure(figsize=(6,3));
# 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',
label='k={0:.1f}'.format(k));
plt.legend();
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).