diff --git a/Introduction.ipynb b/Introduction.ipynb
index bb7e9c2a..5dd1e44b 100644
--- a/Introduction.ipynb
+++ b/Introduction.ipynb
@@ -23,9 +23,9 @@
     "\n",
     "If the water to the left of the dam is deeper than the water to the right (i.e. $h_l > h_r$) then we expect water to flow from left to right in some manner.  In reality we might expect the flow to be somewhat turbulent and multidimensional, but suppose that we could idealize this with a one-dimensional model in which the velocity is purely horizontal and is constant throughout the depth of the fluid, so for $t>0$ the velocity varies only with $x$ and $t$, given by some function $u(x,t)$, and the depth is given by $h(x,t)$.  \n",
     "\n",
-    "In the notebook [Shallow_water.ipynb](Shallow_water.ipynb) we will discuss the one-dimensional **shallow water equations** that can be used to model this idealized situation.  \n",
+    "In [Shallow_water.ipynb](Shallow_water.ipynb) we will discuss the one-dimensional **shallow water equations** that can be used to model this idealized situation.  \n",
     "\n",
-    "The animation below shows the solution to the shallow water equations in the situation just described.  The water is colored red for water that is initially to the left of the dam and blue for water that is initially to the right.  The stripes with different shading are to help you visualize the flow.  Think of the color as a dye that is simply carried along with the water as it flows.  In the notebook [Shallow_water_tracer.ipynb](Shallow_water_tracer.ipynb) we will see that we can in fact model this by adding another equation to the shallow water equations."
+    "The animation below shows the solution to the shallow water equations in the situation just described.  The water is colored red for water that is initially to the left of the dam and blue for water that is initially to the right.  The stripes with different shading are to help you visualize the flow.  Think of the color as a dye that is simply carried along with the water as it flows.  In [Shallow_water_tracer.ipynb](Shallow_water_tracer.ipynb) we will see that we can in fact model this by adding another equation to the shallow water equations."
    ]
   },
   {
@@ -36,6 +36,7 @@
    "source": [
     "%matplotlib inline\n",
     "from exact_solvers import shallow_water\n",
+    "demo_plot = shallow_water.make_demo_plot_function\n",
     "from ipywidgets import widgets, fixed\n",
     "from ipywidgets import interact"
    ]
@@ -43,13 +44,12 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "collapsed": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "interact(shallow_water.make_demo_plot_function(h_l=3., h_r=1., u_l=0., u_r=0),\n",
-    "         t=widgets.FloatSlider(min=0.001,max=0.6,step=0.1,value=0.), fig=fixed(0))"
+    "interact(demo_plot(h_l=3., h_r=1., u_l=0., u_r=0),\n",
+    "         t=widgets.FloatSlider(min=0., max=0.6, step=0.1,\n",
+    "                               value=0.), fig=fixed(0));"
    ]
   },
   {
@@ -114,8 +114,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "interact(shallow_water.make_demo_plot_function(h_l=1., h_r=3., u_l=0., u_r=0),\n",
-    "         t=widgets.FloatSlider(min=0.001,max=0.6,step=0.1,value=0.), fig=fixed(0));"
+    "interact(demo_plot(h_l=1., h_r=3., u_l=0., u_r=0),\n",
+    "         t=widgets.FloatSlider(min=0.0, max=0.6, step=0.1,\n",
+    "                               value=0.), fig=fixed(0));"
    ]
   },
   {
@@ -131,8 +132,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "interact(shallow_water.make_demo_plot_function(h_l=1., h_r=1., u_l=0.8, u_r=-0.8),\n",
-    "         t=widgets.FloatSlider(min=0.001,max=0.6,step=0.1,value=0.), fig=fixed(0));"
+    "interact(demo_plot(h_l=1., h_r=1., u_l=0.8, u_r=-0.8),\n",
+    "         t=widgets.FloatSlider(min=0., max=0.6, step=0.1,\n",
+    "                               value=0.), fig=fixed(0));"
    ]
   },
   {
@@ -148,15 +150,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "interact(shallow_water.make_demo_plot_function(h_l=1., h_r=1., u_l=-0.8, u_r=0.8),\n",
-    "         t=widgets.FloatSlider(min=0.001,max=0.6,step=0.1,value=0.), fig=fixed(0));"
+    "interact(demo_plot(h_l=1., h_r=1., u_l=-0.8, u_r=0.8),\n",
+    "         t=widgets.FloatSlider(min=0.0, max=0.6, step=0.1,\n",
+    "                               value=0.), fig=fixed(0));"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "In the notebook [Shallow_water.ipynb](Shallow_water.ipynb) we will show how to solve the Riemann problem for arbitrary left and right states and in particular how to determine whether each wave is a shock or rarefaction wave."
+    "In [Shallow_water.ipynb](Shallow_water.ipynb) we will show how to solve the Riemann problem for arbitrary left and right states and in particular how to determine whether each wave is a shock or rarefaction wave."
    ]
   },
   {
@@ -218,7 +221,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# A linear example: acoustic waves\n",
+    "## A linear example: acoustic waves\n",
     "\n",
     "We end this introduction with a sample linear system. Acoustic waves in one dimension can be modeled by the system of PDEs\n",
     "\n",
@@ -261,7 +264,11 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "hide"
+    ]
+   },
    "outputs": [],
    "source": [
     "from IPython.core.display import display, HTML\n",
@@ -275,28 +282,23 @@
    "source": [
     "## Phase plane plot for the shallow water equations\n",
     "\n",
-    "The phase plane plot for the nonlinear shallow water equations is more interesting.  Since the eigenvectors of the Jacobian matrix vary with $q$, the curves connecting states separated by a single wave are not straight lines. There is also a distinction between the curves connecting states by a shock wave and those connecting states by a rarefaction wave.  This is discussed in much more detail in the notebook [Shallow_water.ipynb](Shallow_water.ipynb), but here is an interactive view that allows you to see how the solution changes as you move the left and right states around in the phase plane:"
+    "The phase plane plot for the nonlinear shallow water equations is more interesting.  Since the eigenvectors of the Jacobian matrix vary with $q$, the curves connecting states separated by a single wave are not straight lines. There is also a distinction between the curves connecting states by a shock wave and those connecting states by a rarefaction wave.  This is discussed in much more detail in [Shallow_water.ipynb](Shallow_water.ipynb), but here is an interactive view that allows you to see how the solution changes as you move the left and right states around in the phase plane:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "hide"
+    ]
+   },
    "outputs": [],
    "source": [
     "from IPython.core.display import display, HTML\n",
     "app = open('phase_plane_shallow_water_verysmall.html','r').read()\n",
     "display(HTML(app))"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "collapsed": false
-   },
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
@@ -315,7 +317,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython2",
-   "version": "2.7.12"
+   "version": "2.7.13"
   }
  },
  "nbformat": 4,
