%************************************************************
% File:         run_parameterEst_test.m
% Date:         January 8, 2008
% Author:       Jo-Anne Ting
% Description:
%  Creates synthetic data sets and runs the Bayesian
%  de-noising algorithm to estimate the parameter
%  estimates in a high-dimensional linear regression
%  problem where the input data is noisy.
%
%  Compares performances of other algorithms such as:
%  i) Ordinary Least Squares (OLS)
%  ii) Stepwise regression (STEP) (Draper & Smith, 1981)
%  iii) Partial Least Squares (PLS) regression (Wold, 1975)
%  iv) Least Absolute Shrinkage & Selection Operator 
%      (LASSO) (Tibshirani, 1996) 
%  v) Joint-space Factor Analysis for regression 
%     (Massey, 1965)
%
%************************************************************
clear all;
close all;

%------------------------------------------------------------
% 1. Initialize test settings  
%------------------------------------------------------------

% Parameters for data set
%randn('seed',0);
%rand('seed',0);

totalNumTrials = 1;    % number of trials to average over
snrx           = 2;    % signal-to-noise ratio in input data
snry           = 5;    % signal-to-noise ratio in output data

d              = 100;  % input dimensionality of data
r              = 90;   % number of redundant input dimensions
u              = 0;    % number of irrelevant input dimensions
n              = d*50; % number of data samples

% Set indicator variable to 1 if want to run algorithm
OLS_ON         = 1;     % Ordinary Least Squares regression
STEP_ON        = 1;     % Stepwise regression
PLS_ON         = 1;     % Partial Least Squares regression
LASSO_ON       = 1;     % LASSO regression
JFA_ON         = 1;     % Joint-space factor analysis 
BAYES_ON       = 1;     % Bayesian de-noising regression

% For collecting results over multiple trials
if OLS_ON == 1
  totalNMSE_ols_train = []; totalNMSE_ols_test = []; 
end
if STEP_ON == 1
  totalNMSE_step_train = []; totalNMSE_step_test = []; 
end
if PLS_ON == 1
  totalNMSE_pls_train = []; totalNMSE_pls_test = []; 
  num_pls_dim = [];
end
if LASSO_ON == 1
  totalNMSE_lasso_train = []; totalNMSE_lasso_test = []; 
end
if JFA_ON == 1
  totalNMSE_jfa_train = []; totalNMSE_jfa_test = []; 
end
if BAYES_ON == 1
  totalNMSE_bayes_train = []; totalNMSE_bayes_test = []; 
end

%------------------------------------------------------------
% 2. Run algorithms: for each trial, a different data set
%    is created.
%------------------------------------------------------------

