% For more information plese contact:
% Gianluca Meneghello
% gianluca.meneghello@gmail.com
% http://www.mit.edu/~mgl/
%
%
%
% A basic implementation of a multigrid solver for the convection-diffusion
% equation on a unit square. For a quick test, type in Matlab
%
% output = main(2^7+1,10) 
%
% INPUT: 
% - n: the number of mesh point along a mesh direction, it has to be a power of
% two + 1 (e.g., 2^7+1=129). The total number of degrees of freedom is n^2;
% - nits: the number of V-cycle to be used
%
% Both the CS and the FAS algorithms are available, and the used
% algorithm can be selected by setting the variable cntrparams.algorithm to 'CS'
% and 'FAS' respectively. The Poisson equation can be obtained by setting
% cntrparams.adv --- i.e. the coefficient of the convective term --- to zero. A
% purely convective equation can be obtained by setting cntrparam.nu --- i.e.
% the viscosity --- to zero. Other input parameters are clarified in the code.
%
% OUTPUT:
% A single structure, named output, is provided as a result. Any quantity
% present in the main function can be easily added to the output structure.
% The default quatities in the output structure are:
%
% - its: the number of V-cycle iteration [0:nits] 
% - rnormRelaxation: the L2 norm of the residual for the relaxation scheme, as a
% function of the iteration number.
% (without multigrid)
% - rnormTwoGrids: the L2 norm of the residual for the two grid scheme, as a
% function of the iteration number.
% - rnormMG: the L2 norm of the residual for the multigrid scheme, as a
% function of the iteration number
% - x,y,uh,rh,u0: the x,y coordinate, the computed solution, the
% corresponding residual and the exact solution on the finest grid
% - Reh: the mesh-based Reynold number on the finest grid
%
% As a usage example, a plot of the convergence history of the residual for the
% multigrid scheme can be obtained with
%
% plot(output.its,output.rnormMG)
%

%
% the full code is made of the following functions:
%
% === [output] = main(n,nits) ===
% the main function
%
% === [J] = jac(n) === 
% returs the discretization of the convection-diffusion operator on a grid with
% n mesh points on each direction
%
% === [Lp,Lm] = relaxationSetUp(J,n) ===
% sets up the relaxation algorithm by performing the operator splitting and, if
% required, implementing the boundary relaxation
%
% === [U] = RR(u) ===
% restrict the field u to the next coarser grid 
%
% === [u] = II(U) ===
% interpolate the field U to the next finer grid
%
% == [R] = Rop(n,N) ===
% provides the restriction operator between a finer and a corser grid
% characterized by n and N mesh points in each direction respectively. 
% Restriction can the be done as U = R*u and the result is the same as using U =
% RR(u). The implementation of RR is much faster, but this function can be used
% to obtain an explicit description of the restriction operator
%
% === [I] = Iop(N,n) ===
% provides the interpolation operator between a coarser and a finer grid
% characterized by N and n mesh points in each directio respectively.
% The same observation made for Rop applies, with u = II(U) being a much faster
% implenentation
%
% === [uh] = mg(uh,rhs,n) ===
% this function represents the application of a single V-cycle. Both the CS and
% the FAS algorithm are implemented and can be selected by setting the value of
% cntrparams.algorithm to 'CS' or 'FAS' respectively.
%
% === [uh] = twogrids(uh,rhs,n) ===
% the two grid algorithm --- i.e. only two grids are used and the solution is
% computed on the coarsest independently of its size. It can be used for
% comparison with mg.
% 
% === [x,y,u0,rhs] = init(n) ===
% this function initialize the grid point coordinates, the exact solution and
% the corresponding right hand side. WARNING: some choices of boundary condition
% do not return a solution corresponding to the exact one.
% 



