% This program runs tests for the Laplace operator on a square.
% It is adapted from the "p16" function from Trefethen's book
% "Spectral Methods in Matlab"

function main_spectral()

%LOCAL_solve_laplace();
%LOCAL_solve_helmholtz();
%LOCAL_find_eigenfuncs;
LOCAL_laplace_composite;

return

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function LOCAL_solve_laplace()

n        = 20;
N        = n*n;
a        = 2.5;
[L,xvec] = LOCAL_get_L_and_geom(n,a);

[XX1,XX2] = meshgrid(xvec);
xx1       = XX1(:);
xx2       = XX2(:);

INDI      = ones(n-2,1)*(2:(n-1)) + n*(1:(n-2))' * ones(1,n-2);
indi      = INDI(:)';
indb      = LOCAL_complement_vector(indi,N);

u_bound   = cos(xx1(indb)).*exp(xx2(indb));
u_exact   = cos(xx1).*exp(xx2);
u_int     = L(indi,indi)\(L(indi,indb)*(-u_bound));
u         = zeros(N,1);
u(indi)   = u_int;
u(indb)   = u_bound;

fprintf(1,'Max error = %12.5e\n',max(abs(u - u_exact)))
subplot(1,2,1)
mesh(XX1,XX2,reshape(u,n,n));
subplot(1,2,2)
mesh(XX1,XX2,reshape(u - u_exact,n,n))

xxs       = a*[-1.4,1];
u_bound   = log((xx1(indb)-xxs(1)).^2 + (xx2(indb)-xxs(2)).^2);
u_exact   = log((xx1-xxs(1)).^2 + (xx2-xxs(2)).^2);
u_int     = L(indi,indi)\(L(indi,indb)*(-u_bound));
u         = zeros(N,1);
u(indi)   = u_int;
u(indb)   = u_bound;

fprintf(1,'Max error = %12.5e\n',max(abs(u - u_exact)))
subplot(1,2,1)
mesh(XX1,XX2,reshape(u,n,n));
subplot(1,2,2)
mesh(XX1,XX2,reshape(u - u_exact,n,n))

return

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function LOCAL_solve_helmholtz()

n        = 20;
N        = n*n;
a        = 1.2;
[L,xvec] = LOCAL_get_L_and_geom(n,a);

[XX1,XX2] = meshgrid(xvec);
xx1       = XX1(:);
xx2       = XX2(:);

INDI      = ones(n-2,1)*(2:(n-1)) + n*(1:(n-2))' * ones(1,n-2);
indi      = INDI(:)';
indb      = LOCAL_complement_vector(indi,N);

fprintf(1,'\nSolving for u = cos(k1*x1 + k2*x2)\n')
k1        = 3;
k2        = -5;
kh        = sqrt(k1*k1 + k2*k2);
fprintf(1,'   Size of domain = %6.2f wavelengths\n',a*kh/pi)
u_bound   = cos(k1*xx1(indb) + k2*xx2(indb));
u_exact   = cos(k1*xx1 + k2*xx2);
A         = -L - kh*kh*eye(N);
u_int     = A(indi,indi)\(A(indi,indb)*(-u_bound));
u         = zeros(N,1);
u(indi)   = u_int;
u(indb)   = u_bound;
fprintf(1,'Max error = %12.5e\n',max(abs(u - u_exact)))

m         = 3;
kh        = 5;
fprintf(1,'\nSolving for u = cos(m*theta)*besselj(m,kh*r) for m = %d  kh = %6.2f\n',m,kh);
fprintf(1,'   Size of domain = %6.2f wavelengths\n',a*kh/pi)
[tt,rr]   = cart2pol(xx1,xx2);
u_bound   = cos(m*tt(indb)).*besselj(m,kh*rr(indb));
u_exact   = cos(m*tt).*besselj(m,kh*rr);
A         = -L - kh*kh*eye(N);
u_int     = A(indi,indi)\(A(indi,indb)*(-u_bound));
u         = zeros(N,1);
u(indi)   = u_int;
u(indb)   = u_bound;
fprintf(1,'Max error = %12.5e\n',max(abs(u - u_exact)))

