% ModifiedSweepingFGT.m
%
% Computes the fast Gauss transform (FGT) of a 2D data set using the
% modified sweeping algorithm described by Sampath et al. in "Parallel Fast
% Gauss Transform".
%
% Program outputs the error in the FGT, computed with respect to several 
% exact Gauss transforms from random targets and averaged.
% 
% Matt Warnez -- April 2013

% create sources and source positions
N = 6400; % number of source points
d = sqrt(N);
f = rand(N,1); % sources
x1 = repmat(0:1/(d-1):1,1,d); % distribution of sources
x2 = repmat(0:1/(d-1):1,d,1);
x2 = x2(:)';

% numeric parameters (see Sampath et al.)
p = 15;
L = 10;
K = 81; % choose from 3^2, 5^2, 7^2,...
delta = 1/3600; % choose from 1/2^2, 1/3^2, 1/4^2, ...

% compute number of boxes and box dimensions
Bnum = round(1/delta);
Bdim = sqrt(Bnum);
Delt = sqrt(delta);
Kdim = (sqrt(K) - 1)/2;

tic
% wave number setup
k = repmat((-p:p)',2*p+1,1);
k2 = reshape(repmat(-p:p,2*p+1,1),(2*p+1)^2,1);
ktest = sqrt(k.^2+k2.^2) <= p;
k = L/(p*sqrt(delta))*[k(ktest) k2(ktest)];
ksize = size(k,1); k1 = 1i*k(:,1); k2 = 1i*k(:,2);

% precomputation of constants
Gconst = (L/(2*p))^2/pi*exp(-(k(:,1).^2 + k(:,2).^2)*delta/4);
x = k1*Delt; y = k2*Delt;
ex = exp(x); ey = exp(y); exey = -ex.*ey;
ew1 = exp((Kdim + 1)*(x + y)); ew2 = exp(Kdim*(-x + y) + y);
ew3 = exp(Kdim*(x - y) + x); ew4 = exp(-Kdim*(x + y));

% data storage variables
w = zeros(ksize,Bnum); v = w;
I = cell(Bdim); W = I; B = I;
F = zeros(N,1);

% prepare box variables
x1n = x1/Delt; x2n = x2/Delt;
x1n(x1n == 0) = eps; x2n(x2n == 0) = eps;
Bc = cell(Bdim); Bcc = zeros(Bdim);
Bc1 = repmat(((1:Bdim)-0.5)*Delt,1,Bdim);
Bc2 = repmat(((1:Bdim)-0.5)*Delt,Bdim,1);
Bc2 = Bc2(:)'; NN = (1:N)';

% precomputations for each box
for i = 1:Bdim
    for j = 1:Bdim
        
        % distribute sources into boxes, identify box centers
        B{i,j} = NN(logical(((i-1) < x1n(:)).*(x1n(:) <= i).*((j-1) ...
            < x2n(:)).*(x2n(:) <= j)));
        Bc{i,j} = [2*i-1; 2*j-1]*Delt/2;
        
        % construct interaction list
        n = zeros(K,1);
        count = 1;
        for q = -Kdim:Kdim
            for r = -Kdim:Kdim                
                if (i+q)>0 && (j+r)>0 && (i+q)<=Bdim && (j+r)<=Bdim
                    n(count) = i + q + Bdim*(j + r - 1);
                    count = count + 1;
                end                
            end
        end
        I{i,j} = n(n > 0);
        
        % compute offset corners for sweeping algorithm
        www = zeros(4,1);
        if (i+Kdim)<=Bdim
            if (j+Kdim)<=Bdim
                www(4) = i + Kdim + Bdim*(j + Kdim - 1);
            end
            if (j-1-Kdim)>0
                www(2) = i + Kdim + Bdim*(j - Kdim - 2);
            end
        end
        if (i - 1 - Kdim)>0
            if (j - 1 - Kdim)>0
                www(1) = i - 1 - Kdim + Bdim*(j - Kdim - 2);
            end
            if (j + Kdim)<= Bdim
                www(3) = i - 1 - Kdim + Bdim*(j + Kdim - 1);
            end
        end
        www(www == 0) = Bnum + 1;
        W{i,j} = www;
        
    end
end

% form plane-wave expansion (S2W)
for i = 1:Bdim
    for j = 1:Bdim
        
        m = B{i,j};
        w(:,i + Bdim*(j - 1)) = exp(k1*(Bc{i,j}(1) - x1(m)) ...
            + k2*(Bc{i,j}(2) - x2(m)))*f(m);
        
    end
end

% compute initial layer directly (W2La)
v(:,1) = sum(exp(k1*(Bc{1,1}(1) - Bc1(I{1,1})) ...
        + k2*(Bc{1,1}(2) - Bc2(I{1,1}))).*w(:,I{1,1}),2);
for ij = 2:Bdim
    v(:,ij) = sum(exp(k1*(Bc{ij,1}(1) - Bc1(I{ij,1})) ...
        + k2*(Bc{ij,1}(2) - Bc2(I{ij,1}))).*w(:,I{ij,1}),2); 
    v(:,1 + Bdim*(ij-1)) = sum(exp(k1*(Bc{1,ij}(1) - Bc1(I{1,ij})) ...
        + k2*(Bc{1,ij}(2) - Bc2(I{1,ij}))).*w(:,I{1,ij}),2);
end

% sweep over remaining layers (W2Lb)
w = [w zeros(ksize,1)];
for i = 2:Bdim
    for j = 2:Bdim
        
        v(:,i + Bdim*(j - 1)) = v(:,i - 1 + Bdim*(j - 2)).*exey ...
            + v(:,i + Bdim*(j - 2)).*ey + v(:,i - 1 + Bdim*(j - 1)).*ex ...
            + w(:,W{i,j}(1)).*ew1 - w(:,W{i,j}(2)).*ew2 ...
            - w(:,W{i,j}(3)).*ew3 + w(:,W{i,j}(4)).*ew4;
        
    end
end

% evaluate transform (L2T)
vG = bsxfun(@times,Gconst,v);
vR = real(vG); vI = imag(vG);
for i = 1:Bdim
    for j = 1:Bdim
        
        m = B{i,j};
        theta = k(:,1)*(x1(m) - Bc{i,j}(1)) + k(:,2)*(x2(m) - Bc{i,j}(2));
        F(m) = cos(theta)'*vR(:,i + Bdim*(j - 1)) ...
            - sin(theta)'*vI(:,i + Bdim*(j - 1));
        
    end
end
toc

% direct summation at 10 random targets to compute error
i = randi(N,10,1);
Fi = zeros(10,1);
for j = 1:N
    
    Fi = Fi + exp(-((x1(i) - x1(j)).^2 + (x2(i) - x2(j)).^2)/delta)'*f(j);
    
end
disp(['Error is ' num2str(norm(abs(Fi-F(i))./abs(Fi))) '.']);