function [output]= main(n,nits)

  %%%%%%%%%%%%%%%%%%%%%%%%%%
  %%%% PROBLEM CONTROL %%%%%
  %%%%%%%%%%%%%%%%%%%%%%%%%%

  global cntrparams
  % the algorithm used, can be 'CS' or 'FAS'
  cntrparams.algorithm = 'FAS'; 
  % the anisotropy parameter (see the Relaxation section in the Multigrid
  % chapter). A value of epsilon < 1 is correctly treated by the current
  % implementation of linewise Gauss-Seidel relaxation. A value of epsilon >1
  % correspond to stretching in the wrong direction
  cntrparams.epsilon = 1; 
  % control for the linewise Gauss-Seidel relaxation. A value of 0 corresponds
  % to pointwise relaxation, a value of 1 correspond to linewise relaxation.
  cntrparams.LGS = 0;
  % the sweep direction of the relaxation, can be 'F' (forward) or 'B'
  % (backward). A value of 'F' corresponds to a sweep starting in the bottom
  % left corner and ending in the top right corner. A value of 'B' correspond to a
  % sweep in the opposite direction, i.e. starting in the top right corner and
  % ending in the bottom-left corner.
  cntrparams.sweepdirection = 'F'; 
  % the convective field, can be 'const' or 'hiemenz'. The value 'const'
  % correspond to a constant flow field, the value 'hiemenz' correspond to the
  % (unswept) hiemenz flow field.
  cntrparams.advfield = 'hiemenz'; % possible choices are 'const' or 'hiemenz'
  % the angle of the (constant) velocity field with respect to the horizonal
  % axis.
  cntrparams.alph= 45*pi/180; 
  % coefficient for the convection term. For a value of zero the Poisson
  % equation is obtained. For a value of one, the viscosity is the opposite of
  % the Reynolds number. 
  cntrparams.adv = 0;
  % the order of discretization of the advection term. Can be 1 or 2.
  cntrparams.advorder = 1;
  % the viscosity. It corresponds to the inverse of the Reynolds number if
  % cntrparams.adv = 1. If one wants to control the mesh-based Reynolds on the
  % finer grid, cntrparams.nu = 1./(n-1)/Re_h, where Re_h is the mesh-based
  % Reynolds.
  cntrparams.nu = 1; 1./(n-1)/1; 
  % whether a Dirichlet (0) or a homogeneous Neuamnn (1) boundary condition is
  % implemented on each boundary
  cntrparams.neumannS = 0;
  cntrparams.neumannN = 0;
  cntrparams.neumannE = 0;
  cntrparams.neumannW = 0;
  % the number of mesh-lines to be included in the boundary relaxation on each
  % boundary
  cntrparams.brkS = 0;
  cntrparams.brkN = 0;
  cntrparams.brkW = 0;
  cntrparams.brkE = 0;
  cntrparams

  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  %%%% END OF PROBLEM CONTROL   %%%%%
  %%%% YOU SHOULD NOT MODIFY    %%%%%
  %%%% ANYTHING BELOW THIS LINE %%%%%
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
 

  %%%%%%%%%%%%%%%%%%%%%%%%
  %%%% INITIALIZATION %%%%
  %%%%%%%%%%%%%%%%%%%%%%%%
  
  % the total number of degrees of freeedom
  nn = n*n; 
  % the discretized operator  
  Jh = jac(n);  
  % initialization of the rhs and exact solution
  [x,y,u0,rhs] = init(n); 

  %%%%%%%%%%%%%%%%%%%%%%%%%
  %%%% DIRECT SOLUTION %%%%
  %%%%%%%%%%%%%%%%%%%%%%%%%
  tic
  disp('=== Direct solution ===')
  u = u0;
  u = Jh\rhs;
  toc
  output.u = reshape(u,n,n);

  %%%%%%%%%%%%%%%%%%%%%%%%%%%
  %%%% SIMPLE RELAXATION %%%%
  %%%%%%%%%%%%%%%%%%%%%%%%%%%
  tic
  disp(['=== Simple relaxation (',num2str(3*nits),' iterations) ==='])
  [Lp,Lm] = relaxationSetUp(Jh,n);
  output.Lp = Lp;
  output.Lm = Lm;
  u = u0;
  r = Jh*u-rhs;
  rnormRelaxation(1) = norm(r)/nn;
  for cc = 2:3*nits+1
    u = -Lp\(Lm*u-rhs);
    r = Jh*u-rhs;
    rnormRelaxation(cc) = norm(r)/nn;
  end
  toc

  %%%%%%%%%%%%%%%%%%%%%%%%
  %%%% TWO GRID CYCLE %%%%
  %%%%%%%%%%%%%%%%%%%%%%%%
  tic
  disp(['=== Two grid cycle (',num2str(nits),' iterations) ==='])
  uh = u0;
  rh = Jh*uh - rhs;
  rnormTwoGrids(1) = norm(rh)/nn;
  for cc = 2:nits+1
    uh = twogrids(uh,rhs,n);
    rh = Jh*uh - rhs;
    rnormTwoGrids(cc) = norm(rh)/nn;
  end
  toc

  %%%%%%%%%%%%%%%%%%%%%%%%%
  %%%% MULTIGRID CYCLE %%%%
  %%%%%%%%%%%%%%%%%%%%%%%%%
  tic
  disp(['=== Multigrid V-cycle (',num2str(nits),' iterations) ==='])
  uh = u0;
  rh = Jh*uh - rhs;
  rnormMG(1) = norm(rh)/nn;
  for cc = 2:nits+1
    uh = mg(uh,rhs,n);
    rh = Jh*uh - rhs;
    rnormMG(cc) = norm(rh)/nn;
    disp(['V cycle # ', num2str(cc), ' residual is ', num2str(rnormMG(cc))]);
  end
  toc

  %%%%%%%%%%%%%%%%
  %%%% OUTPUT %%%%
  %%%%%%%%%%%%%%%%
  output.its = 0:3*nits
  output.rnormRelaxation = rnormRelaxation;
  output.rnormTwoGrids = [rnormTwoGrids ones(1,nits*3-nits)*nan];
  output.rnormMG = [rnormMG ones(1,nits*3-nits)*nan];
  output.x = x;
  output.y = y;
  output.uh = reshape(uh,n,n);
  output.rh = reshape(rh,n,n);
  output.u0 = reshape(u0,n,n);
  output.Reh = 1./(n-1)/cntrparams.nu;
  dlmwrite('norms.dat',[output.its',output.rnormRelaxation',output.rnormTwoGrids',output.rnormMG'],'delimiter',' ');

