--- a/+rv/constructDiffOps.m	Fri Jan 18 08:58:26 2019 +0100
+++ b/+rv/constructDiffOps.m	Fri Jan 18 09:02:02 2019 +0100
@@ -1,68 +1,74 @@
-function [D_rv, D_flux, DvDt, solutionPenalties, residualPenalties] = constructDiffOps(scheme, grid, order, opSet, waveSpeed, BCs, fluxSplitting)
-    default_arg('fluxSplitting',[]);
-
+function [D_rv, D_flux, DvDt, solutionPenalties, residualPenalties] = constructDiffOps(scheme, g, order, schemeParams, opSet, BCs)
     %% DiffOps for solution vector
-    [D, solutionPenalties] = constructTotalFluxDiffOp(scheme, grid, order, opSet, waveSpeed, BCs, fluxSplitting);
-    D2 = constructSymmetricD2Operator(grid, order, opSet);
-    D_rv = @(v,viscosity)(D + D2(viscosity))*v;
+    [D, solutionPenalties] = constructTotalFluxDiffOp(scheme, g, order, schemeParams, opSet, BCs);
+    D2 = constructSymmetricD2Operator(g, order, opSet);
+    D_rv = @(v,viscosity)(D(v) + D2(v, viscosity));
 
     %% DiffOps for residual viscosity
-    [D_flux, residualPenalties] = constructTotalFluxDiffOp(scheme, grid, max(order-2,2), opSet, waveSpeed, BCs, fluxSplitting);
+    [D_flux, residualPenalties] = constructTotalFluxDiffOp(scheme, g, max(order-2,2), schemeParams, opSet, BCs);
     % DiffOp for flux in residual viscosity. Due to sign conventions of the implemnted schemes, we need to
     % change the sign.
-    D_flux = -D_flux;
-    D_flux = @(v) D_flux*v;
+    D_flux = @(v) -D_flux(v);
     % DiffOp for time derivative in residual viscosity
-    DvDt = @(v)D*v;
+    DvDt = D;
 end
 
-function [D, penalties] = constructTotalFluxDiffOp(scheme, grid, order, opSet, waveSpeed, BCs, fluxSplitting)
-    if isequal(opSet, @sbp.D1Upwind)
-        diffOp = scheme(grid, order, opSet, waveSpeed, fluxSplitting);
-    else
-        diffOp = scheme(grid, order, opSet, waveSpeed);
-    end
+function [D, penalties] = constructTotalFluxDiffOp(scheme, g, order, schemeParams, opSet, BCs)
+    diffOp = scheme(g, order, schemeParams{:}, opSet);
     [D, penalties]  = addClosuresToDiffOp(diffOp, BCs);
 end
 
 function [D, penalties] = addClosuresToDiffOp(diffOp, BCs)
-    D = diffOp.D;
+    if ~isa(diffOp.D, 'function_handle')
+        D = @(v) diffOp.D*v
+    else
+        D = diffOp.D;
+    end
     penalties = cell(size(BCs));
     for i = 1:size(BCs,1)
         for j = 1:size(BCs,2)
             [closure, penalties{i,j}] = diffOp.boundary_condition(BCs{i,j}.boundary, BCs{i,j}.type);
-            D = D + closure;
+            if ~isa(closure, 'function_handle')
+                closure = @(v) closure*v;
+            end
+            D = @(v) D(v) + closure(v);
         end
     end
 end
 
-function D2 = constructSymmetricD2Operator(grid, order, opSet)
-    % TODO: 
-    % Currently only implemented for upwind operators.
-    % Remove this part once the time-dependent D2 operator is implemented for other opSets
-    % or if it is decided that it should only be supported for upwind operators.
-    assert(isequal(opSet,@sbp.D1Upwind))
+function D2 = constructSymmetricD2Operator(g, order, opSet)
+
 
-    m = grid.size();
-    ops = cell(grid.D(),1);
-    I = cell(grid.D(),1);
-    for i = 1:grid.D()
-       lim = {grid.x{i}(1), grid.x{i}(end)};
+    m = g.size();
+    ops = cell(g.D(),1);
+    I = cell(g.D(),1);
+    for i = 1:g.D()
+       lim = {g.x{i}(1), g.x{i}(end)};
        ops{i} = opSet(m(i), lim, order);
        I{i} = speye(m(i));
     end
 
     % TBD: How is this generalized to a loop over dimensions or similar?
-    switch grid.D()
+    switch g.D()
         case 1
+            
             e_r = ops{1}.e_r;
             e_l = ops{1}.e_l;
-            Dm = ops{1}.Dm;
-            Dp = ops{1}.Dp;
             Hi = ops{1}.HI;
             B = e_r*e_r' - e_l*e_l';
-            D2 = @(viscosity) Dm*spdiag(viscosity)*Dp-Hi*(B*spdiag(viscosity)*Dp);
+            if isequal(opSet,@sbp.D1Upwind)
+                Dm = ops{1}.Dm;
+                Dp = ops{1}.Dp;
+                D2 = @(viscosity) Dm*spdiag(viscosity)*Dp-Hi*(B*spdiag(viscosity)*Dp);
+            else
+                D2 = @(viscosity)ops{1}.D2(viscosity);
+            end
         case 2
+            % TODO: 
+            % Currently only implemented for upwind operators.
+            % Remove this part once the time-dependent D2 operator is implemented for other opSets
+            % or if it is decided that it should only be supported for upwind operators.
+            assert(isequal(opSet,@sbp.D1Upwind))
             e_e = kron(ops{1}.e_r,I{2});
             e_w = kron(ops{1}.e_l,I{2});
             Dm_x = kron(ops{1}.Dm,I{2});
@@ -82,4 +88,5 @@
         otherwise
             error('3D not yet implemented')
     end
+    D2 = @(v, viscosity) D2(viscosity)*v;
 end
\ No newline at end of file