%% Using Finite Difference preconditioners to solve the variable diffusivity Helmholtz problem
% This code is a tutorial code for solving a variable coefficient
% Helmholtz (Poisson) problem in a doubly periodic domain. It uses GMRES
% preconditioned with 2nd order finite differences. Sometimes if you have a
% very dense matrix, it can take an extremely long time to converge using
% iterative methods. One way to improve this is to "precondition" your
% matrix by multiplying it by an approximation to the solution so that the
% condition number is smaller and thus will converge faster. Instead of
% solving Ax=b, you solve AP^{-1}y=b for y, and then solve Px=y. Please note
% that this script defines functions at the end, which is only supported by
% MATLAB 2016b or later.
 
%% Construct grid
% Define the various grid parameters.
Nx = 256; % uonal resolution
Ny = 256; % meridional resolution
Lx=3e3;
Ly=3e3;
dx=Lx/Nx;
dy=Ly/Ny; 
x=dx*(1:Nx)';
y=dy*(1:Ny)';
dx2o2 = dx^2/2; 
dy2o2 = dy^2/2;
[xx,yy] = meshgrid(x,y);
 
%% GMRES parameters
% These parameters are used by the GMRES function to tell it to stop
% iterating.
TOL=1e-9; % GMRES convergence tolerance
MAXIT = min(250,Nx); % Upper-bound for # of iterations GMRES can perform.
 
%% FFT related material
% For this problem we will be solving the PDE using spectral methods. For
% the constant diffusivity case see adv_spectralhelmholtz.m. 
%%
% Build the wavenumbers and create the variable diffusivity.
dk = 2*pi/Lx; % Nyquist wavenumbers
dl = 2*pi/Ly;
 
k=[0:Nx/2-1 Nx/2 -Nx/2+1:-1]'*dk;
l=[0:Ny/2-1 Ny/2 -Ny/2+1:-1]'*dl;
[kk,ll]=meshgrid(k,l);
 
l2 = l.^2; % 1D meridional wavenumber arrays for splitting up the 1D problems
il = 1i*l;
 
k2 = k.^2; % 1D uonal wavenumber arrays for differentiating within a sub problem
ik = 1i*k;
 
ll2 = ll.^2; % 2D meridional wavenumber arrays for differentiation purposes 
ill = 1i*ll; 
 
kk2 = kk.^2; % 2D uonal wavenumber arrays for general differentiation purposes
ikk = 1i*kk;
%% Diffusivity
% The key difference for this case (and why we need to use GMRES) is that
% we have a variable diffusivity and so it is no longer trivial to solve
% the matrix problem in Fourier space.
myr=sqrt(4*(xx-0.5*Lx).^2 + (yy-0.5*Ly).^2);
kappa = 100*(1-0.95*exp(-(myr/(0.3*Lx)).^4));
kappax= real(ifft(ikk.*fft(kappa,[],2),[],2));
kappay= real(ifft(ill.*fft(kappa,[],1),[],1));
 
%% Beginning the Pre-conditioner material
% Create the finite difference derivative matrices and make sure the
% boundary conditions are correct.
disp('Building preconditioner...');
tic;
%%
% Build 1st and 2nd derivative matrices from FD (for preconditioner)
disp('Calculating FD matrices...');
FDx=spalloc(Nx,Nx,3*Nx);
FDxx=spalloc(Nx,Nx,3*Nx);
FDy=spalloc(Ny,Ny,3*Ny); 
FDyy=spalloc(Ny,Ny,3*Ny);
 
matx=[1 1 1;-dx 0 dx;dx2o2 0 dx2o2];
maty=[1 1 1;-dy 0 dy;dy2o2 0 dy2o2];
coefs1x=matx\[0;1;0]; 
coefs1y=maty\[0;1;0];
coefs2x=matx\[0;0;1]; 
coefs2y=maty\[0;0;1];
 
for ii=2:Nx-1
    FDx(ii,ii-1:ii+1)=coefs1x;
    FDxx(ii,ii-1:ii+1)=coefs2x;
end
for ii=2:Ny-1
    FDy(ii,ii-1:ii+1)=coefs1y;
    FDyy(ii,ii-1:ii+1)=coefs2y;
