%-------------------------------------------------------------------------------------------------
% anomalous_diff3D_decay_test.m
% 04/2018, David Shirokoff
%
% Description: Computes the solution to the nonlinear diffusion for a
% Gaussian initial data, using new ImEx schemes.
% u_t = Au + Bu + f(x,t), 
% Au ~ sigma d^2u/dx^2,  Bu ~ d/dx(  ( u(x)^gamma - sigma )du/dx)
% Au - implicit, Bu - explicit.
%
% Discretizations: Fourier spectral method.
%
% No inputs. Output: Solution level sets to the screen. 
%--------------------------------------------------------------------------------------------------
function anomalous_diff3D_decay_test
    format long;                        %----------------------------------------------------------
	show_plots = 1;                     % Set to 0 to turn off level set plots.
    option = 1;                         % Allows for modification to allow different
                                        % initial data.  (default = 1).      
                                        %----------------------------------------------------------
                                        % Variable Initializations
    n = 128;                            % number grid points.
    h = 1/n;                            % grid size, domain = [0, 1]
    x = (0:n-1)*h;                      % set grid
    y = (0:n-1)*h;    
    z = (0:n-1)*h;    
    [X, Y, Z] = meshgrid(x, y, z);    
    tf = 2^-4;                          %Fix final time.
    del_t_vec = 2.^[-8:-1:-10]';        %Set dt values to loop over.
    num_t = size(del_t_vec, 1);
    one_step = 1;                       %DECLARE: one_step -- constant variable.    
    gam = 5/3;                          %Set adiabatic constant, (default 5/3 - for ideal gas)
                                        %----------------------------------------------------------
    r_max = 3;                          % Max order being tested. NOTE: changing r_max may change
                                        % the values of (delta, sigma) which will change the error
                                        % constant observed.                                        
    order = 3;                          %Set the scheme order to be used (must be <= r_max)
    
    % Initial data.
    u0 = GetInitialData(X,Y,Z, option);     
    % Get ImEx parameter and splitting parameter.                                        
    [delta sigma] = ImexSplittingParameters(u0, gam, r_max, option);
    
    %Create vectors to store the min and max values of u at t=t_f, for
    %different time steps dt.
    max_vec = zeros(num_t,1); min_vec = zeros(num_t, 1);
    mean_of_u = sum(sum(sum(u0)))*h^3;    
    
    %create and save figure for decay of peak solution
    peakSolutionfig=figure;
    
    %Plot initial data level set
    PlotLevelSet(X,Y,Z, u0, mean_of_u, 0, 0)
    labels = cell(num_t,1);
    for jj=1:num_t
        labels{jj} = ['dt= ' num2str(del_t_vec(jj))];
    end
	%labels(jj) = strcat('dt=', num2str(del_t));
    
    
    %Loop over computations for different del_t values.
    for jj = 1:num_t        
        %Get del_t value, and adjust (slight) for the final time.
        del_t = del_t_vec(jj);
        num_steps = round(tf/del_t);
        true_tf = num_steps*del_t;            

        fprintf('\n dt value = %f \n', del_t);
        
        %Initialize the multistep scheme with u0,...ur in the u_Idata array.
        u_Idata = MultistepInitialization(u0, n, order, delta, sigma, X, Y, Z, h, del_t, gam, option);
        
        %Get ImEx coefficients for order and parameter delta
        [a b c] = GetImExCoefficients(order, delta);
        TStepsNeeded = num_steps - (order - 1)*one_step;    %The first (order - 1) time steps are already computed for initialization. 
                                                            %Adjust # time steps, by the order or scheme.                                                             
        %Call main integration.
        [Uf maxU] = ImEx_TimeIntegrate(option, del_t, TStepsNeeded, X, Y, Z, a, b, c, n, h, u_Idata, sigma, gam, show_plots); 
        
        %Create vector of times corresponding to maxU values
        dt_vect = (1:TStepsNeeded)*del_t + del_t*(order - 1);        
        
        %Add to decay of peak solution plot.
        LW = 'LineWidth'; lw = 3;    
        figure(peakSolutionfig)
        hold on;
        plot(log10(dt_vect), log10( abs(maxU - mean_of_u)), LW, lw);
        set(gca,'FontSize',30)
        legend(labels{1:jj}); 
        xlabel('log_{10} t');
        ylabel('log_{10} Max u - avg u');
        title(sprintf('Max u vs t, Scheme order %d', order));
        grid on; box on;    
        hold off;
    end
