function variable_coefficient_imex_chebyshev
%--------------------------------------------------------------------------
% This function demonstrates the utility of the new ImEx coefficients
% for a variable diffusion equation by performing a covergence test.  
% A Chebyshev grid is used for a spatial discretization of the domain
% [-1, 1].  This code uses the Chebyshev differentiation matrix 
% defined in [N. L. Trefethen, Spectral Methods in MATLAB, (2000)].  
%
% This code time steps the ODE u_t = Amat u + Bmat u + f.
% Modify Amat, Bmat and f for different choices of PDEs and splittings. 
%
% 09/2016, David Shirokoff
%
% Note: The two variables 'order' and 'delta' are needed to defined
% the ImEx coefficients.  
%
%--------------------------------------------------------------------------
    format long;
    tic
    n = 100;    %Number of gridpoints 
                %Time steps are taken to be powers of 2.
    pow2 = -6:-1:-10;  
    del_t_vect = (2.^pow2');    
                %Choose the orders of the scheme you would like to output
                %in the final table by changing the values contained in the
                %'order_vect' vector.  i.e., order_vect = 1:5 will produce 
                %convergence result for orders 1 through 5.                
    order_vect = 1:5;
                %Fix the parameter value for the new ImEx coefficients:
                % 0 < delta <=1 are the allowable values. In practice
                % choose delta as large as possible while still maintaining
                % stability.
    delta = 0.12;
    tf = 1;     %Final integration time.
    test_option = 1; % Specifies the test_option for the test case
                     % test_option = 1: variable diffusion. Add more test
                     % options by modifying the code below. 
                     
    num_order = size(order_vect, 2);
    num_points = size(pow2, 2);
    err_vect = zeros(num_points, num_order);
    [x, x_int] = cheb_grid(n);    
    for ord_counter = 1:num_order
        fprintf('Calculating...Order = %d\n', ord_counter); 
        order = order_vect(ord_counter);
        [a_coeff b_coeff c_coeff] = GetImExCoefficients(order, delta);
        for m = 1:num_points
            del_t = del_t_vect(m);            
            num_steps = round(tf/del_t); 
            true_tf = num_steps*del_t;            
            [Uf, x] = ImEx_TimeIntegrate(test_option, del_t, num_steps, a_coeff, b_coeff, c_coeff, n);        
            [Uexact]=  ExactSolution(test_option, x_int, true_tf);    
            err_vect(m, ord_counter) = max(abs(Uf - Uexact));      
        end
    end 
    cvgrate = zeros(num_points, num_order);
    cvgrate(2:end, :) = log( err_vect(1:end-1, :)./err_vect(2:end, :) )/log(2);    
    fprintf('\nDelta t\t')
    for jj = 1:num_order
        fprintf('\tLinf Err\tRate');
    end
    fprintf('\n');
    for counter = 1:num_points
        fprintf('2^{%2d}\t&', pow2(counter));
        for jj = 1:num_order
            fprintf('\t%1.1e\t &\t%1.1f  &', err_vect(counter, jj), cvgrate(counter, jj));
        end
        fprintf('\n');        
    end
    toc
end

%--------------------------------------------------------------------------
% CHEB compute D = differentiation matrix, x = Chebyshev grid
% Cite for routine:
% [N. L. Trefethen, Spectral Methods in MATLAB, (2000)].  
function [D,x] = cheb(n)
    if n==0, D=0; x=1; return, end
    x = cos(pi*(0:n)/n)';
    c = [2; ones(n-1,1); 2].*(-1).^(0:n)';
    X = repmat(x,1,n+1);
    dX = X-X';
    D = (c*(1./c)')./(dX+(eye(n+1))); % off-diagonal entries
    D = D - diag(sum(D'));
end

%--------------------------------------------------------------------------
% Routine computes the cheb_grid
% [N. L. Trefethen, Spectral Methods in MATLAB, (2000)].  
function [x, x_int] = cheb_grid(n)
    if n==0, D=0; x=1; return, end
    x = cos(pi*(0:n)/n)';
    x_int = x(2:end-1);
end

%--------------------------------------------------------------------------
% Function builds the implicit matrix A and explicit matrix B used
% in the time stepping.
function [Amat, Bmat, x] = Build_matrices(n, test_option)     
    alpha = 2.5;
    [D,x] = cheb(n);
    D2 = D*D;
    Lap = D2(2:end-1, 2:end-1);
    Lsym = (Lap + Lap')/2;
    d_coeff_diag= Var_diffusion(x, test_option);
    d_coeff_mat = diag(d_coeff_diag);
    L = D*d_coeff_mat*D;    
    Amat = alpha*Lsym;
    Bmat = L(2:end-1, 2:end-1) - Amat;
end

%--------------------------------------------------------------------------
% Main routine that time solves the ODE u_t = Amat u + Bmat u + f.
% Integration is done using a fixed time step del_t and evolved for a total
% of num_steps. 
function [Uf, x] = ImEx_TimeIntegrate(test_option, del_t, num_steps, a, b, c, n);        
        
        %Get matrices Amat and Bmat.  Amat is treated implicitly, Bmat
        %explcitly.
        [Amat, Bmat, x] =  Build_matrices(n, test_option);
        x_int = x(2:end-1);        
        order = size(a, 2)-1;
        %Initialize vectors.
        u_old = zeros(n-1, order);  f_old = zeros(n-1, order);
        au_old = zeros(n-1, order); bu_old = zeros(n-1, order);        
        
        % Fix the initial data for the test case for LMM.
        for jj = 1:order
            [u_old(:, jj)]=  ExactSolution(test_option, x_int, -(jj-1)*del_t);           
            bu_old(:, jj) = Bmat*u_old(:, jj);
            au_old(:, jj) = Amat*u_old(:, jj);
            f_old(:, jj) = Forcing(test_option, x_int, -(jj-1)*del_t);
        end           
        
        imp_part = a(1)*eye(n-1) - del_t*c(1)*Amat;
        [Lmat, Umat] = lu(imp_part);
        
        % Main time stepping loop.
        for counter = 1:num_steps        
            rhs = -u_old*a(2:end)' + del_t*( au_old*c(2:end)' + bu_old*b(2:end)' ) + del_t*f_old*b(2:end)';                        
            u = Umat \ (Lmat \ rhs);
            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);
                f_old(:, order - jj+1) = f_old(:, order - jj);
            end                     
            bu_1 = Bmat*u;
            au_1 = Amat*u;            
            u_old(:, 1) = u;
            bu_old(:, 1) = bu_1;
            au_old(:, 1) = au_1;
            f_old(:, 1) =  Forcing(test_option, x_int, counter*del_t);            
        end
        Uf =  u;
end

%--------------------------------------------------------------------------
% Modify the following 3 functions if you wish to include additional 
% test cases. 
%
% Function returns the exact solution at time t for a manufactured solution
% test
function [Uexact] = ExactSolution(test_option, x, t)
    switch test_option
        case 1
            Uexact = sin(20*t)*sin(2*pi*x).*exp(sin(2*pi*x));
    end
end

%-------------------------------------------------------------------------
% This function returns the variable diffusion coefficient in the test case
function [Diff_x] = Var_diffusion(x, test_option)
    switch test_option
        case 1
            Diff_x = 4+3*cos(2*pi*x);
    end
end

%-------------------------------------------------------------------------
% The function computes and returns the forcing for a manufactured solution
% test case.
function [forcing] = Forcing(test_option, x, t)
    switch test_option
        case 1
            forcing = 4*exp(sin(2*pi*x)).*(5*cos(20*t).*sin(2*pi*x) + 4*pi^2*sin(20*t) - 12*pi^2*sin(20*t).*cos(2*pi*x).^2 - 12*pi^2*sin(20*t).*cos(2*pi*x).^3 + 6*pi^2*sin(20*t).*cos(2*pi*x) + 4*pi^2*sin(20*t).*sin(2*pi*x) + 6*pi^2*sin(20*t).*cos(2*pi*x).*sin(2*pi*x) - 4*pi^2*sin(20*t).*cos(2*pi*x).^2.*sin(2*pi*x) - 3*pi^2*sin(20*t).*cos(2*pi*x).^3.*sin(2*pi*x));
    end    
end

%-------------------------------------------------------------------------
% Function computes the new ImEx coefficients for a scheme with order 
% 'order' and parameter 'delta'.
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