% PURPOSE : Demonstrate PF on a simple mixture problem. We only estimate the means of a mixture
%           of two Gaussians, but this can be easily extended. Note that this problem appears a lot
%           when tracking signal such as formants in speech. Here the means are fixed, only to keep the demo
%           simple. As an exercise, one can test the algorithm when the means vary with time.

% AUTHOR   : Nando de Freitas      (jfgf@cs.berkeley.edu)
% DATE     : Feb 2001

clear all; clc; echo off;


% INITIALISATION AND PARAMETERS:
% ==============================

k = 2;                      % Number of mixture components (fixed in this demo).
T = 100;                    % Number of time steps.
N = 50;                     % Number of particles.
resamplingScheme = 1;       % The possible choices are
                            % systematic sampling (2),
                            % residual (1)
                            % and multinomial (3). 
                            % They're all O(N) algorithms. 
delta = .001;               % Drift parameter in the transition prior for mu.
  

% GENERATE THE DATA FROM A MIXTURE OF TWO GAUSSIAN COMPONENTS:
% ===========================================================
sigma  = [.01 .005];       % Fixed variances.
mu     = [0 1];            % Means. These can be made time-varying.
lambda = [.3 .7];          % Fixed mixture weights.
y = zeros(T,1);            % Observations.
muforplot = zeros(T,1);
for t=2:T,
  if rand(1,1) > lambda(1),
    y(t) = mu(2) + sqrtm(sigma(2)).*randn(1,1);
    muforplot(t) = mu(2);
  else
    y(t) = mu(1) + sqrtm(sigma(1)).*randn(1,1);  
    muforplot(t) = mu(1);
  end;
end;  


% PLOT THE GENERATED DATA:
% ========================
figure(1)
clf;
plot(1:T,muforplot,'k',1:T,y,'r','linewidth',2);
ylabel('Data','fontsize',15);
xlabel('Time','fontsize',15);
legend('Generating mean of y','Realisation of y');

%  We will now use a particle filter to try to learn the mixture parameters 
%  mu given the data (y). lambda and sigma are fixed, but the program can be
%  easily extended to estimate these quantities. Watch out for identifiability 
%  issues.


%%%%%%%%%%%%%%%  PERFORM SEQUENTIAL MONTE CARLO  %%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%  ==============================  %%%%%%%%%%%%%%%%%%%%%

% INITIALISATION:
% ==============
muPF = rand(k,T,N);          % These are the particles for the estimate
                             % of mu. Note that there's no need to store
                             % them for all t. We're only doing this to
                             % show you all the nice plots at the end.
muPredPF = rand(k,T,N);      % One-step-ahead predicted values of mu.
w = ones(T,N);               % Importance weights.

disp(' ');
for t=2:T,    
  fprintf('t = %i / %i  \r',t,T);  fprintf('\n')
  
  % PREDICTION STEP:
  % ================ 
  % We use a Gaussian random walk transition prior as proposal.
  for i=1:N,
    muPredPF(:,t,i) = muPF(:,t-1,i) + sqrtm(delta)*randn(k,1);   
  end;

  % EVALUATE IMPORTANCE WEIGHTS:
  % ============================
  % For our choice of proposal, the importance weights are give by:  
  for i=1:N,
    lik =   lambda(1)*inv(sqrt(2*pi*sigma(1))) * exp(-0.5*inv(sigma(1))*((y(t)-muPredPF(1,t,i))^(2))) ...
          + lambda(2)*inv(sqrt(2*pi*sigma(2))) * exp(-0.5*inv(sigma(2))*((y(t)-muPredPF(2,t,i))^(2)));
    w(t,i) = lik;    
  end;  
  w(t,:) = w(t,:)./sum(w(t,:));                % Normalise the weights.
  
  % SELECTION STEP:
  % ===============
  % Here, we give you the choice to try three different types of
  % resampling algorithms. Note that the code for these algorithms
  % applies to any problem!
  if resamplingScheme == 1
    outIndex = residualR(1:N,w(t,:)');        % Residual resampling.
  elseif resamplingScheme == 2
    outIndex = systematicR(1:N,w(t,:)');      % Systematic resampling.
  else  
    outIndex = multinomialR(1:N,w(t,:)');     % Multinomial resampling.  
  end;
  muPF(:,t,:) = muPredPF(:,t,outIndex);     % Keep particles with
                                            % resampled indices.
end;   % End of t loop.

muPF1 = zeros(T,N);
muPF1(:,:) = muPF(1,:,:);
muPF2 = zeros(T,N);
muPF2(:,:) = muPF(2,:,:);


%%%%%%%%%%%%%%%%%%%%%  PLOT THE RESULTS  %%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%  ================  %%%%%%%%%%%%%%%%%%%%%

figure(2)
clf;
plot(1:T,mean(muPF1'),'b',1:T,mean(muPF2'),'r','linewidth',2);
xlabel('Time','fontsize',15);
legend('Estimated \mu_1','Estimated \mu_2');
hold on;
mu1 = mu(1)*ones(T,1);
mu2 = mu(2)*ones(T,1);
plot(1:T,mu1,'k--',1:T,mu2,'k--','linewidth',2);

figure(3)
clf;
subplot(211)
domain = zeros(T,1);
range = zeros(T,1);
if mean(muPF1(1,50,:))>mean(muPF1(2,50,:)) % hack to deal with identifiability.
  thex=[-.5:.05:.5];
else
  thex=[.5:.05:1.5];
end;
hold on
ylabel('Time (t)','fontsize',15)
xlabel('\mu1_t','fontsize',15)
zlabel('p(\mu1_t|y_{1:t})','fontsize',15)
for t=6:5:T,
  [range,domain]=hist(muPF1(t,:),thex);
  waterfall(domain,t,range/sum(range));
end;
view(-30,80);
rotate3d on;
a=get(gca);
set(gca,'ygrid','off');
subplot(212)
domain = zeros(T,1);
range = zeros(T,1);
if mean(muPF1(1,50,:))<mean(muPF1(2,50,:))
  thex=[-.5:.05:.5];
else
  thex=[.5:.05:1.5];
end;
hold on
ylabel('Time (t)','fontsize',15)
xlabel('\mu2_t','fontsize',15)
zlabel('p(\mu2_t|y_{1:t})','fontsize',15)
for t=6:5:T,
  [range,domain]=hist(muPF2(t,:),thex);
  waterfall(domain,t,range/sum(range));
end;
view(-30,80);
rotate3d on;
a=get(gca);
set(gca,'ygrid','off');