end

% function [J] = jac(n)
% returs the discretization of the convection-diffusion operator on a grid with
% n mesh points on each direction
function [J] = jac(n)
  global cntrparams

  epsilon = cntrparams.epsilon;
  alph = cntrparams.alph;
  nu = cntrparams.nu;
  adv = cntrparams.adv;
  advfield = cntrparams.advfield;
  advorder = cntrparams.advorder;
  neumannS = cntrparams.neumannS; 
  neumannN = cntrparams.neumannN; 
  neumannE = cntrparams.neumannE; 
  neumannW = cntrparams.neumannW; 

  nn = n*n;
  h = 1./(n-1);

  % vectors containing the index of the boundary points
  sb = 1:n;
  nb = nn-n+1:nn;
  wb = 1:n:nn;
  eb = n:n:nn;
  bb = [ sb,nb,wb,eb];
  lbb = length(bb);

  % the centered second order discretization of the Laplacian operator. epsilon
  % represents the anisotropy. spdiags will be used to build the sparse matrix,
  % look at 'help spdiags' on Matlab for information on the format used here.
  kz = 0;2*pi/(40*h);
  d0 = 1/h^2*[zeros(nn,1) , epsilon*ones(nn,1), zeros(nn,1),  ones(nn,1),  (-2*(1+epsilon)-kz^2.*h^2).*ones(nn,1),  ones(nn,1),zeros(nn,1),epsilon*ones(nn,1), zeros(nn,1)];

  % the convective the velocity field can be constant, with an angle alph with
  % the horizontal axis, or (unswept) hiemenz. 
  if (strcmpi(advfield,'const'))
    u = ones(nn,1).*cos(alph);
    v = ones(nn,1).*sin(alph);
  elseif (strcmpi(advfield,'hiemenz'))
    x = linspace(0,1,n); y = linspace(0,1,n); [x,y] = meshgrid(x,y);
    u = x; u = u(:); 
    v = 1-y; v = v(:);
    clear x,y;
  else
    error('advfield can be const or hiemenz');
  end

  % the upwinded first or second order discretization of the convective operator
  up = .5*(u+abs(u)); um = -.5*(u-abs(u));
  vp = .5*(v+abs(v)); vm = -.5*(v-abs(v));
  d1 = zeros(nn,9);
  if (advorder==1)
    d1 = 1./(h)*[ 0*vp -1*vp 0*up -1*up (up+vp+um+vm) -1*um 0*um -1*vm 0*vm ];
  elseif (advorder==2)
    d1 = 1./(2*h)*[ 1*vp -2*vp 1*up -2*up 3*(up+vp+um+vm) -2*um 1*um -2*vm 1*vm ];
  else
    error('The discretization order for the convective term must be 1 or 2'); 
  end

  % the discretized operator, including the convective (d1) and diffusive (d0)
  % terms is built using spdiags
  J = spdiags(adv.*d1-nu*d0,[-2*n -n -2 -1 0 1 2 n 2*n],nn,nn);

  % enforce Dirichlet boundary conditions on all boundaries. The indexes of the
  % boundary's points are contained in the vector bb
  J(bb,:) = 0;
  J(bb,bb) = spdiags(ones(lbb,1),0,lbb,lbb);

  % enforce Neumann boundary condition where required
  if (neumannS == 1)
    J(sb,:) = 0;
    for i = 1:n
      J(sb(i),sb(i)+[0 n 2*n]) = -1/(2*h) * [ 3 -2 1];
    end
  end

  if (neumannE == 1)
    J(eb,:) = 0;
    for i = 1:n
      J(eb(i),eb(i)-[0 1 2]) = +1/(2*h) * [ 3 -2 1];
    end
  end

  if (neumannN == 1)
    J(nb,:) = 0;
    for i = 1:n
      J(nb(i),nb(i)-[0 n 2*n]) = +1/(2*h) * [ 3 -2 1];
    end
  end

  if (neumannW == 1)
    J(wb,:) = 0;
    for i = 1:n
      J(wb(i),wb(i)+[0 1 2]) = -1/(2*h) * [ 3 -2 1];
    end
  end

