% Mfun = pc_2grid(A,n,p)
%
% Generates a two-grid preconditioner for a Poisson-type problem
% on an n-by-n-by-n grid. The coarse grid problem has p^3 unknowns;
% large p will improve the convergence, but with a greater setup cost
% and cost per iteration.
%
% Input:
% A - Original matrix (fine grid operator)
% n - dimension of original grid (points per side)
% p - dimension of coarse grid (points per side)
%
% Output:
% Mfun - function to solve with the preconditioner M
%
% Usage:
% Mfun = pc_2grid(A);
% x = pcg(A,b,tol,maxit,Mfun);

function Mfun = pc_2grid(A,n,p)

% Form and factor the projected matrix
P = setup_coarse_grid(n,p);
Aproj = P'*A*P;
idx = symamd(Aproj);
R = chol(Aproj(idx,idx));
P = P(:,idx);

% Get the necessary parameters for a symmetric G-S sweep
L = tril(A);

% Return the function to apply the preconditioner
Mfun = @(r) apply_pc(tril(A),triu(A),A,P,R,r);

function P = setup_coarse_grid(n,p)

% Sample to get coarse mesh nodal points
b = ceil( (n-1)/(p-1) );
I = [0:b:n, n+1];
p = length(I)-2;

% Set up coarse-grid shape functions
N = zeros(n,p);
for j = 1:p
I1 = (I(j)+1):I(j+1);
I2 = I(j+1):(I(j+2)-1);
N(I1,j) = (1:length(I1))/length(I1);
N(I2,j) = (length(I2):-1:1)/length(I2);
end

% Extend to 3D shape functions
N1 = sparse(N);
P = kron(kron(N1,N1),N1);

function x = apply_pc(L,U,A,P,R,r)

% Gauss-Seidel smoother step
x = L\r;
x = x + U\(r-A*x);

% Coarse grid solve
x = x + P*(R\(R'\(P'*(r-A*x))));

% Gauss-Seidel smoother step
x = x + U\(r-A*x);
x = x + L\(r-A*x);