%% 3D Poisson on GPU
%$ -u_xx-u_yy-u_zz=f  
%$ [-L,L]*[-L,L]*[-L,L] with homogeneous Dirichlet B.C..
%$ 2nd order centered difference

%% Xiangxiong Zhang, Purdue University, Oct 2023
% If no GPU device, then do computation on CPU
if gpuDeviceCount("available")<1
    Param.device = 'cpu';
else
    Param.device = 'gpu'; 
    Param.deviceID = 2; % ID=1,2,3,...
end

n = 200; % size of the Laplacian matrix, for 3rd scheme, n must be odd
L = 16; % the domain is [-L, L]*[-L,L]*[-L,L]

%% The computational grid
h = 2*L/(n+1);
x = (-L:h:L)';
x = x(2:end-1);
Length = L;
%% generate the discrete Laplacian
ex = ones(n,1); e = ones(n,1);
K = spdiags([-ex 2*ex -ex], [-1 0 1], n, n)/h^2;
K = full(K);
xx = (0:1/(n+1):1)';
xx = xx(2:end-1);
T = sin(pi*xx*(1:n));
T = T*sqrt(diag(1./sum(T.*T, 1))); % normalize each column of T
lambda = 2-2*cos(pi/(n+1)*(1:n));
lambda = lambda'/h^2;
fprintf('This is a code solving 3D Poison on a grid of size %d by %d by %d \n', n, n,n)
fprintf('scheme is 2nd order centered difference \n') 

switch Param.device
    case 'gpu'
       fprintf('GPU computation: starting to load matrices/data \n') 
       Device = gpuDevice(Param.deviceID);
       T = gpuArray(T); e = gpuArray(e); 
       K = gpuArray(K); x = gpuArray(x);     
end 
Lambda3D = squeeze(tensorprod(lambda, e*e')+tensorprod(e,lambda*e')+tensorprod(e,e*lambda'));
y = x; z = x; 
vx = sin(pi*x/4).^2;
vy = sin(pi*y/8).^4;
vz = sin(pi*z/4).^2;
f = squeeze(tensorprod(vx*vy',vx));
switch Param.device
    case 'gpu'
        wait(Device)
        fprintf('GPU computation: loading finished and GPU computing started \n') 
end
tic;
%% Eigenvector method to solve the linear system by tensor matrix multiplication      
%% T is very special here: T'=T=invT
% the code be the following for general T
% u = tensorprod(f,invT',3,1);
% u = pagemtimes(u,invT');
% u = squeeze(tensorprod(invT,u,2,1));
% u = u./Lambda3D;
% u = tensorprod(u,T',3,1);
% u = pagemtimes(u,T');
% u = squeeze(tensorprod(T,u,2,1));
u = tensorprod(f,T,3,1);
u = pagemtimes(u,T);
u = squeeze(tensorprod(T,u,2,1));
u = u./Lambda3D;
u = tensorprod(u,T,3,1);
u = pagemtimes(u,T);
u = squeeze(tensorprod(T,u,2,1));
 
switch Param.device
    case 'gpu'
        wait(Device);
end
time = toc; 
residue = tensorprod(u,K,3,1)+squeeze(tensorprod(K,u,2,1))+pagemtimes(u,K)-f; 
fprintf('The ell-2 norm residue is %d \n',norm(residue,'fro')*h^1.5)

switch Param.device
    case 'gpu'
        fprintf('The GPU online computation time is %d \n', time)
    case 'cpu'
        fprintf('The CPU online computation time is %d \n', time)
end
