Lecture.14-Higher-Order-Methods-for-DEs.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Higher Order Methods for Solving Differential Equations\n",
    "## Lecture 14"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In a constant magnetic field $\\vec{B}$, particle of electrical charge $q$ moving with velocity $\\vec{v}$ experiences a magnetic force $\\vec{F_M}$ given by\n",
    "\n",
    "\\begin{align}\n",
    "\\vec{F_M} &= q \\vec{v} \\times \\vec{B}\n",
    "\\end{align}\n",
    "\n",
    "Predict the motion of a proton as shown in the figure:\n",
    "\n",
    "![](http://s3.amazonaws.com/answer-board-image/b18ba716-e5f2-4809-8c28-b559175fd1be.bmp)\n",
    "\n",
    "Expanding the cross-product and applying Newton's Second Law leads to the following equations of motion:\n",
    "\n",
    "\\begin{align}\n",
    "m \\frac{d^2x}{dt^2} &= - q v_y B \\\\\n",
    "m \\frac{d^2y}{dt^2} &=   q v_x B \\\\\n",
    "\\;\n",
    "\\end{align}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can write this set of two, second-order differential equations into four, first-order differential equations:\n",
    "\\begin{align}\n",
    "\\frac{dv_x}{dt} &= - q v_y B / m\\\\\n",
    "\\frac{dv_y}{dt} &=   q v_x B / m  \\\\\n",
    "\\frac{dx}{dt} &= v_x \\\\\n",
    "\\frac{dy}{dt} &= v_y \\\\\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's use the following physical parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v0 = 4.7e6    # initial horizontal velocity, m/s\n",
    "q = 1.602e-19 # charge, C\n",
    "m = 1.67e-27  # mass of particle, kg\n",
    "B = 0.35      # magnetic field, T "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A numerical solution using Euler's method is shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_1(tmax = 1e-6, dt = 1e-11):  \n",
    "    \n",
    "    # Allocate memory for arrays and set to zero\n",
    "    N = round(tmax/dt)\n",
    "    x = np.zeros(N)\n",
    "    y = np.zeros(N)\n",
    "    vx = np.zeros(N)\n",
    "    vy = np.zeros(N)\n",
    "    t = np.zeros(N)\n",
    "\n",
    "    # Initial values\n",
    "    x[0] = 0\n",
    "    y[0] = 0\n",
    "    vx[0] = v0\n",
    "    vy[0] = 0\n",
    "    t[0] = 0\n",
    "    \n",
    "    # Calculate the solution\n",
    "    for i in range(N-1):\n",
    "        vx[i+1] = vx[i] - q*vy[i]*B/m*dt\n",
    "        vy[i+1] = vy[i] + q*vx[i]*B/m*dt\n",
    "        \n",
    "        x[i+1]  = x[i] + vx[i]*dt\n",
    "        y[i+1]  = y[i] + vy[i]*dt\n",
    "        \n",
    "        t[i+1]  = t[i] + dt\n",
    "\n",
    "    return x, y, t"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is the solution over 1 $\\mu$s of time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(figsize=(12,6))\n",
    "\n",
    "x, y, t = solve_1(dt=1e-11, tmax=1e-6)\n",
    "\n",
    "# plot the x and y component vs time\n",
    "plt.plot(t*1e6, x*100)\n",
    "plt.plot(t*1e6, y*100)   \n",
    "plt.title('Trajectory of a proton in a magnetic field')\n",
    "plt.ylabel('x or y (cm)')\n",
    "plt.xlabel('t ($\\mu$s)')\n",
    "\n",
    "# plot the trajectory in an inset axis\n",
    "fig.add_axes([0.68, 0.2, 0.2, 0.4])\n",
    "plt.plot(x,y)\n",
    "plt.xlim(-0.18, 0.18)\n",
    "plt.ylim(-0.04, 0.32)\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')\n",
    "plt.xticks([-0.1, 0.0, 0.1])\n",
    "plt.yticks([ 0.0, 0.1, 0.2])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "> In the above cell, change the timestep to `dt = 1e-10` and re-execute the cell.\n",
    "\n",
    "> Change it again to `dt=1e-9` and re-execute.\n",
    "\n",
    "> Change the time step back to `dt=1e-11` and re-execute.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, to simulate only 1 $\\mu$s of model time with Euler's method required 100,000 time steps. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmax = 1e-6\n",
    "dt = 1e-11\n",
    "N = round(tmax/dt)\n",
    "print(N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " For comparison, here is the analytical solution resolved with a decreasing number of points:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(4,1, figsize=(12,8))\n",
    "\n",
    "for i, N in enumerate([500, 250, 100, 50]):\n",
    "    plt.sca(axes[i])\n",
    "    ω = q*B/m\n",
    "    t = np.linspace(0, tmax, N)\n",
    "    x = v0/ω*np.sin(ω*t)\n",
    "    y = -v0/ω*(np.cos(ω*t)-1)\n",
    "\n",
    "    plt.plot(t*1e6, x*100, '.-')\n",
    "    plt.plot(t*1e6, y*100, '.-')\n",
    "    plt.ylabel('N = {}'.format(N))\n",
    "    \n",
    "fig.suptitle('Trajectory of a proton in a magnetic field\\nfor different resolutions')\n",
    "plt.xlabel('t ($\\mu$s)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So even though the solution can be well resolve with only 500, 250, 100, or even 50 points, the numerical solutions required 100,000 points to converge.\n",
    "\n",
    "We can do better (achieve the same accuracy with less computational effort) by improving our numerical method."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Runge-Kutta Methods\n",
    "\n",
    "The goal is to solve\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{dy}{dt} = f(y,t), \\quad y(0) = y_0\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Review Euler's method\n",
    "\n",
    "We can approximate derivative with a **forward difference**\n",
    "\n",
    "\\begin{align}\n",
    " \\frac{dy}{dt} &= \\frac{y(t+\\Delta t) - y(t)}{\\Delta t} + O(\\Delta t) \\\\\n",
    " y(t+\\Delta t) &= y(t) + \\frac{dy}{dt}\\Delta t + O(\\Delta t^2) \n",
    "\\end{align}\n",
    "\n",
    "since $\\frac{dy}{dt} = f(y,t)$ we can write\n",
    "\n",
    "$$  y(t+\\Delta t) = y(t) + f(y,t) \\Delta t + O(\\Delta t^2) $$\n",
    "\n",
    "We say that Euler's method is accurage to *first-order* in $\\Delta t$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Midpoint scheme or Second-Order Runge-Kutta Method\n",
    "\n",
    "Start with a **central difference** representation for a derivative\n",
    "\n",
    "\\begin{align}\n",
    " \\frac{dy}{dt} &= \\frac{y(t+\\Delta t) - y(t-\\Delta t)}{2 \\Delta t} + O(\\Delta t^2) \\\\\n",
    " y(t+\\Delta t) &= y(t-\\Delta t) + 2 \\frac{dy}{dt} \\Delta t + O(\\Delta t^3) \\\\\n",
    "  &= y(t - \\Delta t) + 2 f(y(t), t)\\Delta t +O(\\Delta t^3)\n",
    "\\end{align}\n",
    "\n",
    "$f(y(t), t)$ is evaluated at the midpoint in time between $t - \\Delta t$ and $t+\\Delta t$. \n",
    "\n",
    "We can rewrite this as\n",
    "\n",
    "$$ y(t+\\Delta t) = y(t) + 2 f\\left(y(t + \\Delta t /2), t + \\Delta t/2\\right)\\Delta t +O(\\Delta t^3)$$\n",
    "\n",
    "i.e. $\\Delta t \\rightarrow \\Delta t / 2$, then we can shift $t \\rightarrow t + \\Delta t / 2$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, if we can estimate \n",
    "\n",
    "$$ f\\left(y(t + \\Delta t /2), t + \\Delta t/2\\right) $$\n",
    "\n",
    "we can gain an order in our truncation error. To do this we need to estimate $y(t+\\Delta/2)$.\n",
    "\n",
    "Trick: use an Euler half-step\n",
    "\n",
    "$$\n",
    "y ( t + \\Delta t / 2) = y(t) + f(y(t), t)  (\\Delta t / 2) + O(\\Delta t^2)\n",
    "$$\n",
    "\n",
    "\n",
    "We can show (through Taylor expansion) that this leads to a final estimate of $y(t + \\Delta t)$ that is accurate to *second-order* in $\\Delta t$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Midpoint Algorithm\n",
    "\n",
    "1\\. Estimate slope at $t$\n",
    "\n",
    "$$ s_1 = f(y(t), t) $$\n",
    "\n",
    "2\\. Use $s_1$ to estimate the midpoint between $t$ and $t + \\Delta t$\n",
    "\n",
    "\\begin{align}\n",
    "y^* &= y ( t + \\Delta t) \\\\\n",
    "       &= y(t) + \\frac{\\Delta t}{2} s_1\n",
    "\\end{align}\n",
    " \t\n",
    "\n",
    "3\\. Use $y^*$ to get the the slope at the midpoint\n",
    "\n",
    "$$ s_2 = f( y^*, t +  \\frac{\\Delta t}{2} ) $$\n",
    "\n",
    "4\\. Use $s_2$ to estimate $y(t + \\Delta t)$\n",
    "\n",
    "$$ y(t + \\Delta t) = y(t) + \\Delta t s_2 $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Midpoint Algorithm\n",
    "\n",
    "\\begin{align}\n",
    "s_1 &= f(y_i, t_i) \\\\\n",
    "y^* &= y_i + \\Delta t / 2 s_1 \\\\\n",
    "s_2 &= f(y^*, t_i + \\Delta t /2) \\\\\n",
    "y_{i+1} &= y_i + \\Delta t s_2\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Although the algorithm takes two steps, there is an order gain in accuracy. The error of the scheme is of order  $\\Delta t^2$. This means that decreasing the timestep  $\\Delta t$ by a factor of  2 decreases the error by a factor of 4. \n",
    "\n",
    "It also means that compared to the Euler scheme and for the same time step used the accuracy of the scheme is better."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fourth-Order Runge-Kutta Method\n",
    "\n",
    "A very common scheme is the classical fourth-order Runge-Kutta formula.\n",
    "\n",
    "\t\n",
    "\\begin{align}\n",
    "    s_1 &= f(y_i, t_i,) \\\\\n",
    "    y_{1/2} &= y_i + s_1 \\Delta t \\\\ \n",
    "    s_2 &= f(y_{1/2}, t + \\Delta t / 2) \\\\\n",
    "    y_{1/2}^* &= y_i + s_2 \\Delta t / 2  \\\\\n",
    "    s_3 &= f(y_{1/2}^*, t + \\Delta t / 2) \\\\\n",
    "    y^* &= y_i + s_3 \\Delta t  \\\\\n",
    "    s_4 &= f(y^*, t + \\Delta t) \\\\\n",
    "    y_{i+1} &= y_i  +  1/6 (s_1+2 s_2+2s_3+s_4)\\Delta t\t\\\\\n",
    "\\end{align}\n",
    "\n",
    "\n",
    "The error of the scheme is of order  $\\Delta t^4$. This means that decreasing the timestep  $\\Delta t$ by a factor of 2 decreases the error by a factor of 16. It also means that compared to the Euler or midpoint schemes the accuracy is typically much better."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 1\n",
    "\n",
    "Solve numerically the following ODE\n",
    "\n",
    "$$ \\frac{dy}{dt} = - y $$\n",
    "\n",
    "with initial condition\n",
    "\n",
    "$$ y(0)=1$$\n",
    "\n",
    "using the Euler and midpoint and RK4 methods and compare graphically the numerical solutions to the exact solution. Take \n",
    "$\\Delta t$=0.5 and carry out the simulation up to $t=4$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmax = 4\n",
    "dt = 0.5\n",
    "\n",
    "N = round(tmax/dt)\n",
    "t = np.arange(0, tmax, dt)\n",
    "\n",
    "# Exact solution\n",
    "y_exact = 1 * np.exp(-t)\n",
    "\n",
    "# Euler's method\n",
    "y = np.zeros(N)\n",
    "y[0] = 1\n",
    "for i in range(N-1):\n",
    "    s1 = -y[i]\n",
    "    y[i+1]  = y[i] + s1*dt\n",
    "y_euler = y\n",
    "\n",
    "# Midpoint method\n",
    "y = np.zeros(N)\n",
    "y[0] = 1\n",
    "for i in range(N-1):\n",
    "    s1 = -y[i]\n",
    "    ytmp = y[i] + dt/2 * s1\n",
    "    s2 = -ytmp\n",
    "    y[i+1]  = y[i] + s2*dt\n",
    "y_midpoint = y\n",
    "\n",
    "# RK4 method\n",
    "y = np.zeros(N)\n",
    "y[0] = 1\n",
    "for i in range(N-1):\n",
    "    s1 = -y[i]\n",
    "    y1 = y[i] + s1 * dt/2\n",
    "    s2 = -y1\n",
    "    y2 = y[i] + s2 * dt/2\n",
    "    s3 = -y2\n",
    "    y3 = y[i] + s3 * dt\n",
    "    s4 = -y3\n",
    "    y[i+1] = y[i] + (s1 + 2*s2 + 2*s3 + s4)/6*dt\n",
    "y_RK4 = y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare the solutions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(figsize=(7,7))\n",
    "plt.plot(t, y_exact, 'k-',  label='Exact')\n",
    "plt.plot(t, y_euler, 'o:',  label='Euler')\n",
    "plt.plot(t, y_midpoint, 'o:',  label='Midpoint')\n",
    "plt.plot(t, y_RK4, 'o:',  label='RK4')\n",
    "plt.legend()\n",
    "plt.xlabel('t')\n",
    "plt.ylabel('y')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 2\n",
    "\n",
    "Solve numerically the following ODE\n",
    "\n",
    "$$ \\frac{dy}{dt} = -2 y +3 e^{-4t} $$\n",
    "\n",
    " with initial condition\n",
    " \n",
    "$$ y(0)=1$$ \n",
    "\n",
    " using the Euler and midpoint methods, and compare graphically the numerical solutions to the exact solution. Take \n",
    "$\\Delta t =0.2$ and carry out the simulation up to  $t=3$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmax = 3\n",
    "dt = 0.2\n",
    "\n",
    "N = round(tmax/dt)\n",
    "t = np.arange(0, tmax, dt)\n",
    "\n",
    "# Exact solution\n",
    "y_exact = 1/2* np.exp(-4*t)*(5*np.exp(2*t)-3)\n",
    "\n",
    "# define the right hand side\n",
    "def f(y, t):\n",
    "    return -2*y+3*np.exp(-4*t)\n",
    "\n",
    "# Euler's method\n",
    "y = np.zeros(N)\n",
    "y[0] = 1\n",
    "for i in range(N-1):\n",
    "    s1 = f(y[i], t[i])\n",
    "    y[i+1]  = y[i] + s1*dt\n",
    "y_euler = y\n",
    "\n",
    "# Midpoint scheme\n",
    "y = np.zeros(N)\n",
    "y[0] = 1\n",
    "for i in range(N-1):\n",
    "    s1 = f(y[i], t[i])\n",
    "    ytmp = y[i] + dt/2 * s1\n",
    "    s2 = f(ytmp, t[i])\n",
    "    y[i+1]  = y[i] + s2*dt\n",
    "y_midpoint = y\n",
    "\n",
    "# RK4 method\n",
    "y = np.zeros(N)\n",
    "y[0] = 1\n",
    "for i in range(N-1):\n",
    "    s1 = f(y[i], t[i])\n",
    "    y1 = y[i] + s1 * dt/2\n",
    "    s2 = f(y1, t[i] + dt/2)\n",
    "    y2 = y[i] + s2 * dt/2\n",
    "    s3 = f(y2, t[i] + dt/2)\n",
    "    y3 = y[i] + s3 * dt\n",
    "    s4 = f(y3, t[i] + dt)\n",
    "    y[i+1] = y[i] + (s1 + 2*s2 + 2*s3 + s4)/6*dt\n",
    "y_RK4 = y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(figsize=(7,7))\n",
    "plt.plot(t, y_exact, 'k-',  label='Exact')\n",
    "plt.plot(t, y_euler, 'o:',  label='Euler')\n",
    "plt.plot(t, y_midpoint, 'o:',  label='Midpoint')\n",
    "plt.plot(t, y_RK4, 'o:',  label='RK4')\n",
    "plt.legend()\n",
    "plt.xlabel('t')\n",
    "plt.ylabel('y')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is clear that for the same time step, the fourth-order Runge-Kutta method is much more accurate."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 3\n",
    "Let's return the example from Test 1 on a the trajectory of proton in a magnetic field.\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{dv_x}{dt} &= - q v_y B / m\\\\\n",
    "\\frac{dv_y}{dt} &=   q v_x B / m  \\\\\n",
    "\\frac{dx}{dt} &= v_x \\\\\n",
    "\\frac{dy}{dt} &= v_y \\\\\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_2(tmax = 1e-6, dt = 1e-11):  \n",
    "    \n",
    "    # Allocate memory for arrays and set to zero\n",
    "    N = round(tmax/dt)\n",
    "    x = np.zeros(N)\n",
    "    y = np.zeros(N)\n",
    "    vx = np.zeros(N)\n",
    "    vy = np.zeros(N)\n",
    "    t = np.zeros(N)\n",
    "\n",
    "    # Initial values\n",
    "    x[0] = 0\n",
    "    y[0] = 0\n",
    "    vx[0] = v0\n",
    "    vy[0] = 0\n",
    "    t[0] = 0\n",
    "    \n",
    "    # define a function for the right-hand-side:\n",
    "    def f(vx, vy, x, y, t):\n",
    "        return (-q*vy*B/m, +q*vx*B/m, vx, vy)\n",
    "    \n",
    "    # Calculate the solution using the Midpoint Method:\n",
    "    for i in range(N-1):\n",
    "        s1_vx, s1_vy, s1_x, s1_y = f(vx[i], vy[i], x[i], y[i], t[i])\n",
    "        vx_1 = vx[i] + s1_vx*dt/2\n",
    "        vy_1 = vy[i] + s1_vy*dt/2\n",
    "        x_1  = x[i]  + s1_x*dt/2\n",
    "        y_1  = y[i]  + s1_y*dt/2\n",
    "        \n",
    "        s2_vx, s2_vy, s2_x, s2_y = f(vx_1, vy_1, x_1, y_1, t[i]+dt/2)\n",
    "        vx[i+1] = vx[i] + s2_vx*dt\n",
    "        vy[i+1] = vy[i] + s2_vy*dt\n",
    "        x[i+1]  = x[i]  + s2_x*dt\n",
    "        y[i+1]  = y[i]  + s2_y*dt\n",
    "        \n",
    "        t[i+1]  = t[i] + dt\n",
    "    \n",
    "    return x, y, t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(figsize=(12,6))\n",
    "\n",
    "x, y, t = solve_2(dt=1e-11, tmax=1e-6)\n",
    "\n",
    "# plot the x and y component vs time\n",
    "plt.plot(t*1e6, x*100)\n",
    "plt.plot(t*1e6, y*100)   \n",
    "plt.title('Trajectory of a proton in a magnetic field')\n",
    "plt.ylabel('x or y (cm)')\n",
    "plt.xlabel('t ($\\mu$s)')\n",
    "\n",
    "# plot the trajectory in an inset axis\n",
    "fig.add_axes([0.68, 0.2, 0.2, 0.4])\n",
    "plt.plot(x,y)\n",
    "plt.xlim(-0.18, 0.18)\n",
    "plt.ylim(-0.04, 0.32)\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')\n",
    "plt.xticks([-0.1, 0.0, 0.1])\n",
    "plt.yticks([ 0.0, 0.1, 0.2])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "> Compare the above code with the midpoint scheme with Euler's method. \n",
    "\n",
    "> How large can the time step be made while still getting a convergent solution?\n",
    "\n",
    "> Implement a third solver using a 4th order Runge-Kutta scheme."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_3(tmax = 1e-6, dt = 1e-11):  \n",
    "    \n",
    "    # Allocate memory for arrays and set to zero\n",
    "    N = round(tmax/dt)\n",
    "    x = np.zeros(N)\n",
    "    y = np.zeros(N)\n",
    "    vx = np.zeros(N)\n",
    "    vy = np.zeros(N)\n",
    "    t = np.zeros(N)\n",
    "\n",
    "    # Initial values\n",
    "    x[0] = 0\n",
    "    y[0] = 0\n",
    "    vx[0] = v0\n",
    "    vy[0] = 0\n",
    "    t[0] = 0\n",
    "    \n",
    "    # define a function for the right-hand-side:\n",
    "    def f(vx, vy, x, y, t):\n",
    "        return (-q*vy*B/m, +q*vx*B/m, vx, vy)\n",
    "    \n",
    "    # create an array for the time variable\n",
    "    t = np.arange(0, tmax, dt)\n",
    "    \n",
    "    # Calculate the solution using the RK4 Method:\n",
    "    for i in range(N-1):\n",
    "        # s1 = f(y[i], t[i])\n",
    "        # y1 = y[i] + s1 * dt/2\n",
    "        \n",
    "        \n",
    "        # s2 = f(y1, t[i] + dt/2)\n",
    "        # y2 = y[i] + s2 * dt/2\n",
    "\n",
    "        \n",
    "        # s3 = f(y2, t[i] + dt/2)\n",
    "        # y3 = y[i] + s3 * dt\n",
    "\n",
    "        \n",
    "        # s4 = f(y3, t[i] + dt)\n",
    "        # y[i+1] = y[i] + (s1 + 2*s2 + 2*s3 + s4)/6*dt\n",
    "\n",
    "\n",
    "        t[i+1]  = t[i] + dt\n",
    "    \n",
    "    return x, y, t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(figsize=(12,6))\n",
    "\n",
    "x, y, t = solve_3(dt=1e-11, tmax=1e-6)\n",
    "\n",
    "# plot the x and y component vs time\n",
    "plt.plot(t*1e6, x*100)\n",
    "plt.plot(t*1e6, y*100)   \n",
    "plt.title('Trajectory of a proton in a magnetic field')\n",
    "plt.ylabel('x or y (cm)')\n",
    "plt.xlabel('t ($\\mu$s)')\n",
    "\n",
    "# plot the trajectory in an inset axis\n",
    "fig.add_axes([0.68, 0.2, 0.2, 0.4])\n",
    "plt.plot(x,y)\n",
    "plt.xlim(-0.18, 0.18)\n",
    "plt.ylim(-0.04, 0.32)\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')\n",
    "plt.xticks([-0.1, 0.0, 0.1])\n",
    "plt.yticks([ 0.0, 0.1, 0.2])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solution for the RK4 solver:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_3_soln(tmax = 1e-6, dt = 1e-11):  \n",
    "    \n",
    "    # Allocate memory for arrays and set to zero\n",
    "    N = round(tmax/dt)\n",
    "    x = np.zeros(N)\n",
    "    y = np.zeros(N)\n",
    "    vx = np.zeros(N)\n",
    "    vy = np.zeros(N)\n",
    "    t = np.zeros(N)\n",
    "\n",
    "    # Initial values\n",
    "    x[0] = 0\n",
    "    y[0] = 0\n",
    "    vx[0] = v0\n",
    "    vy[0] = 0\n",
    "    t[0] = 0\n",
    "    \n",
    "    # define a function for the right-hand-side:\n",
    "    def f(vx, vy, x, y, t):\n",
    "        return (-q*vy*B/m, +q*vx*B/m, vx, vy)\n",
    "    \n",
    "    # create an array for the time variable\n",
    "    t = np.arange(0, tmax, dt)\n",
    "    \n",
    "    # Calculate the solution using the RK4 Method:\n",
    "    for i in range(N-1):\n",
    "        # s1 = f(y[i], t[i])\n",
    "        # y1 = y[i] + s1 * dt/2\n",
    "        s1_vx, s1_vy, s1_x, s1_y = f(vx[i], vy[i], x[i], y[i], t[i])\n",
    "        vx_1 = vx[i] + s1_vx*dt/2\n",
    "        vy_1 = vy[i] + s1_vy*dt/2\n",
    "        x_1  = x[i]  + s1_x*dt/2\n",
    "        y_1  = y[i]  + s1_y*dt/2\n",
    "        \n",
    "        # s2 = f(y1, t[i] + dt/2)\n",
    "        # y2 = y[i] + s2 * dt/2\n",
    "        s2_vx, s2_vy, s2_x, s2_y = f(vx_1, vy_1, x_1, y_1, t[i]+dt/2)\n",
    "        vx_2 = vx[i] + s2_vx*dt/2\n",
    "        vy_2 = vy[i] + s2_vy*dt/2\n",
    "        x_2  = x[i]  + s2_x*dt/2\n",
    "        y_2  = y[i]  + s2_y*dt/2\n",
    "        \n",
    "        # s3 = f(y2, t[i] + dt/2)\n",
    "        # y3 = y[i] + s3 * dt\n",
    "        s3_vx, s3_vy, s3_x, s3_y = f(vx_2, vy_2, x_2, y_2, t[i]+dt/2)\n",
    "        vx_3 = vx[i] + s3_vx*dt\n",
    "        vy_3 = vy[i] + s3_vy*dt\n",
    "        x_3  = x[i]  + s3_x*dt\n",
    "        y_3  = y[i]  + s3_y*dt\n",
    "        \n",
    "        # s4 = f(y3, t[i] + dt)\n",
    "        # y[i+1] = y[i] + (s1 + 2*s2 + 2*s3 + s4)/6*dt\n",
    "        s4_vx, s4_vy, s4_x, s4_y = f(vx_3, vy_3, x_3, y_3, t[i]+dt)\n",
    "        vx[i+1] = vx[i] + (s1_vx + 2*s2_vx + 2*s3_vx + s4_vx)/6*dt\n",
    "        vy[i+1] = vy[i] + (s1_vy + 2*s2_vy + 2*s3_vy + s4_vy)/6*dt\n",
    "        x[i+1]  = x[i]  + (s1_x + 2*s2_x + 2*s3_x + s4_x)/6*dt\n",
    "        y[i+1]  = y[i]  + (s1_y + 2*s2_y + 2*s3_y + s4_y)/6*dt\n",
    "\n",
    "        t[i+1]  = t[i] + dt\n",
    "    \n",
    "    return x, y, t"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}