kh        = 5;
xxs       = a*[-1.4,1];
fprintf(1,'\nSolving for u = bessely(kh*|x - xs|)\n')
fprintf(1,'   kh = %6.2f    xs = (%6.2f, %6.2f)\n',kh,xxs(1),xxs(2));
fprintf(1,'   Size of domain = %6.2f wavelengths\n',a*kh/pi)
u_bound   = bessely(0,kh*sqrt((xx1(indb) - xxs(1)).^2 + (xx2(indb) - xxs(2)).^2));
u_exact   = bessely(0,kh*sqrt((xx1 - xxs(1)).^2 + (xx2 - xxs(2)).^2));
A         = -L - kh*kh*eye(N);
u_int     = A(indi,indi)\(A(indi,indb)*(-u_bound));
u         = zeros(N,1);
u(indi)   = u_int;
u(indb)   = u_bound;
fprintf(1,'Max error = %12.5e\n',max(abs(u - u_exact)))

subplot(1,2,1)
mesh(XX1,XX2,reshape(u,n,n));
subplot(1,2,2)
mesh(XX1,XX2,reshape(u - u_exact,n,n))

return

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function LOCAL_find_eigenfuncs()

n        = 20;
N        = n*n;
a        = 2.5;
[L,xvec] = LOCAL_get_L_and_geom(n,a);

[XX1,XX2] = meshgrid(xvec);
xx1       = XX1(:);
xx2       = XX2(:);

INDI      = ones(n-2,1)*(2:(n-1)) + n*(1:(n-2))' * ones(1,n-2);
indi      = INDI(:)';
indb      = LOCAL_complement_vector(indi,N);

A         = -L;

