% A test script for ao/whiten2D
% 
% L. Ferraioli 18-11-08
% 
% $Id: test_ao_whiten2D_2.m,v 1.3 2009/08/26 18:27:06 luigi Exp $
% 

%%

clear all

%% CSD Noise models and TF

[CSD,TF] = get_2D_test_obj_tf_psd();

% get noise coherence
coh = CSD(1,2)./sqrt(abs(CSD(1,1).*CSD(2,2)));
rcoh = real(coh);
rcoh.setName;
icoh = imag(coh);
icoh.setName;

%% Colored Noise generation

% Make white noise

a1 = ao(plist('tsfcn', 'randn(size(t))', 'fs', 10, 'nsecs', 5e5));
a2 = ao(plist('tsfcn', 'randn(size(t))', 'fs', 10, 'nsecs', 5e5));

% Generation from white noise

pl = plist(...
    'csd11', CSD(1,1),...
    'csd12', CSD(1,2),...
    'csd21', CSD(2,1),...
    'csd22', CSD(2,2),...
    'MaxIter', 70, ...
    'PoleType', 2, ...
    'MinOrder', 25, ...
    'MaxOrder', 35, ...
    'Weights', 2, ...
    'Plot', false,...
    'MSEVARTOL', 1e-2,...
    'FITTOL', 1e-4,...
    'UseSym', 0,...
    'Disp', false);

[ac1,ac2] = noisegen2D(a1,a2, pl);

ac1.setYunits('m');
ac2.setYunits('m');

%% Checking generated noise

acxx1 = ac1.psd(plist('Nfft',1e6));
acxx2 = ac2.psd(plist('Nfft',1e6));

acxy = cohere(ac1,ac2,plist('Nfft',5e4));
racxy = real(acxy);
racxy.setName;
iacxy = imag(acxy);
iacxy.setName;

iplot(abs(acxx1),abs(CSD(1,1)))
iplot(abs(acxx2),abs(CSD(2,2)))

iplot(racxy,rcoh,plist('Yscales',{'All','lin'}))
iplot(iacxy,icoh,plist('Yscales',{'All','lin'}))

%%
iplot(abs(acxx1),abs(acxx2))

%% Get Coloring filters

cf11 = ac1.procinfo.find('Filt11');
cf12 = ac1.procinfo.find('Filt12');
cf21 = ac2.procinfo.find('FILT21');
cf22 = ac2.procinfo.find('FILT22');

%% Get Theoretical CSD from coloring filters

rc11 = resp(cf11,plist('f',CSD(1,1).x,'bank','parallel'));
rc12 = resp(cf12,plist('f',CSD(1,1).x,'bank','parallel'));
rc21 = resp(cf21,plist('f',CSD(1,1).x,'bank','parallel'));
rc22 = resp(cf22,plist('f',CSD(1,1).x,'bank','parallel'));


psd11 = rc11.*conj(rc11)+rc12.*conj(rc12);
psd22 = rc21.*conj(rc21)+rc22.*conj(rc22);
psd12 = rc11.*conj(rc21)+rc12.*conj(rc22);
psd21 = conj(psd12);
coher = psd12./sqrt(abs(psd11).*abs(psd22));
rcoher = real(coher);
rcoher.setName;
icoher = imag(coher);
icoher.setName;

% compare with data
iplot(racxy,rcoher,plist('Yscales',{'All','lin'}))
iplot(iacxy,icoher,plist('Yscales',{'All','lin'}))

%% Noise whitening

pl = plist(...
    'csd11', CSD(1,1),...
    'csd12', CSD(1,2),...
    'csd21', CSD(2,1),...
    'csd22', CSD(2,2),...
    'MaxIter', 75, ...
    'PoleType', 3, ...
    'MinOrder', 25, ...
    'MaxOrder', 40, ...
    'Weights', 2, ...
    'Plot', false,...
    'MSEVARTOL', 1e-2,...
    'FITTOL', 1e-3,...
    'UseSym', 0,...
    'Disp', false);

