{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Week 10 Day 2: Runge–Kutta algorithm\n",
    "\n",
    "## Objectives\n",
    "\n",
    "Work through a problem by implementing the RK2 algorithm!\n",
    "\n",
    "First let's rewrite the Euler code from last time, with a few changes to make this more general and adaptable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can rewrite our harmonic motion equation in terms of $\\mathbf{f}(t, \\mathbf{y}) = \\dot{\\mathbf{y}}$, where this is in the standard form:\n",
    "\n",
    "$$\n",
    "\\mathbf{y} =\n",
    "\\left(\n",
    "\\begin{matrix}\n",
    "\\dot{x} \\\\\n",
    "x\n",
    "\\end{matrix}\n",
    "\\right)\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\mathbf{f}(t, \\mathbf{y}) =\n",
    "\\dot{\\mathbf{y}}\n",
    "=\n",
    "\\left(\n",
    "\\begin{matrix}\n",
    "\\ddot{x} \\\\\n",
    "\\dot{x}\n",
    "\\end{matrix}\n",
    "\\right)\n",
    "=\n",
    "\\left(\n",
    "\\begin{matrix}\n",
    "-\\frac{k}{m} x \\\\\n",
    "\\dot{x}\n",
    "\\end{matrix}\n",
    "\\right)\n",
    "=\n",
    "\\left(\n",
    "\\begin{matrix}\n",
    "-\\frac{k}{m} y_1 \\\\\n",
    "y_0\n",
    "\\end{matrix}\n",
    "\\right)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_max = 1  # Size of x max\n",
    "v_0 = 0\n",
    "koverm = 1  # k / m\n",
    "\n",
    "\n",
    "def f(t, y):\n",
    "    \"Y has two elements, x and v\"\n",
    "    return np.array([-koverm * y[1], y[0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def euler_ivp(f, init_y, t):\n",
    "    steps = len(t)\n",
    "    order = len(init_y)  # Number of equations\n",
    "\n",
    "    y = np.empty((steps, order))\n",
    "    y[0] = init_y  # Note that this sets the elements of the first row\n",
    "\n",
    "    for n in range(steps - 1):\n",
    "        h = t[n + 1] - t[n]\n",
    "\n",
    "        # Copute dydt based on *current* position\n",
    "        dydt = f(t[n], y[n])\n",
    "\n",
    "        # Compute next velocity and position\n",
    "        y[n + 1] = y[n] - dydt * h\n",
    "\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = np.linspace(0, 40, 1000 + 1)\n",
    "y = euler_ivp(f, [x_max, v_0], ts)\n",
    "plt.plot(ts, np.cos(ts))\n",
    "plt.plot(ts, y[:, 0], \"--\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Range-Kutta introduction\n",
    "\n",
    "Note that $h =  t_{n+1} - t_n $.\n",
    "\n",
    "$$\n",
    "\\dot{y} = f(t,y)\n",
    "\\\\\n",
    "\\implies y = \\int f(t,y) \\, dt\n",
    "\\\\\n",
    "\\implies y_{n+1} = y_{n} + \\int_{t_n}^{t_{n+1}} f(t,y) \\, dt\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, expand $f$ in a Taylor series around the *midpoint* of the interval:\n",
    "\n",
    "$$\n",
    "f(t,y) \\approx  f(t_{n+\\frac{1}{2}},y_{n+\\frac{1}{2}})\n",
    "       + \\left( t - t_{n+\\frac{1}{2}}\\right)\n",
    "       \\dot{f}(t_{n+\\frac{1}{2}})\n",
    "       + \\mathcal{O}(h^2)\n",
    "$$\n",
    "\n",
    "The second term here is symmetric in the interval, so all we have left is the first term in the integral:\n",
    "\n",
    "$$\n",
    "\\int_{t_n}^{t_{n+1}} f(t,y) \\, dt \\approx\n",
    "    h\\, f(t_{n+\\frac{1}{2}},y_{n+\\frac{1}{2}}) + \\mathcal{O}(h^3)\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Back into the original statement, we get:\n",
    "\n",
    "$$\n",
    "y_{n+1} \\approx\n",
    "\\color{blue}{\n",
    "y_{n}\n",
    "+ h\\, f(t_{n+\\frac{1}{2}},y_{n+\\frac{1}{2}})\n",
    "}\n",
    "+ \\mathcal{O}(h^3)\n",
    "\\tag{rk2}\n",
    "$$\n",
    "\n",
    "We've got one more problem! How do we calculate $f(t_{n+\\frac{1}{2}},y_{n+\\frac{1}{2}})$? We can use the Euler's algorithm that we saw last time:\n",
    "\n",
    "$$\n",
    "y_{n+\\frac{1}{2}}\n",
    "\\approx y_n + \\frac{1}{2} h \\dot{y}\n",
    "= \\color{red}{\n",
    "y_n + \\frac{1}{2} h f(t_{n},y_{n})\n",
    "}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Putting it together, this is our RK2 algorithm:\n",
    "\n",
    "$$\n",
    "\\mathbf{y}_{n+1} \\approx\n",
    "\\color{blue}{\n",
    "\\mathbf{y}_{n}\n",
    "+ \\mathbf{k}_2\n",
    "}\n",
    "\\tag{1.0}\n",
    "$$\n",
    "\n",
    "\n",
    "$$\n",
    "\\mathbf{k}_1 = h \\mathbf{f}(t_n,\\, \\mathbf{y}_n)\n",
    "\\tag{1.1}\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\mathbf{k}_2 = h \\mathbf{f}(t_n + \\frac{h}{2},\\, \\color{red}{\\mathbf{y}_n\n",
    "+ \\frac{\\mathbf{k}_1}{2}})\n",
    "\\tag{1.2}\n",
    "$$\n",
    "\n",
    "Like the book, we've picked up bold face to indicate that we can have a vector of ODEs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rk2_ivp(f, init_y, t):\n",
    "    steps = len(t)\n",
    "    order = len(init_y)\n",
    "\n",
    "    y = np.empty((steps, order))\n",
    "    y[0] = init_y\n",
    "\n",
    "    for n in range(steps - 1):\n",
    "        h = t[n + 1] - t[n]\n",
    "\n",
    "        k1 = h * f(t[n], y[n])  # 1.1\n",
    "        k2 = h * f(t[n] + h / 2, y[n] + k1 / 2)  # 1.2\n",
    "\n",
    "        y[n + 1] = y[n] + k2  # 1.0\n",
    "\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try this with the same grid as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = np.linspace(0, 40, 1000 + 1)\n",
    "y = rk2_ivp(f, [x_max, v_0], ts)\n",
    "plt.plot(ts, np.cos(ts))\n",
    "plt.plot(ts, y[:, 0], \"--\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And, on a coarser grid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = np.linspace(0, 40, 100 + 1)\n",
    "y = rk2_ivp(f, [x_max, v_0], ts)\n",
    "plt.plot(ts, np.cos(ts))\n",
    "plt.plot(ts, y[:, 0], \"--\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can get the RK4 algorithm by keeping another non-zero term in the Taylor series:\n",
    "\n",
    "$$\n",
    "\\mathbf{y}_{n+1} \\approx\n",
    "\\mathbf{y}_{n}\n",
    "+ \\frac{1}{6} (\\mathbf{k}_1 + 2 \\mathbf{k}_2 + 2 \\mathbf{k}_3 + \\mathbf{k}_4 )\n",
    "\\tag{2.0}\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\mathbf{k}_1 = h \\mathbf{f}(t_n,\\, \\mathbf{y}_n)\n",
    "\\tag{2.1}\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\mathbf{k}_2 = h \\mathbf{f}(t_n + \\frac{h}{2},\\,\n",
    "                            \\mathbf{y}_n + \\frac{\\mathrm{k}_1}{2})\n",
    "\\tag{2.2}\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\mathbf{k}_3 = h \\mathbf{f}(t_n + \\frac{h}{2},\\,\n",
    "                            \\mathbf{y}_n + \\frac{\\mathrm{k}_2}{2})\n",
    "\\tag{2.3}\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\mathbf{k}_4 = h \\mathbf{f}(t_n + h,\\,\n",
    "                            \\mathbf{y}_n + \\mathrm{k}_3)\n",
    "\\tag{2.4}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rk4_ivp(f, init_y, t):\n",
    "    steps = len(t)\n",
    "    order = len(init_y)\n",
    "\n",
    "    y = np.empty((steps, order))\n",
    "    y[0] = init_y\n",
    "\n",
    "    for n in range(steps - 1):\n",
    "        h = t[n + 1] - t[n]\n",
    "\n",
    "        k1 = h * f(t[n], y[n])  # 2.1\n",
    "        k2 = h * f(t[n] + h / 2, y[n] + k1 / 2)  # 2.2\n",
    "        k3 = h * f(t[n] + h / 2, y[n] + k2 / 2)  # 2.3\n",
    "        k4 = h * f(t[n] + h, y[n] + k3)  # 2.4\n",
    "\n",
    "        y[n + 1] = y[n] + 1 / 6 * (k1 + 2 * k2 + 2 * k3 + k4)  # 2.0\n",
    "\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try this with the same grid as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = np.linspace(0, 40, 100 + 1)\n",
    "y = rk4_ivp(f, [x_max, v_0], ts)\n",
    "plt.plot(ts, np.cos(ts))\n",
    "plt.plot(ts, y[:, 0], \"--\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Bonus: Performance boost"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "ts = np.linspace(0, 40, 1000 + 1)\n",
    "y = rk4_ivp(f, [x_max, v_0], ts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's JIT both of these functions, and see if we can get a speedup!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba\n",
    "\n",
    "f_jit = numba.njit(f)\n",
    "rk4_ivp_jit = numba.njit(rk4_ivp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "ts = np.linspace(0, 40, 1000 + 1)\n",
    "y = rk4_ivp_jit(f_jit, [x_max, v_0], ts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How's that for almost 0 effort!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RK 4(5)\n",
    "\n",
    "This is the \"adaptive\" step solver from the book. It's a bit more readable, and a logical error was fixed. The definition of the tolerance seems to be different than scipy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rk45_ivp(f, init_y, t_range, tol=1e-8, attempt_steps=20):\n",
    "    order = len(init_y)  # Number of equations\n",
    "\n",
    "    y = [np.array(init_y)]\n",
    "    t = [t_range[0]]\n",
    "    err_sum = 0\n",
    "\n",
    "    # Step size and limits to step size\n",
    "    h = (t_range[1] - t_range[0]) / attempt_steps\n",
    "    hmin = h / 64\n",
    "    hmax = h * 64\n",
    "\n",
    "    while t[-1] < t_range[1]:\n",
    "\n",
    "        # Last step should just be exactly what is needed to finish\n",
    "        if t[-1] + h > t_range[1]:\n",
    "            h = t_range[1] - t[-1]\n",
    "\n",
    "        # Compute k1 - k6 for evaluation and error estimate\n",
    "        k1 = h * f(t[-1], y[-1])\n",
    "        k2 = h * f(t[-1] + h / 4, y[-1] + k1 / 4)\n",
    "        k3 = h * f(t[-1] + 3 * h / 8, y[-1] + 3 * k1 / 32 + 9 * k2 / 32)\n",
    "        k4 = h * f(\n",
    "            t[-1] + 12 * h / 13,\n",
    "            y[-1] + 1932 * k1 / 2197 - 7200 * k2 / 2197 + 7296 * k3 / 2197,\n",
    "        )\n",
    "        k5 = h * f(\n",
    "            t[-1] + h,\n",
    "            y[-1] + 439 * k1 / 216 - 8 * k2 + 3680 * k3 / 513 - 845 * k4 / 4104,\n",
    "        )\n",
    "        k6 = h * f(\n",
    "            t[-1] + h / 2,\n",
    "            y[-1]\n",
    "            + 8 * k1 / 27\n",
    "            + 2 * k2\n",
    "            - 3544 * k3 / 2565\n",
    "            + 1859 * k4 / 4104\n",
    "            - 11 * k5 / 40,\n",
    "        )\n",
    "\n",
    "        # Compute error from higher order RK calculation\n",
    "        err = np.abs(\n",
    "            k1 / 360 - 128 * k3 / 4275 - 2197 * k4 / 75240 + k5 / 50 + 2 * k6 / 55\n",
    "        )\n",
    "\n",
    "        # Compute factor to see if step size should be changed\n",
    "        s = 0 if err[0] == 0 or err[1] == 0 else 0.84 * (tol * h / err[0]) ** 0.25\n",
    "\n",
    "        lower_step = s < 0.75 and h > 2 * hmin\n",
    "        raise_step = s > 1.5 and 2 * h < hmax\n",
    "        no_change = not raise_step and not lower_step\n",
    "\n",
    "        # Accept step and move on\n",
    "        if err[0] < tol or err[1] < tol or no_change:\n",
    "            y.append(\n",
    "                y[-1] + 25 * k1 / 216 + 1408 * k3 / 2565 + 2197 * k4 / 4104 - k5 / 5\n",
    "            )\n",
    "            t.append(t[-1] + h)\n",
    "\n",
    "        # Grow or shrink the step size if needed\n",
    "        if lower_step:\n",
    "            h /= 2\n",
    "        elif raise_step:\n",
    "            h *= 2\n",
    "\n",
    "    return np.array(t), np.array(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = np.linspace(0, 40, 1000 + 1)\n",
    "t, y = rk45_ivp(f, [x_max, v_0], [0, 40], 0.005)\n",
    "plt.plot(ts, np.cos(ts))\n",
    "plt.plot(t, y[:, 0], \".\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare it with the scipy algorithm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.integrate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r45 = scipy.integrate.solve_ivp(f, [0, 40], [x_max, v_0], rtol=0.00001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r45\n",
    "plt.plot(ts, np.cos(ts))\n",
    "plt.plot(r45.t, r45.y[0], \".\")"
   ]
  }
 ],
 "metadata": {
  "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}