Devito-fied first cell of vib_gen

Author: EdCaunt

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

@@ -1,7 +1,7 @@
 {
  "cells": [
   {
-   "cell_type": "raw",
+   "cell_type": "markdown",
    "metadata": {},
    "source": [
     "<!-- Equation labels as ordinary links -->\n",
@@ -510,6 +510,11 @@
    "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",
@@ -519,27 +524,33 @@
     "    '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",
+    "    dt = float(dt)\n",
+    "    b = float(b)\n",
+    "    m = float(m)\n",
     "    Nt = int(round(T/dt))\n",
-    "    u = np.zeros(Nt+1)\n",
-    "    t = np.linspace(0, Nt*dt, Nt+1)\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",
-    "    u[0] = I\n",
     "    if damping == 'linear':\n",
-    "        u[1] = u[0] + dt*V + dt**2/(2*m)*(-b*V - s(u[0]) + F(t[0]))\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",
-    "        u[1] = u[0] + dt*V + \\\n",
-    "               dt**2/(2*m)*(-b*V*abs(V) - s(u[0]) + F(t[0]))\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",
+    "    # First timestep needs to have the backward timestep substituted\n",
+    "    stencil_init = stencil.subs(u.backward, u.forward-2*t.spacing*V)\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",
-    "    for n in range(1, Nt):\n",
-    "        if damping == 'linear':\n",
-    "            u[n+1] = (2*m*u[n] + (b*dt/2 - m)*u[n-1] +\n",
-    "                      dt**2*(F(t[n]) - s(u[n])))/(m + b*dt/2)\n",
-    "        elif damping == 'quadratic':\n",
-    "            u[n+1] = (2*m*u[n] - m*u[n-1] + b*u[n]*abs(u[n] - u[n-1])\n",
-    "                      + dt**2*(F(t[n]) - s(u[n])))/\\\n",
-    "                      (m + b*abs(u[n] - u[n-1]))\n",
-    "    return u, t"
+    "    return u.data, np.linspace(0, Nt*dt, Nt+1)"
    ]
   },
   {
@@ -619,27 +630,6 @@
     "the quadratic polynomial is reproduced by the numerical method (to\n",
     "machine precision).\n",
     "\n",
-    "### Catching bugs\n",
-    "\n",
-    "How good are the constant and quadratic solutions at catching\n",
-    "bugs in the implementation? Let us check that by introducing some bugs.\n",
-    "\n",
-    " * Use `m` instead of `2*m` in the denominator of `u[1]`: code works for constant\n",
-    "   solution, but fails (as it should) for a quadratic one.\n",
-    "\n",
-    " * Use `b*dt` instead of `b*dt/2` in the updating formula for `u[n+1]`\n",
-    "   in case of linear damping: constant and quadratic both fail.\n",
-    "\n",
-    " * Use `F[n+1]` instead of `F[n]` in case of linear or quadratic damping:\n",
-    "   constant solution works, quadratic fails.\n",
-    "\n",
-    "We realize that the constant solution is very useful for catching certain bugs because\n",
-    "of its simplicity (easy to predict what the different terms in the\n",
-    "formula should evaluate to), while the quadratic solution seems\n",
-    "capable of detecting all (?) other kinds of typos in the scheme.\n",
-    "These results demonstrate why we focus so much on exact, simple polynomial\n",
-    "solutions of the numerical schemes in these writings.\n",
-    "\n",
     "<!-- More: classes, cases with pendulum approx u vs sin(u), -->\n",
     "<!-- making UI via parampool -->\n",
     "\n",
@@ -1470,7 +1460,19 @@
    "cell_type": "code",
    "execution_count": 3,
    "metadata": {},
-   "outputs": [],
+   "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",
@@ -1858,7 +1860,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -2191,9 +2193,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "devito",
    "language": "python",
-   "name": "python3"
+   "name": "devito"
   },
   "language_info": {
    "codemirror_mode": {
@@ -2205,7 +2207,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.3"
+   "version": "3.8.1"
   }
  },
  "nbformat": 4,