function FAS
% Recursive implementation of the full approximation scheme (FAS) multigrid
% method for the following nonlinear PDE in two dimensions:
%           lambda^2*Laplacian(u(x,y)) + exp(-u(x,y)) = f(x,y)
% The problem is solved over the square grid B = [0,1]^2 with boundary
% condition u(dB) = 0.
% 
% Program graphs the resulting residual as a function of x and y. For 
% details on FAS, see "A Multigrid Tutorial" By W. L. Briggs, page 98.
%
% Matt Warnez -- March 2013

tic
% problem parameters
lambda = 10;

% numerical paramters
d = 6; % number of final V-cycles
alpha1 = 5; % number of coarsening relaxations per V-cycle
alpha2 = 5; % number of refining relaxations per V-cycle

% numerical setup
N = 2^6;
M = N;
x = (0:N)'/N;
y = (0:M)'/M;

% source function
f = zeros(M+1,N+1);

% precomputations
H = 2*(M^2+N^2);
a1 = 2/3*N^2/H;
a2 = 2/3*M^2/H;
a3 = 1/(lambda^2*H);
a323 = 2*a3/3;
Hl2 = H*lambda^2;
N2l2 = N^2*lambda^2;
M2l2 = M^2*lambda^2;

% initial guess
v = zeros(size(f));
ug = v;

% initial relaxation
for n = 1:alpha1
    for i = 2:N
        for j = 2:M
            K = v(j,i)/3 - a3*exp(-v(j,i))/3 + a1*(v(j,i-1) + v(j,i+1)) ...
                + a2*(v(j-1,i) + v(j+1,i)) - a323*f(j,i);
            
            % Newton's iteration
            dF = (ug(j,i) - a3*exp(-ug(j,i)) - K)/(1 + a3*exp(-ug(j,i)));
            while abs(dF) > 10^-12
                
                ug(j,i) = ug(j,i) - dF;
                dF = (ug(j,i) - a3*exp(-ug(j,i)) - K)/(1 + a3*exp(-ug(j,i)));
                
            end
            ug(j,i) = ug(j,i) - dF;
            
        end
    end
    v = ug;
end

% recursive V-cycles
for k = 1:d, v = FASRecursion(v,f,lambda,N,M,alpha1,alpha2); end

% final relaxation
for n = 1:alpha2
    for i = 2:N
        for j = 2:M
            K = v(j,i)/3 - a3*exp(-v(j,i))/3 + a1*(v(j,i-1) + v(j,i+1)) ...
                    + a2*(v(j-1,i) + v(j+1,i)) - a323*f(j,i);
            
            % Newton's iteration
            dF = (ug(j,i) - a3*exp(-ug(j,i)) - K)/(1 + a3*exp(-ug(j,i)));
            while abs(dF) > 10^-12
                
                ug(j,i) = ug(j,i) - dF;
                dF = (ug(j,i) - a3*exp(-ug(j,i)) - K)/(1 + a3*exp(-ug(j,i)));
                
            end
            ug(j,i) = ug(j,i) - dF;
            
        end
    end
    v = ug;
end

% compute residue
res = zeros(M+1,N+1);
for i = 2:N
    for j = 2:M
        res(j,i) = f(j,i) + Hl2*v(j,i) - N2l2*(v(j,i-1) + v(j,i+1))  ...
            - M2l2*(v(j-1,i) + v(j+1,i)) - exp(-v(j,i));
    end
end
toc

% display
figure;
surf(x,y,res);
xlabel('x'); ylabel('y'); zlabel('Residual');

end

function vh = FASRecursion(vh,fh,lambda,N,M,alpha1,alpha2)

% precomputation
H = 2*(M^2+N^2);
a1 = 2/3*N^2/H;
a2 = 2/3*M^2/H;
a3 = lambda^-2/H;
a323 = 2*a3/3;
ug = vh;

% if coarsest grid, perform additional relaxations
if N > 3 || M > 3, q = 1;
else q = 1000; end