end

% function [Lp,Lm] = relaxationSetUp(J,n);
% sets up the relaxation algorithm by performing the operator splitting and, if
% required, implementing the boundary relaxation
function [Lp,Lm] = relaxationSetUp(J,n);
  global cntrparams;

  LGS = cntrparams.LGS;

  % perform operator splitting by selecting the upper and lower parts of the
  % matrix. LGS diagonals above/below the main diagonal are included in the Lp
  % operator in order to implenent linewise Gauss-Seidel, if required.
  if (cntrparams.sweepdirection == 'F')
    Lp = tril(J,+LGS);
    Lm = triu(J,+1+LGS);
  elseif (cntrparams.sweepdirection == 'B')
    Lp = triu(J,-LGS);
    Lm = tril(J,-1-LGS);
  end

  % boundary relaxation is obtained modifying the splitting. All discrete
  % equations corredponding to the boundary points are included in the Lp
  % operator and are solved implicitly.

  % southern boundary
  brkS = cntrparams.brkS;
  Lp(1:brkS*n,:) = J(1:brkS*n,:);
  Lm(1:brkS*n,:) = 0;

  % northern boundary 
  brkN = cntrparams.brkN;
  Lp(end-brkN*n+1:end,:) = J(end-brkN*n+1:end,:);
  Lm(end-brkN*n+1:end,:) = 0;

  % western boundary
  brkW = cntrparams.brkW;
  brkpoints = zeros(brkW*n,1);
  for i = 1:brkW
    brkpoints((i-1)*n+1:i*n) = i:n:nn;
  end
  Lp(brkpoints,:) = J(brkpoints,:);
  Lm(brkpoints,:) = 0;

  % eastern boundary
  brkE = cntrparams.brkE;
  brkpoints = zeros(brkE*n,1);
  for i = 1:brkE
    brkpoints((i-1)*n+1:i*n) = n-i+1:n:nn;
  end
  Lp(brkpoints,:) = J(brkpoints,:);
  Lm(brkpoints,:) = 0;

