{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Week 10 day 1: ODEs\n",
    "\n",
    "## Continued from previous notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.integrate\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look now at ODEs. Let's take one of the most famous physics examples:\n",
    "\n",
    "$$\n",
    "F_k(x) = m \\ddot{x} = - k x\n",
    "$$\n",
    "\n",
    "This is the equation for the restoring force of a spring - this is an analytically solvable equation with the solution:\n",
    "\n",
    "$$\n",
    "x(t) = x_\\mathrm{max} \\cos\\left(\\sqrt{\\frac{k}{m}} t\\right)\n",
    "$$\n",
    "\n",
    "This is *harmonic motion*. Notice that we had to set some initial conditions here; the value of $x(0)=x_{\\max}$ and $\\dot{x}(0)=0$.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's try to solve this numerically. We'll learn much better ways to do this, but for now, we are just going to use time steps and Euler’s algorithm.\n",
    "\n",
    "<font color=\"blue\">\n",
    "\n",
    "> Note: I learned something today from the IPython release notes. Type `\\Delta` then press tab. Perfect for Python 3!\n",
    "\n",
    "</font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ΔT = 0.01  # timestep\n",
    "steps = 1_000  # Total number of steps\n",
    "x_max = 1  # Size of x max\n",
    "v_0 = 0\n",
    "koverm = 1  # k / m\n",
    "\n",
    "# Compute the final values of t for plotting\n",
    "ts = np.linspace(0, ΔT * (steps - 1), steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = np.empty(steps)\n",
    "vs = np.empty(steps)\n",
    "\n",
    "xs[0] = x_max\n",
    "vs[0] = v_0\n",
    "\n",
    "for i in range(steps - 1):\n",
    "    # Compute a based on *current* position\n",
    "    a = -koverm * xs[i]\n",
    "\n",
    "    # Compute next velocity\n",
    "    vs[i + 1] = vs[i] - a * ΔT\n",
    "\n",
    "    # Compute next position\n",
    "    xs[i + 1] = xs[i] - vs[i] * ΔT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(ts, xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Try it yourself\n",
    "\n",
    "* What happens if you run more steps?\n",
    "* What happens if you increase the step size?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a better solver, we need to rewrite the equation to a standard form, a vector of first order ODEs:\n",
    "\n",
    "$$\n",
    "\\ddot{x} = - \\frac{k}{m} x\n",
    "$$\n",
    "\n",
    "As a series of first order equations by writing\n",
    "\n",
    "$$\n",
    "u = \\left(\n",
    "\\begin{matrix}\n",
    "\\dot{x} \\\\\n",
    "x\n",
    "\\end{matrix}\n",
    "\\right)\n",
    "$$\n",
    "\n",
    "Then,\n",
    "\n",
    "$$\n",
    "\\dot{u} = \\left(\n",
    "\\begin{matrix}\n",
    "\\ddot{x} \\\\\n",
    "\\dot{x}\n",
    "\\end{matrix}\n",
    "\\right) = \\left(\n",
    "\\begin{matrix}\n",
    "- \\frac{k}{m} x \\\\\n",
    "\\dot{x}\n",
    "\\end{matrix}\n",
    "\\right)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(t, y):\n",
    "    \"Y has two elements, x and v\"\n",
    "\n",
    "    return -koverm * y[1], y[0]\n",
    "\n",
    "\n",
    "res = scipy.integrate.solve_ivp(f, [ts[0], ts[-1]], [1, 0], t_eval=ts)\n",
    "print(res.message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(ts, np.cos(ts), label=\"Analytic solution\")\n",
    "plt.plot(ts, xs, label=\"Hand solution\")\n",
    "plt.plot(res.t, res.y[0], label=\"RK4(5) solution\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "compclass",
   "language": "python",
   "name": "compclass"
  },
  "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": 4
}