end

%------------------------------------------------------------------
% Function: MultistepInitialization
%
% Description: Given the initial data u0, computes the initializations
% u_1, ... u_r for a multistep scheme.  Currently it is only coded for
% orders 1 -- 3.
% Remark -- for clarity this routine was not coded using
% recursion, however, is set up to enable a recursive routine for
% modifications to orders 4, and 5.
%
% Inputs: del_t -- time step,
% order -- scheme order
% delta, sigma -- ImEx parameter and splitting parameter.
% (X,Y,Z) -- spatial grid point values.
% n -- spatial grid number
% gam -- adiabatic constant (default 5/3)
%
% Outputs: u_Idata -- a 4d array, containing multistep initialization 
% functions u0, u1, u2. 
%------------------------------------------------------------------
function u_Idata = MultistepInitialization(u0, n, order, delta, sigma, X, Y, Z, h, del_t, gam, option)
    %Initialization for multistep methods: populate initial data array. 3d function x order of the scheme.
    u_Idata = zeros(n, n, n, order);    
    one_step = 1;
    
    %Choose the order of scheme:
    switch order
        case 1
            %u_Idata(:,:,:, j) corresponds to data at time: t = (order - j)*dt;
            u_Idata(:,:,:, 1) = u0; 
                
        case 2
            %----------------------------------------------------------
            %Computes the initial condition at 
            %t=dt, using one first order step (second order LTE).         
            r = 1;
            [a b c] = GetImExCoefficients(r, delta);
            u_Idata(:,:,:, 1) = u0;
            [Uf maxU] = ImEx_TimeIntegrate(option, del_t, one_step, X, Y, Z, a, b, c, n, h, u_Idata, sigma, gam, 0); 
            u_Idata(:,:,:,2) = u0;
            u_Idata(:,:,:,1) = Uf;
                
        case 3                
            %----------------------------------------------------------
            %STEP 1: Compute u_1 using 2nd order scheme with step size
            %del_t_small
   
            %fix step size
            del_t_small = 0.1*del_t^(3/2);              
            num_small_steps = ceil(del_t/del_t_small);
            del_t_small = del_t/num_small_steps;
            
            %Compute u(del_t_small) using one 1st order step
            %(error in u(del_t_small) is second order del_t_small^2)
            r = 1;
            [a b c] = GetImExCoefficients(r, delta);
            u_Idata(:,:,:,1) = u0;
            [U_del_t_small maxU] = ImEx_TimeIntegrate(option, del_t_small, one_step, X, Y, Z, a, b, c, n, h, u_Idata, sigma, gam, 0); 
    
            %Initialize the 2nd order scheme, at t = 0 and t = del_t_small. 
            u_Idata(:,:,:,1) = U_del_t_small;   
            u_Idata(:,:,:,2) = u0;
            
            %Get 2nd order coefficients. 
            r = 2;          
            [a b c] = GetImExCoefficients(r, delta);
                                
            %set the number of steps
            numsteps = num_small_steps-one_step;
                
            %Integrate: 2nd order scheme to get u(del_t), using del_t_small 
            [Uf maxU] = ImEx_TimeIntegrate(option, del_t_small, numsteps, X, Y, Z, a, b, c, n, h, u_Idata, sigma, gam, 0); 
                            
            %----------------------------------------------------------
            %STEP 2: Compute u_2 using one 2nd order step of size del_t
                
            %Initialize u_Idata for a 2nd order scheme, with u_0, u_1: 
            u_Idata(:,:,:,1) = Uf;   
            u_Idata(:,:,:,2) = u0;            
            [Uf maxU] = ImEx_TimeIntegrate(option, del_t, one_step, X, Y, Z, a, b, c, n, h, u_Idata, sigma, gam, 0); 
            
            %----------------------------------------------------------
            %Populate the array of initial data with u_0, u_1, u_2
            %(needed for 3rd order)
            u_Idata(:,:,:,3) = u0;
            u_Idata(:,:,:,2) = u_Idata(:,:,:,1);
            u_Idata(:,:,:,1) = Uf;            
	end