end

% function U = RR(u)
% restrict the field u to the next coarser grid 
function U = RR(u)
  n = sqrt(length(u));
  N = ceil(n/2);
  utmp = reshape(u,n,n);
  Utmp = zeros(N,N);

  Utmp(2:N-1,2:N-1) = 0.25*utmp(3:2:n-2,3:2:n-2) + ...
    0.125*( utmp(2:2:n-3,3:2:n-2) + utmp(4:2:n-1,3:2:n-2) ...
    + utmp(3:2:n-2,2:2:n-3) + utmp(3:2:n-2,4:2:n-1) ) + ...
    0.0625*( utmp(2:2:n-3,2:2:n-3) + utmp(2:2:n-3,4:2:n-1) ...
    + utmp(4:2:n-1,2:2:n-3) + utmp(4:2:n-1,4:2:n-1) );

  Utmp(2:N-1,1) = 0.5*utmp(3:2:n-2,1) + 0.25*( utmp(2:2:n-3,1) + utmp(4:2:n-1,1) );
  Utmp(2:N-1,N) = 0.5*utmp(3:2:n-1,n) + 0.25*( utmp(2:2:n-3,n) + utmp(4:2:n-1,n) );
  Utmp(1,2:N-1) = 0.5*utmp(1,3:2:n-1) + 0.25*( utmp(1,2:2:n-3) + utmp(1,4:2:n-1) );
  Utmp(N,2:N-1) = 0.5*utmp(n,3:2:n-1) + 0.25*( utmp(n,2:2:n-3) + utmp(n,4:2:n-1) );

  Utmp(1,1) = utmp(1,1); 
  Utmp(1,end) = utmp(1,end); 
  Utmp(end,1) = utmp(end,1); 
  Utmp(end,end) = utmp(end,end);

  U = reshape(Utmp,N*N,1);

end

% function u = II(U)
% interpolate the field U to the next finer grid
function u = II(U)
  N = sqrt(length(U));
  n = 2*N-1;
  Utmp = reshape(U,N,N);
  utmp = zeros(n,n);

  utmp(1:2:n,1:2:n) = Utmp(1:N,1:N);
  utmp(2:2:n-1,1:2:n) = ( Utmp(1:N-1,1:N) + Utmp(2:N,1:N) ) / 2.;
  utmp(1:2:n,2:2:n-1) = ( Utmp(1:N,1:N-1) + Utmp(1:N,2:N) ) / 2.;
  utmp(2:2:n-1,2:2:n-1) = ( Utmp(1:N-1,1:N-1) + Utmp(1:N-1,2:N) + Utmp(2:N,1:N-1) + Utmp(2:N,2:N) ) / 4.;
  u = reshape(utmp,n*n,1);

end


