%***************************************************************
% Function: fapbf.m
% Date:     March 2005
% Author:   Stefan Schaal 
% Description:
%   Implements Joint-space factor analysis for regression
%   with hidden varibles, z, for high dimensional inputs.
%
%
% Inputs: unless a required variable, a "[]" can be given to 
%         receive default initializtion.
%
% X     - input data matrix (must be given)
% Y     - output data matrix (must be given)
% cyc   - maximum number of cycles of EM (default 100)
% tol   - termination tolerance (prop change in likelihood) 
%         (default 1.e-6)
% wx    - initial guess of the factor loadings for the inputs 
%         (default ones())
% wz    - initial guess of the regression vector 
%         (default randn())
% psix  - initial guess of the noise variance in the inputs 
%         (default rand())
% psiz  - initial guess of the noise variance in the hidden 
%         variables (default rand())
% sigma2- initial guess of the variance of the regression 
%         output (default 1)
%
% Outputs:
%
% wx         - factor loadings for inputs
% wz         - factor loadings for hidden z 
% psix       - noise variance of inputs
% psiz       - noise variance of hidden variables
% sigma2     - variance of output
% LL         - log likelihood curve
% mx         - the mean vector of the input data
% my         - the mean vector of the output data
% beta_true  - the the regression vector for noise free data
% beta_noise - the the regression vector for the given noise 
%              level
%
% The generative model of the data is thus:
%            x = wx.*t+N(0,PHIx)
%            z = wz.*t+N(0,PHIz)
%            y = sum(z)+N(0,sigma2)
%
% Iterates until a proportional change < tol in the log 
% likelihood or cyc steps of EM.
%
%***************************************************************
function [wx, wz, psix, psiz, sigma2, LL, mx, my, ...
          beta_true, beta_noise]=...
          fapbf(X, Y, cyc, tol, wx, wz, psix, psiz, sigma2)

% Check that we have the inputs needed
%--------------------------------------
if (nargin<2),
  error('X and Y must be given to the function');
end;

