--- a/+time/+rk/General.m	Wed Jan 09 12:14:30 2019 +0100
+++ b/+time/+rk/General.m	Wed Jan 09 22:57:13 2019 +0100
@@ -6,18 +6,17 @@
         v       % Solution vector
         n       % Time level
         scheme  % The scheme used for the time stepping, e.g rk4, rk6 etc.
-        coeffs  % Butcher tableau coefficients
+        bt
         V       % All stage approximations in most recent time step
         K       % All stage rates in most recent time step
     end
 
 
     methods
-        % Timesteps v_t = F(v,t), using the specified ButcherTableau
+        % Timesteps v_t = F(t,v), using the specified ButcherTableau
         % from t = t0 with timestep dt and initial conditions v(0) = v0
-        function obj = General(F, dt, t0, v0, bt, discreteData)
+        function obj = General(F, dt, t0, v0, bt)
             assertType(bt, 'time.rk.ButcherTableau')
-            default_arg('discreteData', []);
             obj.F = F;
             obj.dt = dt;
             obj.t = t0;
@@ -25,13 +24,7 @@
             obj.n = 0;
 
             assert(bt.isExplicit())
-            obj.coeffs = struct('s',bt.nStages,'a',bt.a(:,1:end-1),'b',bt.b,'c',bt.c);
-
-            if isempty(discreteData)
-                obj.scheme = @(v,t,n) time.rk.rungekutta(v, t, dt, F, obj.coeffs);
-            else
-                obj.scheme = @(v,t,n) time.rk.rungekuttaDiscreteData(v, t, dt, F, obj.coeffs, discreteData, n);
-            end
+            obj.bt = bt;
         end
 
         % v: Current solution
@@ -39,16 +32,74 @@
         % V: All stage approximations in most recent time step
         % K: All stage rates in most recent time step
         % T: Time points (corresponding to V and K) in most recent time step
-        function [v,t,V,T,K] = getV(obj)
+        function [v,t] = getV(obj)
             v = obj.v;
             t = obj.t;
-            V = obj.V;
-            K = obj.K;
-            T = obj.t + obj.dt*obj.coeffs.b;
         end
 
         function obj = step(obj)
-            [obj.v, obj.V, obj.K] = obj.scheme(obj.v, obj.t, obj.n);
+            s = obj.bt.nStages();
+            a = obj.bt.a;
+            b = obj.bt.b;
+            c = obj.bt.c;
+
+            % Compute rates K
+            K = zeros(length(v), s);
+            for i = 1:s
+                V_i = obj.v;
+                for j = 1:i-1
+                    V_i = V_i + dt*a(i,j)*K(:,j);
+                end
+                K(:,i) = F(t+dt*c(i), V_i);
+            end
+
+            % Compute updated solution
+            v_next = v;
+            for i = 1:s
+                v_next = v_next + dt*b(i)*K(:,i);
+            end
+
+            obj.v = v_next;
+            obj.t = obj.t + obj.dt;
+            obj.n = obj.n + 1;
+        end
+
+        % TBD: Method name
+        % TBD: Parameter name
+        %
+        % Takes a regular step but with discreteRates(:,i) added to RHS for stage i.
+        %  v_t = F(t,v) + discreteRates(:, ...)
+        %
+        % Also returns the stage approximations (V) and stage rates (K).
+        function [v,t, V, K] = stepWithDiscreteData(obj, discreteRates)
+            s = obj.bt.nStages();
+            a = obj.bt.a;
+            b = obj.bt.b;
+            c = obj.bt.c;
+
+            % Compute rates K and stage approximations V
+            K = zeros(length(v), s);
+            V = zeros(length(v), s);
+            for i = 1:s
+                V_i = obj.v;
+                for j = 1:i-1
+                    V_i = V_i + dt*a(i,j)*K(:,j);
+                end
+
+                K_i = F(t+dt*c(i), V_i);
+                K_i = K_i + discreteRates(:,i);
+
+                V(:,i) = V_i;
+                K(:,i) = K_i;
+            end
+
+            % Compute updated updated solution
+            v_next = v;
+            for i = 1:s
+                v_next = v_next + dt*b(i)*K(:,i);
+            end
+
+            obj.v = v_next;
             obj.t = obj.t + obj.dt;
             obj.n = obj.n + 1;
         end
@@ -57,6 +108,7 @@
         % in the time interval [t0, tEnd].
         % The time-step obj.dt is assumed to be aligned with [t0, tEnd] already.
         function tvec = timePoints(obj, t0, tEnd)
+            % TBD: Should this be implemented here or somewhere else?
             N = round( (tEnd-t0)/obj.dt );
             tvec = zeros(N*obj.s, 1);
             s = obj.coeffs.s;
@@ -71,6 +123,7 @@
         % in time interval [t0, tEnd], substage points included.
         % The time-step obj.dt is assumed to be aligned with [t0, tEnd] already.
         function weights = quadWeights(obj, t0, tEnd)
+            % TBD: Should this be implemented here or somewhere else?
             N = round( (tEnd-t0)/obj.dt );
             b = obj.coeffs.b;
             weights = repmat(b', N, 1);