function [dataout] = ft_denoise_prewhiten(cfg, datain, noise)

% FT_DENOISE_PREWHITEN applies a spatial prewhitening operation to the data using the

% inverse noise covariance matrix. The consequence is that all channels are expressed

% in singnal-to-noise units, causing different channel types to be comparable. This

% ensures equal weighting in source estimation on data with different channel types.

%

% Use as

% dataout = ft_denoise_prewhiten(cfg, datain, noise)

% where the datain is the original data from FT_PREPROCESSING and

% noise should contain the estimated noise covariance from

% FT_TIMELOCKANALYSIS.

%

% The configuration structure can contain

% cfg.channel = cell-array, see FT_CHANNELSELECTION (default = 'all')

% cfg.split = cell-array of channel types between which covariance is split, it can also be 'all' or 'no'

% cfg.lambda = scalar, or string, regularization parameter for the inverse

% cfg.kappa = scalar, truncation parameter for the inverse

%

% The channel selection relates to the channels that are pre-whitened using the same

% selection of channels in the noise covariance. All channels present in the input

% data structure will be present in the output, including trigger and other auxiliary

% channels.

%

% See also FT_DENOISE_SYNTHETIC, FT_DENOISE_PCA, FT_DENOISE_DSSP, FT_DENOISE_TSP

% Copyright (C) 2018-2019, Robert Oostenveld and Jan-Mathijs Schoffelen

%

% This file is part of FieldTrip, see http://www.fieldtriptoolbox.org

% for the documentation and details.

%

% FieldTrip is free software: you can redistribute it and/or modify

% it under the terms of the GNU General Public License as published by

% the Free Software Foundation, either version 3 of the License, or

% (at your option) any later version.

%

% FieldTrip is distributed in the hope that it will be useful,

% but WITHOUT ANY WARRANTY; without even the implied warranty of

% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the

% GNU General Public License for more details.

%

% You should have received a copy of the GNU General Public License

% along with FieldTrip. If not, see <http://www.gnu.org/licenses/>.

%

% $Id$

% these are used by the ft_preamble/ft_postamble function and scripts

ft_revision = '$Id$';

ft_nargin = nargin;

ft_nargout = nargout;

ft_defaults

ft_preamble init

ft_preamble debug

ft_preamble loadvar datain

ft_preamble provenance datain

% the ft_abort variable is set to true or false in ft_preamble_init

if ft_abort

% do not continue function execution in case the outputfile is present and the user indicated to keep it

return

end

% ensure that the input data is correct, the next line is needed for a

% attempt correct detection of the data chanunit (with a hdr-field it fails

% for meggrad data)

if isfield(datain, 'hdr'), datain = rmfield(datain, 'hdr'); end

% check if the input data is valid for this function

datain = ft_checkdata(datain, 'datatype', {'raw' 'timelock' 'freq'}, 'haschantype', 'yes', 'haschanunit', 'yes');

noise = ft_checkdata(noise, 'datatype', { 'timelock' 'freq'}, 'haschantype', 'yes', 'haschanunit', 'yes');

% check if the input cfg is valid for this function

cfg = ft_checkconfig(cfg, 'forbidden', {'channels'}); % prevent accidental typos, see issue 1729

% set the defaults

cfg.channel = ft_getopt(cfg, 'channel', 'all');

cfg.split = ft_getopt(cfg, 'split', 'all');

cfg.lambda = ft_getopt(cfg, 'lambda', 0);

cfg.kappa = ft_getopt(cfg, 'kappa', []);

cfg.tol = ft_getopt(cfg, 'tol', []);

cfg.realflag = ft_getopt(cfg, 'realflag', true); % for complex-valued crsspctrm

cfg.invmethod = ft_getopt(cfg, 'invmethod', 'tikhonov');

dtype_datain = ft_datatype(datain);

% check for allowed input combinations

switch dtype_datain

case 'raw'

assert(ft_datatype(noise, 'timelock'), 'noise data should be of datatype ''timelock''');

case 'timelock'

assert(ft_datatype(noise, 'timelock'), 'noise data should be of datatype ''timelock''');

case 'freq'

if ft_datatype(noise, 'freq')

% this is only allowed if both structures have the same singleton frequency

assert(numel(noise.freq)==1 && numel(datain.freq)==1 && isequal(noise.freq,datain.freq), ...

'with both datain and noise of datatype ''freq'', only singleton and equal frequency bins are allowed');

elseif ft_datatype(noise, 'timelock')