% Pre-process data to be mean zero with unit variance
%----------------------------------------------------
[n,d] = size(X);
mx    = mean(X,1)';
my    = mean(Y);
X     = X-repmat(mx',n,1);
Y     = Y-my;
stdX  = std(X,1)';
stdY  = std(Y,1);
X     = X./repmat(stdX',n,1);
Y     = Y/stdY;


% Initialization of parameters 
%-------------------------------------------------------
% Default: 
% Aassume all signal in x and y is noise, and wx & wz are
% essentially zero

% Check user-supplied initialization
if (~exist('cyc') | isempty(cyc)),
  cyc = 100;
end;
if (~exist('tol') | isempty(tol)),
 tol = 1.e-6;
end;
if (~exist('wx') | isempty(wx)),
  wx = 1+randn(d,1)*0.1;
end;
if (~exist('wz') | isempty(wz)),
  wz = (X'*Y);
  wz = wz/norm(wz);
end;
if (~exist('psix') | isempty(psix)),
  psix = var(X)'; 
end;
if (~exist('psiz') | isempty(psiz)),
  psiz = 1+randn(d,1)*0.1;
end;
if (~exist('sigma2') | isempty(sigma2)),
  sigma2 = var(Y);
end;

% Initialize some handy variables and start iterating
loglik           = -1.e10;
LL               = [];
const            = d/2*log(2*pi);
count_violations = 0;
var_thres        = 1.e-4;

% EM Algorithm
%------------------------------------------------------------
for i=1:cyc;
  % some handy pre-computations
  one_over_sigma2 = 1./sigma2;
  one_over_psiz   = 1./psiz;
  one_over_psix   = 1./psix;
  wz_over_psiz    = wz.*one_over_psiz;
  wx_over_psix    = wx.*one_over_psix;

  r = ones(d,1)+wx.*wx_over_psix;
  p = psiz+wz.^2./r;
  s = wz_over_psiz./(r+wz.*wz_over_psiz);
  sum_p = sum(p);
  p_over_sigm2_plus_sum_p = p./(sigma2+sum_p);

  %---------------------------------------------------------
  % E-step
  %---------------------------------------------------------

  % An efficient computation of the posterior mean of Z and T
  %---------------------------------------------------------
  T1 = Y*((p-p_over_sigm2_plus_sum_p*sum_p)*one_over_sigma2)';
  T2 = repmat((p.*s.*wx_over_psix)',n,1).*X;
  T2 = T2 - sum(T2,2)*p_over_sigm2_plus_sum_p';
  T3 = T1.*repmat(s',n,1);
  T4 = repmat((wx_over_psix./(r+wz.*wz_over_psiz))',n,1).*X ...
       + T2.*repmat(s',n,1);
  E_Z  = T1+T2;
  E_T  = T3+T4;
  
  % Compute the necessary expectations
  %---------------------------------------------------------
  E_ZmTm = repmat( ...
        ((p-p.*p_over_sigm2_plus_sum_p).*s)',n,1) ...
        + E_Z.*E_T;
  E_Tm2 = repmat((1./(r+wz.*wz_over_psiz)+...
        (p-p.*p_over_sigm2_plus_sum_p).*(s.^2))',n,1) ...
        + E_T.^2;
  E_Zm2 = repmat((p-p.*p_over_sigm2_plus_sum_p)',n,1) ...
        + E_Z.^2;
  E_1ZZ1 = sum_p*(1-sum(p_over_sigm2_plus_sum_p)) ...
        + sum(E_Z,2).^2;


  % Monitor the lower bound of complete log likelihood
  %---------------------------------------------------------
  if mod(i,10)==1,
    % Compute posterior covariance of Z and T
    Omega_inv = [ones(d)./sigma2+diag(1./psiz) ...
        -diag(wz./psiz);...
        -diag(wz./psiz) ... 
         diag(1+(wx.^2)./psix+wz.*wz./psiz)];

    oldloglik = loglik;
    loglik = ...
     - n/2*log(sigma2) ...
     - 0.5/sigma2*sum(Y.^2-2*Y.*sum(E_Z,2)+E_1ZZ1) ...
     + sum(-n/2*log(psiz) - 0.5./psiz.*...
       sum(E_Zm2 - 2*repmat(wz',n,1).*E_ZmTm ...
          + repmat((wz.^2)',n,1).*E_Tm2,1)') ...
     + sum(-n/2*log(psix) - 0.5./psix.*...
       sum(X.^2 - 2*repmat(wx',n,1).*X.*E_T ...
          + repmat((wx.^2)',n,1).*E_Tm2,1)')...
     - 0.5*sum(sum(E_Tm2))...
     - 0.5*log(det(Omega_inv));

    if (abs(loglik-oldloglik)/abs(loglik) < tol | ...
            ~finite(loglik))
      disp( sprintf(...
        'Likelihood ratio tolerance reached: (%g)-(%g) < %g',...
         loglik,oldloglik,tol));
      break;
    end;
  end

  % Debugging statements
  if (mod(i,500)==1),
    disp(sprintf('Iter %4i: lik %g',i,loglik));
  end;
  LL(i,:) = loglik;

  %---------------------------------------------------------
  % M-step
  %---------------------------------------------------------
  wz     = (sum(E_ZmTm,1)./sum(E_Tm2,1))';
  wx     = (sum(X.*E_T,1)./sum(E_Tm2,1))';
  sigma2 = mean(Y.^2-2*Y.*sum(E_Z,2)+E_1ZZ1,1);
  psiz   = mean(E_Zm2-2*repmat(wz',n,1).*E_ZmTm + ...
           repmat((wz.^2)',n,1).*E_Tm2,1)';
  psix   = mean(X.^2-2*repmat(wx',n,1).*X.*E_T + ...
           repmat((wx.^2)',n,1).*E_Tm2,1)';

  % limit the minimum of a variance
  if (sigma2 < var_thres)
    sigma2 = var_thres;
  end
  inds = find(psiz < var_thres);
  psiz(inds) = var_thres;
  inds = find(psix < var_thres);
  psix(inds) = var_thres;

end % for i=1:cyc

% Some summary statements
if (count_violations > 0),
  fprintf('--> %d violations in the loglik increase\n', count_violations);
end
fprintf('Iter %4i: lik %g \n',i,loglik);

% Fix all parameters due to unit variance scaling
%------------------------------------------------------------
wx     = wx.*stdX;
wz     = wz.*stdY;
psix   = psix.*(stdX.^2);
psiz   = psiz*(stdY^2);
sigma2 = sigma2*(stdY^2);

% Compute the true regression estimate
%------------------------------------------------------------
temp1 = 1./((wx.^2)./psix + (wz.^2)./psiz + 1);
temp2 = ones(d,1)*ones(1,d) ./ sigma2 + diag(1./psiz) -...
        diag(((wz./psiz).^2).*temp1);
temp3 = ones(d,1)*ones(1,d) ./ sigma2 + diag(1./psiz);

% Compute the true regression and noise-optimal regression 
% coefficients 
beta_true = sigma2 / ...
            (sigma2 - ones(1,d)*inv(temp3)*ones(d,1))*...
            ones(1,d)*inv(temp3)*diag(wz./wx./psiz);
beta_true = beta_true';

beta_noise = sigma2 / ...
            (sigma2 - ones(1,d)*inv(temp2)*ones(d,1))*...
            ones(1,d)*inv(temp2) * ...
            diag((wz./psiz).*temp1.*(wx./psix));
beta_noise = beta_noise';
