{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# 3.8 A Nonlinear conservation law: shallow water equation in 2D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "We consider the shallow water equations as an example of a nonlinear conservation law, i.e. we consider\n",
    "\n",
    "$$\n",
    "  \\partial_t \\mathbf{U} + \\operatorname{div}(\\mathbf{F} (\\mathbf{U} )) = 0 \\qquad in \\qquad \\Omega \\times[0,T],\n",
    "$$\n",
    "\n",
    "with \n",
    "\n",
    "$$\n",
    "\\mathbf{U} = (h, \\mathbf{w}) = (h, hu, hv) = (\\mathbf{u}_1, \\mathbf{u}_2, \\mathbf{u}_3)\n",
    "$$\n",
    "\n",
    "and \n",
    "\n",
    "$$\n",
    "  \\mathbf{F}(\\mathbf{U}) = \\left( \\begin{array}{c@{}c@{\\qquad}c@{}c} h u & & h v & \\\\ h u^2 &+ \\frac12 g h^2 & h u v & \\\\ h u v &  & h v^2 & + \\frac12 g h^2 \\end{array} \\right) \n",
    "  = \\left( \\begin{array}{cc} \\mathbf{u}_2 & \\mathbf{u}_3 \\\\ \\frac{\\mathbf{u}_2^2}{\\mathbf{u}_1} + \\frac12 g \\mathbf{u}_1^2 & \\frac{\\mathbf{u}_2 \\mathbf{u}_3}{\\mathbf{u}_1} \\\\ \\frac{\\mathbf{u}_2 \\mathbf{u}_3}{\\mathbf{u}_1} & \\frac{\\mathbf{u}_3^2}{\\mathbf{u}_1} + \\frac12 g \\mathbf{u}_1^2 \\end{array} \\right)  \n",
    "  = \\left( \\begin{array}{cc} h \\mathbf{w}^T \\\\ h \\mathbf{w} \\otimes \\mathbf{w} + \\frac12 g h^2 \\mathbf{I} \\end{array} \\right)\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "## Jacobian of the flux for shallow water:\n",
    "\n",
    "\\begin{align*}\n",
    "\\mathbf{A}_1 & = \\left(\\begin{array}{ccc} 0 & 1 & 0 \\\\ - \\frac{\\mathbf{u}_2^2}{\\mathbf{u}_1^2} + g \\mathbf{u}_1 & 2 \\frac{\\mathbf{u}_2}{\\mathbf{u}_1} & 0 \\\\ - \\frac{\\mathbf{u}_2 \\mathbf{u}_3}{\\mathbf{u}_1^2} & \\frac{\\mathbf{u}_3}{\\mathbf{u}_1} & \\frac{\\mathbf{u}_2}{\\mathbf{u}_1} \\end{array} \\right) = \\left( \\begin{array}{ccc} 0 & 1 & 0 \\\\ - u^2 + g h & 2 u & 0 \\\\ - u v & v & u \\end{array} \\right) \n",
    "\\end{align*}\n",
    "\n",
    "\\begin{align*}\n",
    "\\mathbf{A}_2 & = \\left( \\begin{array}{ccc} 0 & 0 & 1 \\\\ - \\frac{\\mathbf{u}_2\\mathbf{u}_3}{\\mathbf{u}_1^2} & \\frac{\\mathbf{u}_3}{\\mathbf{u}_1} & \\frac{\\mathbf{u}_2}{\\mathbf{u}_1} \\\\ - \\frac{\\mathbf{u}_3^2}{\\mathbf{u}_1^2} + g \\mathbf{u}_1  & 0 & 2\\frac{\\mathbf{u}_3}{\\mathbf{u}_1} \\end{array} \\right) = \\left( \\begin{array}{ccc} 0 & 0 & 1 \\\\ - uv & v & u \\\\ - v^2 + gh & 0 & 2 v \\end{array} \\right)\n",
    "\\end{align*}\n",
    "\n",
    "\\begin{align*}\n",
    "\\mathbf{A}(\\mathbf{u}, \\mathbf{n}) & = n_1 \\mathbf{A}_1 + n_2 \\mathbf{A}_2 = \\left( \\begin{array}{ccc} 0 & n_1 & n_2 \\\\\n",
    "- u \\alpha - g h n_1  & \\alpha + u n_1 & u n_2 \\\\\n",
    "- v \\alpha - g h n_2  & v n_1 & \\alpha + v n_2 \\end{array} \\right), \\quad \\text{ with } \\alpha = (\\mathbf{u}, \\mathbf{n}),\n",
    "\\end{align*}\n",
    "\n",
    "$$\n",
    "\\text{spectrum:} \\qquad \\rho( \\mathbf{A}(\\mathbf{u}, \\mathbf{n}) ) = \\{ \\alpha + \\sqrt{g h}, \\alpha, \\alpha - \\sqrt{g h} \\} \n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### The dam break problem (geometry)\n",
    "\n",
    "![geomsketch](geomsketch.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "# geometry description (including boundary/mat names)\n",
    "from ngsolve import * \n",
    "from netgen.occ import *\n",
    "from ngsolve.webgui import * \n",
    "wp = WorkPlane()\n",
    "wp.MoveTo(-12,-5).LineTo(-3,-5).NameVertex(\"fillet\")\n",
    "wp.LineTo(-3,-1).NameVertex(\"fillet\")\n",
    "wp.LineTo(0,-1).LineTo(0,0).LineTo(-12,0).Close()\n",
    "geo = wp.Face()\n",
    "geo = geo.MakeFillet(list(set(geo.vertices[\"fillet\"])),2)\n",
    "geo = geo + geo.Mirror(Axis((0,0,0),X)).Reversed()\n",
    "geo = Glue([geo,geo.Mirror(Axis((0,0,0),Y)).Reversed()])\n",
    "geo.faces.Min(X).name=\"upperlevel\"\n",
    "geo.faces.Max(X).name=\"lowerlevel\"\n",
    "geo.edges.name = \"wall\"\n",
    "geo.edges.Nearest((0,0)).name = \"dam\"\n",
    "#Draw(geo)\n",
    "geo = OCCGeometry(geo,dim=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "from ngsolve import *\n",
    "from ngsolve.webgui import Draw\n",
    "mesh = Mesh(geo.GenerateMesh(maxh=2))\n",
    "mesh.Curve(3)\n",
    "Draw(mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Vectorial (`dim=3`) approximation space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "order = 2\n",
    "fes = L2(mesh,order=order)**3\n",
    "#fes = L2(mesh,order=order,dim=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### initial and boundary conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "U,V = fes.TnT() # \"Trial\" and \"Test\" function\n",
    "h, hu, hv = U\n",
    "\n",
    "# initial conditions\n",
    "h0mat = {\"upperlevel\" : 10, \"lowerlevel\" : 2}\n",
    "U0 = CF((mesh.MaterialCF(h0mat),0,0))\n",
    "\n",
    "# boundary conditions\n",
    "hbndreg = mesh.BoundaryCF({\"wall\" : h, \"dam\" : 0})\n",
    "hubndreg = mesh.BoundaryCF({\"wall\" : -hu, \"dam\" : 0})\n",
    "hvbndreg = mesh.BoundaryCF({\"wall\" : -hv, \"dam\" : 0})\n",
    "\n",
    "Ubnd = CF((hbndreg,hubndreg,hvbndreg))\n",
    "\n",
    "# constant for gravitational force\n",
    "g=1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Flux definition and numerical flux choice (Lax-Friedrich)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "$$\n",
    "  \\mathbf{F}(\\mathbf{U}) = \\left( \\begin{array}{c@{}c@{\\qquad}c@{}c} h u & & h v & \\\\ h u^2 &+ \\frac12 g h^2 & h u v & \\\\ h u v &  & h v^2 & + \\frac12 g h^2 \\end{array} \\right) \n",
    "  = \\left( \\begin{array}{cc} \\mathbf{u}_2 & \\mathbf{u}_3 \\\\ \\frac{\\mathbf{u}_2^2}{\\mathbf{u}_1} + \\frac12 g \\mathbf{u}_1^2 & \\frac{\\mathbf{u}_2 \\mathbf{u}_3}{\\mathbf{u}_1} \\\\ \\frac{\\mathbf{u}_2 \\mathbf{u}_3}{\\mathbf{u}_1} & \\frac{\\mathbf{u}_3^2}{\\mathbf{u}_1} + \\frac12 g \\mathbf{u}_1^2 \\end{array} \\right)  \n",
    "  = \\left( \\begin{array}{cc} h \\mathbf{w}^T \\\\ h \\mathbf{w} \\otimes \\mathbf{w} + \\frac12 g h^2 \\mathbf{I} \\end{array} \\right)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "def F(U):\n",
    "    h, hvx, hvy = U\n",
    "    vx = hvx/h\n",
    "    vy = hvy/h\n",
    "    return CF(((hvx,hvy),\n",
    "               (hvx*vx + 0.5*g*h**2, hvx*vy),\n",
    "               (hvx*vy, hvy*vy + 0.5*g*h**2)),dims=(3,2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "$$\n",
    "  \\hat{\\mathbf{F}}(\\mathbf{U}) \\cdot \\mathbf{n} = \\frac{(\\mathbf{F}(\\mathbf{U}^l)+\\mathbf{F}(\\mathbf{U}^r)}{2} \\cdot \\mathbf{n} + \\frac{|\\lambda|}{2} (\\mathbf{U}^l - \\mathbf{U}^r), \\quad (\\cdot^l: \\text{current element}, \\quad \\cdot^r: \\text{neighbor element})\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     3
    ],
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "n = specialcf.normal(mesh.dim)\n",
    "def Max(u,v):\n",
    "    return IfPos(u-v,u,v)\n",
    "def Fmax(A,B): # max. characteristic speed:\n",
    "    ha, hua, hva = A\n",
    "    hb, hub, hvb = B\n",
    "    vnorma = sqrt(hua**2+hva**2)/ha\n",
    "    vnormb = sqrt(hub**2+hvb**2)/hb\n",
    "    return Max(vnorma+sqrt(g*A[0]),vnormb+sqrt(g*B[0]))\n",
    "\n",
    "def Fhatn(U): # numerical flux\n",
    "    Uhat = U.Other(bnd=Ubnd)\n",
    "    return (0.5*F(U)+0.5*F(Uhat))*n + Fmax(U,Uhat)/2*(U-Uhat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### DG formulation\n",
    "\n",
    "We recall that a `BilinearForm` is allowed to be nonlinear in the first argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "def DGBilinearForm(fes,F,Fhatn,Ubnd):\n",
    "    a = BilinearForm(fes, nonassemble=True)    \n",
    "    a += - InnerProduct(F(U),Grad(V)) * dx \n",
    "    a += InnerProduct(Fhatn(U),V) * dx(element_boundary=True)\n",
    "    return a\n",
    "\n",
    "a = DGBilinearForm(fes,F,Fhatn,Ubnd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "#### Simple fix to deal with shocks: artificial diffusion:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "from DGdiffusion import AddArtificialDiffusion\n",
    "\n",
    "artvisc = Parameter(1.0)\n",
    "if order > 0:\n",
    "    AddArtificialDiffusion(a,Ubnd,artvisc,compile=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Visualization of solution quantities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "gfu = GridFunction(fes)\n",
    "gfh, gfhu, gfhv = gfu\n",
    "gfvu = gfhu/gfh\n",
    "gfvv = gfhv/gfh\n",
    "momentum = CF((gfhu,gfhv))\n",
    "velocity = CF((gfvu,gfvv))\n",
    "gfu.Set(U0) \n",
    "scenes = [ \\\n",
    "    Draw(momentum,mesh,\"mom\"),\n",
    "    Draw(velocity,mesh,\"vel\"),\n",
    "    Draw(gfh,mesh,\"h\") ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Explicit Euler time stepping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     1,
     9,
     16,
     18,
     23
    ],
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "def TimeLoop(a,gfu,dt,T,nsamplings=100,scenes=scenes,multidim_draw=False,md_nsamplings=20):\n",
    "    if multidim_draw:\n",
    "        gfu_t = GridFunction(gfu.space,multidim=True)\n",
    "    #gfu.Set(U0)\n",
    "    res = gfu.vec.CreateVector()\n",
    "    fes = a.space\n",
    "    t = 0; i = 0\n",
    "    nsteps = int(ceil(T/dt))    \n",
    "    invma = fes.Mass(1).Inverse() @ a.mat\n",
    "    if multidim_draw:\n",
    "        gfu_t.AddMultiDimComponent(gfu.vec)\n",
    "    with TaskManager():\n",
    "        while t <= T - 0.5*dt:\n",
    "            res = invma * gfu.vec\n",
    "            gfu.vec.data -= dt * res\n",
    "            t += dt\n",
    "            if (i+1) % int(nsteps/nsamplings) == 0:\n",
    "                for s in scenes: s.Redraw()     \n",
    "            if multidim_draw and (i+1) % int(nsteps/md_nsamplings) == 0:\n",
    "                gfu_t.AddMultiDimComponent(gfu.vec)\n",
    "            i+=1\n",
    "            print(\"\\rt = {:.10}\".format(t),end=\"\")\n",
    "    for s in scenes: s.Redraw()\n",
    "    if multidim_draw:\n",
    "        return gfu_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "gfu.Set(U0)\n",
    "with TaskManager():\n",
    "    gfu_t = TimeLoop(a,gfu,dt=0.0025,T=15,multidim_draw=True,md_nsamplings=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "run_control": {
     "marked": false
    },
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "Draw(gfu_t.components[0],mesh,\"h\",interpolate_multidim=True,animate=True, \n",
    "     deformation=True, settings = {\"camera\" : \n",
    "                                   {\"transformations\" : \n",
    "                                    [{ \"type\": \"rotateX\", \"angle\": -20},\n",
    "                                     { \"type\": \"rotateZ\", \"angle\": 0}]}},\n",
    "     min=2, max=5, autoscale=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "**Tasks**\n",
    "\n",
    "You may play around with this example. \n",
    "\n",
    "* change the artificial diffusion parameter: How does it influence the solution\n",
    "* change boundary conditions: left boundary -> fixed height and non-reflecting\n",
    "* change the initial heights\n",
    "* introduce a (circular) obstacle below the dam\n",
    "* Generate output to create a video: To this end take a look at the final unit of this section."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.10"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}