% function [R] = Rop(n,N)
% provides the restriction operator between a finer and a corser grid
% characterized by n and N mesh points in each direction respectively. 
% Restriction can the be done as U = R*u and the result is the same as using U =
% RR(u). The implementation of RR is much faster, but this function can be used
% to obtain an explicit description of the restriction operator
function [R] = Rop(n,N)
  nn = n*n;
  NN = N*N;

  % SW S SE W C E NW N E
  st = [-n-1 -n -n+1  -1 0 1 n-1 n n+1];

  R = sparse([],[],[],NN,nn,nn*9);
  di = 1/16*[ 1 , 2 , 1. , 2 , 4 , 2 , 1 , 2 , 1] ; % interior
  dh = 1/4*[ 0 , 0 , 0. , 1 , 2 , 1 , 0 , 0 , 0] ; % horizontal boundaries
  dv = 1/4*[ 0 , 1 , 0. , 0 , 2 , 0 , 0 , 1 , 0] ; % vertical boundaries
  dc = [ 0 , 0 , 0. , 0 , 1 , 0 , 0 , 0 , 0] ; % corners

  for J = 2:N-1
    for I = 2:N-1
      j = 2*J-1;
      i = 2*I-1;
      CC = (J-1)*N+I;
      cc = (j-1)*n+i;
      R(CC,st+cc)  = di;
    end
  end

  for J = 1:N-1:N
    for I = 2:N-1
      j = 2*J-1;
      i = 2*I-1;
      CC = (J-1)*N+I;
      cc = (j-1)*n+i;
      R(CC,:) = spdiags(dh,st+cc-1,1,nn);
    end
  end

  for J = 2:N-1
    for I = 1:N-1:N
      j = 2*J-1;
      i = 2*I-1;
      CC = (J-1)*N+I;
      cc = (j-1)*n+i;
      R(CC,:) = spdiags(dv,st+cc-1,1,nn);
    end
  end

  for J = 1:N-1:N
    for I = 1:N-1:N
      j = 2*J-1;
      i = 2*I-1;
      CC = (J-1)*N+I;
      cc = (j-1)*n+i;
      R(CC,:) = spdiags(dc,st+cc-1,1,nn);
    end
  end

end


% function [Iout] = Iop(N,n)
% provides the interpolation operator between a coarser and a finer grid
% characterized by N and n mesh points in each directio respectively.
% The same observation made for Rop applies, with u = II(U) being a much faster
% implenentation
function [Iout] = Iop(N,n)
  nn = n*n;
  NN = N*N;

  % SW S SE W C E NW N E
  st = [-n-1 -n -n+1  -1 0 1 n-1 n n+1];

  Iout = sparse([],[],[],nn,NN,nn.*9);
  di = 1/4*[ 1 , 2 , 1. , 2 , 4 , 2 , 1 , 2 , 1] ; % interior
  dh = 1/2*[ 0 , 0 , 0. , 1 , 2 , 1 , 0 , 0 , 0] ; % horizontal boundaries
  dv = 1/2*[ 0 , 1 , 0. , 0 , 2 , 0 , 0 , 1 , 0] ; % vertical boundaries
  dc = [ 0 , 0 , 0. , 0 , 1 , 0 , 0 , 0 , 0] ; % corners

  % injection
  for J = 1:N
    for I = 1:N
      j = 2*J-1;
      i = 2*I-1;
      CC = (J-1)*N+I;
      cc = (j-1)*n+i;
      Iout(cc,CC) = 1.; 
    end
  end

  % linear interpolation
  for J = 1:N
    for I = 1:N-1
      j = 2*J-1;
      i = 2*I;
      CC = (J-1)*N+I;
      cc = (j-1)*n+i;
      %Iout(cc,:) = spdiags([0.5,0.5],[CC,CC+1],NN,1);
      Iout(cc,CC) = 0.5;
      Iout(cc,CC+1) = 0.5;
    end
  end

  for J = 1:N-1
    for I = 1:N
      j = 2*J;
      i = 2*I-1;
      CC = (J-1)*N+I;
      cc = (j-1)*n+i;
      Iout(cc,CC) = 0.5;
      Iout(cc,CC+N) = 0.5;
    end
  end

  % bilinear interpolation
  for J = 1:N-1
    for I = 1:N-1
      j = 2*J;
      i = 2*I;
      CC = (J-1)*N+I;
      cc = (j-1)*n+i;
      Iout(cc,CC) = 0.25;
      Iout(cc,CC+1) = 0.25;
      Iout(cc,CC+N) = 0.25;
      Iout(cc,CC+N+1) = 0.25;
    end
  end
end


