%***********************************************************
% File: 	test.m
% Date: 	July 26, 2007 
% Author: 	Jo-Anne Ting
% Description: 
%   Creates synthetic data sets and runs
%   Variational Bayesian Least Squares (VBLS) 
%   algorithm as written in Ting et al. (2005). 
%
%************************************************************
close all;
clear all;


%------------------------------------------------------------
% Set values of indicator variables for test settings 
%------------------------------------------------------------
numTrials = 1;


% Set indicator variable to 1 if want to run algorithm 
VBLS_ON   = 1;
RR_ON     = 1;
STEP_ON   = 1;
PLS_ON    = 1;
LASSO_ON  = 1;

if VBLS_ON == 1
  totalTrainingNMSE_vbls = [];
  totalPredictionNMSE_vbls = [];
end
if RR_ON == 1
  totalTrainingNMSE_rr = [];
  totalPredictionNMSE_rr = [];
end
if STEP_ON == 1
  totalTrainingNMSE_step = [];
  totalPredictionNMSE_step = [];
end
if PLS_ON == 1
  totalTrainingNMSE_pls = [];
  totalPredictionNMSE_pls = [];
end
if LASSO_ON == 1
  totalTrainingNMSE_lasso = [];
  totalPredictionNMSE_lasso = [];
end



% Repeat this for a certain number of trials
for tt=1:numTrials

%------------------------------------------------------------
% 1. Create a synthetic dataset
%
% Create a synthetic data set with redundant and irrelevant 
% dimensions. In the example below, we create a dataset with 
% 100-dimensional inputs and a scalar output, having 10 
% relevant input dimensions, r redundant input dimensions 
% (for example, input data is simply all zeros) and u irrelevant 
% input dimensions (i.e. the corresponding regression coefficient 
% for that dimension is 0). The training dataset has noise in the 
% outputs. The test dataset has the same characteristics as the 
% training dataset, except it has no noise in the outputs.
%------------------------------------------------------------

% Data characteristics
%---------------------
numSamples       = 1000;  % Number of samples
numTestSamples   = 20;    % Number of test samples
numDim           = 100;   % Number of input dimensions
r                = 30;    % Number of redundant input dimensions
u                = 60;    % Number of irrelevant input dimensions
numRelevDim      = numDim - r - u;

% Set the noise level in the outputs
r2               = 0.9;   % coefficient of determination

% Create noisy training data (noise in outputs only)
%-----------------------------------------------------

% First, create a rotation matrix
T                = rand(numRelevDim);  
T                = diag(abs(randn(numRelevDim,1))+1)*orth(T); 

% Create mean-zero training data
coreTrainX       = randn(numSamples, numRelevDim)*T;
coreTrainX       = coreTrainX - repmat(mean(coreTrainX,1), numSamples, 1);
trueB            = randn(numRelevDim,1)*100;
trainY           = coreTrainX*trueB;
varNoise         = var(trainY)*(1/r2 - 1);
trainY           = trainY + randn(numSamples,1)*sqrt(varNoise);
trainY           = trainY - mean(trainY);

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

% Apply the rotation matrix to inputs so that redundant dimensions
% contain linear combinations of noisy relevant data
trainX = [coreTrainX zeros(numSamples, r) randn(numSamples, u)]*R;

% Create noiseless test data
%----------------------------

% Create mean-zero noiseless test data
coreTestX         = randn(numTestSamples, numRelevDim)*T;
coreTestX         = coreTestX -repmat(mean(coreTestX, 1), numTestSamples, 1);
testY             = coreTestX*trueB; 
testY             = testY - mean(testY);

% Use the same rotation matrix as for training input to create correlations
testX              = [coreTestX zeros(numTestSamples, r) ...
                      randn(numTestSamples, u)]*R;

% Pre-process training data to be mean zero and variance 1
%---------------------------------------------------------
trainYP = (trainY - mean(trainY)) ./ std(trainY);
trainXP = trainX - repmat(mean(trainX, 1), numSamples, 1);
stdtrainX = std(trainX, 1);
for m=1:numDim
  if stdtrainX(m) ~= 0
    trainXP(:,m) = trainXP(:,m) ./ repmat(stdtrainX(m), numSamples, 1);
  end
end

