{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Shallow water equations with a tracer\n",
    "\n",
    "We can augment the shallow water equations discussed in [Shallow_water.ipynb](Shallow_water.ipynb) with a tracer $\\phi(x,t)$ that measures the concentration of a tracer that is advected with the fluid motion (and that has no influence on the fluid dynamics).  If $\\phi$ is measured in units of mass per unit volume (which is really per unit area of the vertical cross section in this one-dimensional example) then the mass per unit length along the $x$ axis is given by $h(x,t)\\phi(x,t)$.  The quantity $\\phi$ satisfies the variable coefficient advection equation in advective (non-conservative) form:\n",
    "$$\n",
    "\\phi_t(x,t) + u(x,t)\\phi_x(x,t) = 0.\n",
    "$$\n",
    "This is also called the \"color equation\", since we can think of $\\phi$ as measuring the concentration of a dye that changes the color of the water.  We will use this interpretation in the plots below.  By setting the intitial conditions $\\phi(x,0)$ to be piecewise constant with different values corresponding to different colors, we can visualize the motion of the fluid better.   We will use two shades of blue for the water that is initially to the left of a dam at $x=0$ and two lighter shades of blue for the water that is initially to the right.  \n",
    "\n",
    "The quantity $h\\phi$ satisfies the conservative form of the advection equation,\n",
    "$$\n",
    "(h\\phi)_t + (uh\\phi)_x = 0.\n",
    "$$\n",
    "This can be derived by combining the color equation with the conservation of mass equation $h_t +(hu)_x  = 0$.  Since $h\\phi$ measures the concentration of dye per unit length in $x$, and we assume molecules of dye are not created or destroyed, it makes sense that this is the conserved quantity.\n",
    "\n",
    "The full system of equations in conservation form is thus:\n",
    "$$\n",
    "\\begin{split}\n",
    "h_t + (hu)_x &=0\\\\\n",
    "(hu)_t + \\left(hu^2 + \\frac 1 2 gh^2\\right)_x &=0\\\\\n",
    "(h\\phi)_t + (uh\\phi)_x &= 0\n",
    "\\end{split}\n",
    "$$\n",
    "\n",
    "The Riemann solution for this system has 3 waves.  The wave speeds are the eigenvalues of the Jacobian matrix, and are $u-c,~u,~u+c$ where $c = \\sqrt{gh}$ is the gravity wave speed.  The 1-wave and 3-wave are the nonlinear shallow water waves, which can be computed as in the shallow water equations using the first two equations of the system alone (since the tracer does not affect the fluid dynamics).  The 2-wave is a contact discontinuity (linearly degenerate wave) moving with the fluid velocity and carrying only a jump in the tracer concentration $\\phi$.  The characteristic speed for this wave is the fluid velocity, which is constant across this wave, and so characteristics travel parallel to the wave on either side.  The nonlinear waves can each be either a shock wave or rarefaction wave depending on the initial data, as explained in [Shallow_water.ipynb](Shallow_water.ipynb).  Below we consider a dam break problem in which case there is one of each."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "context = 'notebook'\n",
    "from ipywidgets import widgets, fixed\n",
    "if context == 'notebook':\n",
    "    from ipywidgets import interact\n",
    "elif context == 'html':\n",
    "    from utils.jsanimate_widgets import interact\n",
    "elif context == 'pdf':\n",
    "    from utils.snapshot_widgets import interact\n",
    "import seaborn as sns\n",
    "sns.set_style('white',{'legend.frameon':'True'});\n",
    "\n",
    "from exact_solvers import shallow_water\n",
    "demo_plot = shallow_water.make_demo_plot_function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dam break problem\n",
    "\n",
    "Here we plot the solution to the dam break Riemann problem for the shallow water equations with the addition of a tracer.  The tracer is now also take to be piecewise constant, with light blue representing the value $\\phi_l$ taken in the water the left of the dam, while dark blue represents the value $\\phi_r$ in the water on the right side of the dam.\n",
    "\n",
    "The structure of the depth and velocity are exactly the same as seen and discussed in detail in [Shallow_water.ipynb](Shallow_water.ipynb), and the value of $\\phi$ is constant across the 1-rarefaction and 3-shock waves.    The discontinuity in $\\phi$ shows the propagation of the 2-wave in the Riemann solution, the contact discontinuity across which $\\phi$ is discontinuous while  $h$ and $u$ are continuous. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interact(demo_plot(h_l=3., h_r=1., u_l=0., u_r=0, stripes=False),\n",
    "         t=widgets.FloatSlider(min=0., max=0.6, step=0.1,\n",
    "                               value=0.), fig=fixed(0));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add more discussion\n",
    "\n",
    "Show the eigenvectors and discuss linearly degenerate family.\n",
    "\n",
    "Also include plots in the x-t and phase plane."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Old version\n",
    "\n",
    "Why use pyclaw here since the exact solver already includes tracer stripes?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from clawpack import pyclaw\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import animation\n",
    "from clawpack.visclaw.JSAnimation import IPython_display\n",
    "from clawpack.riemann import shallow_roe_tracer_1D\n",
    "depth = 0\n",
    "momentum = 1\n",
    "tracer = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clawpack solution of the dam break problem\n",
    "\n",
    "In this notebook we use the finite volume methods in Clawpack to illustrate the solution to a Riemann problem, since we take initial data for the tracer $\\phi$ that is not simply piecewise constant.\n",
    "\n",
    "This allows us to also impose solid wall boundary conditions at the left and right boundaries of the domain and observe the resulting reflection and interaction of the rarefaction wave and shock wave that arise from the Riemann problem.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = pyclaw.ClawSolver1D(shallow_roe_tracer_1D)\n",
    "solver.num_eqn = 3\n",
    "solver.num_waves = 3\n",
    "solver.bc_lower[0] = pyclaw.BC.wall\n",
    "solver.bc_upper[0] = pyclaw.BC.wall\n",
    "x = pyclaw.Dimension(-1.0,1.0,2000,name='x')\n",
    "domain = pyclaw.Domain(x)\n",
    "state = pyclaw.State(domain,solver.num_eqn)\n",
    "\n",
    "state.problem_data['grav'] = 1.0\n",
    "\n",
    "grid = state.grid\n",
    "xc = grid.p_centers[0]\n",
    "\n",
    "hl = 1.\n",
    "hr = 1./3\n",
    "ul = 0.\n",
    "ur = 0.\n",
    "\n",
    "state.q[depth,:] = hl*(xc<=0) + hr*(xc>0)\n",
    "state.q[momentum,:] = hl*ul*(xc<=0) + hr*ur*(xc>0)\n",
    "state.q[tracer,:] = xc\n",
    "\n",
    "claw = pyclaw.Controller()\n",
    "claw.tfinal = 5\n",
    "claw.solution = pyclaw.Solution(state,domain)\n",
    "claw.solver = solver\n",
    "claw.keep_copy = True\n",
    "claw.num_output_times = 50\n",
    "claw.verbosity = 0\n",
    "\n",
    "claw.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot surface height.  Fill area below surface with alternating dark\n",
    "# and light stripes that advect with the tracer.\n",
    "# Red stripes for tracer starting on the left and blue for the right.\n",
    "# This code is written to include the possibility of a non-zero bathymetry.\n",
    "# This code should probably be moved to a file in the future.\n",
    "\n",
    "fig = plt.figure(figsize=[12,6])\n",
    "ax = fig.add_subplot(111)\n",
    "fills = []\n",
    "frame = claw.frames[0]\n",
    "h = frame.q[0,:]\n",
    "b = 0*h\n",
    "surface = h+b\n",
    "tracer = frame.q[2,:]\n",
    "\n",
    "x, = frame.state.grid.p_centers    \n",
    "\n",
    "line, = ax.plot(x, surface,'-k',linewidth=3)\n",
    "\n",
    "fills = {'cornflowerblue': None,\n",
    "         'blue': None,\n",
    "         'salmon': None,\n",
    "         'red': None}\n",
    "colors = fills.keys()\n",
    "\n",
    "def set_stripe_regions(tracer):\n",
    "    # Designate areas for each color of stripe\n",
    "    stripes = {}\n",
    "    stripes['cornflowerblue'] = (tracer>=0)\n",
    "    stripes['blue'] = (tracer%0.1>=0.05)*(tracer>=0)\n",
    "    stripes['salmon'] = (tracer<=0)\n",
    "    stripes['red'] = (tracer%0.1>=0.05)*(tracer<=0)\n",
    "    return stripes\n",
    "\n",
    "stripes = set_stripe_regions(tracer)\n",
    "\n",
    "for color in colors:\n",
    "    fills[color] = ax.fill_between(x,b,surface,facecolor=color,where=stripes[color],alpha=0.5)\n",
    "\n",
    "plt.xlabel('$x$'); plt.ylabel('depth ($h$)'); plt.axis('scaled')\n",
    "ax.set_xlim(-1,1); ax.set_ylim(0,1.2)\n",
    "\n",
    "def fplot(frame_number):\n",
    "    # Remove old fill_between plots\n",
    "    for color in colors:\n",
    "        fills[color].remove()\n",
    "        \n",
    "    frame = claw.frames[frame_number]\n",
    "    h = frame.q[0,:]\n",
    "    b = 0*h\n",
    "    tracer = frame.q[2,:]\n",
    "    surface = h+b\n",
    "    line.set_data(x,surface)\n",
    "    stripes = set_stripe_regions(tracer)\n",
    "    for color in colors:\n",
    "        fills[color] = ax.fill_between(x,b,surface,facecolor=color,where=stripes[color],alpha=0.5)\n",
    "    return line,\n",
    "\n",
    "animation.FuncAnimation(fig, fplot, frames=len(claw.frames), interval=200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this animation, note that both the mass of water and quantity of dye in any vertical stripe of a given color is conserved, and that the width of each stripe varies inversely with the depth of water in the stripe.  The edges of the stripes are always moving at the fluid velocity."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}