function [uh] = mg(uh,rhs,n)

  global cntrparams
  nn = n*n;
  LGS = cntrparams.LGS;
  if (n <= 5)

    Jh = jac(n);
    uh = Jh\rhs;

  else

    N = (n-1)/2+1;
    % the follwing two lines can be uncommented if an explicit representation of
    % the restriction and iterpolation operator are required, but the
    % implentation of the Rop and Iop functions is quite slow
    %R = Rop(n,N);
    %I = Iop(N,n);
    Jh = jac(n);
    JH = jac(N);

    [Lp,Lm] = relaxationSetUp(Jh,n);

    % relaxation: relaxation is peformed by "inverting" the Lp operator. This is
    % more costly than a real Gauss-Seidel iteration, but clearer from a
    % notation point of view and, for this reason, is used here. 
    uh = -Lp\(Lm*uh-rhs);
    uh = -Lp\(Lm*uh-rhs);

    if ( strcmpi(cntrparams.algorithm,'FAS') )  % the FAS algorithm

      % restriction
      tau = JH*(RR(uh)) - RR(Jh*uh);
      rhsH = RR(rhs) + tau;

      % multigrid on coarser grids
      uHold = RR(uh);
      uH = mg(uHold,rhsH,N);

      % interpolation
      corrH = uH - uHold;
      uh = uh + II(corrH);

    elseif ( strcmpi(cntrparams.algorithm,'CS') ) % the CS algorithm

      rh = rhs - Jh*uh;
      rhsH = RR(rh);
      corrH = mg(zeros(size(rhsH)),rhsH,N); 
      uh = uh + II(corrH);

    else
      error('cntrparams.algorithm must be FAS or CS');
    end

    % relaxation
    uh = -Lp\(Lm*uh-rhs);

  end

end

function [uh] = twogrids(uh,rhs,n)

  global cntrparams;
  nn = n*n;
  LGS = cntrparams.LGS;

  N = (n-1)/2+1;

  %R = Rop(n,N);
  %I = Iop(N,n);
  Jh = jac(n);
  JH = jac(N);

  [Lp,Lm] = relaxationSetUp(Jh,n);

  % relaxation
  uh = -Lp\(Lm*uh-rhs);
  uh = -Lp\(Lm*uh-rhs);

  if ( strcmpi(cntrparams.algorithm,'FAS') )  % the FAS algorithm

    % restriction
    tau = JH*RR(uh) - RR(Jh*uh);
    rhsH = RR(rhs) + tau;

    % solve on coarse grid
    uH = JH\rhsH;

    % interpolation
    corrH = uH - RR(uh);
    uh = uh + II(corrH);

  elseif ( strcmpi(cntrparams.algorithm,'CS') ) % the CS algorithm

    rh = rhs - Jh*uh;
    corrH = JH\(RR(rh));
    uh = uh + II(corrH);

  else
    error('cntrparams.algorithm must be FAS or CS');
  end

  % relaxation
  uh = -Lp\(Lm*uh-rhs);

end

function [x,y,u0,rhs] = init(n)
  global cntrparams
  % initial condition, exact solution and rhs, from PETER's book
  x = linspace(0,1,n); y = linspace(0,1,n); [x,y] = meshgrid(x,y);
  nu = cntrparams.nu; adv = cntrparams.adv; alph = cntrparams.alph;

  u0 = (x.^2 - x.^4) .* (y.^4-y.^2);   u0(2:end-1,2:end-1) = rand(n-2,n-2); u0 = u0(:);
  %u0 = (x.^2 - x.^4) .* (y.^4-y.^2);    u0 = u0(:);

  rhs = + adv.*( cos(alph).*(2.*x-4.*x.^3).*(y.^4-y.^2) + sin(alph).*(x.^2-x.^4).*(4.*y.^3-2.*y) )  - nu.*( (2-12.*x.^2).*(y.^4-y.^2) + (x.^2-x.^4).*(12.*y.^2-2) );
  rhs(:,1) = 0; rhs(:,end)=0; rhs(1,:) = 0; rhs(end,:) = 0; rhs = rhs(:); 
end


