fdm-devito-notebooks/01_vib/vib_gen.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "mu^{\\prime\\prime} + f(u^{\\prime}) + s(u) = F(t),\\quad u(0)=I,\\ u^{\\prime}(0)=V,\\ t\\in (0,T]\n",
    "\\thinspace .\n",
    "\\label{vib:ode2} \\tag{1}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have also included a possibly nonzero initial value for $u^{\\prime}(0)$.\n",
    "The parameters $m$, $f(u^{\\prime})$, $s(u)$, $F(t)$, $I$, $V$, and $T$ are\n",
    "input data.\n",
    "\n",
    "There are two main types of damping (friction) forces: linear $f(u^{\\prime})=bu$, or\n",
    "quadratic $f(u^{\\prime})=bu^{\\prime}|u^{\\prime}|$. Spring systems often feature linear\n",
    "damping, while air resistance usually gives rise to quadratic damping.\n",
    "Spring forces are often linear: $s(u)=cu$, but nonlinear versions\n",
    "are also common, the most famous is the gravity force on a pendulum\n",
    "that acts as a spring with $s(u)\\sim \\sin(u)$.\n",
    "\n",
    "\n",
    "## A centered scheme for linear damping\n",
    "<div id=\"vib:ode2:fdm:flin\"></div>\n",
    "\n",
    "Sampling ([1](#vib:ode2)) at a mesh point $t_n$, replacing\n",
    "$u^{\\prime\\prime}(t_n)$ by $[D_tD_tu]^n$, and $u^{\\prime}(t_n)$ by $[D_{2t}u]^n$ results\n",
    "in the discretization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto1\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "[mD_tD_t u + f(D_{2t}u) + s(u) = F]^n,\n",
    "\\label{_auto1} \\tag{2}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which written out means"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:step3b\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "m\\frac{u^{n+1}-2u^n + u^{n-1}}{\\Delta t^2}\n",
    "+ f(\\frac{u^{n+1}-u^{n-1}}{2\\Delta t}) + s(u^n) = F^n,\n",
    "\\label{vib:ode2:step3b} \\tag{3}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where $F^n$ as usual means $F(t)$ evaluated at $t=t_n$.\n",
    "Solving ([3](#vib:ode2:step3b)) with respect to the unknown\n",
    "$u^{n+1}$ gives a problem: the $u^{n+1}$ inside the $f$ function\n",
    "makes the equation *nonlinear* unless $f(u^{\\prime})$ is a linear function,\n",
    "$f(u^{\\prime})=bu^{\\prime}$. For now we shall assume that $f$ is linear in $u^{\\prime}$.\n",
    "Then"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:step3b2\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "m\\frac{u^{n+1}-2u^n + u^{n-1}}{\\Delta t^2}\n",
    "+ b\\frac{u^{n+1}-u^{n-1}}{2\\Delta t} + s(u^n) = F^n,\n",
    "\\label{vib:ode2:step3b2} \\tag{4}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which gives an explicit formula for $u$ at each\n",
    "new time level:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:step4\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "u^{n+1} = (2mu^n + (\\frac{b}{2}\\Delta t - m)u^{n-1} +\n",
    "\\Delta t^2(F^n - s(u^n)))(m + \\frac{b}{2}\\Delta t)^{-1}\n",
    "\\label{vib:ode2:step4} \\tag{5}\n",
    "\\thinspace .\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the first time step we need to discretize $u^{\\prime}(0)=V$\n",
    "as $[D_{2t}u = V]^0$ and combine\n",
    "with ([5](#vib:ode2:step4)) for $n=0$. The discretized initial condition\n",
    "leads to"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:ic:du\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "u^{-1} = u^{1} - 2\\Delta t V,\n",
    "\\label{vib:ode2:ic:du} \\tag{6}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which inserted in ([5](#vib:ode2:step4)) for $n=0$ gives an equation\n",
    "that can be solved for\n",
    "$u^1$:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:step4b\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "u^1 = u^0 + \\Delta t\\, V\n",
    "+ \\frac{\\Delta t^2}{2m}(-bV - s(u^0) + F^0)\n",
    "\\thinspace .\n",
    "\\label{vib:ode2:step4b} \\tag{7}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A centered scheme for quadratic damping\n",
    "<div id=\"vib:ode2:fdm:fquad\"></div>\n",
    "\n",
    "When $f(u^{\\prime})=bu^{\\prime}|u^{\\prime}|$, we get a quadratic equation for $u^{n+1}$\n",
    "in ([3](#vib:ode2:step3b)). This equation can be straightforwardly\n",
    "solved by the well-known formula for the roots of a quadratic equation.\n",
    "However, we can also avoid the nonlinearity by introducing\n",
    "an approximation with an error of order no higher than what we\n",
    "already have from replacing derivatives with finite differences.\n",
    "\n",
    "\n",
    "We start with ([1](#vib:ode2)) and only replace\n",
    "$u^{\\prime\\prime}$ by $D_tD_tu$, resulting in"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:quad:idea1\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "[mD_tD_t u + bu^{\\prime}|u^{\\prime}| + s(u) = F]^n\\thinspace .\n",
    "\\label{vib:ode2:quad:idea1} \\tag{8}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, $u^{\\prime}|u^{\\prime}|$ is to be computed at time $t_n$. The idea\n",
    "is now to introduce\n",
    "a *geometric mean*, defined by"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "(w^2)^n \\approx w^{n-\\frac{1}{2}}w^{n+\\frac{1}{2}},\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "for some quantity $w$ depending on time. The error in the geometric mean\n",
    "approximation is $\\Oof{\\Delta t^2}$, the same as in the\n",
    "approximation $u^{\\prime\\prime}\\approx D_tD_tu$. With $w=u^{\\prime}$ it follows\n",
    "that"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "[u^{\\prime}|u^{\\prime}|]^n \\approx u^{\\prime}(t_{n+\\frac{1}{2}})|u^{\\prime}(t_{n-\\frac{1}{2}})|\\thinspace .\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step is to approximate\n",
    "$u^{\\prime}$ at $t_{n\\pm 1/2}$, and fortunately a centered difference\n",
    "fits perfectly into the formulas since it involves $u$ values at\n",
    "the mesh points only. With the approximations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:quad:idea2\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "u^{\\prime}(t_{n+1/2})\\approx [D_t u]^{n+\\frac{1}{2}},\\quad u^{\\prime}(t_{n-1/2})\\approx [D_t u]^{n-\\frac{1}{2}},\n",
    "\\label{vib:ode2:quad:idea2} \\tag{9}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "we get"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto2\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "[u^{\\prime}|u^{\\prime}|]^n\n",
    "\\approx [D_tu]^{n+\\frac{1}{2}}|[D_tu]^{n-\\frac{1}{2}}| = \\frac{u^{n+1}-u^n}{\\Delta t}\n",
    "\\frac{|u^n-u^{n-1}|}{\\Delta t}\n",
    "\\thinspace .\n",
    "\\label{_auto2} \\tag{10}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The counterpart to ([3](#vib:ode2:step3b)) is then"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:step3b:quad\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "m\\frac{u^{n+1}-2u^n + u^{n-1}}{\\Delta t^2}\n",
    "+ b\\frac{u^{n+1}-u^n}{\\Delta t}\\frac{|u^n-u^{n-1}|}{\\Delta t}\n",
    "+ s(u^n) = F^n,\n",
    "\\label{vib:ode2:step3b:quad} \\tag{11}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which is linear in the unknown $u^{n+1}$. Therefore, we can easily solve\n",
    "([11](#vib:ode2:step3b:quad))\n",
    "with respect to $u^{n+1}$ and achieve the explicit updating formula"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "u^{n+1} =  \\left( m + b|u^n-u^{n-1}|\\right)^{-1}\\times \\nonumber\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:step4:quad\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    " \\qquad \\left(2m u^n - mu^{n-1} + bu^n|u^n-u^{n-1}| + \\Delta t^2 (F^n - s(u^n))\n",
    "\\right)\n",
    "\\thinspace .\n",
    "\\label{vib:ode2:step4:quad} \\tag{12}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Make exercise to solve complicated u^1 equation with Bisection/Newton -->\n",
    "\n",
    "In the derivation of a special equation for the first\n",
    "time step we run into some trouble: inserting ([6](#vib:ode2:ic:du))\n",
    "in ([12](#vib:ode2:step4:quad)) for $n=0$ results in a complicated nonlinear\n",
    "equation for $u^1$. By thinking differently about the problem we can\n",
    "easily get away with the nonlinearity again. We have for $n=0$ that\n",
    "$b[u^{\\prime}|u^{\\prime}|]^0 = bV|V|$. Using this value in ([8](#vib:ode2:quad:idea1))\n",
    "gives"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto3\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "[mD_tD_t u + bV|V| + s(u) = F]^0\n",
    "\\thinspace .\n",
    "\\label{_auto3} \\tag{13}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Writing this equation out and using ([6](#vib:ode2:ic:du)) results in the\n",
    "special equation for the first time step:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:step4b:quad\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "u^1 = u^0 + \\Delta t V + \\frac{\\Delta t^2}{2m}\\left(-bV|V| - s(u^0) + F^0\\right)\n",
    "\\thinspace .\n",
    "\\label{vib:ode2:step4b:quad} \\tag{14}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A forward-backward discretization of the quadratic damping term\n",
    "\n",
    "The previous section first proposed to discretize the quadratic\n",
    "damping term $|u^{\\prime}|u^{\\prime}$ using centered differences:\n",
    "$[|D_{2t}|D_{2t}u]^n$. As this gives rise to a nonlinearity in\n",
    "$u^{n+1}$, it was instead proposed to use a geometric mean combined\n",
    "with centered differences.  But there are other alternatives. To get\n",
    "rid of the nonlinearity in $[|D_{2t}|D_{2t}u]^n$, one can think\n",
    "differently: apply a backward difference to $|u^{\\prime}|$, such that\n",
    "the term involves known values, and apply a forward difference to\n",
    "$u^{\\prime}$ to make the term linear in the unknown $u^{n+1}$. With\n",
    "mathematics,"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:nonlin:fbdiff\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "[\\beta |u^{\\prime}|u^{\\prime}]^n \\approx \\beta |[D_t^-u]^n|[D_t^+ u]^n =\n",
    "\\beta\\left\\vert\\frac{u^n-u^{n-1}}{\\Delta t}\\right\\vert\n",
    "\\frac{u^{n+1}-u^n}{\\Delta t}\\thinspace .\n",
    "\\label{vib:ode2:nonlin:fbdiff} \\tag{15}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The forward and backward differences both have an error proportional\n",
    "to $\\Delta t$ so one may think the discretization above leads to\n",
    "a first-order scheme.\n",
    "However, by looking at the formulas, we realize that the forward-backward\n",
    "differences in ([15](#vib:ode2:nonlin:fbdiff))\n",
    "result in exactly the same scheme as in\n",
    "([11](#vib:ode2:step3b:quad)) where we\n",
    "used a geometric mean and centered differences and committed errors\n",
    "of size $\\Oof{\\Delta t^2}$. Therefore, the forward-backward\n",
    "differences in ([15](#vib:ode2:nonlin:fbdiff))\n",
    "act in a symmetric way and actually produce a second-order\n",
    "accurate discretization of the quadratic damping term.\n",
    "\n",
    "\n",
    "## Implementation\n",
    "<div id=\"vib:ode2:solver\"></div>\n",
    "\n",
    "\n",
    "The algorithm arising from the methods in the sections [A centered scheme for linear damping](#vib:ode2:fdm:flin)\n",
    "and [A centered scheme for quadratic damping](#vib:ode2:fdm:fquad) is very similar to the undamped case in\n",
    "the section [vib:ode1:fdm](#vib:ode1:fdm). The difference is\n",
    "basically a question of different formulas for $u^1$ and\n",
    "$u^{n+1}$. This is actually quite remarkable. The equation\n",
    "([1](#vib:ode2)) is normally impossible to solve by pen and paper, but\n",
    "possible for some special choices of $F$, $s$, and $f$. On the\n",
    "contrary, the complexity of the\n",
    "nonlinear generalized model ([1](#vib:ode2)) versus the\n",
    "simple undamped model is not a big deal when we solve the\n",
    "problem numerically!\n",
    "\n",
    "The computational algorithm takes the form\n",
    "\n",
    "1. $u^0=I$\n",
    "\n",
    "2. compute $u^1$ from ([7](#vib:ode2:step4b)) if linear\n",
    "    damping or ([14](#vib:ode2:step4b:quad)) if quadratic damping\n",
    "\n",
    "3. for $n=1,2,\\ldots,N_t-1$:\n",
    "\n",
    "a. compute $u^{n+1}$ from ([5](#vib:ode2:step4)) if linear\n",
    "      damping or ([12](#vib:ode2:step4:quad)) if quadratic damping\n",
    "\n",
    "\n",
    "Modifying the `solver` function for the undamped case is fairly\n",
    "easy, the big difference being many more terms and if tests on\n",
    "the type of damping:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import sympy as sp\n",
    "from devito import Dimension, Constant, TimeFunction, Eq, solve, Operator\n",
    "\n",
    "\n",
    "def solver(I, V, m, b, s, F, dt, T, damping='linear'):\n",
    "    \"\"\"\n",
    "    Solve m*u'' + f(u') + s(u) = F(t) for t in (0,T],\n",
    "    u(0)=I and u'(0)=V,\n",
    "    by a central finite difference method with time step dt.\n",
    "    If damping is 'linear', f(u')=b*u, while if damping is\n",
    "    'quadratic', f(u')=b*u'*abs(u').\n",
    "    F(t) and s(u) are Python functions.\n",
    "    \"\"\"\n",
    "    dt = float(dt)\n",
    "    b = float(b)\n",
    "    m = float(m)\n",
    "    Nt = int(round(T/dt))\n",
    "    t = Dimension('t', spacing=Constant('h_t'))\n",
    "\n",
    "    u = TimeFunction(name='u', dimensions=(t,),\n",
    "                     shape=(Nt+1,), space_order=2)\n",
    "\n",
    "    u.data[0] = I\n",
    "\n",
    "    if damping == 'linear':\n",
    "        # dtc for central difference (default for time is forward, 1st order)\n",
    "        eqn = m*u.dt2 + b*u.dtc + s(u) - F(u)\n",
    "        stencil = Eq(u.forward, solve(eqn, u.forward))\n",
    "    elif damping == 'quadratic':\n",
    "        # fd_order set as backward derivative used is 1st order\n",
    "        eqn = m*u.dt2 + b*u.dt*sp.Abs(u.dtl(fd_order=1)) + s(u) - F(u)\n",
    "        stencil = Eq(u.forward, solve(eqn, u.forward))\n",
    "\n",
    "    # First timestep needs to have the backward timestep substituted\n",
    "    # Has to be done to the equation otherwise the stencil will have\n",
    "    # forward timestep on both sides\n",
    "    # FIXME: Doesn't look like you can do subs or solve on anything inside an Abs\n",
    "    eqn_init = eqn.subs(u.backward, u.forward-2*t.spacing*V)\n",
    "    stencil_init = Eq(u.forward, solve(eqn_init, u.forward))\n",
    "    # stencil_init = stencil.subs(u.backward, u.forward-2*t.spacing*V)\n",
    "\n",
    "    op_init = Operator(stencil_init, name='first_timestep')\n",
    "    op = Operator(stencil, name='main_loop')\n",
    "    op_init.apply(h_t=dt, t_M=1)\n",
    "    op.apply(h_t=dt, t_m=1, t_M=Nt-1)\n",
    "\n",
    "    return u.data, np.linspace(0, Nt*dt, Nt+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The complete code resides in the file [`vib.py`](${src_vib}/vib.py).\n",
    "\n",
    "## Verification\n",
    "<div id=\"vib:ode2:verify\"></div>\n",
    "\n",
    "### Constant solution\n",
    "\n",
    "For debugging and initial verification, a constant solution is often\n",
    "very useful. We choose $\\uex(t)=I$, which implies $V=0$.\n",
    "Inserted in the ODE, we get\n",
    "$F(t)=s(I)$ for any choice of $f$. Since the discrete derivative\n",
    "of a constant vanishes (in particular, $[D_{2t}I]^n=0$,\n",
    "$[D_tI]^n=0$, and $[D_tD_t I]^n=0$), the constant solution also fulfills\n",
    "the discrete equations. The constant should therefore be reproduced\n",
    "to machine precision. The function `test_constant` in `vib.py`\n",
    "implements this test.\n",
    "\n",
    "### Linear solution\n",
    "\n",
    "Now we choose a linear solution: $\\uex = ct + d$. The initial condition\n",
    "$u(0)=I$ implies $d=I$, and $u^{\\prime}(0)=V$ forces $c$ to be $V$.\n",
    "Inserting $\\uex=Vt+I$ in the ODE with linear damping results in"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "0 + bV + s(Vt+I) = F(t),\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "while quadratic damping requires the source term"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "0 + b|V|V + s(Vt+I) = F(t)\\thinspace .\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the finite difference approximations used to compute $u^{\\prime}$ all\n",
    "are exact for a linear function, it turns out that the linear $\\uex$\n",
    "is also a solution of the discrete equations.\n",
    "[vib:exer:verify:gen:linear](#vib:exer:verify:gen:linear) asks you to carry out\n",
    "all the details.\n",
    "\n",
    "### Quadratic solution\n",
    "\n",
    "Choosing $\\uex = bt^2 + Vt + I$, with $b$ arbitrary,\n",
    "fulfills the initial conditions and\n",
    "fits the ODE if $F$ is adjusted properly. The solution also solves\n",
    "the discrete equations with linear damping. However, this quadratic\n",
    "polynomial in $t$ does not fulfill the discrete equations in case\n",
    "of quadratic damping, because the geometric mean used in the approximation\n",
    "of this term introduces an error.\n",
    "Doing [vib:exer:verify:gen:linear](#vib:exer:verify:gen:linear) will reveal\n",
    "the details. One can fit $F^n$ in the discrete equations such that\n",
    "the quadratic polynomial is reproduced by the numerical method (to\n",
    "machine precision).\n",
    "\n",
    "<!-- More: classes, cases with pendulum approx u vs sin(u), -->\n",
    "<!-- making UI via parampool -->\n",
    "\n",
    "## Visualization\n",
    "<div id=\"vib:ode2:viz\"></div>\n",
    "\n",
    "The functions for visualizations differ significantly from\n",
    "those in the undamped case in the `vib_undamped.py` program because,\n",
    "in the present general case, we do not have an exact solution to\n",
    "include in the plots. Moreover, we have no good estimate of\n",
    "the periods of the oscillations as there will be one period\n",
    "determined by the system parameters, essentially the\n",
    "approximate frequency $\\sqrt{s'(0)/m}$ for linear $s$ and small damping,\n",
    "and one period dictated by $F(t)$ in case the excitation is periodic.\n",
    "This is, however,\n",
    "nothing that the program can depend on or make use of.\n",
    "Therefore, the user has to specify $T$ and the window width\n",
    "to get a plot that moves with the graph and shows\n",
    "the most recent parts of it in long time simulations.\n",
    "\n",
    "The `vib.py` code\n",
    "contains several functions for analyzing the time series signal\n",
    "and for visualizing the solutions.\n",
    "\n",
    "## User interface\n",
    "<div id=\"vib:ode2:ui\"></div>\n",
    "\n",
    "\n",
    "The `main` function is changed substantially from\n",
    "the `vib_undamped.py` code, since we need to\n",
    "specify the new data $c$, $s(u)$, and $F(t)$.  In addition, we must\n",
    "set $T$ and the plot window width (instead of the number of periods we\n",
    "want to simulate as in `vib_undamped.py`). To figure out whether we\n",
    "can use one plot for the whole time series or if we should follow the\n",
    "most recent part of $u$, we can use the `plot_empricial_freq_and_amplitude`\n",
    "function's estimate of the number of local maxima. This number is now\n",
    "returned from the function and used in `main` to decide on the\n",
    "visualization technique."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "def main():\n",
    "    import argparse\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--I', type=float, default=1.0)\n",
    "    parser.add_argument('--V', type=float, default=0.0)\n",
    "    parser.add_argument('--m', type=float, default=1.0)\n",
    "    parser.add_argument('--c', type=float, default=0.0)\n",
    "    parser.add_argument('--s', type=str, default='u')\n",
    "    parser.add_argument('--F', type=str, default='0')\n",
    "    parser.add_argument('--dt', type=float, default=0.05)\n",
    "    parser.add_argument('--T', type=float, default=140)\n",
    "    parser.add_argument('--damping', type=str, default='linear')\n",
    "    parser.add_argument('--window_width', type=float, default=30)\n",
    "    parser.add_argument('--savefig', action='store_true')\n",
    "    a = parser.parse_args()\n",
    "    from scitools.std import StringFunction\n",
    "    s = StringFunction(a.s, independent_variable='u')\n",
    "    F = StringFunction(a.F, independent_variable='t')\n",
    "    I, V, m, c, dt, T, window_width, savefig, damping = \\\n",
    "       a.I, a.V, a.m, a.c, a.dt, a.T, a.window_width, a.savefig, \\\n",
    "       a.damping\n",
    "\n",
    "    u, t = solver(I, V, m, c, s, F, dt, T)\n",
    "    num_periods = empirical_freq_and_amplitude(u, t)\n",
    "    if num_periods <= 15:\n",
    "        figure()\n",
    "        visualize(u, t)\n",
    "    else:\n",
    "        visualize_front(u, t, window_width, savefig)\n",
    "    show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The program `vib.py` contains\n",
    "the above code snippets and can solve the model problem\n",
    "([1](#vib:ode2)). As a demo of `vib.py`, we consider the case\n",
    "$I=1$, $V=0$, $m=1$, $c=0.03$, $s(u)=\\sin(u)$, $F(t)=3\\cos(4t)$,\n",
    "$\\Delta t = 0.05$, and $T=140$. The relevant command to run is"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "        Terminal> python vib.py --s 'sin(u)' --F '3*cos(4*t)' --c 0.03\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This results in a [moving window following the function](${doc_notes}/vib/html//mov-vib/vib_generalized_dt0.05/index.html) on the screen.\n",
    "[Figure](#vib:ode2:fig:demo) shows a part of the time series.\n",
    "\n",
    "<!-- dom:FIGURE: [fig-vib/vib_gen_demo.png, width=600 frac=1.0] Damped oscillator excited by a sinusoidal function. <div id=\"vib:ode2:fig:demo\"></div> -->\n",
    "<!-- begin figure -->\n",
    "<div id=\"vib:ode2:fig:demo\"></div>\n",
    "\n",
    "<p>Damped oscillator excited by a sinusoidal function.</p>\n",
    "<img src=\"fig-vib/vib_gen_demo.png\" width=600>\n",
    "\n",
    "<!-- end figure -->\n",
    "\n",
    "\n",
    "## The Euler-Cromer scheme for the generalized model\n",
    "<div id=\"vib:ode2:EulerCromer\"></div>\n",
    "\n",
    "The ideas of the Euler-Cromer method from the section [vib:model2x2:EulerCromer](#vib:model2x2:EulerCromer)\n",
    "carry over to the generalized model. We write ([1](#vib:ode2))\n",
    "as two equations for $u$ and $v=u^{\\prime}$. The first equation is taken as the\n",
    "one with $v^{\\prime}$ on the left-hand side:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:EulerCromer:veq\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "v^{\\prime} = \\frac{1}{m}(F(t)-s(u)-f(v)),\n",
    "\\label{vib:ode2:EulerCromer:veq} \\tag{16}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:EulerCromer:ueq\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "u^{\\prime} = v\\thinspace .\n",
    "\\label{vib:ode2:EulerCromer:ueq} \\tag{17}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, the idea is to step ([16](#vib:ode2:EulerCromer:veq)) forward using\n",
    "a standard Forward Euler method, while we update $u$ from\n",
    "([17](#vib:ode2:EulerCromer:ueq)) with a Backward Euler method,\n",
    "utilizing the recent, computed $v^{n+1}$ value. In detail,"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:EulerCromer:dveq0a\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "\\frac{v^{n+1}-v^n}{\\Delta t} = \\frac{1}{m}(F(t_n)-s(u^n)-f(v^n)),\n",
    "\\label{vib:ode2:EulerCromer:dveq0a} \\tag{18}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:EulerCromer:dueq0a\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "\\frac{u^{n+1}-u^n}{\\Delta t} = v^{n+1},\n",
    "\\label{vib:ode2:EulerCromer:dueq0a} \\tag{19}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "resulting in the explicit scheme"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:EulerCromer:dveq\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "v^{n+1} = v^n + \\Delta t\\frac{1}{m}(F(t_n)-s(u^n)-f(v^n)),\n",
    "\\label{vib:ode2:EulerCromer:dveq} \\tag{20}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:EulerCromer:dueq0\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "u^{n+1} = u^n + \\Delta t\\,v^{n+1}\\thinspace .\n",
    "\\label{vib:ode2:EulerCromer:dueq0} \\tag{21}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We immediately note one very favorable feature of this scheme: all the\n",
    "nonlinearities in $s(u)$ and $f(v)$ are evaluated at a previous time\n",
    "level. This makes the Euler-Cromer method easier to apply and\n",
    "hence much more convenient than the centered scheme for the second-order\n",
    "ODE ([1](#vib:ode2)).\n",
    "\n",
    "The initial conditions are trivially set as"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto4\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "v^0 = V,\n",
    "\\label{_auto4} \\tag{22}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto5\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "u^0 = I\\thinspace .\n",
    "\\label{_auto5} \\tag{23}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[hpl 1: odespy for the generalized problem]\n",
    "\n",
    "<!-- Discussion of which method is best: -->\n",
    "<!-- <http://gamedev.stackexchange.com/questions/33694/pros-and-cons-of-different-integrators> -->\n",
    "\n",
    "## The Stoermer-Verlet algorithm for the generalized model\n",
    "<div id=\"vib:model2x2:gen:StormerVerlet\"></div>\n",
    "\n",
    "We can easily apply the ideas from the section [vib:model2x2:StormerVerlet](#vib:model2x2:StormerVerlet) to\n",
    "extend that method to the generalized model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\begin{align*}\n",
    "v^{\\prime} &= \\frac{1}{m}(F(t)-s(u)-f(v)),\\\\ \n",
    "u^{\\prime} &= v\\thinspace .\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, since the scheme is essentially centered differences for\n",
    "the ODE system on a staggered mesh, we do not go into detail here,\n",
    "but refer to the section [A staggered Euler-Cromer scheme for a generalized model](#vib:ode2:staggered).\n",
    "\n",
    "## A staggered Euler-Cromer scheme for a generalized model\n",
    "<div id=\"vib:ode2:staggered\"></div>\n",
    "\n",
    "The more general model for vibration problems,"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto6\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "mu'' + f(u') + s(u) = F(t),\\quad u(0)=I,\\ u'(0)=V,\\ t\\in (0,T],\n",
    "\\label{_auto6} \\tag{24}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "can be rewritten as a first-order ODE system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:staggered:veq\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "v' = m^{-1}\\left(F(t) - f(v) - s(u)\\right),\n",
    "\\label{vib:ode2:staggered:veq} \\tag{25}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:staggered:ueq\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "u' = v\\thinspace .\n",
    "\\label{vib:ode2:staggered:ueq} \\tag{26}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "mathcal{I}_t is natural to introduce a staggered mesh (see the section [vib:model2x2:staggered](#vib:model2x2:staggered)) and seek $u$ at mesh points $t_n$ (the numerical value is\n",
    "denoted by $u^n$) and $v$ between mesh points at $t_{n+1/2}$ (the numerical\n",
    "value is denoted by $v^{n+\\frac{1}{2}}$).\n",
    "A centered difference approximation to ([26](#vib:ode2:staggered:ueq))-([25](#vib:ode2:staggered:veq)) can then be written in operator notation as"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:staggered:dveq\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "\\lbrack D_tv = m^{-1}\\left(F(t) - f(v) - s(u)\\right)\\rbrack^n,\n",
    "\\label{vib:ode2:staggered:dveq} \\tag{27}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:staggered:dueq\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "\\lbrack D_t u = v\\rbrack^{n+\\frac{1}{2}}\\thinspace .\n",
    "\\label{vib:ode2:staggered:dueq} \\tag{28}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Written out,"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:staggered:dveq2\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "\\frac{v^{n+\\frac{1}{2}} - v^{n-\\frac{1}{2}}}{\\Delta t}\n",
    "= m^{-1}\\left(F^n - f(v^n) - s(u^n)\\right),\n",
    "\\label{vib:ode2:staggered:dveq2} \\tag{29}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:staggered:dueq2\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "\\frac{u^n - u^{n-1}}{\\Delta t} = v^{n+\\frac{1}{2}}\\thinspace .\n",
    "\\label{vib:ode2:staggered:dueq2} \\tag{30}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With linear damping, $f(v)=bv$, we can use an arithmetic mean\n",
    "for $f(v^n)$: $f(v^n)\\approx = \\frac{1}{2}(f(v^{n-\\frac{1}{2}}) +\n",
    "f(v^{n+\\frac{1}{2}}))$. The system\n",
    "([29](#vib:ode2:staggered:dveq2))-([30](#vib:ode2:staggered:dueq2))\n",
    "can then be solved with respect to the unknowns $u^n$ and $v^{n+\\frac{1}{2}}$:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:staggered:v:scheme:lin\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "v^{n+\\frac{1}{2}} = \\left(1 + \\frac{b}{2m}\\Delta t\\right)^{-1}\\left(\n",
    "v^{n-\\frac{1}{2}} + {\\Delta t}\n",
    "m^{-1}\\left(F^n - {\\frac{1}{2}}f(v^{n-\\frac{1}{2}}) - s(u^n)\\right)\\right),\n",
    "\\label{vib:ode2:staggered:v:scheme:lin} \\tag{31}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:staggered:u:scheme:lin\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "u^n  = u^{n-1} + {\\Delta t}v^{n-\\frac{1}{2}}\\thinspace .\n",
    "\\label{vib:ode2:staggered:u:scheme:lin} \\tag{32}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In case of quadratic damping, $f(v)=b|v|v$, we can use a geometric mean:\n",
    "$f(v^n)\\approx b|v^{n-\\frac{1}{2}}|v^{n+\\frac{1}{2}}$. Inserting this approximation\n",
    "in ([29](#vib:ode2:staggered:dveq2))-([30](#vib:ode2:staggered:dueq2)) and\n",
    "solving for the unknowns $u^n$ and $v^{n+\\frac{1}{2}}$ results in"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:staggered:v:scheme:quad\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "v^{n+\\frac{1}{2}} = (1 + \\frac{b}{m}|v^{n-\\frac{1}{2}}|\\Delta t)^{-1}\\left(\n",
    "v^{n-\\frac{1}{2}} + {\\Delta t}\n",
    "m^{-1}\\left(F^n - s(u^n)\\right)\\right),\n",
    "\\label{vib:ode2:staggered:v:scheme:quad} \\tag{33}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:staggered:u:scheme:quad\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "u^n  = u^{n-1} + {\\Delta t}v^{n-\\frac{1}{2}}\\thinspace .\n",
    "\\label{vib:ode2:staggered:u:scheme:quad} \\tag{34}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The initial conditions are derived at the end of\n",
    "the section [vib:model2x2:staggered](#vib:model2x2:staggered):"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:staggered:u02\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "u^0 = I,\n",
    "\\label{vib:ode2:staggered:u02} \\tag{35}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"vib:ode2:staggered:v02\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "v^\\frac{1}{2} = V - \\frac{1}{2}\\Delta t\\omega^2I\n",
    "\\label{vib:ode2:staggered:v02} \\tag{36}\\thinspace .\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[hpl 2: Must deal with two things: how to set initial conditions on $v$, and\n",
    "how to make phase-space plots with $v$ at integer mesh points. See\n",
    "article.]\n",
    "\n",
    "## The PEFRL 4th-order accurate algorithm\n",
    "<div id=\"vib:ode2:PEFRL\"></div>\n",
    "\n",
    "A variant of the Euler-Cromer type of algorithm, which provides an\n",
    "error $\\Oof{\\Delta t^4}$ if $f(v)=0$, is called PEFRL\n",
    "[[PEFRL_2002]](#PEFRL_2002). This algorithm is very well suited for integrating\n",
    "dynamic systems (especially those without damping) over very long time\n",
    "periods. Define"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "g(u,v) = \\frac{1}{m}(F(t)-s(u)-f(v))\\thinspace .\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The algorithm is explicit and features these steps:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto7\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "u^{n+1,1} = u^n       + \\xi\\Delta t v^n,\n",
    "\\label{_auto7} \\tag{37}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto8\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "v^{n+1,1} = v^n       + \\frac{1}{2}(1-2\\lambda)\\Delta t g(u^{n+1,1},v^n),\n",
    "\\label{_auto8} \\tag{38}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto9\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "u^{n+1,2} = u^{n+1,1} + \\chi\\Delta t v^{n+1,1},\n",
    "\\label{_auto9} \\tag{39}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto10\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "v^{n+1,2} = v^{n+1,1} + \\lambda\\Delta t g(u^{n+1,2}, v^{n+1,1}),\n",
    "\\label{_auto10} \\tag{40}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto11\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "u^{n+1,3} = u^{n+1,2} + (1-2(\\chi + \\xi))\\Delta t v^{n+1,2},\n",
    "\\label{_auto11} \\tag{41}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto12\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "v^{n+1,3} = v^{n+1,2} + \\lambda\\Delta t g(u^{n+1,3}, v^{n+1,2}),\n",
    "\\label{_auto12} \\tag{42}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto13\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "u^{n+1,4} = u^{n+1,3} + \\chi\\Delta t v^{n+1,3},\n",
    "\\label{_auto13} \\tag{43}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto14\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "v^{n+1}   = v^{n+1,3} + \\frac{1}{2}(1-2\\lambda)\\Delta t g(u^{n+1,4},v^{n+1,3}),\n",
    "\\label{_auto14} \\tag{44}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto15\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "u^{n+1}   = u^{n+1,4} + \\xi\\Delta t v^{n+1}\n",
    "\\label{_auto15} \\tag{45}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Compare with eq 22 in article -->\n",
    "The parameters $\\xi$, $\\lambda$, and $\\xi$ have the values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto16\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "\\xi = 0.1786178958448091,\n",
    "\\label{_auto16} \\tag{46}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto17\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "\\lambda = -0.2123418310626054,\n",
    "\\label{_auto17} \\tag{47}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Equation labels as ordinary links -->\n",
    "<div id=\"_auto18\"></div>\n",
    "\n",
    "$$\n",
    "\\begin{equation}  \n",
    "\\chi = -0.06626458266981849\n",
    "\\label{_auto18} \\tag{48}\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercises and Problems\n",
    "\n",
    "\n",
    "\n",
    "<!-- --- begin exercise --- -->\n",
    "\n",
    "## Exercise 1: Implement the solver via classes\n",
    "<div id=\"vib:exer:gen:class\"></div>\n",
    "\n",
    "Reimplement the `vib.py` program using a class `Problem` to hold all\n",
    "the physical parameters of the problem, a class `Solver` to hold the\n",
    "numerical parameters and compute the solution, and a class\n",
    "`Visualizer` to display the solution.\n",
    "\n",
    "<!-- --- begin hint in exercise --- -->\n",
    "\n",
    "**Hint.**\n",
    "Use the ideas and examples from Sections 5.5.1 and 5.5.2 \n",
    "in [[Langtangen_decay]](#Langtangen_decay).  More specifically, make a superclass\n",
    "`Problem` for holding the scalar physical parameters of a problem and\n",
    "let subclasses implement the $s(u)$ and $F(t)$ functions as methods.\n",
    "Try to call up as much existing functionality in `vib.py` as possible.\n",
    "\n",
    "<!-- --- end hint in exercise --- -->\n",
    "\n",
    "\n",
    "<!-- --- begin solution of exercise --- -->\n",
    "**Solution.**\n",
    "The complete code looks like this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'scitools'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-3-b591f37272a6>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mnumpy\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0;32mimport\u001b[0m \u001b[0mscitools\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstd\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mplt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0msympy\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0msym\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mvib\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0msolver\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mvib_solver\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'scitools'"
     ]
    }
   ],
   "source": [
    "# Reimplementation of vib.py using classes\n",
    "\n",
    "import numpy as np\n",
    "import scitools.std as plt\n",
    "import sympy as sym\n",
    "from vib import solver as vib_solver\n",
    "from vib import visualize as vib_visualize\n",
    "from vib import visualize_front as vib_visualize_front\n",
    "from vib import visualize_front_ascii as vib_visualize_front_ascii\n",
    "from vib import plot_empirical_freq_and_amplitude as \\\n",
    "                            vib_plot_empirical_freq_and_amplitude\n",
    "\n",
    "class Vibration:\n",
    "    '''\n",
    "    Problem: m*u'' + f(u') + s(u) = F(t) for t in (0,T],\n",
    "    u(0)=I and u'(0)=V. The problem is solved\n",
    "    by a central finite difference method with time step dt.\n",
    "    If damping is 'linear', f(u')=b*u, while if damping is\n",
    "    'quadratic', f(u')=b*u'*abs(u'). Zero damping is achieved\n",
    "    with b=0. F(t) and s(u) are Python functions.\n",
    "    '''\n",
    "    def __init__(self, I=1, V=0, m=1, b=0, damping='linear'):\n",
    "        self.I = I; self.V = V; self.m = m; self.b=b;\n",
    "        self.damping = damping\n",
    "    def s(self, u):\n",
    "        return u\n",
    "    def F(self, t):\n",
    "        '''Driving force. Zero implies free oscillations'''\n",
    "        return 0\n",
    "\n",
    "class Free_vibrations(Vibration):\n",
    "    '''F(t) = 0'''\n",
    "    def __init__(self, s=None, I=1, V=0, m=1, b=0, damping='linear'):\n",
    "        Vibration.__init__(self, I=I, V=V, m=m, b=b, damping=damping)\n",
    "        if s != None:\n",
    "            self.s = s\n",
    "\n",
    "class Forced_vibrations(Vibration):\n",
    "    '''F(t)! = 0'''\n",
    "    def __init__(self, F, s=None, I=1, V=0, m=1, b=0,\n",
    "                 damping='linear'):\n",
    "        Vibration.__init__(self, I=I, V=V, m=m, b=b,\n",
    "                           damping=damping)\n",
    "        if s != None:\n",
    "            self.s = s\n",
    "        self.F = F\n",
    "\n",
    "class Solver:\n",
    "    def __init__(self, dt=0.05, T=20):\n",
    "        self.dt = dt; self.T = T\n",
    "\n",
    "    def solve(self, problem):\n",
    "        self.u, self.t = vib_solver(\n",
    "            problem.I, problem.V,\n",
    "            problem.m, problem.b,\n",
    "            problem.s, problem.F,\n",
    "            self.dt, self.T, problem.damping)\n",
    "\n",
    "class Visualizer:\n",
    "    def __init__(self, problem, solver, window_width, savefig):\n",
    "        self.problem = problem; self.solver = solver\n",
    "        self.window_width = window_width; self.savefig = savefig\n",
    "    def visualize(self):\n",
    "        u = self.solver.u; t = self.solver.t   # short forms\n",
    "        num_periods = vib_plot_empirical_freq_and_amplitude(u, t)\n",
    "        if num_periods <= 40:\n",
    "            plt.figure()\n",
    "            vib_visualize(u, t)\n",
    "        else:\n",
    "            vib_visualize_front(u, t, self.window_width, self.savefig)\n",
    "            vib_visualize_front_ascii(u, t)\n",
    "        plt.show()\n",
    "\n",
    "def main():\n",
    "    # Note: the reading of parameter values would better be done\n",
    "    # from each relevant class, i.e. class Problem should read I, V,\n",
    "    #  etc., while class Solver should read dt and T, and so on.\n",
    "    # Consult, e.g., Langtangen: \"A Primer on Scientific Programming\",\n",
    "    # App E.\n",
    "    import argparse\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--I', type=float, default=1.0)\n",
    "    parser.add_argument('--V', type=float, default=0.0)\n",
    "    parser.add_argument('--m', type=float, default=1.0)\n",
    "    parser.add_argument('--b', type=float, default=0.0)\n",
    "    parser.add_argument('--s', type=str, default=None)\n",
    "    parser.add_argument('--F', type=str, default='0')\n",
    "    parser.add_argument('--dt', type=float, default=0.05)\n",
    "    parser.add_argument('--T', type=float, default=20)\n",
    "    parser.add_argument('--window_width', type=float, default=30.,\n",
    "                        help='Number of periods in a window')\n",
    "    parser.add_argument('--damping', type=str, default='linear')\n",
    "    parser.add_argument('--savefig', action='store_true')\n",
    "    # Hack to allow --SCITOOLS options\n",
    "    # (scitools.std reads this argument at import)\n",
    "    parser.add_argument('--SCITOOLS_easyviz_backend',\n",
    "                        default='matplotlib')\n",
    "    a = parser.parse_args()\n",
    "\n",
    "    from scitools.std import StringFunction\n",
    "    if a.s != None:\n",
    "        s = StringFunction(a.s, independent_variable='u')\n",
    "    else:\n",
    "        s = None\n",
    "    F = StringFunction(a.F, independent_variable='t')\n",
    "\n",
    "    if a.F == '0':  # free vibrations\n",
    "        problem = Free_vibrations(s=s, I=a.I, V=a.V, m=a.m, b=a.b,\n",
    "                                  damping=a.damping)\n",
    "    else:   # forced vibrations\n",
    "        problem = Forced_vibrations(lambda t: np.sin(t),\n",
    "                                    s=s, I=a.I, V=a.V,\n",
    "                                    m=a.m, b=a.b, damping=a.damping)\n",
    "\n",
    "    solver = Solver(dt=a.dt, T=a.T)\n",
    "    solver.solve(problem)\n",
    "\n",
    "    visualizer = Visualizer(problem, solver,\n",
    "                            a.window_width, a.savefig)\n",
    "    visualizer.visualize()\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- --- end solution of exercise --- -->\n",
    "Filename: `vib_class`.\n",
    "\n",
    "<!-- --- end exercise --- -->\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "<!-- --- begin exercise --- -->\n",
    "\n",
    "## Problem 2: Use a backward difference for the damping term\n",
    "<div id=\"vib:exer:quad:damping:bw\"></div>\n",
    "\n",
    "As an alternative to discretizing the damping terms $\\beta u^{\\prime}$ and\n",
    "$\\beta |u^{\\prime}|u^{\\prime}$ by centered differences, we may apply\n",
    "backward differences:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\begin{align*}\n",
    "[u^{\\prime}]^n &\\approx [D_t^-u]^n,\\\\ \n",
    "[|u^{\\prime}|u^{\\prime}]^n &\\approx [|D_t^-u|D_t^-u]^n\\\\ \n",
    "&= |[D_t^-u]^n|[D_t^-u]^n\\thinspace .\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The advantage of the backward difference is that the damping term is\n",
    "evaluated using known values $u^n$ and $u^{n-1}$ only.\n",
    "Extend the [`vib.py`](${src_vib}/vib.py) code with a scheme based\n",
    "on using backward differences in the damping terms. Add statements\n",
    "to compare the original approach with centered difference and the\n",
    "new idea launched in this exercise. Perform numerical experiments\n",
    "to investigate how much accuracy that is lost by using the backward\n",
    "differences.\n",
    "\n",
    "\n",
    "<!-- --- begin solution of exercise --- -->\n",
    "**Solution.**\n",
    "The new discretization approach of the linear and quadratic damping\n",
    "terms calls for new derivations of the updating formulas (for $u$) in\n",
    "the solver. Since backward difference approximations will be used for\n",
    "the damping term, we may also use this approximation for the initial\n",
    "condition on $u^{\\prime}(0)$ without deteriorating the convergence\n",
    "rate any further. Note that introducing backward difference\n",
    "approximations for the damping term make our computational schemes\n",
    "first order, as opposed to the original second order schemes which\n",
    "used central difference approximations also for the damping terms. The\n",
    "motivation for also using a backward difference approximation for the\n",
    "initial condition on $u^{\\prime}(0)$, is simply that the computational\n",
    "schemes get much simpler.\n",
    "\n",
    "With linear damping, the new discretized form of the equation reads"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "m\\frac{u^{n+1}-2u^n+u^{n-1}}{\\Delta t^2} + b\\frac{u^n - u^{n-1}}{\\Delta t} + s(u^n) = F^n,\n",
    "\\nonumber\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which gives us"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "u^{n+1} = \\left(2-\\frac{\\Delta t b}{m}\\right)u^n + \\frac{\\Delta t^2}{m}\\left(F^n - s(u^n)\\right) + \\left(\\frac{\\Delta t b}{m} - 1\\right)u^{n-1}.\n",
    "\\nonumber\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With $n=0$, the updating formula becomes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "u^1 = \\left(2-\\frac{\\Delta t b}{m}\\right)u^0 + \\frac{\\Delta t^2}{m}\\left(F^0 - s(u^0)\\right) + \\left(\\frac{\\Delta t b}{m} - 1\\right)u^{-1},\n",
    "\\nonumber\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which requires some further elaboration because of the unknown $u^{-1}$. We handle this by discretizing the initial condition $u^{\\prime}(0) = V$ by a backward difference approximation as"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\frac{u^0 - u^{-1}}{\\Delta t} = V,\n",
    "\\nonumber\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which implies that"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "u^{-1} = u^0 - \\Delta t V.\n",
    "\\nonumber\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inserting this expression for $u^{-1}$ in the updating formula for $u^{n+1}$, and simplifying,\n",
    "gives us the following special formula for the first time step:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "u^1 = u^0 + \\Delta t V + \\frac{\\Delta t^2}{m}\\left(-bV - s(u^0) + F^0\\right).\n",
    "\\nonumber\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Switching to quadratic damping, the new discretized form of the equations becomes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "m\\frac{u^{n+1}-2u^n+u^{n-1}}{\\Delta t^2} + b|\\frac{u^n - u^{n-1}}{\\Delta t}|\\frac{u^n - u^{n-1}}{\\Delta t} + s(u^n) = F^n,\n",
    "\\nonumber\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which leads to"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "u^{n+1} = 2u^n - u^{n-1} - \\frac{b}{m}|u^n - u^{n-1}|(u^n - u^{n-1}) + \\frac{\\Delta t^2}{m}\\left(F^n - s(u^n)\\right).\n",
    "\\nonumber\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With $n=0$, this updating formula becomes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "u^1 = 2u^0 - u^{-1} - \\frac{b}{m}|u^0 - u^{-1}|(u^0 - u^{-1}) + \\frac{\\Delta t^2}{m}\\left(F^0 - s(u^0)\\right).\n",
    "\\nonumber\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we handle the unknown $u^{-1}$ via the same expression as above, which be derived from a backward difference approximation to the initial condition on the derivative. Inserting this expression for $u^{-1}$ and simplifying, gives the special updating formula for $u^1$ as"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "u^1 = u^0 + \\Delta t V + \\frac{\\Delta t^2}{m}\\left(-b|V|V - s(u^0) + F^0\\right).\n",
    "\\nonumber\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We implement these new computational schemes in a new solver function\n",
    "`solver_bwdamping`, so that the discrete solution for $u$ can be found\n",
    "by both the original and the new solver. The difference between the\n",
    "two different solutions is then visualized in the same way as the\n",
    "original solution in `main`.\n",
    "\n",
    "The convergence rates computed in `test_mms` demonstrates that our\n",
    "scheme now is a first order scheme, as $r$ is seen to approach 1.0\n",
    "with decreasing $\\Delta t$.\n",
    "\n",
    "Both solvers reproduce a constant solution exactly (within machine\n",
    "precision), whereas sinusoidal and quadratic solutions differ, as\n",
    "should be expected after comparing the schemes. Pointing out the\n",
    "\"best\" approach is difficult: the backward differences yield a much\n",
    "simpler mathematical problem to be solved, while the more complicated\n",
    "method converges faster and gives more accuracy for the same cost. On\n",
    "the other hand, the backward differences can yield any reasonable\n",
    "accuracy by lowering $\\Delta t$, and the results are obtained within a\n",
    "few seconds on a laptop.\n",
    "\n",
    "Here is the complete computer code, arising from copying `vib.py` and changing\n",
    "the functions that have to be changed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "#import matplotlib.pyplot as plt\n",
    "import scitools.std as plt\n",
    "\n",
    "def solver_bwdamping(I, V, m, b, s, F, dt, T, damping='linear'):\n",
    "    \"\"\"\n",
    "    Solve m*u'' + f(u') + s(u) = F(t) for t in (0,T],\n",
    "    u(0)=I and u'(0)=V. All terms except damping is discretized\n",
    "    by a central finite difference method with time step dt.\n",
    "    The damping term is discretized by a backward diff. approx.,\n",
    "    as is the init.cond. u'(0). If damping is 'linear', f(u')=b*u,\n",
    "    while if damping is 'quadratic', f(u')=b*u'*abs(u').\n",
    "    F(t) and s(u) are Python functions.\n",
    "    \"\"\"\n",
    "    dt = float(dt); b = float(b); m = float(m) # avoid integer div.\n",
    "    Nt = int(round(T/dt))\n",
    "    u = np.zeros(Nt+1)\n",
    "    t = np.linspace(0, Nt*dt, Nt+1)\n",
    "\n",
    "    u_original = np.zeros(Nt+1); u_original[0] = I  # for testing\n",
    "\n",
    "    u[0] = I\n",
    "    if damping == 'linear':\n",
    "        u[1] = u[0] + dt*V + dt**2/m*(-b*V - s(u[0]) + F(t[0]))\n",
    "    elif damping == 'quadratic':\n",
    "        u[1] = u[0] + dt*V + \\\n",
    "               dt**2/m*(-b*V*abs(V) - s(u[0]) + F(t[0]))\n",
    "    for n in range(1, Nt):\n",
    "        if damping == 'linear':\n",
    "            u[n+1] = (2 - dt*b/m)*u[n] + dt**2/m*(- s(u[n]) + \\\n",
    "                                F(t[n])) + (dt*b/m - 1)*u[n-1]\n",
    "        elif damping == 'quadratic':\n",
    "            u[n+1] = 2*u[n] - u[n-1] - b/m*abs(u[n] - \\\n",
    "            u[n-1])*(u[n] - u[n-1]) + dt**2/m*(-s(u[n]) + F(t[n]))\n",
    "    return u, t\n",
    "\n",
    "\n",
    "import sympy as sym\n",
    "\n",
    "def test_constant():\n",
    "    \"\"\"Verify a constant solution.\"\"\"\n",
    "    u_exact = lambda t: I\n",
    "    I = 1.2; V = 0; m = 2; b = 0.9\n",
    "    w = 1.5\n",
    "    s = lambda u: w**2*u\n",
    "    F = lambda t: w**2*u_exact(t)\n",
    "    dt = 0.2\n",
    "    T = 2\n",
    "    #u, t = solver(I, V, m, b, s, F, dt, T, 'linear')\n",
    "    u, t = solver_bwdamping(I, V, m, b, s, F, dt, T, 'linear')\n",
    "    difference = np.abs(u_exact(t) - u).max()\n",
    "    print difference\n",
    "    tol = 1E-13\n",
    "    assert difference < tol\n",
    "\n",
    "    #u, t = solver(I, V, m, b, s, F, dt, T, 'quadratic')\n",
    "    u, t = solver_bwdamping(I, V, m, b, s, F, dt, T, 'quadratic')\n",
    "    difference = np.abs(u_exact(t) - u).max()\n",
    "    print difference\n",
    "    assert difference < tol\n",
    "\n",
    "def lhs_eq(t, m, b, s, u, damping='linear'):\n",
    "    \"\"\"Return lhs of differential equation as sympy expression.\"\"\"\n",
    "    v = sym.diff(u, t)\n",
    "    if damping == 'linear':\n",
    "        return m*sym.diff(u, t, t) + b*v + s(u)\n",
    "    else:\n",
    "        return m*sym.diff(u, t, t) + b*v*sym.Abs(v) + s(u)\n",
    "\n",
    "def test_quadratic():\n",
    "    \"\"\"Verify a quadratic solution.\"\"\"\n",
    "    I = 1.2; V = 3; m = 2; b = 0.9\n",
    "    s = lambda u: 4*u\n",
    "    t = sym.Symbol('t')\n",
    "    dt = 0.2\n",
    "    T = 2\n",
    "\n",
    "    q = 2  # arbitrary constant\n",
    "    u_exact = I + V*t + q*t**2\n",
    "    F = sym.lambdify(t, lhs_eq(t, m, b, s, u_exact, 'linear'))\n",
    "    u_exact = sym.lambdify(t, u_exact, modules='numpy')\n",
    "    #u1, t1 = solver(I, V, m, b, s, F, dt, T, 'linear')\n",
    "    u1, t1 = solver_bwdamping(I, V, m, b, s, F, dt, T, 'linear')\n",
    "    diff = np.abs(u_exact(t1) - u1).max()\n",
    "    print diff\n",
    "    tol = 1E-13\n",
    "    #assert diff < tol\n",
    "\n",
    "    # In the quadratic damping case, u_exact must be linear\n",
    "    # in order to exactly recover this solution\n",
    "    u_exact = I + V*t\n",
    "    F = sym.lambdify(t, lhs_eq(t, m, b, s, u_exact, 'quadratic'))\n",
    "    u_exact = sym.lambdify(t, u_exact, modules='numpy')\n",
    "    #u2, t2 = solver(I, V, m, b, s, F, dt, T, 'quadratic')\n",
    "    u2, t2 = solver_bwdamping(I, V, m, b, s, F, dt, T, 'quadratic')\n",
    "    diff = np.abs(u_exact(t2) - u2).max()\n",
    "    print diff\n",
    "    assert diff < tol\n",
    "\n",
    "def test_sinusoidal():\n",
    "    \"\"\"Verify a numerically exact sinusoidal solution when b=F=0.\"\"\"\n",
    "    from math import asin\n",
    "\n",
    "    def u_exact(t):\n",
    "        w_numerical = 2/dt*np.arcsin(w*dt/2)\n",
    "        return I*np.cos(w_numerical*t)\n",
    "\n",
    "    I = 1.2; V = 0; m = 2; b = 0\n",
    "    w = 1.5  # fix the frequency\n",
    "    s = lambda u: m*w**2*u\n",
    "    F = lambda t: 0\n",
    "    dt = 0.2\n",
    "    T = 6\n",
    "    #u, t = solver(I, V, m, b, s, F, dt, T, 'linear')\n",
    "    u, t = solver_bwdamping(I, V, m, b, s, F, dt, T, 'linear')\n",
    "    diff = np.abs(u_exact(t) - u).max()\n",
    "    print diff\n",
    "    tol = 1E-14\n",
    "    #assert diff < tol\n",
    "\n",
    "    #u, t = solver(I, V, m, b, s, F, dt, T, 'quadratic')\n",
    "    u, t = solver_bwdamping(I, V, m, b, s, F, dt, T, 'quadratic')\n",
    "    print diff\n",
    "    diff = np.abs(u_exact(t) - u).max()\n",
    "    assert diff < tol\n",
    "\n",
    "def test_mms():\n",
    "    \"\"\"Use method of manufactured solutions.\"\"\"\n",
    "    m = 4.; b = 1\n",
    "    w = 1.5\n",
    "    t = sym.Symbol('t')\n",
    "    u_exact = 3*sym.exp(-0.2*t)*sym.cos(1.2*t)\n",
    "    I = u_exact.subs(t, 0).evalf()\n",
    "    V = sym.diff(u_exact, t).subs(t, 0).evalf()\n",
    "    u_exact_py = sym.lambdify(t, u_exact, modules='numpy')\n",
    "    s = lambda u: u**3\n",
    "    dt = 0.2\n",
    "    T = 6\n",
    "    errors_linear = []\n",
    "    errors_quadratic = []\n",
    "    # Run grid refinements and compute exact error\n",
    "    for i in range(5):\n",
    "        F_formula = lhs_eq(t, m, b, s, u_exact, 'linear')\n",
    "        F = sym.lambdify(t, F_formula)\n",
    "        #u1, t1 = solver(I, V, m, b, s, F, dt, T, 'linear')\n",
    "        u1, t1 = solver_bwdamping(I, V, m, b, s,\n",
    "                                  F, dt, T, 'linear')\n",
    "        error = np.sqrt(np.sum((u_exact_py(t1) - u1)**2)*dt)\n",
    "        errors_linear.append((dt, error))\n",
    "\n",
    "        F_formula = lhs_eq(t, m, b, s, u_exact, 'quadratic')\n",
    "        #print sym.latex(F_formula, mode='plain')\n",
    "        F = sym.lambdify(t, F_formula)\n",
    "        #u2, t2 = solver(I, V, m, b, s, F, dt, T, 'quadratic')\n",
    "        u2, t2 = solver_bwdamping(I, V, m, b, s,\n",
    "                                  F, dt, T, 'quadratic')\n",
    "        error = np.sqrt(np.sum((u_exact_py(t2) - u2)**2)*dt)\n",
    "        errors_quadratic.append((dt, error))\n",
    "        dt /= 2\n",
    "    # Estimate convergence rates\n",
    "    tol = 0.05\n",
    "    for errors in errors_linear, errors_quadratic:\n",
    "        for i in range(1, len(errors)):\n",
    "            dt, error = errors[i]\n",
    "            dt_1, error_1 = errors[i-1]\n",
    "            r = np.log(error/error_1)/np.log(dt/dt_1)\n",
    "        # check r for final simulation with (final and) smallest dt\n",
    "        # note that the method now is 1st order, i.e. r should\n",
    "        # approach 1.0\n",
    "        print r\n",
    "        assert abs(r - 1.0) < tol\n",
    "\n",
    "import os, sys\n",
    "sys.path.insert(0, os.path.join(os.pardir, 'src-vib'))\n",
    "from vib import (plot_empirical_freq_and_amplitude,\n",
    "                 visualize_front, visualize_front_ascii,\n",
    "                 minmax, periods, amplitudes,\n",
    "                 solver as solver2)\n",
    "\n",
    "def visualize(list_of_curves, legends, title='', filename='tmp'):\n",
    "    \"\"\"Plot list of curves: (u, t).\"\"\"\n",
    "    for u, t in list_of_curves:\n",
    "        plt.plot(t, u)\n",
    "        plt.hold('on')\n",
    "    plt.legend(legends)\n",
    "    plt.xlabel('t')\n",
    "    plt.ylabel('u')\n",
    "    plt.title(title)\n",
    "    plt.savefig(filename + '.png')\n",
    "    plt.savefig(filename + '.pdf')\n",
    "    plt.show()\n",
    "\n",
    "def main():\n",
    "    import argparse\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('--I', type=float, default=1.0)\n",
    "    parser.add_argument('--V', type=float, default=0.0)\n",
    "    parser.add_argument('--m', type=float, default=1.0)\n",
    "    parser.add_argument('--b', type=float, default=0.0)\n",
    "    parser.add_argument('--s', type=str, default='4*pi**2*u')\n",
    "    parser.add_argument('--F', type=str, default='0')\n",
    "    parser.add_argument('--dt', type=float, default=0.05)\n",
    "    parser.add_argument('--T', type=float, default=20)\n",
    "    parser.add_argument('--window_width', type=float, default=30.,\n",
    "                        help='Number of periods in a window')\n",
    "    parser.add_argument('--damping', type=str, default='linear')\n",
    "    parser.add_argument('--savefig', action='store_true')\n",
    "    # Hack to allow --SCITOOLS options\n",
    "    # (scitools.std reads this argument at import)\n",
    "    parser.add_argument('--SCITOOLS_easyviz_backend',\n",
    "                        default='matplotlib')\n",
    "    a = parser.parse_args()\n",
    "    from scitools.std import StringFunction\n",
    "    s = StringFunction(a.s, independent_variable='u')\n",
    "    F = StringFunction(a.F, independent_variable='t')\n",
    "    I, V, m, b, dt, T, window_width, savefig, damping = \\\n",
    "       a.I, a.V, a.m, a.b, a.dt, a.T, a.window_width, a.savefig, \\\n",
    "       a.damping\n",
    "\n",
    "    # compute u by both methods and then visualize the difference\n",
    "    u, t = solver2(I, V, m, b, s, F, dt, T, damping)\n",
    "    u_bw, _ = solver_bwdamping(I, V, m, b, s, F, dt, T, damping)\n",
    "    u_diff = u - u_bw\n",
    "\n",
    "    num_periods = plot_empirical_freq_and_amplitude(u_diff, t)\n",
    "    if num_periods <= 40:\n",
    "        plt.figure()\n",
    "        legends = ['1st-2nd order method',\n",
    "                   '2nd order method',\n",
    "                   '1st order method']\n",
    "        visualize([(u_diff, t), (u, t), (u_bw, t)], legends)\n",
    "    else:\n",
    "        visualize_front(u_diff, t, window_width, savefig)\n",
    "        #visualize_front_ascii(u_diff, t)\n",
    "    plt.show()\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()\n",
    "    #test_constant()\n",
    "    #test_sinusoidal()\n",
    "    #test_mms()\n",
    "    #test_quadratic()\n",
    "    raw_input()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a comparison of standard method (2nd order) and backward differences for damping (1st order) for 10 (left) and 40 (right) time steps per period:\n",
    "\n",
    "<!-- dom:FIGURE: [fig-vib/vib_gen_bwdamping.png, width=800 frac=1] -->\n",
    "<!-- begin figure -->\n",
    "\n",
    "<p></p>\n",
    "<img src=\"fig-vib/vib_gen_bwdamping.png\" width=800>\n",
    "\n",
    "<!-- end figure -->\n",
    "\n",
    "\n",
    "<!-- --- end solution of exercise --- -->\n",
    "Filename: `vib_gen_bwdamping`.\n",
    "\n",
    "<!-- --- end exercise --- -->\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "<!-- --- begin exercise --- -->\n",
    "\n",
    "## Exercise 3: Use the forward-backward scheme with quadratic damping\n",
    "<div id=\"vib:exer:quad:damping:fwbw\"></div>\n",
    "\n",
    "We consider the generalized model with quadratic damping, expressed\n",
    "as a system of two first-order equations as in the section [A staggered Euler-Cromer scheme for a generalized model](#vib:ode2:staggered):"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\begin{align*}\n",
    "u^{\\prime} &= v,\\\\ \n",
    "v' &= \\frac{1}{m}\\left( F(t) - \\beta |v|v - s(u)\\right)\\thinspace .\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, contrary to what is done in the section [A staggered Euler-Cromer scheme for a generalized model](#vib:ode2:staggered),\n",
    "we want to apply the idea of a forward-backward discretization:\n",
    "$u$ is marched forward by a one-sided Forward Euler scheme applied\n",
    "to the first equation, and\n",
    "thereafter $v$ can be marched forward by a Backward Euler scheme in the\n",
    "second\n",
    "% if BOOK == \"book\":\n",
    "equation, see in the section [vib:model2x2:EulerCromer](#vib:model2x2:EulerCromer).\n",
    "% else:\n",
    "equation.\n",
    "% endif\n",
    "Express the idea in operator notation and write out the\n",
    "scheme. Unfortunately, the backward difference for the $v$ equation\n",
    "creates a nonlinearity $|v^{n+1}|v^{n+1}$.  To linearize this\n",
    "nonlinearity, use the known value $v^n$ inside the absolute value\n",
    "factor, i.e., $|v^{n+1}|v^{n+1}\\approx |v^n|v^{n+1}$.  Show that the\n",
    "resulting scheme is equivalent to the one in the section [A staggered Euler-Cromer scheme for a generalized model](#vib:ode2:staggered) for some time level $n\\geq 1$.\n",
    "\n",
    "What we learn from this exercise is that the first-order differences\n",
    "and the linearization trick play together in \"the right way\" such that\n",
    "the scheme is as good as when we (in the section [A staggered Euler-Cromer scheme for a generalized model](#vib:ode2:staggered))\n",
    "carefully apply centered differences and a geometric mean on a\n",
    "staggered mesh to achieve second-order accuracy.\n",
    "% if BOOK == \"book\":\n",
    "There is a\n",
    "difference in the handling of the initial conditions, though, as\n",
    "explained at the end of the section [vib:model2x2:EulerCromer](#vib:model2x2:EulerCromer).\n",
    "% endif\n",
    "Filename: `vib_gen_bwdamping`.\n",
    "\n",
    "<!-- --- end exercise --- -->"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "devito",
   "language": "python",
   "name": "devito"
  },
  "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}