for numTrials = 1:totalNumTrials

  % Create a synthetic data set 
  %----------------------------------------------------------
  % Create correlated data
  T = rand(d-r-u);
  T = diag(abs(randn(d- r-u, 1)) + 1)*orth(T); 
 
  % Create input (X) and output (Y) data 
  X = [randn(n, d-r-u)*T zeros(n, r) randn(n, u)];
  
  % Irrelevant and redundant dimensions in b are set to 0 
  b = [(1:d - u - r)'; zeros(u + r, 1)];
  Y = X * b;

  % Apply noise to inputs X and outputs Y
  stdx = std(X)'*sqrt(1/snrx);
  stdy = std(Y,1)*sqrt(1/snry);

  % Add noise to relevant input dimensions in the training data
  X = X + [randn(n, d-r-u) zeros(n, r+u)] .* ...
      repmat(stdx', n, 1);
  Y = Y + randn(n, 1)*stdy;

  % Create the rotation matrix, R, over redundant dimensions
  for i=1:r
    for j=1:d-r- u
      A(j,i) = rand(1);
    end
    totalSum = sum(A(1:d-r- u, i));
    A(1:d-r-u,i) = A(1:d-r-u, i) ./ totalSum;
  end
  R = eye(d);
  if r ~= 0
    R(1:d-r-u, d-r-u+1:d-u) = A;
  end

  % Apply rotation matrix R to inputs so that redundant 
  % dimensions contain linear combinations of noisy 
  % relevant data
  X = X * R;

  % Check condition number & rank
  %cond(X'*X)
  %rank(X'*X)

  % Create noise-less test data
  % (input test data is rotated, with no noise)
  Xt = [randn(n, d-r-u)*T zeros(n, r) randn(n, u)];
  Yt = Xt * b;
  Xt = Xt * R;


  % Run OLS regression 
  %----------------------------------------------------------
  if OLS_ON == 1
  b_ols = pinv(X'*X)*X'*Y;
  mx = mean(X,1)';
  my = mean(Y);

  % Calculate nMSE on training and test data
  Yp_ols_train = (X - repmat(mx',n,1))*b_ols + my;
  nMSE_ols_train = mean((Yp_ols_train - Y).^2) / var(Y);
  Yp_ols_test = (Xt - repmat(mx',n,1))*b_ols + my;
  nMSE_ols_test = mean((Yp_ols_test - Yt).^2) / var(Yt);
  end

  % Run stepwise regression 
  %----------------------------------------------------------
  if STEP_ON == 1
  [b_step, r2_step, Yp_step] = stepwise_regression(X, Y, Xt);
  
  % Calculate nMSE on training and test data
  nMSE_step_train = 1 - r2_step;
  nMSE_step_test = mean((Yp_step - Yt).^2) / var(Yt); 
  end

  % Run PLS regression 
  %----------------------------------------------------------
  if PLS_ON == 1
  % Do a search by progressively increasing the number of
  % projections until r2 stops increasing
  last_r2cross_pls = 0;
  last_r2_pls = 0;
  for k=1:d
    [b_pls, U, W, mx, my, r2_pls, r2cross_pls] = ...
      pls_v1(X, Y, k);
    if (r2cross_pls < last_r2cross_pls)
      k = k-1;
      break;
    end
    last_r2cross_pls = r2cross_pls; 
  end
  % Store the number of projections chosen
  num_pls_dim = [k; num_pls_dim];
  [b_pls, U, W, mx, my, r2_pls, r2cross_pls, Yp_pls] = ...
      pls_v1(X, Y, k, Xt);

  % Calculate nMSE on training and test data
  nMSE_pls_train = 1 - r2_pls;
  nMSE_pls_test = mean((Yp_pls - Yt).^2) / var(Yt);
  end

  % Run LASSO regression 
  %----------------------------------------------------------
  if LASSO_ON == 1
  [b_lasso, msr] = arrfit(X, Y, 0.01, 1e-4);

  % Calculate nMSE on training and test data
  Yp_lasso_train = X*b_lasso;
  nMSE_lasso_train = mean((Yp_lasso_train - Y).^2) / var(Y);
  Yp_lasso = Xt*b_lasso;
  nMSE_lasso_test = mean((Yp_lasso - Yt).^2) / var(Yt);
  end

  % Run Joint-space Factor Analysis for regression 
  %----------------------------------------------------------
  if JFA_ON == 1
  iter = 2000;
  threshold = 1e-18;
  [wx_jfa, wz_jfa, psix_jfa, psiz_jfa, sigma2_jfa, ...
   LL_jfa, mx, my, bt_jfa, bn_jfa]=...
       fapbf(X, Y, iter, threshold, [], [], [], []);

  % Calculate nMSE on training and test data
  Yp_jfa_train = (X - repmat(mx',n, 1))*bt_jfa + my;
  nMSE_jfa_train = mean((Yp_jfa_train - Y).^2) / var(Y);
  Yp_jfa = (Xt - repmat(mx', n, 1))*bt_jfa + my;
  nMSE_jfa_test = mean((Yp_jfa - Yt).^2) / var(Yt);
  end

  % Run Bayesian de-noising regression 
  %----------------------------------------------------------
  if BAYES_ON == 1
  options.noise = 1;           % initial output noise variance
  options.threshold = 1e-18;   % threshold for convergence
  options.numIter = 5000;      % max number of EM iterations
  options.printFreq = 500;     % frequency of debugging prints
  options.clipDims = 1;        % clip irrelevant dimensions
  noisyBLR_result = noisyBLR(X, Y, options);  

  % Calculate nMSE on training and test data
  Yp_bayes_train = (X - repmat(noisyBLR_result.meanX', n, 1)) ...
                   * noisyBLR_result.beta_true ...
                   + noisyBLR_result.meanY;
  nMSE_bayes_train = mean((Yp_bayes_train - Y).^2) / var(Y);
  Yp_bayes = (Xt - repmat(noisyBLR_result.meanX', n, 1)) ...
                   * noisyBLR_result.beta_true ...
                   + noisyBLR_result.meanY;
  nMSE_bayes_test = mean((Yp_bayes - Yt).^2) / var(Yt);
  end



  % Store results in "totalNMSE" vectors 
  %----------------------------------------------------------
  if OLS_ON == 1
  totalNMSE_ols_train = [totalNMSE_ols_train; nMSE_ols_train]; 
  totalNMSE_ols_test = [totalNMSE_ols_test; nMSE_ols_test]; 
  end
  if STEP_ON == 1
  totalNMSE_step_train = [totalNMSE_step_train; nMSE_step_train];
  totalNMSE_step_test = [totalNMSE_step_test; nMSE_step_test];
  end
  if PLS_ON == 1
  totalNMSE_pls_train = [totalNMSE_pls_train; nMSE_pls_train];
  totalNMSE_pls_test = [totalNMSE_pls_test; nMSE_pls_test];
  end
  if LASSO_ON == 1
  totalNMSE_lasso_train = [totalNMSE_lasso_test; nMSE_lasso_train]; 
  totalNMSE_lasso_test = [totalNMSE_lasso_test; nMSE_lasso_test]; 
  end
  if JFA_ON == 1
  totalNMSE_jfa_train = [totalNMSE_jfa_train; nMSE_jfa_train]; 
  totalNMSE_jfa_test = [totalNMSE_jfa_test; nMSE_jfa_test]; 
  end
  if BAYES_ON == 1
  totalNMSE_bayes_train = [totalNMSE_bayes_train; nMSE_bayes_train]; 
  totalNMSE_bayes_test = [totalNMSE_bayes_test; nMSE_bayes_test]; 
  end

end  % numTrials


%------------------------------------------------------------
% 3. Display average training & test nMSE results 
%------------------------------------------------------------

% Average training nMSE values
disp('---------------------------');
disp('Average training nMSE values:');
if OLS_ON == 1
  disp(sprintf('OLS-training-nmse: %f', ...
       mean(totalNMSE_ols_train) ));
end
if STEP_ON == 1
  disp(sprintf('STEP-training-nmse: %f', ...
       mean(totalNMSE_step_train) )); 
end
if PLS_ON == 1
  disp(sprintf('PLS-training-nmse: %f', ...
        mean(totalNMSE_pls_train) )); 
end
if LASSO_ON == 1
  disp(sprintf('LASSO-training-nmse: %f', ...
       mean(totalNMSE_lasso_train) ));
end
if JFA_ON == 1
  disp(sprintf('JFA-training-nmse: %f', ...
       mean(totalNMSE_jfa_train) )); 
end
if BAYES_ON == 1
  disp(sprintf('BAYES-training-nmse: %f', ...
       mean(totalNMSE_bayes_train) ));
end

% Average test nMSE values (on noise-free test data)
disp('---------------------------');
disp('Average test nMSE values:');
if OLS_ON == 1
  disp(sprintf('OLS-test-nmse: %f', ...
       mean(totalNMSE_ols_test) ));    
end
if STEP_ON == 1
  disp(sprintf('STEP-test-nmse: %f', ...
       mean(totalNMSE_step_test) ));
end 
if PLS_ON == 1
  disp(sprintf('PLS-test-nmse: %f', ...
       mean(totalNMSE_pls_test) )); 
end
if LASSO_ON == 1
  disp(sprintf('LASSO-test-nmse: %f', ...
       mean(totalNMSE_lasso_test) )); 
end 
if JFA_ON == 1
  disp(sprintf('JFA-test-nmse: %f', ...
       mean(totalNMSE_jfa_test) )); 
end
if BAYES_ON == 1
  disp(sprintf('BAYES-test-nmse: %f', ...
       mean(totalNMSE_bayes_test) )); 
end