% initial relaxation
for n = 1:alpha1*q
    for i = 2:N
        for j = 2:M
            K = vh(j,i)/3 - a3*exp(-vh(j,i))/3 + a1*(vh(j,i-1) + vh(j,i+1)) ...
                + a2*(vh(j-1,i) + vh(j+1,i)) - a323*fh(j,i);
            
            % Newton's iteration
            dF = (ug(j,i) - a3*exp(-ug(j,i)) - K)/(1 + a3*exp(-ug(j,i)));
            while abs(dF) > 10^-12                
                ug(j,i) = ug(j,i) - dF;
                dF = (ug(j,i) - a3*exp(-ug(j,i)) - K)/(1 + a3*exp(-ug(j,i)));                
            end
            ug(j,i) = ug(j,i) - dF;
            
        end
    end
    vh = ug;
end

% if not coarsest grid
if q == 1
    
    % precomputation
    Hl2 = H*lambda^2;
    N2l2 = N^2*lambda^2;
    M2l2 = M^2*lambda^2;
    Ahvh = zeros(M+1,N+1);
    
    % compute residual
    for i = 2:N
        for j = 2:M
            Ahvh(j,i) = N2l2*(vh(j,i-1) + vh(j,i+1)) - Hl2*vh(j,i) ...
                + M2l2*(vh(j-1,i) + vh(j+1,i)) + exp(-vh(j,i));
        end
    end
    r2h = restrict(fh - Ahvh,N,M); % residual on restricted grid
    
    % compute nonlinear operation on coarser grid
    Hl2 = (M^2+N^2)*lambda^2/2;
    N2l2 = (N*lambda)^2/4;
    M2l2 = (M*lambda)^2/4;
    f2h = zeros(M/2+1,N/2+1);
    v2h = restrict(vh,N,M);
    for i = 2:N/2
        for j = 2:M/2
            f2h(j,i) = N2l2*(v2h(j,i-1) + v2h(j,i+1)) - Hl2*v2h(j,i) ...
                + M2l2*(v2h(j-1,i) + v2h(j+1,i)) + exp(-v2h(j,i)) + r2h(j,i);
        end
    end
    
    % call V-cycle on coarser grid
    u2h = FASRecursion(v2h,f2h,lambda,N/2,M/2,alpha1,alpha2);
    
    % correct solution on finer grid
    vh = vh + prolong(u2h-v2h,N,M);
    
    % final relaxation
    ug = vh;
    for n = 1:alpha2
        for i = 2:N
            for j = 2:M
                K = vh(j,i)/3 - a3*exp(-vh(j,i))/3 + a1*(vh(j,i-1) ...
                    + vh(j,i+1)) + a2*(vh(j-1,i) + vh(j+1,i)) - a323*fh(j,i);
                
                % Newton's iteration
                dF = (ug(j,i) - a3*exp(-ug(j,i)) - K)/(1 + a3*exp(-ug(j,i)));
                while abs(dF) > 10^-12
                    ug(j,i) = ug(j,i) - dF;
                    dF = (ug(j,i) - a3*exp(-ug(j,i)) - K)/(1 + a3*exp(-ug(j,i)));
                end
                ug(j,i) = ug(j,i) - dF;
            end
        end
        vh = ug;
    end
    
end
    
end

function u2h = restrict(uh,N,M)
% interpolative grid restriction

u2h = zeros(M/2+1,N/2+1);

for i = 3:2:N-1
    for j = 3:2:M-1
        
        u2h((j+1)/2,(i+1)/2) = uh(j,i)/4 + (uh(j,i-1) + uh(j,i+1) ...
            + uh(j-1,i) + uh(j+1,i))/8 + (uh(j-1,i-1) + uh(j-1,i+1) ...
            + uh(j+1,i-1) + uh(j+1,i+1))/16;
        
    end
end    

end

function uh = prolong(u2h,N,M)
% interpolative grid prolongation

uh = zeros(M+1,N+1);

for i = 2:1:N/2+1
    for j = 2:1:M/2+1
        
        uh(2*j-1,2*i-1) = u2h(j,i);
        uh(2*j-2,2*i-1) = (u2h(j,i) + u2h(j-1,i))/2;
        uh(2*j-1,2*i-2) = (u2h(j,i) + u2h(j,i-1))/2;
        uh(2*j-2,2*i-2) = (u2h(j,i) + u2h(j-1,i) + u2h(j,i-1) + u2h(j-1,i-1))/4;
        
    end
end    

end