{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Runge–Kutta algorithm\n",
    "\n",
    "Let's try a ODE solver, the Runge-Kutta algorithm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numba\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's setup an ODE function to solve. We can write our ODE as a system of linear first order ODE equations:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The harmonic motion equation can be written 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": {
    "tags": []
   },
   "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": "markdown",
   "metadata": {},
   "source": [
    "Now let's do the most basic, stupid thing possible: The Euler algorithm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "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",
    "        # Compute 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": "markdown",
   "metadata": {},
   "source": [
    "We can see, when compared to the actual solution, that the Euler very quickly is destroyed by errors at each step (in this case, the frequency is not too bad, though):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "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",
    "We've picked up bold face to indicate that we can have a vector of ODEs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "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": {
    "tags": []
   },
   "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": {
    "tags": []
   },
   "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": {
    "tags": []
   },
   "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 course grid as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "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": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "ts = np.linspace(0, 40, 1000 + 1)\n",
    "y = rk4_ivp(f, [x_max, v_0], ts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "f_jit = numba.njit(f)\n",
    "rk4_ivp_jit = numba.njit(rk4_ivp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "ts = np.linspace(0, 40, 1000 + 1)\n",
    "y = rk4_ivp_jit(f_jit, np.array([x_max, v_0]), ts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can inspect the types if you'd like to add them after running once:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_jit.inspect_types()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## See also:\n",
    "\n",
    "* [CompClass: RK](https://nbviewer.jupyter.org/github/henryiii/compclass/blob/master/classes/week10/2_rk.ipynb)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:performance-minicourse] *",
   "language": "python",
   "name": "conda-env-performance-minicourse-py"
  },
  "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}