end
%%
%Employ periodicity at the ends
FDx(1, [Nx 1 2]) = coefs1x; 
FDx(Nx, [Nx-1 Nx 1]) = coefs1x;
FDxx(1, [Nx 1 2]) = coefs2x; 
FDxx(Nx, [Nx-1 Nx 1]) = coefs2x;
FDy(1, [Ny 1 2]) = coefs1y; 
FDy(Ny, [Ny-1 Ny 1]) = coefs1y;
FDyy(1, [Ny 1 2]) = coefs2y; 
FDyy(Ny, [Ny-1 Ny 1]) = coefs2y;
%%
% Now build big matrices
Ix = speye(Nx,Nx); 
Iy = speye(Ny,Ny); 
%% Construct the Preconditioner
% Do this in pieces so we don't run out of memory.
disp('Constructing actual preconditioner...');
kappadiag = spdiags(kappa(:),0,Nx*Ny,Nx*Ny);
bigFDxx = kron(FDxx,Iy);
precon = kappadiag*bigFDxx;
clear bigFDxx
%% 
bigFDyy=kron(Ix,FDyy);
precon = precon + kappadiag*bigFDyy;
clear bigFDyy kappadiag
%%
kappaxdiag = spdiags(kappax(:),0,Nx*Ny,Nx*Ny);
bigFDx=kron(FDx,Iy);
precon = precon + kappaxdiag*bigFDx;
clear bigFDx kappaxdiag
%%
kappaydiag = spdiags(kappay(:),0,Nx*Ny,Nx*Ny);
bigFDy=kron(Ix,FDy);
precon = precon + kappaydiag*bigFDy;
clear bigFDy kappaydiag
clear Ix Iy FDx FDxx FDy FDyy coefs1x coefs2x coefs1y coefs2y
disp('done');
%% Compute the LU factorization
tic;
[ll,uu,pp,qq]=lu(precon);
toc;
disp('done');
%% Solve the Problem using GMRES
% Define the RHS and solve the Helmholtz problem using GMRES. GMRES takes
% the inputs: the matrix A, the vector b, restart number [], tolerance of
% the method TOL, maximum number of iterations MAXIT, and the
% preconditioner matrix M.
% A does not have to be a matrix explicitly, here we have used an anonymous
% function that calls the karhelm function so that its output is the
% matrix. This allows more parameters to be passed to the karnhelm function
% than is traditionally allowed. A similar procedure is used for the
% preconditioner matrix.
RHShelm = sin(6*pi*xx/Lx).*sin(4*pi*yy/Ly);
disp('Helmholtz inversion without preconditioner');
tic;
[u00]=gmres(@(foo)func_helmholtz(foo,kappax,kappay,kappa,ill,ikk,Nx,Ny),RHShelm(:),[],TOL,MAXIT);
toc;
disp('Helmholtz inversion with incomplete LU preconditioner...');
tic;
[u0]=gmres(@(foo)func_helmholtz(foo,kappax,kappay,kappa,ill,ikk,Nx,Ny),RHShelm(:),[],TOL,MAXIT,@(bar)func_preconLU(bar,ll,uu,pp,qq));
toc;
disp('Zero diffusivity helmholtz inversion without preconditioner');
tic;
[u0const]=gmres(@(foo)func_helmholtz(foo,zeros(size(kappa)),zeros(size(kappa)),mean(kappa(:))*ones(size(kappa)),ill,ikk,Nx,Ny),RHShelm(:),[],TOL,MAXIT);
toc;
%% Plot the Results
u0=reshape(u0,Ny,Nx); % the preconditioned solution
u00=reshape(u00,Ny,Nx); % the not preconditioned solution
u0const=reshape(u0const,Ny,Nx); % the constant diffusivity solution
ud=u0-u0const; % the difference between constant and non-constant diffusivity
udo = u0-u00; % the difference between preconditioned and not
maxdiff=max(abs(ud(:)))/max(abs(u0(:)))
maxdiff0 = max(abs(udo(:)))/max(abs(u0(:)))
figure(1)
clf
subplot(2,2,1)
pcolor(xx,yy,u0),shading flat; colorbar;
title('preconditioned var diff')
subplot(2,2,2)
pcolor(xx,yy,u0const),shading flat; colorbar;
title('const diff')
subplot(2,2,3)
pcolor(xx,yy,u0-u0const),shading flat; colorbar;
title('var - const')
subplot(2,2,4)
pcolor(xx,yy,u0-u00),shading flat; colorbar;
title('precon - no precon')
figure(2)
clf
pcolor(xx,yy,kappa),shading flat; colorbar;
title('diffusivity')

%% Helmholtz Operator
% This function returns the results of the Helmholtz operator taking in the
% variable u, the diffusivity kappa, its derivatives kappax and kappay, i
% times the wavenumber vector squared, ikk and ill, and the grid sizes Nx
% and Ny.
function result = func_helmholtz(u,kappax,kappay,kappa,ill,ikk,Nx,Ny)
    u = reshape(u,Ny,Nx);
    %% Compute the derivaitves
    ux = real(ifft(ikk.*fft(u,[],2),[],2));
    uxx = real(ifft(ikk.*fft(ux,[],2),[],2));
    uy = real(ifft(ill.*fft(u,[],1),[],1));
    uyy = real(ifft(ill.*fft(uy,[],1),[],1));
    %% Compute the operator
    result = kappax.*ux + kappa.*uxx + kappay.*uy + kappa.*uyy - u;
    %% Return as a vector
    result = result(:);
end

%% Precondition using LU
% Pre-condition using stored sparse LU factors. (with permutation matrices P and Q)
function approxans = func_preconLU(rhs,L,U,P,Q)
    approxans = Q*(U\(L\(P*rhs)));
end