[aw1,aw2] = whiten2D(ac1,ac2, pl);

aw1.setYunits('m');
aw2.setYunits('m');

%% Checking whitened noise

awxx1 = aw1.psd(plist('Nfft',1e6));
awxx2 = aw2.psd(plist('Nfft',1e6));
%%
iplot(awxx1,awxx2)

%% Get coherence

awxy = cohere(aw1,aw2,plist('Nfft',6e4));
rawxy = real(awxy);
rawxy.setName;
iawxy = imag(awxy);
iawxy.setName;

%% Check whitened data coherence
racxy.setName('a_xy');
iacxy.setName('a_xy');
rawxy.setName('wa_xy');
iawxy.setName('wa_xy');
iplot(racxy.split(plist('split_type','frequencies','frequencies',[4e-4 5])),rcoh.split(plist('split_type','frequencies','frequencies',[1e-4 5])),rawxy.split(plist('split_type','frequencies','frequencies',[4e-4 5])),plist('Yscales',{'All','lin'},'Ylabels',{'All','Re(coh)'}))
iplot(iacxy.split(plist('split_type','frequencies','frequencies',[4e-4 5])),icoh.split(plist('split_type','frequencies','frequencies',[1e-4 5])),iawxy.split(plist('split_type','frequencies','frequencies',[4e-4 5])),plist('Yscales',{'All','lin'},'Ylabels',{'All','Im(coh)'}))

%% Get whitening filters

wf11 = aw1.procinfo.find('Filt11');
wf12 = aw1.procinfo.find('Filt12');
wf21 = aw2.procinfo.find('FILT21');
wf22 = aw2.procinfo.find('FILT22');

% filters respose
rw11 = resp(wf11,plist('f',CSD(1,1).x,'bank','parallel'));
rw12 = resp(wf12,plist('f',CSD(1,1).x,'bank','parallel'));
rw21 = resp(wf21,plist('f',CSD(1,1).x,'bank','parallel'));
rw22 = resp(wf22,plist('f',CSD(1,1).x,'bank','parallel'));

%% Build expected csd and coherence of whitened data

wp11 = conj(rw11).*psd11+psd12.*conj(rw12);
wp12 = psd11.*conj(rw21)+psd12.*conj(rw22);
wp21 = psd21.*conj(rw11)+psd22.*conj(rw12);
wp22 = psd21.*conj(rw21)+psd22.*conj(rw22);

u11 = rw11.*wp11+rw12.*wp21;
u11.setName;
u12 = rw11.*wp12+rw12.*wp22;
u12.setName;
u21 = rw21.*wp11+rw22.*wp21;
u21.setName;
u22 = rw21.*wp12+rw22.*wp22;
u22.setName;
wcoher = u12./sqrt(abs(u11).*abs(u22));
rwcoher = real(wcoher);
rwcoher.setName;
iwcoher = imag(wcoher);
iwcoher.setName;

% compare with data
iplot(rawxy,rwcoher,plist('Yscales',{'All','lin'}))
iplot(iawxy,iwcoher,plist('Yscales',{'All','lin'}))

%% Coherence of starting data

axy = cohere(a1,a2,plist('Nfft',1e4));
raxy = real(axy);
raxy.setName;
iaxy = imag(axy);
iaxy.setName;

% compare with whitened
iplot(rawxy,raxy,plist('Yscales',{'All','lin'}))
iplot(iawxy,iaxy,plist('Yscales',{'All','lin'}))

%% Comparing Theoretical data

iplot(rcoher,rwcoher,plist('Yscales',{'All','lin'}))
iplot(icoher,iwcoher,plist('Yscales',{'All','lin'}))

iplot(rcoh,rwcoher,plist('Yscales',{'All','lin'}))
iplot(icoh,iwcoher,plist('Yscales',{'All','lin'}))



















% END