%%% This program generates three sets of data drawn from three different
%%% multivariate normal distributions. The data is collected in the three
%%% matrices A, B, C, which each hold "m" columns - one column for each
%%% sample. This script also contains code that estimates the mean and
%%% variances for each test case.
%%%
%%% P.G. Martinsson, April 25, 2016

clear 
rng('default')
rng(0)

%%% Set parameters for matrix A:
m     = 1000;    % Number of samples drawn.
mu    = [0;0];   % Vector of means.
ss    = [3;1];   % Each covariances.
u1    = [1;2];   % Direction of first principal component.
u2    = [-2;1];  % Direction of second principal component.
[U,~] = qr([u1,u2]); % Generate ON basis by doing G-S on [u1,u2].

%%% Draw the m samples:
A  = mu*ones(1,m) + U*diag(ss)*randn(2,m);

%%% Estimate the means:
mu_est = mean(A,2);

%%% Estimate the covariance matrix:
AA = A - mu_est*ones(1,m);
SS = (1/(m-1))*(AA*AA');

%%% Compute the spectrum of the covariance matrix:
[U_est,D] = eig(SS);
ss_est    = sqrt(diag(D));
[~,ind]   = sort(-ss_est);
ss_est    = ss_est(ind);
U_est     = U_est(:,ind);

%%% Plot the results
figure(1)
U_scaled     = U*diag(ss);
U_est_scaled = U_est*diag(ss_est);
han = ...
plot(A(1,:),A(2,:),'g.',...
     mu(1)+[0,0;U_scaled(1,:)],mu(2)+[0,0;U_scaled(2,:)],'b',...
     mu(1)-[0,0;U_scaled(1,:)],mu(2)-[0,0;U_scaled(2,:)],'b',...
     mu_est(1)+[0,0;U_est_scaled(1,:)],mu_est(2)+[0,0;U_est_scaled(2,:)],'r',...
     mu_est(1)-[0,0;U_est_scaled(1,:)],mu_est(2)-[0,0;U_est_scaled(2,:)],'r',...
     'LineWidth',3);
legend(han([1,2,6]),'samples','exact distribution','estimated distribution')
axis equal

fprintf(1,'=== Matrix A\n')
fprintf(1,'Exact     means     = '); fprintf(1,'%15.10f  ',mu);     fprintf(1,'\n');
fprintf(1,'Estimated means     = '); fprintf(1,'%15.10f  ',mu_est); fprintf(1,'\n');
fprintf(1,'Exact     variances = '); fprintf(1,'%15.10f  ',ss);     fprintf(1,'\n');
fprintf(1,'Estimated variances = '); fprintf(1,'%15.10f  ',ss_est); fprintf(1,'\n');

%%% Matrix B
m  = 1000;
mu = [1;1;1];
u1 = [1;1;2];
u2 = [1;-1;0];
[U,~] = qr([u1,u2]);
ss = [3;1;0.01];

%%% Draw the m samples:
B  = mu*ones(1,m) + U*diag(ss)*randn(3,m);

%%% Estimate the means:
mu_est = mean(B,2);

%%% Estimate the covariance matrix:
BB = B - mu_est*ones(1,m);
SS = (1/(m-1))*(BB*BB');

%%% Compute the spectrum of the covariance matrix:
[U_est,D] = eig(SS);
ss_est    = sqrt(diag(D));
[~,ind]   = sort(-ss_est);
ss_est    = ss_est(ind);
U_est     = U_est(:,ind);

%%% Plot the results
figure(2)
U_scaled     = U*diag(ss);
U_est_scaled = U_est*diag(ss_est);
han = ...
plot3(B(1,:),B(2,:),B(3,:),'g.',...
      mu(1)+[0,0,0;U_scaled(1,:)],mu(2)+[0,0,0;U_scaled(2,:)],mu(3)+[0,0,0;U_scaled(3,:)],'b',...
      mu(1)-[0,0,0;U_scaled(1,:)],mu(2)-[0,0,0;U_scaled(2,:)],mu(3)-[0,0,0;U_scaled(3,:)],'b',...
      mu_est(1)+[0,0,0;U_est_scaled(1,:)],mu_est(2)+[0,0,0;U_est_scaled(2,:)],mu_est(3)+[0,0,0;U_est_scaled(3,:)],'r',...
      mu_est(1)-[0,0,0;U_est_scaled(1,:)],mu_est(2)-[0,0,0;U_est_scaled(2,:)],mu_est(3)-[0,0,0;U_est_scaled(3,:)],'r',...
      'LineWidth',3);
legend(han([1,2,8]),'samples','exact distribution','estimated distribution')
axis equal

fprintf(1,'=== Matrix B\n')
fprintf(1,'Exact     means     = '); fprintf(1,'%15.10f  ',mu);     fprintf(1,'\n');
fprintf(1,'Estimated means     = '); fprintf(1,'%15.10f  ',mu_est); fprintf(1,'\n');
fprintf(1,'Exact     variances = '); fprintf(1,'%15.10f  ',ss);     fprintf(1,'\n');
fprintf(1,'Estimated variances = '); fprintf(1,'%15.10f  ',ss_est); fprintf(1,'\n');

%%% Matrix C
m  = 10000;
n  = 5;
mu = rand(n,1);
[U,~,~] = qr(randn(n),0);
ss = [3;3;1;0.01*rand(n-3,1)];
gg = randn(n,m);
C  = mu*ones(1,m) + U*diag(ss)*gg;

%%% Estimate the means:
mu_est = mean(C,2);

%%% Estimate the covariance matrix:
CC = C - mu_est*ones(1,m);
SS = (1/(m-1))*(CC*CC');

%%% Compute the spectrum of the covariance matrix:
[U_est,D] = eig(SS);
ss_est    = sqrt(diag(D));
[~,ind]   = sort(-ss_est);
ss_est    = ss_est(ind);
U_est     = U_est(:,ind);

fprintf(1,'=== Matrix C\n')
fprintf(1,'Exact     means     = '); fprintf(1,'%15.10f  ',mu);     fprintf(1,'\n');
fprintf(1,'Estimated means     = '); fprintf(1,'%15.10f  ',mu_est); fprintf(1,'\n');
fprintf(1,'Exact     variances = '); fprintf(1,'%15.10f  ',ss);     fprintf(1,'\n');
fprintf(1,'Estimated variances = '); fprintf(1,'%15.10f  ',ss_est); fprintf(1,'\n');

save testmatrices.mat A B C