% EM example for SPM course, April 2004, W. Penny

% Numerically, EM shows greater benefit over OLS with larger V eg. 100
V=20; % number of voxels

% Number of samples per voxel
N=10;
disp(sprintf('%d observations per voxel'));

% Fix true hyperparameters
true_alpha=1;
true_beta=0.1+rand(V,1)*0.9;

% Sample effect sizes
true_mu=randn(V,1)*(1/true_alpha);
for i=1:V,
    y(i,:)=true_mu(i)+randn(1,N)*(1/true_beta(i));
end

% overwrite these from disk ?
from_disk=1;
if from_disk
    load em_data
end
save em_data y true_mu true_beta


% Get LS/ML estimates
for i=1:V,
        ls_mu(i,:)=(1/N)*sum(y(i,:));
        ls_beta(i)=N/sum((y(i,:)-ls_mu(i)).^2);
end
ls_error=sum((true_mu-ls_mu).^2);
max_mu=max([ls_mu;true_mu])+0.5;
min_mu=min([ls_mu;true_mu])-0.5;

% Plot true effect sizes
figure
plot(true_mu,'ro','LineWidth',2);
hold on
plot([0 V+1],[0 0],'k-','LineWidth',2);
%ylabel('w_i','FontSize',18);
xlabel('i','FontSize',18);
%a=axis;a(2)=a(2)+1;axis(a);
axis([0 V+1 min_mu max_mu]);
%set(gca,'FontSize',18);

% Plot data variances
% figure
% plot(1./true_beta,'rx');
% title('True observation noise variances');
% xlabel('Voxel Number');
% a=axis;a(2)=a(2)+1;axis(a);

% Plot Data
figure
hold on
for i=1:V,
    %plot(i,y(i,:),'k.','LineWidth',3);
    plot(i,y(i,:),'k.','MarkerSize',15);
end

xlabel('i','FontSize',18);
ylabel('y_i','FontSize',18);
a=axis;a(2)=a(2)+1;axis(a);
hold on
plot([0 V+1],[0 0],'k-','LineWidth',2);

% Plot Data
figure
hold on
for i=1:V,
    %plot(i,y(i,:),'k.','LineWidth',3);
    plot(i,y(i,:),'k.','MarkerSize',15);
end
hold on
plot(ls_mu,'x','Markersize',10,'LineWidth',2);
plot([0 V+1],[0 0],'k-','LineWidth',2);


xlabel('i','FontSize',18);
ylabel('y_i','FontSize',18);
a=axis;a(2)=a(2)+1;axis(a);
hold on
plot([0 V+1],[0 0],'k-','LineWidth',2);

% Plot OLS/ML estimates
figure
plot(true_mu,'ro','LineWidth',2);
hold on
plot(ls_mu,'x','Markersize',10,'LineWidth',2);

axis([0 V+1 min_mu max_mu]);
str1=sprintf('True');
str2=sprintf('OLS');
%legend(str1,str2);
%title(sprintf('OLS estimates, error=%1.2f','FontSize',18));
%set(gca,'FontSize',18);
xlabel('i','FontSize',18);
%ylabel('\hat{\mu}_i');
%ylabel('$\hat{w}_i$','FontSize',18);
hold on
plot([0 V+1],[0 0],'k-','LineWidth',2);

% Now run EM
max_its=10;
alpha=0;
gamma=ones(V,1);
mu=zeros(V,1);
beta=zeros(V,1);

figure
for it=1:max_its,
    
    err=sqrt((mu-true_mu)'*(mu-true_mu)/20);
    disp(sprintf('Error SD=%1.2f',err));
    
    for i=1:V,
        mu(i)=(1/N)*gamma(i)*sum(y(i,:));
        beta(i)=(N-gamma(i))/sum((y(i,:)-mu(i)).^2);
        gamma(i)=N*beta(i)/(N*beta(i)+alpha);
    end
    alpha=sum(gamma)/sum(mu.^2)
    error=sum((true_mu-mu).^2);
    
    clf;
    plot(true_mu,'ro','LineWidth',2);
    hold on
    plot(mu,'x','Markersize',10,'LineWidth',2);
    hold on
    plot([0 V+1],[0 0],'k-','LineWidth',2);
    axis([0 V+1 min_mu max_mu]);
    xlabel('i','FontSize',18);
    %ylabel('Effect Size','FontSize',18);
    str1=sprintf('True');
    str2=sprintf('Bayes');
    
    disp(sprintf('Iteration %d',it));
    %legend(str1,str2);
    %title(sprintf('Bayes estimates, It=%d, error=%1.2f',it,error));
    %set(gca,'FontSize',18);
    
    % Compute population mean - as described in hierarchical models
    % chapter - under the separable models section
    xi=ones(N,1);
    beta_mu=0;num_mu=0;   
    for i=1:V,
        d=inv((1/alpha)*xi*xi'+(1/beta(i))*eye(N));
        beta_mu=beta_mu+xi'*d*xi;      
        num_mu=xi'*d*y(i,:)';
    end
    mu_pop=num_mu/beta_mu
    
    pause
end
 
disp(sprintf('True alpha=%1.2f, Estimated alpha=%1.2f',true_alpha,alpha));