%***************************************************************
% Function: noisyBLR.m
% Date:     January 2006
% Author:   Jo-Anne Ting
% Description:
%   Implements the Bayesian Linear Regression algorithm for high
%   dimensional data with input noise as written in Ting, 
%   D'Souza & Schaal  (ICML 2006).
%
% Inputs: 
%
%  (where N = number of data samples, 
%         d = dimensionality of input data)
%
%  x                 : training input (Nxd matrix)
%  y                 : training output (Nx1 matrix)
%  options           : data structure for algorithm settings 
%  options.noise     : initial value for output data noise 
%                      variance
%  options.threshold : tolerance for termination of algorithm
%                     (proportional change in likelihood)
%  options.numIter   : maximum number of EM iterations
%                      desired
%  options.printFreq : frequency of print statements 
%                      (for debugging purposes)
%  options.clipDims  : if 1, then clip dimensions by
%                      setting corresponding regression
%                      coefficient to 0
%
% Outputs:
%
%  A data structure "result" containing the final converged
%  values of:           
%  result.E_alpha   : mean of the precision variables
%  result.a_alpha   : parameter a of the Gamma distribution
%  result.b_alpha   : parameter b of the Gamma distribution
%  result.E_wz      : mean of wz 
%  result.E_wx      : mean of wx 
%  result.sigma_wz  : variance of wz
%  result.sigma_wx  : variance of wx
%  result.sigma2    : noise variance of y
%  result.psiz      : noise variance of z
%  result.psix      : noise variance of x
%  result.meanX     : mean of training input x
%  result.meanY     : mean of training input y 
%  result.stdX      : std deviation of training input x
%  result.stdY      : std deviation of training input y 
%  result.beta_true : true regression coefficient estimate
%  result.beta_noise: noisy regression coefficient estimate 
%  result.LL        : lower bound to incomplete log likelihood
%
%***************************************************************
function [result] = noisyBLR(x, y, options);

% Check that we have the inputs needed
%--------------------------------------
if nargin < 2
  error('Insufficient inputs to vbls function');
end

if (~exist('options')|isempty(options))
  % Default settings for EM algorithm
  options.noise     = 1;
  options.threshold = 1e-12;
  options.numIter   = 5000;
  options.printFreq = 100;
end
if (isempty(options.noise))
  options.noise = 1;
end
if (isempty(options.threshold))
  options.threshold = 1e-12;
end
if (isempty(options.numIter))
  options.numIter = 5000;
end
if (isempty(options.printFreq))
  options.printFreq = 100;
end

[N,d] = size(x); 


