%
% OCF: Orthogonal Connectivity Factorization.
%      Given a set of connectivity matrices (e.g. correlations),
%      finds pairs of variable sets (e.g. brain areas) 
%      whose connectivity is changing as strongly as possible in that set.
%
% Call function as:
%   [W,V,TC]=OCF(CONNTENSOR,[NUMBERCOMPS],[further options]);
%
% INPUT:
%
% CONNTENSOR:  Connectivity matrices stacked as a tensor.
%              CONNTENSOR(:,:,i) is the i-th connectivity matrix,
%              E.g. in i-th time window, i-th subject, or i-th condition.
%
% NUMBERCOMPS (optional): Number of components to be computed.
%              If no value given, a default of 5 is used.
%
% OUTPUT:
%
% Each row of W together with the corresponding row of V
%    contain one pair of spatial patterns estimated.
%
% The rows of TC contain the corresponding values of connectivities
%    for each index i in CONNTENSOR (e.g. time courses of connectivities)
%
%
% FURTHER OPTIONS:
% 'finetuning': try to improve matrix PCA results by constrained PCA method
% 'robust': robust version of constrained PCA algorithm
% 'diagnostics': plot various diagnostics on screen
% 'silent': suppress output on terminal
%
% Aapo Hyvarinen, version of 20 Jan 2016
%


function [W,V,TC]=ocf(connTensor,varargin)


%SET OPTIONS
%%%%%%%%%%%%%

%Get number of components
if length(varargin)>0, %If given in function call, read it
  numberComps=varargin{1}; 
  if ~isnumeric(numberComps), 
     error('First input option must be number of components');
  end
else
  disp('Number of components not given, set to 5 by default');
  numberComps=5; 
end
%Set default values for all other options. 
%(Descriptions below say what happens by default when variables set to zero.)
%Main options:
finetuning=0; %No likelihood-based finetuning, just matrix PCA
robust=0; %Use squaring-based objective, not robust abs-based objective
diagnostics=0; %No diagnostics plots
silent=0; %Do not suppress output on terminal
%Some undocumented options: (rarely useful, for developers only)
random=0; %Finetuning initialized by matrix PCA, not randomly
doubledefl=0; %Deflate away only a one-dimensional subspace for component
usemeans=0; %Subtract means before matrix PCA
%Read options from function call
options=varargin;
for optionind=1:length(options)
  if ~(optionind==1 & isnumeric(varargin{optionind})) %skip numbercomps arg
    switch options{optionind}
     case 'finetuning', finetuning=1;
     case 'robust', robust=1; 
     case 'diagnostics', diagnostics=1;
     case 'silent', silent=1;
     case 'doubledefl', doubledefl=1; 
     case 'random', random=1;
     case 'usemeans', usemeans=1;
     case '', ; %Allow empty option to simplify input of options
     otherwise error(['Unknown option: "',options{optionind},'"'])
    end
  end
end



%DO ORTHOGONAL CONNECTIVITY FACTORIZATION 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if ~silent, disp('--Launching OCF computation--'), end

%Clear output variables
clear W V TC

%Get number of channels and matrices (windows)
channels=size(connTensor,1);
windows=size(connTensor,3);

%Collapse connectivity tensor into matrix, 
%i.e. vectorize each connectivity matrix
connTensor=reshape(connTensor,channels*channels,windows);