end


%------------------------------------------------------------------
% Function: GetInitialData
%
% Description: Returns the initial data at time t = 0.
%
% Inputs: option -- specify which test is being run (default 1)
% (X,Y,Z) -- spatial grid point values.
%
% Outputs: u0 -- initial data.
%------------------------------------------------------------------
function u0 = GetInitialData(X,Y,Z, option)
    switch option
        case 1
            %Guassian initial data
            V = (X-0.5).^2 + (Y-0.5).^2 + (Z-0.5).^2;   
            u0 = 1 + exp(-V/0.15^2);
    end
end


%------------------------------------------------------------------
% Function: GetImExParameters
%
% Description: Computes the ImEx parameter delta, and splitting
% parameter sigma for a given manufactured solution option number.
%
% Inputs: option -- specify which test is being run (default 2)
% gam -- adiabatic constant (default 5/3).
% r_max -- output coefficients (delta,sigma) stabilize schemes with 
% orders 1 <= r <= r_max. 
%
% Outputs: (delta, sigma) parameters. 
%------------------------------------------------------------------
function [delta sigma] = ImexSplittingParameters(u0, gam, r, option)
        switch option
            case 1
                eta = 0.1;                  %Gap parameter, 0 < eta < 1 (default = 0.1)
                                            %eta closer to 0 gives better error constant but less robust.                
                coeff = u0.^gam;            %Estimate/compute max and min values of the solution.
                dmin = min(min(min(coeff)));
                dmax = max(max(max(coeff))); 
                kap = dmin/dmax*(1-eta);
                cr = cos(pi/r)^(-r);
                sigma = dmin*(1-1/2*eta)*( 1 + cr )/(1 + kap*cr);
                delta = min(2 - 2*( (1-kap)/(1+kap*cr ) )^(1/r), 1);    
        end
end


%------------------------------------------------------------------
% Function: PlotLevelSet
%
% Description: Plots the 3d level set of phi = isovalue.
%
% Inputs: phi -- the 3d scalar function with values at (X,Y,Z).
% (X,Y,Z) -- spatial grid point values.
% isovalue -- the level set value to be plotted.
%
% Outputs: 3d plot to the screen.
%------------------------------------------------------------------
function PlotLevelSet(X,Y,Z, phi, isovalue, t, order)
    figure
    p = patch(isosurface(X,Y,Z,phi,isovalue));
    set(p, 'EdgeColor', 'none');
    set(p, 'FaceColor', 'blue');
    
    daspect([1,1,1])
    view(3); 
    xmin = X(1,1,1);    xmax = X(1, end, 1);
    ymin = Y(1,1,1);    ymax = Y(end, 1,1);
    zmin = Z(1,1,1);    zmax = Z(1,1,end);
    
    axis([xmin xmax ymin ymax zmin zmax])
    LW = 'linewidth'; MS = 'markersize'; ms = 5; lw = 3;
    isonormals(X,Y,Z,phi,p)
    set(p, MS, ms, LW, lw)
    grid on, box on, axis square
    set(gca,'FontSize',20)
    
    if (order > 0)
        title(sprintf('order = %d, \t t=%0.5f',order, t))
    else
        title(sprintf('\t t=%0.5f', t))
    end
    xlabel('x'), ylabel('y'), zlabel('z')
    camlight
    lighting gouraud
    rotate3d on
end



