function [THETA ASSIGN] = EM(Data,Means,Cov,runs)

%EM
%Expectation-maximization algorithm for maximum-likelihood estimation of mixture proportions 
%and group memberships for an unknown sample where groups have known
%parameters (from training dataset).
%Designed to use continuous, normally-distributed data, like otolith
%microchemistry data

%%This code authored by 
%%J. Wilson White
%%Dept. of Ecology, Evolution, and Marine Biology
%%Univ. California, Santa Barbara, 93106
%%w_white@lifesci.ucsb.edu
%%May 2006

%%Algorithm based on:
%%Millar RB (1987) Maximum likelihood estimation of mixed stock fishery
%%composition. Can J Fish Aquat Sci 44:583-590


%INPUT: 
% Matrix 'Data': rows are individuals, columns are predictor variables
% Matrix 'Means': means of source populations.  Predictor variables are in
% columns (same order as in 'Data'); one row for each population
% Array 'Cov': covariance matrices for each population.  Format is (r,c,k),
% where r = c = no. variables, k = no. populations, and (:,:,1) is the 
% covariance matrix for pop. 1, (:,:,2) is the covariance matrix for pop. 2, etc.  
% Scalar 'runs': number of algorithm iterations.  Must be > 100,
% recommended > 1000.

%OUTPUT: 
% Vector 'THETA' = probability of membership in each population (mixing proportion)
% Vector 'ASSIGN' = maximum-likelihood group assignments

C = length(Means(1,:)); %number of variables
K = length(Means(:,1)); %number of groups
N = length(Data(:,1)); %number of individuals
Theta(1,1:K) = 1/K; %uninformative prior on Theta

for i = 2:runs %begin algorithm
    ProbK = 0;
    for j = 1:N
        Prob = 0;
        for k = 1:K
            Prob(k) =  Theta(i-1,k)*mvnpdf(Data(j,:),Means(k,:),Cov(:,:,k)); %likelihood that indiv j is in group k
        end
        for k = 1:K
        ProbK(j,k) = Prob(k)/sum(Prob(:)); %Scaled probability that j is in k
        Assign(j,k,i) = ProbK(j,k); %hold assignment probabilities
        end
    end
    for k = 1:K
        Theta(i,k) = sum(ProbK(:,k))/N; %estimate mixing proportions
        if Theta(i,k) <= 1E-10
            Theta(i,k) = 0; %threshold for convenience
        end
    end
end %end algorithm
plot(Theta(:,1)) %plot raw Theta for group 1 to monitor convergence 

%if convergence doesn't happen in < 100 steps, change burnin below
burnin = 100;
%truncate burnin from Theta and Assign
Theta = Theta(burnin:end,:); 
Assign = Assign(:,:,burnin:end);

%Result is mean of post-covergence iterations
for k = 1:K
    THETA(k) = mean(Theta(:,k)); 
    for j = 1:N
    AssignMeans(j,k) = mean(Assign(j,k,:));
    end
end

for i = 1:N
    for k = 1:K
        if AssignMeans(i,k) == max(AssignMeans(i,:)); %assign individuals to group with highest probability
            ASSIGN(i) = k;
        end
    end
end
    
            
    