#!/usr/bin/env python
# coding: utf-8

# > This is one of the 100 recipes of the [IPython Cookbook](http://ipython-books.github.io/), 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
get_ipython().run_line_magic('matplotlib', 'inline')


# 2. We define a few parameters appearing in our model.

# In[ ]:


m = 1.  # particle's mass
k = 1.  # drag coefficient
g = 9.81  # gravity acceleration


# 3. 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.


# 4. 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]


# 3. Now, we simulate the system for different values of $k$. We use the SciPy function `odeint`, defined in the `scipy.integrate` package.

# In[ ]:


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](http://ipython-books.github.io/), by [Cyrille Rossant](http://cyrille.rossant.net), Packt Publishing, 2014 (500 pages).