% this is OK

end

otherwise

ft_error('unsupported input data');

end

% select channels and trials of interest, by default this will select all channels and trials

tmpcfg = keepfields(cfg, {'trials', 'channel', 'tolerance', 'showcallinfo', 'trackcallinfo', 'trackusage', 'trackdatainfo', 'trackmeminfo', 'tracktimeinfo', 'checksize'});

datain = ft_selectdata(tmpcfg, datain);

noise = ft_selectdata(tmpcfg, noise);

% restore the provenance information

[cfg, datain] = rollback_provenance(cfg, datain);

[cfg, noise] = rollback_provenance(cfg, noise);

if ft_datatype(noise, 'timelock')

if ~isfield(noise, 'cov')

ft_error('noise covariance is not present');

else

noisecov = noise.cov;

end

elseif ft_datatype(noise, 'freq')

if ~isfield(noise, 'crsspctrm')

ft_error('noise cross-spectrum is not present');

else

if istrue(cfg.realflag)

noisecov = real(noise.crsspctrm);

else

noisecov = noise.crsspctrm;

end

end

end

% determine whether it is EEG and/or MEG data

hasgrad = isfield(datain, 'grad');

haselec = isfield(datain, 'elec');

hasopto = isfield(datain, 'opto');

if isequal(cfg.split, 'no')

chantype = {};

elseif isequal(cfg.split, 'all')

chantype = unique(noise.chantype);

else

chantype = cfg.split;

end

% zero out the off-diagonal elements for the specified channel types

if numel(chantype)>0

invnoise = zeros(size(noisecov));

tra = zeros(size(noisecov));

for i=1:numel(chantype)

sel = strcmp(noise.chantype, chantype{i});

%noisecov(sel,~sel) = 0;

%noisecov(~sel,sel) = 0;

invnoise(sel,sel) = ft_inv(noisecov(sel,sel), 'lambda', cfg.lambda, 'kappa', cfg.kappa, 'tolerance', cfg.tol, 'method', cfg.invmethod);

[U,S,V] = svd(invnoise(sel,sel), 'econ');

diagS = diag(S)./numel(chantype);

selS = 1:rank(invnoise(sel,sel));

tra(sel,sel) = U(:,selS)*diag(sqrt(diagS(selS)))*U(:,selS)';

end

%invnoise = ft_inv(noisecov, 'lambda', cfg.lambda, 'kappa', cfg.kappa, 'tolerance', cfg.tol, 'method', cfg.invmethod);

else

% invert the noise covariance matrix

invnoise = ft_inv(noisecov, 'lambda', cfg.lambda, 'kappa', cfg.kappa, 'tolerance', cfg.tol, 'method', cfg.invmethod);

[U,S,V] = svd(invnoise,'econ');

diagS = diag(S);

%sel = diagS./diagS(1)>1e-12;

sel = 1:rank(invnoise);

% the prewhitening projection first rotates to orthogonal channels,

% then scales, and then rotates the channels back to (more or less)

% their original MEG-channel representation

tra = U(:,sel)*diag(sqrt(diagS(sel)))*U(:,sel)';

end

prewhiten = [];

prewhiten.tra = tra;

prewhiten.labelold = noise.label;

prewhiten.labelnew = noise.label;

prewhiten.chantypeold = noise.chantype;

prewhiten.chantypenew = noise.chantype;

prewhiten.chanunitold = noise.chanunit;

prewhiten.chanunitnew = repmat({'snr'}, size(noise.chantype));

% apply the projection to the data

dataout = ft_apply_montage(removefields(datain, {'grad', 'elec', 'opto'}), prewhiten, 'keepunused', 'yes');

if hasgrad

% the gradiometer structure needs to be updated to ensure that the forward model remains consistent with the data

dataout.grad = ft_apply_montage(datain.grad, prewhiten, 'balancename', 'prewhiten');

end

if haselec

% the electrode structure needs to be updated to ensure that the forward model remains consistent

dataout.elec = ft_apply_montage(datain.elec, prewhiten, 'balancename', 'prewhiten');

end

if hasopto

% the electrode structure needs to be updated to ensure that the forward model remains consistent

dataout.opto = ft_apply_montage(datain.opto, prewhiten, 'balancename', 'prewhiten');

end

% do the general cleanup and bookkeeping at the end of the function

ft_postamble debug

ft_postamble previous datain

ft_postamble provenance dataout

ft_postamble history dataout

ft_postamble savevar dataout