%--------------------------------------------------------------------------
% Function: ImEx_TimeIntegrate
%
% Description: Uses ImEx multistep methods to integrate to the final 
% time.
%
% Outputs: Uf -- solution u at grid points x and final time 
% maxU -- the maximum value max_{x,y,z} u(x,y,z, t) at every time step.
%--------------------------------------------------------------------------
function [Uf maxU] = ImEx_TimeIntegrate(option, del_t, num_steps, X, Y, Z, a, b, c, n, h, u_Idata, diff_coeffa, gam, show_plots)
        
        % Build Fourier wavenumbers.
        L = 1;
        ikx = 1i*(mod((1:n) - ceil(n/2+1),n) - floor(n/2))*(2*pi)/L;
        ikx(n/2 + 1) = 1i*(n/2)*(2*pi)/L;   ikx(1) = 0;
        iky = 1i*(mod((1:n) - ceil(n/2+1),n) - floor(n/2))*(2*pi)/L;
        iky(n/2 + 1) = 1i*(n/2)*(2*pi)/L;    iky(1) = 0;
        ikz = 1i*(mod((1:n) - ceil(n/2+1),n) - floor(n/2))*(2*pi)/L;
        ikz(n/2 + 1) = 1i*(n/2)*(2*pi)/L;    ikz(1) = 0;
    
        kxyz = zeros(n, n, n);          %Build Laplacian multipliers
        ikx_mat = zeros(n, n, n);
        iky_mat = zeros(n, n, n);
        ikz_mat = zeros(n, n, n);
        
        for jj = 1:n
            for kk = 1:n
                for ll = 1:n
                    kxyz(jj, kk, ll) = iky(jj)^2 + ikx(kk)^2 + ikz(ll)^2;
                
                    ikx_mat(jj, kk, ll) = ikx(kk);
                    iky_mat(jj, kk, ll) = iky(jj);
                    ikz_mat(jj, kk, ll) = ikz(ll);
                end
            end
        end
        
        % Compute initial data: u_n, Au_n and Bu_n, for n = 1, 2.., r-1. 
        % We use the exact initial data for the manufactured solution.
        order = size(a, 2)-1;
        u_old = zeros(n, n, n, order);
        bu_old = zeros(n, n, n, order);
        au_old = zeros(n, n, n, order);
        f_old = zeros(n, n, n, order);        
        for jj = 1:order
            %Evaluate at time: t = (order-jj)*del_t;
            u_old(:,:,:, jj) = u_Idata(:,:,:,jj);
            [diff_coeffb] = GetDiffusionCoeff(option, X, Y, Z, u_old(:, :, :, jj), diff_coeffa, gam);
            
            % Compute the Bu terms
            utemp_hat = fftn(u_old(:, :, :, jj) );
            ux_temp = ifftn(ikx_mat.*utemp_hat);
            uy_temp = ifftn(iky_mat.*utemp_hat);
            uz_temp = ifftn(ikz_mat.*utemp_hat);
            vx = ux_temp.*diff_coeffb; 
            vy = uy_temp.*diff_coeffb;
            vz = uz_temp.*diff_coeffb;
            divhat = fftn(vx).*ikx_mat + fftn(vy).*iky_mat + fftn(vz).*ikz_mat;
            bu_old(:, :, :, jj) = real(ifftn(divhat)); 
            
            % Compute the Au terms
            au_old(:, :, :, jj) = diff_coeffa*real(ifftn(kxyz.*utemp_hat));                    
        end                    
        imp_part = a(1)*ones(n, n, n) - diff_coeffa*del_t*c(1)*kxyz;
        strobe = floor(num_steps/4);
        maxU = zeros(num_steps, 1);    
        
        % Main time loop -- integrate num_steps.
        for counter = 1:num_steps    
        	
            % compute new right hand side.
            rhs = zeros(n,n,n);
            for ind = 1:order
                rhs = rhs - u_old(:, :, :, ind)*a(1+ind) + del_t*( au_old(:,:, :, ind)*c(1+ind) + bu_old(:, :, :, ind)*b(1+ind) );            
            end
            
            % solve for new u_{n+1} via Fourier transform.
            u = real(ifftn(fftn(rhs)./imp_part));
            
            % update stored vectors, one time step.
            for jj = 1:order-1                
                u_old(:, :, :, order - jj + 1) = u_old(:, :, :, order - jj);
                bu_old(:, :, :, order - jj+1) = bu_old(:, :, :, order - jj);
                au_old(:, :, :, order - jj+1) = au_old(:, :, :, order - jj);
                
            end
            [diff_coeffb] = GetDiffusionCoeff(option, X, Y, Z, u, diff_coeffa, gam);
            
            % For efficiency, compute and store A u_{n+1}, B u_{n+1}, and
            % f(t_{n+1})
            utemp_hat = fftn(u);
            ux_temp = ifftn(ikx_mat.*utemp_hat);
            uy_temp = ifftn(iky_mat.*utemp_hat);
            uz_temp = ifftn(ikz_mat.*utemp_hat);
            vx = ux_temp.*diff_coeffb; 
            vy = uy_temp.*diff_coeffb;
            vz = uz_temp.*diff_coeffb;
            divhat = fftn(vx).*ikx_mat + fftn(vy).*iky_mat + fftn(vz).*ikz_mat;
            bu_1 = real(ifftn(divhat)); 
            au_1 = diff_coeffa*real(ifftn(kxyz.*utemp_hat));
           
            u_old(:, :, :, 1) = u;
            bu_old(:, :, :, 1) = bu_1;
            au_old(:, :, :, 1) = au_1;
            
            %Store the maxU value.
            maxU(counter) = max(max(max(u)));

            %Plot solution level set.
            if (mod(counter, strobe) == 0 && show_plots == 1)
                mean_of_u = sum(sum(sum(u_Idata(:,:,:, order))))*h^3;
                t = (counter + order - 1)*del_t;        %Compute time t -- shifted
                                                        %by the number of
                                                        %initialization
                                                        %steps for
                                                        %multistep schemes
                PlotLevelSet(X,Y,Z, u, mean_of_u, t, order);
            end    
        end
        Uf =  u;        
