% Function to compute a spectral density matrix estimate of a three-component 
% color image using the Welch method of averaging modified periodograms 
% adapted to vector data following Brillinger.
% Follow description in section 9.6 of MDSP book.
% Eric Dubois 2018-04-23
function WelchSDM = WelchSDM_est(A,N)
% A(:,:,3) is a three-component color image in double format.
% The three components are set in the calling program.
% N x N is the block size. N should be a power of 2. Usually N=64
% WelchSDM(N,N,3,3) is the N x N spectral density estimate on a linear scale for
% each of the nine spectral densities in the spectral density matrix.
% Since the spectral density matrix at each frequency is Hermitian, only
% six spectral densities need to be computed but for simplicity and since 
% the computational complexity is low, all nine are computed.
% The spectral densities are computed at normalized frequencies 
% ux=-1/2:1/N:1/2-(1/N); vx=ux;
%
% Remove the mean from each component of the image. A non-zero mean implies
% a Dirac delta at zero frequency. Az is the new image with zero mean.
I = size(A,1);
J = size(A,2);
prodsize=I*J;
for C=1:3
    mean(C)=sum(sum(A(:,:,C)))/prodsize;
    Az(:,:,C)=A(:,:,C)-mean(C);
end
% Compute the 2D separable window (we use a 4-term Blackman-Harris 
% (or Harris-Nutall) Window)
n=0:1:N-1;
wd=zeros(N,1);
a0=0.40217;
a1=0.49703;
a2=0.09892;
a3=0.00183;
wd(n+1)=a0-a1*cos(2*pi*n/N)+a2*cos(2*pi*2*n/N)-a3*cos(2*pi*3*n/N);
W=wd*wd';
V = sum(sum(W.*W));  % total power of the window
% normalize the window so that the total power is K = N^2.
W = N*W/sqrt(V);
%
% Compute the spectral estimate using an x and y overlap of N/4
% Note that this overlap could be set as a parameter of the algorithm
PDS=zeros(N,N,3,3);
shx=N/4;     % x overlap
shy=N/4;     % y overlap
% Number of blocks in x and y direction
Mx=floor((I-N)/(N-shx));
My=floor((J-N)/(N-shy));

for i=0:1:Mx
    for j=0:1:My
        %Compute the three windowed components and their DFT
        for C=1:3
            BLK(:,:,C) = W.*Az(i*(N-shx)+1:i*(N-shx)+N,j*(N-shy)+1:j*(N-shy)+N,C);
            BLK_DFT(:,:,C) = fft2(BLK(:,:,C));
        end
        %Compute the nine spectral estimates for block i,j (ignoring
        %Hermitian symmetry)
        for C1=1:3
            for C2=1:3
                PDS(:,:,C1,C2) = PDS(:,:,C1,C2) + BLK_DFT(:,:,C1).*conj(BLK_DFT(:,:,C2))/(N^2);
            end
        end
    end
end
%Compute the average. The number of blocks is L = (Mx+1)*(My+1)
PDS=PDS/((Mx+1)*(My+1)); 
%Shuffle the FFT blocks to put the origin at the center
for C1=1:3
    for C2=1:3
        WelchSDM(:,:,C1,C2) = [PDS(N/2+1:N,N/2+1:N,C1,C2) PDS(N/2+1:N,1:N/2,C1,C2);...
            PDS(1:N/2,N/2+1:N,C1,C2) PDS(1:N/2,1:N/2,C1,C2)];
    end
end