% function [lam,mu,psi] = faEM(x,K,its,minpsi)
%
% Performs factor analysis using K components for the N-dimensional
% data x. its iterations of EM are used and the sensor variances are
% not allowed to fall below minpsi (if minpsi is not given, its default
% value is 0). The parameters are randomly initialized using the mean
% and variance of each input.
%
% Input:
%
%   x(:,t) = the N-dimensional training vector for the tth training case
%   K = number of factors to use
%   its = number of iterations of EM to apply
%   minpsi = minimum variance of sensor noise (default: 0)
%  
% Output:
%
%   lam = N x K factor loading matrix
%   mu = mean of training data
%   psi = vector of sensor variances
%

function [lam,mu,psi] = faEM(x,K,its,minpsi);

if nargin==3 minpsi = 0; end;
N = size(x,1); T = size(x,2);

% Initialize the parameters
mu = mean(x,2);
vr = std(x,[],2).^2;
lam = randn(N,K).*(sqrt(vr/K)*ones(1,K));
psi = (vr>=minpsi).*vr*2 + (vr<minpsi)*minpsi;

% Precompute some stuff that doesn't change
xmu = x-mu*ones(1,T);
Dxmxm = mean(xmu.*xmu,2);

% Do its iterations of EM
for i=1:its
  % Precompute some useful stuff that doesn't change in this iteration
  ipsilam = ((1./psi)*ones(1,K)).*lam;
  omega = (lam'*ipsilam+eye(K,K))^-1;
  Ez = omega*(ipsilam'*xmu);
    
  % Accumulate the sufficient statistics (E-step)
  xmEz = zeros(N,K); Ezz = zeros(K,K);
  for t=1:T
    xmEz = xmEz + xmu(:,t)*Ez(:,t)';
    Ezz = Ezz + Ez(:,t)*Ez(:,t)';
  end;
  xmEz = xmEz/T;
  Ezz = omega+Ezz/T;

  % Update the parameters (M-step)
  lam = xmEz*Ezz^-1;
  psi = Dxmxm - sum(lam.*xmEz,2)/T;
%  psi = sum((xmu-lam*Ez).^2,2)+sum((lam*omega).*lam,2);
end;