end


%--------------------------------------------------------------------------
% Function: GetDiffusionCoeff
%
% Description: Computes the nonlinear diffusion coefficient.
%
% Inputs: option -- manufactured solution option (default 2)
% (X,Y,Z) -- allows for variable spatial dependence to be incorporated 
% u_in -- function value
% gam -- adiabatic constant (default 5/3)
% sigma -- splitting parameter
%
% Outputs: diff_coeffb -- diffusion coefficient used in the B (explicit)
% term of the ImEx scheme.
%--------------------------------------------------------------------------
function [diff_coeffb] = GetDiffusionCoeff(option, X, Y, Z, u_in, sigma, gam)
    switch option
        case 1
            diff_coeffb = u_in.^gam - sigma;         
    end
end


%------------------------------------------------------------------
% Function: GetImExCoefficients
%
% Description: Inputs the scheme order and ImEx parameter delta, and 
% returns the ImEx time-stepping coefficients. For order r, there are
% r + 1 coefficients (for each, a_coeff, b_coeff, and c_coeff). 
%
% Inputs: order (integer 1-5), delta (float, 0 < delta <= 1)
% Outputs: ImEx coefficients, a_coeff, b_coeff, c_coeff
%
%------------------------------------------------------------------
function [a_coeff b_coeff c_coeff] = GetImExCoefficients(order, delta)
    switch order
        case 1
            a_coeff = [delta -delta];
            b_coeff = [0 delta];
            c_coeff = [1 delta-1];

        case 2
            a_coeff = [(2*delta - delta^2/2) (2*delta^2 - 4*delta) (2*delta-(3*delta^2)/2)];
            b_coeff = [0 2*delta (delta - 1)^2 - 1];
            c_coeff = [1 (2*delta - 2) (delta - 1)^2];
            
        case 3
            a_coeff =[(delta^3/3 - (3*delta^2)/2 + 3*delta), -((3*delta^3)/2 - (15*delta^2)/2 + 9*delta), ...
                (3*delta^3 - (21*delta^2)/2 + 9*delta), (-3*delta + (9*delta^2)/2 - (11*delta^3)/6) ];
            b_coeff =[ 0, 3*delta,  ((delta - 1)^2 + (2*delta - 2)*(delta - 1) - 3), (delta - 1)^3 + 1];
            c_coeff = [1, (3*delta - 3), ((delta - 1)^2 + (2*delta - 2)*(delta - 1)), (delta - 1)^3];
            
        case 4 
            a_coeff = [(- delta^4/4 + (4*delta^3)/3 - 3*delta^2 + 4*delta), -(-(4*delta^4)/3 + (22*delta^3)/3 - ...
                18*delta^2 + 16*delta), (- 3*delta^4 + 18*delta^3 - 36*delta^2 + 24*delta), ...
                -(- 4*delta^4 + (58*delta^3)/3 - 30*delta^2 + 16*delta), 4*delta - 9*delta^2 + (22*delta^3)/3 - (25*delta^4)/12 ];
            b_coeff = [ 0, 4*delta, ((delta - 1)^2 + (2*delta - 2)*(delta - 1) + (3*delta - 3)*(delta - 1) - 6),  ...
                ((delta - 1)*((delta - 1)^2 + (2*delta - 2)*(delta - 1)) + (delta - 1)^3 + 4), (delta - 1)^4  - 1];
            c_coeff = [1, (4*delta - 4), ((delta - 1)^2 + (2*delta - 2)*(delta - 1) + (3*delta - 3)*(delta - 1)), ...
                ((delta - 1)*((delta - 1)^2 + (2*delta - 2)*(delta - 1)) + (delta - 1)^3), (delta - 1)^4]; 
            
        case 5
           a_coeff =[(delta^5/5 - (5*delta^4)/4 + (10*delta^3)/3 - 5*delta^2 + 5*delta), -((5*delta^5)/4 - ...
               (95*delta^4)/12 + (65*delta^3)/3 - 35*delta^2 + 25*delta), ((10*delta^5)/3 - (65*delta^4)/3 + (190*delta^3)/3 ...
               - 90*delta^2 + 50*delta), -(5*delta^5 - 35*delta^4 + (280*delta^3)/3 - 110*delta^2 + 50*delta), ...
               (5*delta^5 - (365*delta^4)/12 + (200*delta^3)/3 - 65*delta^2 + 25*delta), 15*delta^2 - (55*delta^3)/3 ...
               + (125*delta^4)/12 - (137*delta^5)/60 - 5*delta];
           b_coeff = [0, 5*delta, ((delta - 1)^2 + (2*delta - 2)*(delta - 1) + (3*delta - 3)*(delta - 1) + ...
               (4*delta - 4)*(delta - 1) - 10),  ((delta - 1)*((delta - 1)^2 + (2*delta - 2)*(delta - 1)) + ...
               (delta - 1)^3 + (delta - 1)*((delta - 1)^2 + (2*delta - 2)*(delta - 1) + (3*delta - 3)*(delta - 1)) + 10), ...
               ((delta - 1)^4 + ((delta - 1)*((delta - 1)^2 + (2*delta - 2)*(delta - 1)) + (delta - 1)^3)*(delta - 1) - 5), 1 + (delta - 1)^5];
           c_coeff = [1, (5*delta - 5), ((delta - 1)^2 + (2*delta - 2)*(delta - 1) + (3*delta - 3)*(delta - 1) + ...
               (4*delta - 4)*(delta - 1)), ((delta - 1)*((delta - 1)^2 + (2*delta - 2)*(delta - 1)) + (delta - 1)^3 + ...
               (delta - 1)*((delta - 1)^2 + (2*delta - 2)*(delta - 1) + (3*delta - 3)*(delta - 1))), ((delta - 1)^4 + ...
               ((delta - 1)*((delta - 1)^2 + (2*delta - 2)*(delta - 1)) + (delta - 1)^3)*(delta - 1)), (delta - 1)^5];           
    end    
end