--- a/+rv/ResidualViscosity.m	Thu Dec 06 17:03:22 2018 +0100
+++ b/+rv/ResidualViscosity.m	Fri Dec 07 13:11:53 2018 +0100
@@ -1,12 +1,13 @@
-% class describing the viscosity 
 classdef ResidualViscosity < handle
     properties
         D % Diff op approximating the gradient of the flux f(u)
         waveSpeed % Wave speed at each grid point, e.g f'(u). %TBD: Better naming?
-        Cmax % Constant controling magnitude of upwind dissipation
-        Cres % Constant controling magnitude residual dissipation
-        h % Length scale used for scaling the viscosity.
+        Cmax % Constant controlling relative amount of upwind dissipation
+        Cres % Constant controling relative amount of upwind dissipation
+        h % Length scale used for scaling the viscosity. Typically grid spacing.
         viscosity % Stores the computed viscosity.
+        normalization % Function used to normalize the residual such that it is amplified in the
+                      % shocks.
 
         % Convenice (for verification and plotting) TBD: Decide on if it should be kept.
         u_t % Stores the latest approximated time derivative of the solution.
@@ -15,13 +16,10 @@
     end
 
     methods
-        % TODO: - Consider passing residual normalization as a function handle.
-        %         or choosing a type of normalization on construction.
-        %         Could for example be 1, norm((v-mean(v),inf) or normInfNeighborhood(v)
-        %         but working
-        %       - Decide on how to treat waveSpeed. It would be nice to just pass a constant value without
-        %         wrapping it in a function.
-        function obj = ResidualViscosity(D, waveSpeed, Cmax, Cres, h, N)
+        %  TBD: Decide on how to treat waveSpeed. It would be nice to just pass a constant value without
+        %       wrapping it in a function.
+        function obj = ResidualViscosity(D, waveSpeed, Cmax, Cres, h, N, normalization)
+            default_arg('normalization',@(v)norm(v-mean(v),inf));
             obj.D = D;
             obj.waveSpeed = waveSpeed;
             obj.h = h;
@@ -31,22 +29,14 @@
             obj.u_t = zeros(N,1);
             obj.grad_f = zeros(N,1);
             obj.residual = zeros(N,1);
+            obj.normalization = normalization;
         end
 
         function obj = update(obj, v, dvdt)
             obj.u_t = dvdt;
             obj.grad_f = obj.D(v);
             obj.residual = obj.u_t + obj.grad_f;
-            %obj.viscosity = min(obj.Cmax*obj.h*abs(obj.waveSpeed(v)), obj.Cres*obj.h^2*abs(obj.residual)/norm(v-mean(v),inf));
-            obj.viscosity = obj.smoothen(obj.Cres*obj.h^2*abs(obj.residual)/norm(v-mean(v),inf));
-
-        end
-
-        function smoothendVector = smoothen(obj, vector)
-            smoothendVector = vector;
-            for i = 2:length(vector)-1
-                smoothendVector(i) = (1/6)*(vector(i-1) + 4*vector(i) + vector(i+1));
-            end
+            obj.viscosity = min(obj.Cmax*obj.h*abs(obj.waveSpeed(v)), obj.Cres*obj.h^2*abs(obj.residual)/obj.normalization(v));
         end
 
         function [residual, u_t, grad_f] = getResidual(obj)
@@ -59,16 +49,4 @@
             viscosity = obj.viscosity;
         end
     end
-    % Remove or fix. Should be able to handle values close to zero. Should work in 2d and 3d.
-    methods (Static)
-        function R_norm = normInfNeighborhood(v)
-            n = length(v);
-            R_norm = zeros(n,1);
-            R_norm(1,1) = norm(v(1:3), inf);
-            R_norm(n,1) = norm(v(n-3:n), inf);
-            for i = 2:n-1
-                R_norm(i,1) = norm(v(i-1:i+1), inf);
-            end
-        end        
-    end
 end
\ No newline at end of file