% Pre-process data to be mean zero with unit variance
%----------------------------------------------------
meanX = mean(x, 1)';
meanY = mean(y);
stdX = std(x, 1)';
stdY = std(y, 1);
y = (y - meanY) ./ stdY;
x = x - repmat(meanX', N, 1);
for i=1:d
  if stdX(i) ~= 0
    x(:,i) = x(:,i) ./ repmat(stdX(i)', N, 1);
  end
end


% Initialization of parameters  (No changes needed here) 
%-------------------------------------------------------
a_alpha0     = 1e-8*ones(d, 1);  				
b_alpha0     = 1e-8*ones(d, 1);  	
a_alpha      = a_alpha0;
b_alpha      = b_alpha0;
E_alpha      = a_alpha ./ b_alpha;	

% Initialize wz with Partial Least Squares projections
E_wz         = (x'*y);
E_wz         = E_wz / norm(E_wz);
E_wx         = 1 + randn(d, 1)*0.1;
sigma2_wz    = 1 ./ E_alpha;  
sigma2_wx    = 1 ./ E_alpha;   
E_wz2        = sigma2_wz + E_wz.^2;
E_wx2        = sigma2_wx + E_wx.^2;

loglik       = -1.e10;  % Initial log likelihood value
LL           = [];      % keeps track of log likelihood value
% Limit variance measures to a minimum to avoid numerical errors 
var_thres    = 1e-4;   

% Initialize noise variances
sigma2       = var(y);	
psix         = var(x)';
psiz         = 1+randn(d,1)*0.1;   
% Make sure psix is not 0
inds = find(psix == 0);
psix(inds) = var_thres;


% EM Algorithm
%------------------------------------------------------------
for numIter = 1:options.numIter

  % Some handy pre-computations
  %-----------------------------
  one_over_sigma2 = 1 ./ sigma2;
  one_over_psiz = 1 ./ psiz;
  one_over_psix = 1 ./ psix;
  wz_over_psiz  = E_wz .* one_over_psiz;
  wx_over_psix  = E_wx .* one_over_psix;
  wz2_over_psiz = E_wz2 .* one_over_psiz;
  wx2_over_psix = E_wx2 .* one_over_psix;

  r = ones(d,1) + wx2_over_psix;
  p = psiz + (E_wz.^2) ./ (r + one_over_psiz.*sigma2_wz);
  s = wz_over_psiz ./ (r + wz2_over_psiz);
  sum_p = sum(p);
  p_over_sigm2_plus_sum_p = p ./ (sigma2+sum_p);


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

  % An efficient computation of the 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+wz2_over_psiz))', N, 1).*x ...
       + T2.*repmat(s', N, 1);
  E_Z = T1 + T2;
  E_T = T3 + T4;

  % Computer other expectations needed
  %----------------------------------------------------------
  E_ZmTm = repmat( ...
           ((p-p.*p_over_sigm2_plus_sum_p).*s)', N, 1 ) ...
           + E_Z.*E_T;
  E_Tm2 = repmat( (1./(r+wz2_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;
 
  % Calculate mean of precision variable
  %----------------------------------------------------------
  a_alpha = a_alpha0 + 1; 
  b_alpha = b_alpha0 + 0.5*(E_wz2 + E_wx2);
  E_alpha = a_alpha ./ b_alpha;

  % Calculate mean & variances of wz 
  %----------------------------------------------------------
  sigma2_wz = 1./(E_alpha + sum(E_Tm2, 1)'./psiz);
  inds = find(sigma2_wz < var_thres);
  sigma2_wz(inds) = var_thres;
  E_wz = sum(E_ZmTm, 1)' ./ (E_alpha.*psiz + sum(E_Tm2, 1)');
  E_wz2 = E_wz.^2 + sigma2_wz;

  % Calculate mean & variances of wx
  %----------------------------------------------------------
  sigma2_wx = 1./(E_alpha+ sum(E_Tm2, 1)'./psix);
  inds = find(sigma2_wx < var_thres);
  sigma2_wx(inds) = var_thres;
  E_wx = sum(E_T.*x, 1)' ./ (E_alpha.*psix + sum(E_Tm2, 1)');
  E_wx2 = E_wx.^2 + sigma2_wx;

  % Monitor the lower bound of complete log likelihood
  %---------------------------------------------------------
  if mod(numIter, options.printFreq) == 1
     % Computer posterior covariance of Z and T
     Omega_inv = [ones(d)./sigma2 + diag(1./psiz) ...
                  -diag(E_wz./psiz); ...
                  -diag(E_wz./psiz) ...
                  diag(1 + E_wx2./psix + E_wz2./psiz)];

     oldloglik = loglik;
     loglik = ...
       - N/2*log(sigma2) ...
       - 0.5/sigma2*sum(y.^2 - 2*y.*sum(E_Z,2) + E_1ZZ1) ...
       - N/2*sum(log(psiz)) ...
       - sum(0.5./psiz .* ...
         sum(E_Zm2 - 2*repmat(E_wz', N, 1).*E_ZmTm ...
             + repmat(E_wz2', N, 1).*E_Tm2, 1)') ...
       - N/2*sum(log(psix)) ...
       - sum(0.5./psix .* ...
         sum(x.^2 - 2*repmat(E_wx', N, 1).*x.*E_T ...
             + repmat(E_wx2', N, 1).*E_Tm2, 1)') ...
       - 0.5*sum(sum(E_Tm2)) ...
       - sum(0.5*E_alpha.*E_wz2) ...
       - sum(0.5*E_alpha.*E_wx2) ...
       - sum(a_alpha0.*log(b_alpha)) ...
       - sum(E_alpha .* b_alpha0) ...
       + 0.5*sum(log(sigma2_wz)) ...
       + 0.5*sum(log(sigma2_wx)) ...
       - 0.5*log(det(Omega_inv)) ...
       - sum(log(b_alpha));

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

  % Debugging statements
  if mod(numIter, options.printFreq) == 1
    disp(sprintf('Iter %i : loglik = %g', numIter, loglik));
    LL(numIter,:) = loglik;
  end
 
  %---------------------------------------------------------
  % M-step
  %---------------------------------------------------------
  
  sigma2 = mean(y.^2 - 2*y.*sum(E_Z, 2) + E_1ZZ1, 1); 
  psiz = mean(E_Zm2 - 2*repmat(E_wz', N, 1).*E_ZmTm ...
         + repmat(E_wz2', N, 1).*E_Tm2, 1)';
  psix = mean(x.^2 - 2*repmat(E_wx', N, 1).*x.*E_T ...
         + repmat(E_wx2', N, 1).*E_Tm2, 1)'; 

  % Limit the variance parameters to a minimum
  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 numIter = 1:options.numIter

% Some summary statements
fprintf('Iter %4i: loglik %g \n', numIter, loglik);


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

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

% True regression estimate (for noiseless test data)
beta_true = sigma2 / ...
           (sigma2 - ones(1,d)*inv(temp3)*ones(d,1)) * ...
           ones(1,d) * inv(temp3) * diag(E_wz./E_wx./psiz);  
beta_true = beta_true';

% Regression estimate for noisy test data
beta_noise = sigma2 / ...
            (sigma2 - ones(1,d)*inv(temp2)*ones(d,1)) *...
            ones(1,d) * inv(temp2) * ...
            diag( (E_wz./psiz).*temp1.*(E_wx./psix) );
beta_noise = beta_noise';

% ARD Determination: clip off dimension if it is too large
%------------------------------------------------------------
if options.clipDims
  % Use mean/variance of t-distributed wz to get prob. of
  % significance and to clip dimensions
  for i=1:d 
    [h,p] = ztest(E_wz(i,1), 0, sqrt(sigma2_wz(i,1)));
    prob(i,1) = p;
  end
  % Value of wz(d) is very likely 0 with prob > 5% 
  inds = find(prob > 0.05);
  beta_true(inds) = 0;
  beta_noise(inds) = 0;  
end

% Fill the output data structure 
%------------------------------------------------------------
result.E_alpha    = E_alpha;
result.a_alpha    = a_alpha;
result.b_alpha    = b_alpha;
result.E_wz       = E_wz;
result.E_wx       = E_wx;
result.sigma2_wz  = sigma2_wz;
result.sigma2_wx  = sigma2_wx;
result.sigma2     = sigma2;
result.psiz       = psiz;
result.psix       = psix;
result.meanX      = meanX;
result.meanY      = meanY;
result.stdX       = stdX;
result.stdY       = stdY;
result.beta_true  = beta_true;
result.beta_noise = beta_noise;
result.LL         = LL;
