%************************************************************
% File:         mixtureModel.m
% Date:         September 4, 2006
% Author:       Jo-Anne Ting
% Description:
%  Mixture Model for globally linear regression 
%  (assume data is generated from a Gaussian distribution
%  and outliers are generated from a uniform distribution). 
%  To simulate a uniform distribution, we use a Gaussian
%  distribution with a very wide, fixed variance.
%
% Inputs:
%
%  (where N = number of data samples,
%         d = dimensionality of input data)
%
%  X               : training input (Nxd matrix)
%  Y               : training output (Nx1 matrix)
%  n_iter          : max number of EM iterations
%  tol             : tolerance for termination of
%                    algorithm (proportional change in
%                    likelihood)
%
% Outputs:
%
% A data structure "s" containing the final converged
% values of:
%
%  s.alpha1       : probability that samples are inliers
%  s.alpha2       : probability that samples are outliers
%  s.mu1          : mean of "inlier" Gaussian component
%                   (capturing inliers)
%  s.Sigma1       : covariance of "inlier" Gaussian component
%  s.mu2          : mean of "outlier" Gaussian component
%                   (capturing outliers)
%  s.Sigma2       : covarianace of outlier Gaussian 
%                   component  
%  s.z1           : membership weights for data cloud
%  s.z2           : membership weights for outlier component
%
%************************************************************
function [s] = mixtureModel(X, Y, n_iter, tol)

X1  = [X Y];
[N,d] = size(X1);

% Setup of priors
%-------------------------------------------------------
s.alpha1 = 0.5;
s.alpha2 = 1-s.alpha1;
s.mu1    = zeros(d,1); 
s.Sigma1 = 1e3*eye(d); 
s.mu2    = zeros(d,1);
% Fix variance to be very large (for noise)
s.Sigma2 = 1e4*eye(d);  
% To avoid computational problems
SMALL    = 1e-6;

% EM-Iterations
%-------------------------------------------------------
for i=1:n_iter,

  % E-step
  %-----------------------------------------------------
  % Calculate membership weights ( Nx1 vectors)
  z1_num = s.alpha1 / ...
    (sqrt((2*pi)^d)*sqrt(det(s.Sigma1 + eye(d)*SMALL))).* ...
     exp( -0.5*(X1 - repmat(s.mu1', N, 1)) * ...
          inv(s.Sigma1 + SMALL*eye(d)) * ...
         (X1 - repmat(s.mu1', N, 1))'...
        );
  z2_num = s.alpha2 / ...
     (sqrt((2*pi)^d)*sqrt(det(s.Sigma2))).* ...
       exp( -0.5*(X1 - repmat(s.mu2', N, 1)) * ...
            inv(s.Sigma2)*...
            (X1 - repmat(s.mu2', N, 1))'...
          );
  z1_num = diag(z1_num);
  z2_num = diag(z2_num);
  s.z1 = z1_num ./ (z1_num + z2_num);
  s.z2 = z2_num ./ (z1_num + z2_num);

  % M-step
  %-----------------------------------------------------
  s.alpha1 = sum(s.z1)/N; 
  s.alpha2 = sum(s.z2)/N; 
  s.Sigma1 = ( (X1 - repmat(s.mu1', N, 1))'*...
               ((X1 - repmat(s.mu1', N, 1)).* ...
               repmat(s.z1,1,d)) ...
             )./ sum(s.z1);
  s.mu1 = (sum(X1.*repmat(s.z1,1,d),1) ./ ...
          sum(s.z1) )';  
  s.mu2 = (sum(X1.*repmat(s.z2,1,d),1) ./ ...
          (sum(s.z2) + SMALL) )'; 

end  % for i=1:n_iter
