%mgm_2d.m --- Linear Multigrid method for a 2D model PDE on [0,1]^2
%         div(D * grad(u)) + c(x,y)*u = d(x,y)  with D=D(x,y)=(a,b), c<0,
%Usage:                                          and u|Gamma = e(x,y)
%      [u r]=mgm_2d(m,n,levs, gam1,gam2,nu, maxit,tol,iout, fmg,u_init);  
%                        for the solution u and residual r on level=levs
%  e.g.   u=mgm_2d(13,17,2, 2,1,2, 10,1.0e-3);  
%where    m,n = m x n defines the finest grid [2^J1-1 x 2^J2-1] say 31x65
%        levs = the number of levels requested say 2 for 2-grid etc
%      gam1/2 = the relaxation steps in pre-restriction/post-interpolation
%          nu = the number of cycling pattern say 2 (for W) or 1 (for V)
%       maxit : the number of maximum MG iterations say 3
%         tol : the requested tolerance (achieved within `maxit')
%        iout : level of intermediate output (0 - default no  1 - active)
%         fmg : 1 (default ON) 0 for OFF -- full MG for initial guess u
%      u_init : initial approximation of u e.g. u=ones(m,n) [NB iprob=2 now]
%        u, r : output the solution u on T_J=T_levs / the residual history r
%(c) K Chen (2004) University of Liverpool, UK
function [u,res]=mgm(m,n,levi, gam1,gam2,nu, maxit,tol,iout,fmg,u_init)
         if nargin<6, help mgm_2d, return;end
         global iprob
         iprob=2; %=============== Problem switch (may add more)
%%%%%%%%%%%%%%% Declare Global Variables to save on dummy variables %%%%%%%
        global hx hy nnx nny levs  
        levs = levi; t0=cputime;
         if nargin<7, maxit=2; end
         if nargin<8,   tol=1.0e-3; end
         if nargin<9,  iout=0; end
         if nargin<10, fmg=1; end
         if nargin<11, u_init=zeros(m,n); end
 %Notation:  Uj is the solution u on level j -------------------------------
 %        :COFj is the sparse matrix "A" on level j
 %        :RHSj is the right hand side "d" on level j
 %        :  Rj is the residual r on level j
 %        : BCj is the boundary condition for j (2 x (mj + nj))
 %    L_J :  The Finest level (J=levs) with L_1 = The Coarsest level 
%%%%%%%%%%%%%%% Initial set-up Start %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        global hx hy nnx nny levs   space iter  
        global U BC COF RHS R 
        mn=min(m+1,n+1); levm=floor(log(mn)/log(2)); 
        if levs>levm | rem(m+1,2^(levs-1))>0 | rem(n+1,2^(levs-1))>0 | ...
         (m+1)/2^(levs-1)<1|(n+1)/2^(levs-1)<1| nu<1 |gam1<1|gam2<1|...
         ~rem(min(m,n),2)
         help mgm_2d,m,n,levs,levm
         disp('                  CHECK the first 6 input parameters!')
         error('Either too many levels requested or Wrong m x n!');return
        end
        space='                                          ';
        U=[];RHS=[];COF=[];R=[];BC=[]; hx=[];hy=[]; nnx=[];nny=[];
disp(['|____ ' mfilename ' _______Linear Multigrid Method_________________|']);
fprintf('\t\tUsing %d levels, gamma_1=%d,gamma_2=%d,%d-cycling\n',...
        levs,gam1,gam2,nu) 
%_____________________________Make variables
          maxi=floor(log(min(m,n))/log(10)); nus=[]; 
for i=1:levs  % NB  levs=2 for 2-grid
   nnx(levs-i+1)=(m+1)/2^(i-1) - 1; % nnx(1)=coarest,  nnx(levs)=finest
   nny(levs-i+1)=(n+1)/2^(i-1) - 1; % Grid points
end %-------------------- -----------------------
for i=1:levs  % NB  levs=2 for 2-grid
   hx(i)=1/(nnx(i)+1);  hy(i)=1/(nny(i)+1);    % Stepsize
    id = floor(log(i)/log(10)); % spacing for %d below i.e. names
      s=sprintf('U%d%s',i,space(1:maxi-id));     U=[U;s];
      s=sprintf('RHS%d%s',i,space(1:maxi-id)); RHS=[RHS;s];
      s=sprintf('COF%d%s',i,space(1:maxi-id)); COF=[COF;s];
      s=sprintf('R%d%s',i,space(1:maxi-id));     R=[R;s];
      s=sprintf('BC%d%s',i,space(1:maxi-id));   BC=[BC;s];
end % Make enough arrays
for i=1:levs % Declare U1 R1 U2 R2 U3 R3 etc as global quantities
 eval(['global ' U(i,:) ' ' BC(i,:) ' ' COF(i,:) ' ' RHS(i,:) ' ' R(i,:) ])
 fprintf('\t T%2d :  grid %3d X %d\n', i,nnx(i),nny(i))
end %%%%%%%%%%%% Initial set-up Completed %%%%%%%%%%%%%%%%%%%%%%%%%%%

  Get_A_rhs_R_BC;  % <===== Get hold of all coefficinet matrices and BCs

%%%%%%%%%%%%%%% Coefficients A/RHS/R set %%%%%% A*u=RHS %%%%%%%%%%%%%%%%%%%%%

disp(['|____ Coarest Lev: T1  and Finest Lev: T' num2str(levs) ' ____|'])

  iter=1; eval(['U' num2str(levs) '=u_init;']); % only used if fmg=0
  if fmg==1
  alev=1; Goto_finest(1,alev,gam1,gam2,nu, iout); % T1 to the finest T_J
  end
%_______________________________________________________________________|
r_fine=sprintf('R%d=residual(COF%d,U%d,RHS%d); err2=norm(R%d,''fro'');',...
         levs,levs,levs,levs,levs);
eval(r_fine); erratio=-1; res(iter)=err2;
 if iout>1 
  eval(['u=U' num2str(levs) ';']); 
  figure; subplot(121),mesh(u); eval(['r=R' num2str(levs) ';']);%residual/T_J
  title(['FMG u at MG step=' num2str(iter) ' size=' num2str(size(u))]); 
          subplot(122); mesh(r); title('Current residual (in place)')
 end

%=======================================================================|
while (err2>tol & erratio<1.00 & iter<maxit) % Global while

  alev=levs; Goto_finest(0,alev,gam1,gam2,nu, iout); % MG cycle T_J to T_J 

  err0=err2; eval(r_fine) %++++++++++++++++ Update the residual on T_J

  erratio=err2/max(err0,eps); % Check convergence rate via erratio 
if iter<10 | iter>maxit-4 | iout>0
fprintf('\t residual(%3d) = %e  Ratio:%5.2f\n',iter,err2,erratio)
end
   iter=iter+1; res(iter)=err2;
 if iout>1 
  eval(['u=U' num2str(levs) ';']); % Assign for output "u=u_J"
  figure; subplot(121),mesh(u); eval(['r=R' num2str(levs) ';']);%residual/T_J
  title([' u at MG step=' num2str(iter) ' size=' num2str(size(u))]); 
          subplot(122); mesh(r); title('Current residual (in place)')
 end
end % Global while 
%-----------------------------------------------------------------------| 
  fprintf('%s: Total Iterations=%2d Err=%5.1e (final rate %5.2f) CPU=%4.1f\n',...
           mfilename,iter, err2, erratio,cputime-t0)
 if iout<=1 
  eval(['u=U' num2str(levs) ';']); % Assign for output "u=u_J"
  figure; subplot(121),mesh(u); eval(['r=R' num2str(levs) ';']);%residual/T_J
  title(['Final u after MG steps=' num2str(iter) ' size=' num2str(size(u))]); 
          subplot(122); mesh(r); title('Final residual (in place)')
  c=zlabel(['residual_2 in 2\_norm = ' num2str(err2)]); set(c,'color','r')
 end
return %+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++MG END 
function uc = restrict(u,lev); %==========================================[1]
        global hx hy nnx nny levs  
   range_x=2:2:nnx(lev); range_y=2:2:nny(lev); uc=[];
  for j=range_y
      jc=floor(j/2); % The full weighting (FW)
      for i=range_x
      ic=floor(i/2);
      uc(ic,jc) = ( ...
          u(i-1,j-1)+u(i+1,j+1)+u(i-1,j+1)+u(i+1,j-1) +...
          2*(u(i,j-1)+u(i,j+1)+u(i-1,j)+u(i+1,j)) +...
          4*u(i,j) )/16;
   end %j
  end %i
return
function v = interpo(uc, Type,BC); % BC.4 pieces ==========================[2]
% X: coarse grid points  from  o: fine grids  [restriction]  above
% X: coarse grid points  to  o: fine grids    {interpolation} here
%       +...b...B---b---B---b---B---b---+  ^
%    7  b...o...o---o---o---o---o---o---b  |
% (3 6  B...o...X---o---X---o---X---o---B  |
%    5  b---o---o---o---o---o---o---o---b  y
% (2 4  B---o---X---o---X---o---X---o---B  |
%    3  b---o---o---o---o---o---o---o---b  |
% (1 2  B---o---X---o---X---o---X---o---B  |
%    1  b---o---o---o---o---o---o---o---b  | Fine
%       +---b---B---b---B---b---B---b---+  |
%           1   2   3   4   5   6   7 -----+-x---->
%   (Coarse    (1      (2      (3          | 
%===Extend the BC for easier loops
   [m n] = size(uc);  %Type=0: Coarse grid correction / 1: u at BC used
    if Type==1 % Mainly used for FMG case --------
     uc = [BC.bottom(:) uc BC.top(:)]; %horizontal lexio-ordering
      left = [given_bc(0,0) BC.left given_bc(0,1)];
     right = [given_bc(1,0) BC.right given_bc(1,1)];
     uc = [left; uc; right];  % done all 4 sides
    else       % Mainly for normal MG ------------
     uc = [zeros(m,1) uc zeros(m,1)];
     uc = [zeros(1,n+2); uc; zeros(1,n+2)];
    end % correction of 0 at Dirichlet BCs ------- 
for i=2:m+2   %------------------ x top/left box
    i2 = 2*(i-1); %[i2-1, i2, i2+1]
   for j=1:n+1  %------------------ y
    j2 = 2*j;     %[j2-1, j2, j2+1]
     v(i2-1,j2-1) = uc(i-1,j); % coinciding points
     v(i2-1,j2+1) = uc(i-1,j+1);
     v(i2+1,j2-1) = uc(i,j);
     v(i2+1,j2+1) = uc(i,j+1);
     v(i2  ,j2) = ( uc(i,j)+uc(i,j+1) + uc(i-1,j) ...
                   +uc(i-1,j+1) )/4; %- center point
     v(i2-1,j2  ) = (uc(i-1,j)+uc(i-1,j+1))/2; %left
     v(i2-1,j2-1) = (uc(i-1,j)+uc(i  ,j  ))/2; %botto
     v(i2+1,j2  ) = (uc(i  ,j)+uc(i  ,j+1))/2; %right
     v(i2  ,j2+1) = (uc(i-1,j+1)+uc(i,j+1))/2; %top
   end %j
end %i
    v = v(2:end-1,2:end-1);
return
function u_out = smooth(u,lev, steps); %==================================[3]
      global hx hy %_______|  A smoothing step using the GS method
      global U BC COF RHS R 
  eval([' global ' COF(lev,:) ' ' RHS(lev,:) ])
  eval([' rhs = ' RHS(lev,:) '; A =' COF(lev,:) ';' ])
      [m n]=size(u); h1=hx(lev);h2=hy(lev);
  for iter=1:steps %------------Gauss Seidel----------
  for jg=1:n 
   for ig=1:m
      Sum=0; row = (jg-1)*m + ig;
      for j=max(1,jg-1):min(jg+1,n) % local ranges j
        if j==jg, is=max(1,ig-1):min(ig+1,m); else
                  is=ig; end %------Enforce 5-point knowledge
      for i=is                      % local ranges i
        col = (j-1)*m + i;
       if i~=ig | j~=jg   % exclude (ig,jg) for GS 
        Sum = Sum + A(row,col)*u(i,j);
       end %if
      end %j
      end %i
        Sum = rhs(ig,jg)-Sum; u(ig,jg) = Sum/A(row,row); % GS step
   end %jg
  end %ig
 end %iter
   u_out=u;
return 
function Goto_finest(Type,lev_now, gam1,gam2,nu, iout); %==================[4]
     global hx hy nnx nny levs  % From input alev to the finest lev=1
     global U BC COF RHS R check%check(1:levs) for recording \nu cycling info
     global alev iter space
     alev=lev_now; % to pass on alev via "global" (alternatively in output)
     for i=1:levs
      eval([' global ' U(i,:) ' ' BC(i,:) ' ' RHS(i,:) ' ' R(i,:) ])
     end
% On entry:  alev=1 --> The FMG : cycling / restriction + interpolation
%         :              up to level J (the finest grid) to finish
%         :  alev=J --> Start a complete MG cycle (finest J -> back to J)
  if levs<=1&Type==0, return; end %Allow uni-grid to return (done already)
  if iout>0 & iter<2
   fprintf('--- ----------- %s %2d\t [ENTER the MGM cycle from level %d]\n',...
                space(1:alev*2),alev, alev) % Down
  end       % check = IGAM : the cycling-pattern counters
     a_pass=0;  check(1:levs)=0; %alev to pass on via "global"
while (alev < levs) | a_pass<1  % Up loop   () () ()..................
                       a_pass=1;  % One off pass for normal MG cycle
  if check(alev)<nu
   if check(alev)==0&Type==0&alev<levs %Set the Correction v = 0
     s=sprintf('%s=sparse(%d,%d);',U(alev,:),nnx(alev),nny(alev));eval(s)
   end
   Goto_coarest(gam1,iout); % Up to the coarest level = "T1"
  else
     if alev<levs, check(alev)=0; end
  end 
     if levs<=1, return; end %% Allow uni-grid
%________Interpolate____________________ Coarse alev => fine alev+1 _____
  s=sprintf('v = interpo(%s,Type,%s); ',U(alev,:),BC(alev,:)); eval(s)
%________Set for next level if any ______________________________________
  if alev<levs
    alev = alev+1;  % Advance the level counter
    if iout>0 & iter<2
     fprintf('Lev %2d Interpol %s %2d\n',levs,space(1:alev*2), alev) % Up
    end
  s=sprintf('%s = v + %s;',U(alev,:),U(alev,:)); eval(s) %add correction
  end %if

%________Post-interpolation smoothing step_______________________________
  s=sprintf('%s=smooth(%s,%d,gam2);',U(alev,:),U(alev,:),alev); eval(s)
  check(alev)=check(alev)+1;  % Advance cycling counter at Interpo
end%while alev_______done.......% Up loop   () () ()..................
return
function k = Goto_coarest(gam1,iout); %====================================[5]
     global hx hy nnx nny levs  %levs=\ell+1 and on G_\ell: n_\ell=2
     global U BC COF RHS R 
     global alev iter space
 for i=1:levs
 eval(['global ' U(i,:) ' ' BC(i,:) ' ' R(i,:) ' ' RHS(i,:) ' ' COF(i,:)])
 end
while alev>1     %-------------------- Down loop   () () () if not yet T1 yet
  %________Pre-restriction smoothing
    s=sprintf('%s = smooth(%s,%d, gam1); ',U(alev,:),U(alev,:),alev);eval(s)
    s=sprintf('%s = residual(%s,%s,%s);',R(alev,:),COF(alev,:),...
                U(alev,:),RHS(alev,:)); eval(s)
  %________Restrict  u -> uc_________ RHS_c = R_f^c*Res_f
     s=sprintf('%s = restrict(%s,alev); ',RHS(alev-1,:),R(alev,:)); eval(s)
  if iout>0 & iter<2
   fprintf('Lev %2d Restrict %s %2d\n',levs,space(1:alev*2), alev) % Down
  end
    alev = alev-1; %________Adjust the level counter
end % while alev %-------------------- Down loop   () () () %REACHED T1
  alev=1;  % The EXACT solver on the coarest level T1 ----------
  s=sprintf('[ m n] = size(%s); ', RHS(alev,:));eval(s)
  s=sprintf('ue = %s \\ %s(:);',COF(alev,:),RHS(alev,:));eval(s)
  s=sprintf('%s = reshape(ue,m,n);', U(alev,:));eval(s)
  if iout>0 & iter<2
   fprintf('Lev %2d ExactSoL %s %2d\n',levs,space(1:alev*2), alev) % Down
  end
return
function r=residual(A,u,rhs) %=============================================[6]
        [m n] = size(u);r=[];
   for jg=1:n  %%% Identical steps : r=reshape(rhs,m*n,1)-A*reshape(u,m*n,1);
    for ig=1:m %%% to these 2 lines: r=reshape(r,m,n); 
      Sum=rhs(ig,jg); row = (jg-1)*m + ig;
      for j=max(1,jg-1):min(jg+1,n) % local ranges j
        if j==jg, is=max(1,ig-1):min(ig+1,m); else
                  is=ig; end %------Enforce 5-point knowledge
      for i=is                      % local ranges i
        col = (j-1)*m + i;
        Sum = Sum - A(row,col)*u(i,j);
      end %i
      end %j
        r(ig,jg) = Sum;
   end %ig
  end %jg
return
function a=coef_a(x,y) %==========================================PROB
    a = 1+x^2/100;
return
function b=coef_b(x,y) %==========================================PROB
    b = 1-y^2/100;
return
function c=coef_c(x,y) %==========================================PROB
    c = -x^2/100-y^2/100;
return
function d=coef_d(x,y) %==========================================PROB
   global iprob
   if iprob==1
    d = -5/3+(-8+(776809/64800+(11/135+(-19/150+...
      (1/75-y/100)*y)*y)*y)*y)*y+(6-3/400*y^2+...
      (-1871/64800+(1/675+(1/120+(1/75-y/100)*y)*y)*y+...
      (-y^2/100+21/400+(3/200-x/100)*x)*x)*x)*x;
   else
    d = (6960+(-14400+(0.1436519E7/100+(144-144*y)*y)*y)*y+...
      (-14400+(0.1443479E7/100+(-144+144*x)*x)*x)*x)*...
      exp(-60*( (x-1/2)^2+(y-1/2)^2));
   end
return 
function u=given_bc(x,y) %=========================================PROB
   global iprob
   if iprob==1
    u = (x-1/2)^3+(y-1/3)^4;
   else
    u = exp(-60*( (x-1/2)^2+(y-1/2)^2 ));
   end
return 
%_______|  problem+geometry dependent !!!   (One-off setting initially) ====
function Get_A_rhs_R_BC %===============================================PROB
      global hx hy nnx nny levs  
      global U BC COF RHS R 
for i=1:levs % Declare U1 R1 U2 R2 U3 R3 etc as global quantities
 eval(['global ' U(i,:) ' ' BC(i,:) ' ' COF(i,:) ' ' RHS(i,:) ' ' R(i,:) ])
end
for i=1:levs  % Geometry-dependent !! -------------Set Boundary conditions
 bc=[]; %______________________________________________
 for j=1:nny(i) % left
  bc(j) = given_bc(0,j*hy(i)); % Geometry-dependent !!!
 end
  s=sprintf('%s.left = bc; ', BC(i,:)); eval(s) 
 bc=[]; %______________________________________________
 for j=1:nnx(i) % bottom
  bc(j) = given_bc(j*hx(i),0); % Geometry-dependent !!!
 end
  s=sprintf('%s.bottom = bc; ', BC(i,:)); eval(s) 
 bc=[]; %______________________________________________
 for j=1:nny(i) % right
  bc(j) = given_bc(1,j*hy(i)); % Geometry-dependent !!!
 end
  s=sprintf('%s.right = bc; ', BC(i,:)); eval(s) 
 bc=[]; %______________________________________________
 for j=1:nnx(i) % top
  bc(j) = given_bc(j*hy(i),1); % Geometry-dependent !!!
 end
  s=sprintf('%s.top = bc; ', BC(i,:)); eval(s) 
end
%%%%%%%%%%%%%%% Initialising B.C. completed  %%%%%%%%%%%%%%%%%%%
 for lev=1:levs %------------Set coefficients on all levels-----
      eval([' bc = ' BC(lev,:) ';'])
      m=nnx(lev);n=nny(lev); h1=hx(lev);h2=hy(lev);
      coeffs=sparse(m*n,m*n); % matrix on lev -- horiz lexico ordering
         rhs=sparse(m,n);   % rhs on lev
   for jg=1:n 
      y = jg*hy(lev);
      for ig=1:m
      x=ig*hx(lev);
        row = (jg-1)*m + ig;
        a=coef_a(x,y)/h1^2; c=coef_c(x,y); 
        b=coef_b(x,y)/h2^2; d=coef_d(x,y);

      for j=max(1,jg-1):min(jg+1,n) % local ranges j
        if j==jg, is=max(1,ig-1):min(ig+1,m); else
                  is=ig; end %------Enforce 5-point knowledge
      for i=is                      % local ranges i
        col = (j-1)*m + i;
        if ig==i & jg==j
        coeffs(row,col) = -2*a -2*b +c;
        elseif ig==i 
        coeffs(row,col) = a;
        elseif jg==j
        coeffs(row,col) = b;
        end %if
      end %i
      end %j
%===========Set the RHS d_ij with BC updates === Vertical ordering !!!!
        if ig==1 %-------Update rhs at BC
          rhs(ig,jg) = d - a*bc.left(jg);
        elseif ig==m
          rhs(ig,jg) = d - a*bc.right(jg);
        elseif jg==1
          rhs(ig,jg) = d - b*bc.bottom(ig);
        elseif jg==n
          rhs(ig,jg) = d - b*bc.top(ig);
        else
          rhs(ig,jg) = d;  % The middle guys
        end 
   end %ig
  end %jg
  s=sprintf('%s=rhs; %s=coeffs;',RHS(lev,:),COF(lev,:));eval(s)
  eval([ R(lev,:) ' = rhs;']) % True residual if u=0
  eval([ U(lev,:) ' = sparse(m,n);']) % Set u=0 for now
 end %lev
return %===============================================END of PROB PART
