diff --git a/week6/ex1.py b/week6/ex1.py
index e044ca3..11479e4 100755
--- a/week6/ex1.py
+++ b/week6/ex1.py
@@ -2,30 +2,23 @@
 
 import numpy as np
 
-def pdeHyperbolic(a, x, t, f, g = None, dtype=np.float64):
+def pdeHyperbolic(a, x, t, f, g, dtype=np.float64):
     """ Solve a Hyperbolic Partial Differential using finite differences. """
     
     m = len(x) # Amount of objects to track
     n = len(t) # Length of Time Vector
 
     # Determine stepsizes
-    h = 1
-    if m > 1:
-        h = x[0] - x[1]
-
-    k = 1
-    if n > 1:
-        k = t[0] - t[1]
+    h = x[0] - x[1]
+    k = t[0] - t[1]
 
     λ_sq = (a*k/h)**2
 
-
     # Create array to hold the solution
     w = np.zeros((n,m), dtype=dtype)
 
     # Create finite difference matrix 
     A = np.diag(m*[2*(1 - λ_sq)], k=0) + np.diag((m-1)*[λ_sq], k=-1) + np.diag((m-1)*[λ_sq], k=1)
-    print(A)
 
     # Initialise first two timesteps
     w[0] = f(x, t[0])
@@ -38,16 +31,13 @@ def pdeHyperbolic(a, x, t, f, g = None, dtype=np.float64):
     return w
 
 
-
-
-
-def test_pdeHyperbolic_case1(m=1e2, n=1e3,l=1, T=1):
+def test_pdeHyperbolic_case1(x_steps=1e2, t_steps=1e2, max_x=1, max_t=1):
     
     a = 1 # from the Schroedinger Equation
 
     # Setup spatial and time grids
-    x = np.linspace(0, l, m)
-    t = np.linspace(0, T, n)
+    x = np.linspace(0, max_x, x_steps)
+    t = np.linspace(0, max_t, t_steps)
 
     # Boundary conditions
     def f(x,t):
@@ -60,43 +50,18 @@ def test_pdeHyperbolic_case1(m=1e2, n=1e3,l=1, T=1):
     sol = pdeHyperbolic(a, x, t, f, g)
 
     # Plot it with the exact solution
-    exact_sol = lambda x,t: np.sin(2*np.pi*x)*(np.cos(2*np.pi*t) + np.sin(2*np.pi*t))
-    t_high_res = np.linspace(0, T, n * 1e4)
+    exact_f = lambda x,t: np.sin(2*np.pi*x)*(np.cos(2*np.pi*t) + np.sin(2*np.pi*t))
 
-    from matplotlib import pyplot
-    from matplotlib import animation as anim
-    fig, _ = pyplot.subplots()
+    plot_animation(x, sol, func=exact_f)
 
-    if False:
-        for _, x_i in enumerate(x):
-            pyplot.plot(t_high_res, exact_sol(x_i, t_high_res), label="Exact")
-        pyplot.plot(t, sol[:,1], label="iter")
-    
-        pyplot.grid()
-        pyplot.xlabel("t")
-        pyplot.ylabel("w")
 
-    else:
-        def animate(i):
-            pyplot.clf()
-            pyplot.ylim(-2, 2)
-            pyplot.grid()
-            pyplot.plot(x, exact_sol(x, t[i]), label="exact")
-            pyplot.plot(x, sol[i,:], label="iter")
-    
-        frames = np.arange(1, n, dtype=np.int)
-        myAnim = anim.FuncAnimation(fig, animate, frames, interval = 10 )
-
-    pyplot.legend()
-    pyplot.show()
-
-def test_pdeHyperbolic_case2(m=200, n=400, l=1, T=1):
+def test_pdeHyperbolic_case2(x_steps=2e2, t_steps=4e2, max_x=1, max_t=1):
     
     a = 1 # from the Schroedinger Equation
 
     # Setup spatial and time grids
-    x = np.linspace(0, l, m)
-    t = np.linspace(0, T, n)
+    x = np.linspace(0, max_x, x_steps)
+    t = np.linspace(0, max_t, t_steps)
 
     # Boundary conditions
     def f(x,t):
@@ -108,24 +73,37 @@ def test_pdeHyperbolic_case2(m=200, n=400, l=1, T=1):
     # Solve it
     sol = pdeHyperbolic(a, x, t, f, g)
 
+    plot_animation(x, sol)
+
+
+
+def plot_animation(x, sol, func=None, interval = 10):
     from matplotlib import pyplot
     from matplotlib import animation as anim
     fig, _ = pyplot.subplots()
-
+    
     def animate(i):
         pyplot.clf()
         pyplot.grid()
-        pyplot.ylim(-1.5,1.5)
+        pyplot.ylim(-1.5, 1.5)
+        pyplot.title('t = {}/{}'.format(i, len(sol)))
+        #if func is not None:
+        #    pyplot.plot(x, func(x, sol[i]), label='exact')
         pyplot.plot(x, sol[i,:], label="iter")
- 
-    frames = np.arange(1, n)
-    myAnim = anim.FuncAnimation(fig, animate, frames, interval = 5 )
 
-    pyplot.legend()
-    pyplot.show()
+    frames = np.arange(0, len(sol))
+
+    try:
+        myAnim = anim.FuncAnimation(fig, animate, frames, interval = interval )
+
+        pyplot.legend()
+        pyplot.show()
+    except AttributeError:
+        # This final error is fugly
+        pass
 
 
 
 if __name__ == "__main__":
-    #test_pdeHyperbolic_case1()
-    test_pdeHyperbolic_case2()
+    test_pdeHyperbolic_case1()
+    #test_pdeHyperbolic_case2()