% Pre-process test data too
testYP = (testY - mean(trainY)) ./ std(trainY);
testXP = testX - repmat(mean(trainX, 1), numTestSamples, 1);
for m=1:numDim
  if stdtrainX(m) ~= 0
     testXP(:, m) = testXP(:, m) ./ repmat(stdtrainX(m), numTestSamples, 1);
  else
  end
end



%------------------------------------------------------------
%                        RUN ALGORITHMS 
%------------------------------------------------------------

%------------------------------------------------------------
% 2. Run vbls
%------------------------------------------------------------

% Algorithm settings that may be modified by user: 
if VBLS_ON == 1
  options.noise         = 1;        % initial output noise variance
  options.threshold     = 1e-10;     % threshold for convergence
  options.numIterations = 10000;    % max number of EM iterations
  vbls_result           = vbls(trainX, trainY, options);
end

%------------------------------------------------------------
% 3. Run ridge regression
%------------------------------------------------------------
if RR_ON == 1
  %ridge_b = pinv(trainX'*trainX)*trainX'*trainY;
  ridge_b = inv(trainX'*trainX + eye(numDim)*1e-10)*trainX'*trainY;
end

%------------------------------------------------------------
% 4. Run stepwise regression 
%------------------------------------------------------------
if STEP_ON == 1
  [step.bb, step.se, step.pval, step.inmodel, step.stats] = ...
      stepwisefit(trainX, trainY, 'display', 'off');
  step_b = [step.bb.*step.inmodel'; step.stats.intercept]; 
end

%------------------------------------------------------------
% 5. Run Partial Least Squares regression on normalized data 
%------------------------------------------------------------
if PLS_ON == 1

  % Find optimal number of projections using loo-cv (i.e. r2cross)
  pls.lowest_r2cross = 0; 
  for m=1:numDim
    [pls.beta, pls.U, pls.W, pls.mx, pls.my, pls.r2, pls.r2cross, ...
     pls_yPred] = pls_v1(trainXP, trainYP, m, testXP);
     if m == 1
       pls.lowest_r2cross = pls.r2cross;
     else
        if (pls.r2cross < pls.lowest_r2cross)
           pls.lowest_r2cross = pls.r2cross;
        end
     end
  end
  pls.m = m;
  [pls.beta, pls.U, pls.W, pls.mx, pls.my, pls.r2, pls.r2cross, ...
   pls_yPred] = pls_v1(trainXP, trainYP, pls.m, testXP);
end


%------------------------------------------------------------
% 6. Run LASSO regression on normalized data
%------------------------------------------------------------
if LASSO_ON == 1
  [lasso_b, lasso.msr] = arrfit(trainXP, trainYP, 5, 1e-4);
  % Count number of relevant inputs
  lasso.indices = find(lasso_b ~= 0);
  lasso.neurons = lasso.indices; 
end

%------------------------------------------------------------
% 7. Calculate prediction (test) and training error 
%------------------------------------------------------------

disp('----------------------------------------------------');
disp(sprintf('Training and Test nMSE values for Trial %d', tt));

% Print training error 
%-----------------------
if VBLS_ON == 1
  vbls_yTrain = trainXP*vbls_result.b_mean;
  vbls_trainnmse = mean((trainYP - vbls_yTrain).^2) / var(trainYP);
  disp(sprintf('Training nmse for VBLS is %f', vbls_trainnmse));
end
if RR_ON == 1
  rr_yTrain = trainX*ridge_b;
  rr_trainnmse = mean((trainY - rr_yTrain).^2) / var(trainY);
  disp(sprintf('Training nmse for Ridge Regression is %f', rr_trainnmse));
end
if STEP_ON == 1
  step_yTrain = [trainX ones(numSamples,1)] *step_b;
  step_trainnmse = mean((trainY - step_yTrain).^2) / var(trainY);
  disp(sprintf('Training nmse for Stepwise Regression is %f', step_trainnmse));
end
if PLS_ON == 1
  pls_trainnmse = 1 - pls.r2;
  disp(sprintf('Training nmse for PLS Regression is %f', pls_trainnmse));
end
if LASSO_ON == 1
  lasso_yTrain = trainXP*lasso_b;
  lasso_trainnmse = mean((trainYP - lasso_yTrain).^2) / var(trainYP); 
  disp(sprintf('Training nmse for LASSO Regression is %f', lasso_trainnmse));
end

% Print prediction error on noiseless test data
%-----------------------------------------------
if VBLS_ON == 1
  vbls_yPred = testXP*vbls_result.b_mean;
  vbls_testnmse = mean((testYP - vbls_yPred).^2) / var(testYP);
  disp(sprintf('Prediction nmse for VBLS is %f', vbls_testnmse));
end
if RR_ON == 1
  rr_yPred = testX*ridge_b;
  rr_testnmse = mean((testY - rr_yPred).^2) / var(testY);
  disp(sprintf('Prediction nmse for Ridge Regression is %f', rr_testnmse));
end
if STEP_ON == 1
  step_yPred = [testX ones(numTestSamples,1)] *step_b;
  step_testnmse = mean((testY - step_yPred).^2) / var(testY);
  disp(sprintf('Prediction nmse for Stepwise Regression is %f', step_testnmse));
end
if PLS_ON == 1
   pls_testnmse = mean((testYP - pls_yPred).^2) ./ var(testYP);
   pls.numDim = pls.m;
   pls.neuronData = pls.U*pls.U'*ones(numDim,1);  
  disp(sprintf('Prediction nmse for PLS Regression is %f', pls_testnmse));
end
if LASSO_ON == 1
  lasso_yPred = testXP*lasso_b;
  lasso_testnmse = mean((testYP - lasso_yPred).^2) / var(testYP); 
  disp(sprintf('Prediction nmse for LASSO Regression is %f', lasso_testnmse));
end


%------------------------------------------------------------
% 8. Store prediction and training error results
%------------------------------------------------------------
if VBLS_ON == 1
  totalTrainingNMSE_vbls = [totalTrainingNMSE_vbls; vbls_trainnmse];
  totalPredictionNMSE_vbls = [totalPredictionNMSE_vbls; vbls_testnmse];
end
if RR_ON == 1
  totalTrainingNMSE_rr = [totalTrainingNMSE_rr; rr_trainnmse];
  totalPredictionNMSE_rr = [totalPredictionNMSE_rr; rr_testnmse];
end
if STEP_ON == 1
  totalTrainingNMSE_step = [totalTrainingNMSE_step; step_trainnmse];
  totalPredictionNMSE_step = [totalPredictionNMSE_step; step_testnmse];
end
if PLS_ON == 1
  totalTrainingNMSE_pls = [totalTrainingNMSE_pls; pls_trainnmse];
  totalPredictionNMSE_pls = [totalPredictionNMSE_pls; pls_testnmse];
end
if LASSO_ON == 1
  totalTrainingNMSE_lasso = [totalTrainingNMSE_lasso; lasso_trainnmse];
  totalPredictionNMSE_lasso = [totalPredictionNMSE_lasso; lasso_testnmse];
end


end  % end of for tt=1:numTrials





%------------------------------------------------------------
% 9. Display averaged error results
%------------------------------------------------------------
disp('====================================================')
disp(sprintf('Average Training and Test nMSE Results for r2 = %g, r = %d, u = %d', ...
     r2, r, u))

if VBLS_ON == 1
   disp(sprintf('Average training nmse for VBLS is %f', mean(totalTrainingNMSE_vbls)));
end
if RR_ON == 1
   disp(sprintf('Average training nmse for RR is %f', mean(totalTrainingNMSE_rr)));
end
if STEP_ON == 1
   disp(sprintf('Average training nmse for STEP is %f', mean(totalTrainingNMSE_step)));
end
if PLS_ON == 1
   disp(sprintf('Average training nmse for PLS is %f', mean(totalTrainingNMSE_pls)));
end
if LASSO_ON == 1
   disp(sprintf('Average training nmse for LASSO is %f', mean(totalTrainingNMSE_lasso)));
end

if VBLS_ON == 1
   disp(sprintf('Average prediction nmse for VBLS is %f', mean(totalPredictionNMSE_vbls)));
end
if RR_ON == 1
   disp(sprintf('Average prediction nmse for RR is %f', mean(totalPredictionNMSE_rr)));
end
if STEP_ON == 1
   disp(sprintf('Average prediction nmse for STEP is %f', mean(totalPredictionNMSE_step)));
end
if PLS_ON == 1
   disp(sprintf('Average prediction nmse for PLS is %f', mean(totalPredictionNMSE_pls)));
end
if LASSO_ON == 1
   disp(sprintf('Average prediction nmse for LASSO is %f', mean(totalPredictionNMSE_lasso)));
end