figure(1)
subplot(1,1,1)
hold off
AINT          = A(indi,indi);
[V,D]         = eig(AINT);
ll            = diag(D);
[ll,ind]      = sort(ll);
U_MAT         = zeros(N,length(indi));
U_MAT(indi,:) = V(:,ind);
[K1,K2] = meshgrid(1:20);
kk = K1.*K1 + K2.*K2;
ll_exact = (pi/(2*a))^2 * sort(kk(:)');
for i = 1:N
  fprintf(1,'Eigenfunction nr %3i: computed lambda(i) = %15.12f \n',i,ll(i))
  fprintf(1,'                      exact    lambda(i) = %15.12f \n',ll_exact(i))
  mesh(XX1,XX2,reshape(U_MAT(:,i),n,n))
  title(sprintf('Eigenfunction nr %d = %15.12f   (exact=%15.12f)',...
                i,ll(i),ll_exact(i)))
  keyboard
end

return

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function LOCAL_laplace_composite()

n         = 20;
n1        = 2*n+1;
n2        = n+1;
N         = n1*n2;
a         = 0.5;
[L,xvec]  = LOCAL_get_L_and_geom(n+1,a);
x1side    = [a+xvec',3*a+xvec(2:end)'];
x2side    = a+xvec';
[XX1,XX2] = meshgrid(x1side,x2side);
xx1       = XX1(:);
xx2       = XX2(:);
xx        = [reshape(XX1,1,numel(XX1));...
             reshape(XX2,1,numel(XX2))];
     
%%% Set up index vectors for box 1:
ind1_own   = 1:((n+1)*n);
J1_own     = 1:((n+1)*n);
ind1_share = (n+1)*n + (1:(n+1));
J1_share   = (n+1)*n + (1:(n+1));

%%% Set up index vectors for box 2:
ind2_own   = (n+1)*(n+1) + (1:((n+1)*n));
J2_own     = (n+1) + (1:((n+1)*n));
ind2_share = (n+1)*n + (1:(n+1));
J2_share   = 1:(n+1);

A = zeros(N);
A(ind1_own,  [ind1_own,ind1_share]) = L(J1_own,  [J1_own,J1_share]);
A(ind1_share,[ind1_own,ind1_share]) = L(J1_share,[J1_own,J1_share]);
A(ind2_own,  [ind2_own,ind2_share]) = L(J2_own,  [J2_own,J2_share]);

%%% Check how well A does at differentation:
%figure(1)
%rsq = xx1.*xx1 + xx2.*xx2;
%u   = exp(-rsq);
%Lu_exact = (4*rsq - 4).*exp(-rsq);
%subplot(1,2,1)
%mesh(XX1,XX2,reshape(A*u,n2,n1))
%subplot(1,2,2)
%mesh(XX1,XX2,reshape(Lu_exact - A*u,n2,n1))
%fprintf(1,'Max error in differentation = %12.5e\n',max(abs(A*u - Lu_exact)))

%%% Solve a Laplace problem:
figure(2)
hmin = xvec(2) - xvec(1);  % minimal grid spacing.
dd   = min([abs(xx1');...
            abs(xx1'-4*a);...
            abs(xx2');...
            abs(xx2'-2*a)]);
indb = find(dd < (0.5*hmin));
indi = LOCAL_complement_vector(indb,N);
xxs       = a*[-1.4,1];
%FUNC      = inline('xx1.*xx1 - xx2.*xx2','xx1','xx2');
%FUNC      = inline('cos(xx1).*exp(xx2)','xx1','xx2');
FUNC      = inline('log((xx1-0.5).^2 + (xx2+1).^2)','xx1','xx2');
u_exact   = FUNC(xx1,xx2);
u_bound   = FUNC(xx1(indb),xx2(indb));
u_int     = A(indi,indi)\(-A(indi,indb)*u_bound);
u         = zeros(N,1);
u(indb)   = u_bound;
u(indi)   = u_int;
disp(max(abs(u - u_exact)))
keyboard
u_bound   = log((xx1(indb)-xxs(1)).^2 + (xx2(indb)-xxs(2)).^2);
u_exact   = log((xx1-xxs(1)).^2 + (xx2-xxs(2)).^2);
u_int     = A(indi,indi)\(A(indi,indb)*(-u_bound));
u         = zeros(N,1);
u(indi)   = u_int;
u(indb)   = u_bound;
disp(norm(u_exact - u))

keyboard


INDI      = ones(n-1,1)*(2:(n-1)) + n*(1:(n-2))' * ones(1,n-2);
indi      = INDI(:)';
indb      = LOCAL_complement_vector(indi,N);

u_bound   = cos(xx1(indb)).*exp(xx2(indb));
u_exact   = cos(xx1).*exp(xx2);
u_int     = L(indi,indi)\(L(indi,indb)*(-u_bound));
u         = zeros(N,1);
u(indi)   = u_int;
u(indb)   = u_bound;

fprintf(1,'Max error = %12.5e\n',max(abs(u - u_exact)))
subplot(1,2,1)
mesh(XX1,XX2,reshape(u,n,n));
subplot(1,2,2)
mesh(XX1,XX2,reshape(u - u_exact,n,n))
fprintf(1,'Max error = %12.5e\n',max(abs(u - u_exact)))
subplot(1,2,1)
mesh(XX1,XX2,reshape(u,n,n));
subplot(1,2,2)
mesh(XX1,XX2,reshape(u - u_exact,n,n))

return

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [L,xvec] = LOCAL_get_L_and_geom(N_side,a)
[D,xvec] = cheb(N_side-1);
xvec     = a*xvec(end:(-1):1);
D        = (1/a)*D;
%y = x;
%[xx,yy] = meshgrid(x(2:N),y(2:N));
%xx = xx(:); yy = yy(:);       % stretch 2D grids to 1D vectors
%f = 10*sin(8*xx.*(yy-1));
D2 = D^2; 
I  = eye(N_side);
L  = kron(I,D2) + kron(D2,I);
return

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function indc = LOCAL_complement_vector(indi,N)
indtmp = 1:N;
indtmp(indi) = 2*N*ones(1,length(indi));
indtmp = sort(indtmp);
indc = indtmp(1:(N - length(indi)));
return