%Remove means of connectivities (i.e.  global connectivities)
if ~usemeans %...unless disabled by user (usually bad idea)
  connTensor_mean=mean(connTensor')';
  connTensor=connTensor-connTensor_mean*ones(1,size(connTensor,2));
else
  connTensor_mean=zeros(size(connTensor,1),1);
end

%Decide if matrix too big for ordinary PCA in matrix space
if min(channels^2,windows)>1000, 
  matrixtoobig=1;
  fprintf('Dimension too big for matrix PCA, using iterative algorithm only\n');
else 
  matrixtoobig=0; 
end

%To compute more than one component, deflation is used.
%The following matrix will contain the tensor from which 
%previously found components have been deflated away
connTensor_deflated=connTensor;

%LOOP: COMPUTE MATRIX PRINCIPAL COMPONENTS ONE-BY-ONE (incl. finetuning)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

for component=1:numberComps

  if ~silent, fprintf('Component %i: ',component); end

  %Do matrix PCA unless either 
  % a) finetuning with random initial point chosen, or
  % b) matrix too big
  
  if ~(random & finetuning) & ~matrixtoobig

    %COMPUTE ONE PRINCIPAL COMPONENT OF CONNECTIVITY MATRICES 
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

    if ~silent, fprintf(' Doing matrix PCA... '); end
 
    %Calculate one dominant eigenvalue and eigenvector
    %of the connectivity covariance matrix (really big matrix)

    %Choose in which direction to do EVDS, depending on dimension of matrix
    %We don't use matlab's svd(s) function since it is slow and unstable
    if channels^2<windows
        %EIGS of covariance matrix of connectivities
        [matrixpc_vec,Sing]=eigs(connTensor_deflated*connTensor_deflated',1);
    else
        %EIGS of "covariance matrix" of connectivities transposed
        [Ematrixtrans,Sing]=eigs(connTensor_deflated'*connTensor_deflated,1);
        %With transformation needed:
        matrixpc_vec=connTensor_deflated*...
                      Ematrixtrans*diag(1./sqrt(diag(Sing)+eps));
    end

    %Reshape matrix principal component to matrix form
    matrixpc_mat=reshape(matrixpc_vec,[channels,channels]);
    %Compute time course of matrix principal component 
    %("real" needed in case of rounding errors)
    tc_matrixpca=real(matrixpc_vec'*connTensor_deflated);

    %LOW-RANK APPROXIMATION OF OBTAINED MATRIX PRINCIPAL COMPONENT
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

    %Do EVD on matrix principal component
    [E,D]=eig(matrixpc_mat);
    D=real(D); %just in case of numerical errors, make real-valued
    %Compute two-rank approximation based on EVD
    [Dsorted,Dord]=sort(diag(D),'descend');
    %choose indices to use: now first and last since sorted without abs
    %i.e. largest and smallest eigenvalues regardless of sign
    ind1=Dord(1);  ind2=Dord(length(Dord));
    %Finally, get the two eigenvectors used in two-rank approximation
    emax=E(:,ind1); emin=E(:,ind2);
    %(Time course computed only after possible finetuning)

  else
    % Optional (or forced for too big data): 
    % Instead of matrix PCA, do constrained PCA with random initialization 
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

    if ~silent, fprintf(' Initializing randomly... '); end
      %Default: Take two orthogonal, real-valued vectors
      emax=randn(channels,1); 
      emax=emax/norm(emax);
      emin=randn(channels,1);
      emin=emin-emin'*emax*emax;
      emin=emin/norm(emin);
      %Create hypothetical time course needed for iterative algorithm
      matrixpc_vec=emax*emax'-emin*emin';
      matrixpc_vec=matrixpc_vec(:);
      tc_matrixpca=real(matrixpc_vec'*connTensor_deflated);

    end %End of section where we either do matrix PCA 
        %or find random initialization for finetuning


  %TRY TO IMPROVE MATRIX-LEVEL PC BY FINETUNING (=CONSTRAINED PCA)
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

  if finetuning 
    if ~silent, fprintf(' Running iterative algorithm... '); end
    %Initial points directly obtained as emax,emin from matrix PCA above.
    %First, save matrix PCA results in case finetuning fails
    emax_initial=emax; emin_initial=emin;
    %Maximum number of iterations
    maxiter=1000;
    %Initialize vector of objective function values for tracking progress
    obj=zeros(1,maxiter);
    %Open window for showing progress if diagnostics option
    if diagnostics, figure(6), end
    %Initialize weight vector "r" used in relaxed optimization problem
    %based on time course of matrix PCA
    r_vec=tc_matrixpca; if robust, r_vec=sign(r_vec); end
    k=0; %iteration counter initialized
    converged=0; %convergence flag set to zero (not converged)
    while ~converged
      k=k+1; %increment iteration counter
      emax_old=emax; emin_old=emin; %for checking convergence
 
      %ALTERNATING VARIABLES OPTIMIZATION METHOD
      %What is denoted by a,b in the paper is emin, emax here.
      %Compute reweighted connectivity matrix, denoted sum_t r_t C_t in paper
      Cweighted=reshape(connTensor_deflated*r_vec',[channels,channels]);
      %Do EVD on it and find largest and smallest eigenvalues
      [Ealt,Dalt]=eig(Cweighted);
      [minval,minind]=min(diag(Dalt));  [maxval,maxind]=max(diag(Dalt));
      emin=Ealt(:,minind);  emax=Ealt(:,maxind);
      %Recompute optimal weights (r_vec)
      patt_mat=emax*emax'-emin*emin';
      r_vec=patt_mat(:)'*connTensor_deflated;  
      %Compute objective function, simple formula based on r_vec
      obj(k)=sum(abs(r_vec).^(2-robust));
      %If robust option, modify optimal weights (must be after obj computation)
      if robust, r_vec=sign(r_vec); end 

      %Decide if end of iterations. Must have k>1 or criterion does
      %not work. Correction on 20 Jan 2016: here we must take into
      %account the possibility that the signs of eigenvectors flip in
      %subsequent iterations.
      if (k>1) & (k==maxiter|2-abs(emax'*emax_old)-abs(emin'*emin_old)<1e-8)
           converged=1; end
      %Plot progress of optimization if requested
      if diagnostics & (mod(k,10)==2 | converged), 
             plot(obj(1:k)/obj(1)), drawnow, end
    end
    %Check if finetuning failed (hardly ever happens with this algorithm)
    if obj(k)<obj(1), 
        fprintf(' Finetuning failed, reverting to initial point! ')
        emax=emax_initial; emin=emin_initial; 
    end
    %Report convergence result
    if ~silent, 
      if k==maxiter, 
          fprintf(' Max number of iterations reached without convergence!');
      else
          fprintf(' Converged in %i iterations.',k);
      end
    end
  end %of finetuning (constrained PCA) section




  %FINALIZE RESULTS FOR THIS COMPONENT
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

  %Compute spatial patterns (sum/diff) based on eigenvectors (or fine-tuned)
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

  %Basic computation of spatial patterns as sums and diffs of eigenvectors
  %where the eigenvectors may be "finetuned" corresponding to a,b in the paper
  w=(emax+emin)/sqrt(2); 
  v=(emax-emin)/sqrt(2);

  %Sign of spatial patterns arbitrary, so make their signs mainly positive
  w=w*sign(sum(w.^3));
  v=v*sign(sum(v.^3));

  %Store spatial patterns in output matrices
  W(component,:)=w;
  V(component,:)=v;

  %Compute time course of component and store in output matrix
  %Note: mean added back to the connectivity tensor here
  for t=1:windows
     connmat=reshape(connTensor_deflated(:,t)+connTensor_mean...
                     ,[channels,channels]);
     TC(component,t)=real(emax'*connmat*emax-emin'*connmat*emin);
            %"real" needed just in case of rounding errors
  end

  %Deflate component just estimated from the connectivivity tensor
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  %Compute the matrix form of the component just found
  patt_mat=emax*emax'-emin*emin';
  %Vectorize and normalize it
  patt_mat=patt_mat(:)/norm(patt_mat(:));
  %Project it away from all the connectivity matrices
  connTensor_deflated=connTensor_deflated...
                -patt_mat*(patt_mat'*connTensor_deflated); 
  %Optional: project away a related component to make deflation stronger
  if doubledefl
    patt_mat_alt=emax*emax'+emin*emin';
    patt_mat_alt=patt_mat_alt(:)/norm(patt_mat_alt(:));
    connTensor_deflated=connTensor_deflated...
                -patt_mat_alt*(patt_mat_alt'*connTensor_deflated); 
  end


  %Optional: PLOT DIAGNOSTICS 
  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

  if diagnostics & ~random & ~matrixtoobig
    
    %SHOW HOW GOOD RANK-ONE AND RANK-TWO APPROXIMATIONS ARE
    figure(1), clf
    dD=diag(D);
    bar(dD(Dord));
    title('Eigenspectrum of matrix obtained by matrix PCA')
    figure(2), clf
    imagesc(matrixpc_mat);colorbar
    title('Matrix principal component')
    figure(3), clf
    Cdiff_appr2=D(ind1,ind1)*E(:,ind1)*E(:,ind1)'...
               +D(ind2,ind2)*E(:,ind2)*E(:,ind2)';
    imagesc(Cdiff_appr2);colorbar
    title('Two-rank approx of matrix PC')
    %PLOT TIME COURSES 
    figure(4), clf
    %Add mean back to matrix PC time course
       tc_matrixpca_withmean=tc_matrixpca+...
            real(matrixpc_vec'*connTensor_mean*ones(1,size(connTensor,2)));
    plot(1:windows,TC(component,:)/std(TC(component,:)),'b',...
        1:windows,sign(TC(component,:)*tc_matrixpca_withmean')*...
         tc_matrixpca_withmean/std(tc_matrixpca_withmean),'g')
    title('Time course: green=matrixPCA, blue=2rank(output)')
    if component<numberComps, fprintf(' Press any key...'); pause, end
  end %of plotting diagnostics section

  if ~silent, fprintf('\n'); end

end %of loop through matrix principal components ("for component")

if ~silent, fprintf('Done.\n'); end