diff --git a/exact_solvers/shallow_water.py b/exact_solvers/shallow_water.py
index 2f3aed49..c1c7075d 100644
--- a/exact_solvers/shallow_water.py
+++ b/exact_solvers/shallow_water.py
@@ -427,7 +427,13 @@ def reval_rho_u(x):
     num_vars = len(primitive_variables)
 
     def plot_shallow_water_demo(t=0.5, fig=0):
-        q = np.array(reval(x/t))
+        if t == 0:
+            q = np.zeros((2,len(x)))
+            q[0,:] = q_l[0]*(x<=0) + q_r[0]*(x>0)
+            q[1,:] = q_l[1]*(x<=0) + q_r[1]*(x>0)
+        else:
+            q = np.array(reval(x/t))
+
         if t<0.02:
             q[1] = np.where(x<0, q_l[1], q_r[1])
 
diff --git a/riemann.tplx b/riemann.tplx
index 4551a57a..f48aad5c 100644
--- a/riemann.tplx
+++ b/riemann.tplx
@@ -30,7 +30,6 @@
 
 ((* block title *))
 \title{The Riemann Problem for Hyperbolic PDEs: Theory and Approximate Solvers}
-\DefineVerbatimEnvironment{Highlighting}{Verbatim}{commandchars=\\\{\},fontsize=\small}
 ((* endblock title *))
 
 ((* block maketitle *))
@@ -47,3 +46,20 @@
 ((( super() )))
 ((*- endif -*))
 ((*- endblock any_cell -*))
+
+((* block input scoped *))
+    ((( add_my_prompt(cell.source | highlight_code(strip_verbatim=True), cell, 'In ', 'incolor') )))
+((* endblock input *))
+
+% Purpose: Renders an output/input prompt
+((* macro add_my_prompt(text, cell, prompt, prompt_color) -*))
+    ((*- if cell.execution_count is defined -*))
+    ((*- set execution_count = "" ~ (cell.execution_count | replace(None, " ")) -*))
+    ((*- else -*))
+    ((*- set execution_count = " " -*))
+    ((*- endif -*))
+    ((*- set indention =  " " * (execution_count | length + 7) -*))
+\begin{Verbatim}[fontsize=\small,commandchars=\\\{\}]
+((( text | add_prompts(first='{\color{' ~ prompt_color ~ '}' ~ prompt ~ '[{\\color{' ~ prompt_color ~ '}' ~ execution_count ~ '}]:} ', cont=indention) )))
+\end{Verbatim}
+((*- endmacro *))