--- a/+time/+cdiff/cdiff.m	Mon Jan 07 15:01:26 2019 +0100
+++ /dev/null	Thu Jan 01 00:00:00 1970 +0000
@@ -1,16 +0,0 @@
-% Takes a step of
-%   v_tt = Dv+Ev_t+S
-%
-%   1/k^2 * (v_next - 2v + v_prev) = Dv  + E 1/(2k)(v_next - v_prev) + S
-%
-function [v_next, v] = cdiff(v, v_prev, k, D, E, S)
-    %   1/k^2 * (v_next - 2v + v_prev) = Dv  + E 1/(2k)(v_next - v_prev) + S
-    %       ekv to
-    %   A v_next = B v + C v_prev + S
-    I = speye(size(D));
-    A =  1/k^2 * I - 1/(2*k)*E;
-    B =  2/k^2 * I + D;
-    C = -1/k^2 * I - 1/(2*k)*E;
-
-    v_next = A\(B*v + C*v_prev + S);
-end
\ No newline at end of file

--- a/+time/Cdiff.m	Mon Jan 07 15:01:26 2019 +0100
+++ b/+time/Cdiff.m	Mon Jan 07 15:56:35 2019 +0100
@@ -1,8 +1,8 @@
 classdef Cdiff < time.Timestepper
     properties
-        D
-        E
-        S
+        A, B, C
+        AA, BB, CC
+        G
         k
         t
         v
@@ -12,25 +12,35 @@
 
 
     methods
-        % Solves u_tt = Du + Eu_t + S
-        % D, E, S can either all be constants or all be function handles,
-        % They can also be omitted by setting them equal to the empty matrix.
-        % Cdiff(D, E, S, k, t0, n0, v, v_prev)
-        function obj = Cdiff(D, E, S, k, t0, n0, v, v_prev)
-            m = length(v);
-            default_arg('E',sparse(m,m));
-            default_arg('S',sparse(m,1));
+        % Solves
+        %   A*v_tt + B*v_t + C*v = G(t)
+        %   v(t0) = v0
+        %   v_t(t0) = v0t
+        % starting at time t0 with timestep k
+        % Using
+        % A*Dp*Dm*v_n + B*D0*v_n + C*v_n = G(t_n)
+        function obj = Cdiff(A, B, C, G, v0, v0t, k, t0)
+            m = length(v0);
+            default_arg('A', speye(m));
+            default_arg('B', sparse(m,m));
+            default_arg('G', @(t) sparse(m,1));
+            default_arg('t0', 0);
 
-            obj.D = D;
-            obj.E = E;
-            obj.S = S;
+            obj.A = A;
+            obj.B = B;
+            obj.C = C;
+            obj.G = G;
 
+            % Rewrite as AA*v_(n+1) + BB*v_n + CC*v_(n-1) = G(t_n)
+            obj.AA = A/k^2 + B/2/k;
+            obj.BB = -2*A/k^2 + C;
+            obj.CC = A/k^2 - B/2/k;
 
             obj.k = k;
-            obj.t = t0;
-            obj.n = n0;
-            obj.v = v;
-            obj.v_prev = v_prev;
+            obj.v_prev = v0;
+            obj.v = v0 + k*v0t;
+            obj.t = t0+k;
+            obj.n = 1;
         end
 
         function [v,t] = getV(obj)
@@ -43,8 +53,19 @@
             t = obj.t;
         end
 
+        function E = getEnergy(obj)
+            v  = obj.v;
+            vp = obj.v_prev;
+            vt = (obj.v - obj.v_prev)/obj.k;
+
+            E = vt'*obj.A*vt + v'*obj.C*vp;
+        end
+
         function obj = step(obj)
-            [obj.v, obj.v_prev] = time.cdiff.cdiff(obj.v, obj.v_prev, obj.k, obj.D, obj.E, obj.S);
+            v_next = obj.AA\(-obj.BB*obj.v - obj.CC*obj.v_prev + obj.G(obj.t));
+
+            obj.v_prev = obj.v;
+            obj.v      = v_next;
             obj.t = obj.t + obj.k;
             obj.n = obj.n + 1;
         end