% EM_Cluster.m
% This code was written exclusively for Example 13.6 where the EM algorithm
% for Gaussian mixture models was applied to the Iris data set in the context of
% data clustering. 
% Inputs:
% X0: Input data in form of a D by N matrix with N the number of samples and D
%     the number of features of a sample. We assume N > D.
% K:  The number of Gaussian components.
% c0: Initial mixing coefficients in form of a K-dimensional column vector.
% s1: Initial random state for random permutation of input samples.
% s2: Initial random state for (initial) K-means++ clustering. 
% T: Number of iterations for (initial) K-means++ clustering.
% t: Convergence tolerance for the EM algorithm. 
% Output:
% X: randomly permutated input data.
% id: an N-dimensional vector whose ith component is an integer being the
%     index of the ith sample of X in the original input X0.
% C: a K-component array of matrices with ith component matrix containing
%    the samples of the ith cluster. 
% idc: a K-component array of vectors with ith component vector containing
%      the index of the ith cluster in the original input X0.
% Sf: a K-component array of matrices whose ith component is an estimate of the
%     covariance of the ith component Gaussian.
% Mf: a K-column matrix whose ith column is an estimate of the mean of the ith
%     component Gaussian.
% cf: a K-dimensional column vector whose whose ith component is an
%     estimate of the ith mixing coefficient. 
% Gf: final values of the responsibilities in form of a K by N matrix whose
%     (k,n)th component if gamma_k,n.
% Nf: a profile of the Nk values after each iteration in form of a K by M
%     matrix, where M denotes the number of E-M iterations.
% Er: profile of convergence tolerance versus E-M iterations. 
% Mc0: confusion matrix for the initial K-means clustering.
% Mc: confusion matrix for the EM-based clustering.
% Written by W.-S. Lu, University of Victoria.
% Last updated: August 24, 2025. 
% Example: 
% load D_iris
% X0 = D_iris(1:4,:);
% c0 = [0.2 0.4 0.4]';
% [X,id,C,idc,Sf,Mf,cf,Gf,Nf,Er,Mc0,Mc] = EM_Cluster(X0,3,c0,17,10,12,0.016);
function [X,id,C,idc,Sf,Mf,cf,Gf,Nf,Er,Mc0,Mc] = EM_Cluster(X0,K,c0,s1,s2,T,t)
% Prepare X S0 M0 C0 Ctrue using K-means++
[D,N] = size(X0);
rand('state',s1)
id = randperm(N);
X = X0(:,id);
% Initial Kmeans clustering
rand('state',s2)
p = randperm(N);
c1 = X(:,p(1));
Ctr = c1;
id0 = 1:1:N;
idw = setdiff(id0,p(1));
Xw = X(:,idw);
for k = 1:K-1
    dw = [];
    for i = 1:N-k
        xi = Xw(:,i);
        qi = [];
        for j = 1:k
            qi = [qi norm(Ctr(:,j)-xi)^2];
        end
        dw = [dw min(qi)];
    end
    wt = dw/sum(dw);
    qk = Xw*wt(:);
    dk = zeros(1,N-k);
    for i = 1:N-k
        dk(i) = norm(Xw(:,i)-qk);
    end
    [~,idk] = min(dk);
    xk = Xw(:,idk);
    Ctr = [Ctr xk];
    idw = setdiff(1:1:N-k,idk);
    Xwnew = Xw(:,idw);
    Xw = Xwnew;
end
ki = 0;
J = [];
while ki < T
    idx0 = zeros(1,N);
    R = zeros(N,K);
    jw = 0;
    for i = 1:N
        xi = X(:,i);
        di = zeros(1,K);
        for k = 1:K
            di(k) = norm(xi - Ctr(:,k))^2;
        end 
        [mi,idi] = min(di);
        jw = jw + mi;
        R(i,idi) = 1;
        idx0(i) = idi;
    end
    J = [J jw];
    for k = 1:K
        Ctr(:,k) = (X*R(:,k))/sum(R(:,k));
    end
    ki = ki + 1;
end
% Prepare confusion matrix for initial Kmeans clustering
in1 = find(idx0 == 1); 
in2 = find(idx0 == 2); 
in3 = find(idx0 == 3);
z1 = id(in1);
z2 = id(in2);
z3 = id(in3);
t1 = zeros(1,length(z1)); 
t2 = zeros(1,length(z2));
t3 = zeros(1,length(z3));
for i = 1:length(z1)
    if z1(i) >= 1 & z1(i) <= 50
        t1(i) = 1;
    elseif z1(i) >= 51 & z1(i) <= 100
        t1(i) = 2;
    else t1(i) = 3;
    end
end
for i = 1:length(z2)
    if z2(i) >= 1 & z2(i) <= 50
        t2(i) = 1;
    elseif z2(i) >= 51 & z2(i) <= 100
        t2(i) = 2;
    else t2(i) = 3;
    end
end
for i = 1:length(z3)
    if z3(i) >= 1 & z3(i) <= 50
        t3(i) = 1;
    elseif z3(i) >= 51 & z3(i) <= 100
        t3(i) = 2;
    else t3(i) = 3;
    end
end
Mc0 = zeros(3,3);
Mc0(1,1) = length(find(t1==1));
Mc0(1,2) = length(find(t1==2));
Mc0(1,3) = length(find(t1==3));
Mc0(3,1) = length(find(t2==1));
Mc0(3,2) = length(find(t2==2));
Mc0(3,3) = length(find(t2==3));
Mc0(2,1) = length(find(t3==1));
Mc0(2,2) = length(find(t3==2));
Mc0(2,3) = length(find(t3==3));
disp('Confusion matrix:')
Mc0
disp('Accuracy in percent:')
tx0 = sum(sum(Mc0));
tt0 = trace(Mc0);
100*(tt0/tx0)
disp('Rate of dis-clustering in percent:')
100*((tx0-tt0)/tx0)
% Prepare initial parameters
X1 = []; X2 = []; X3 = [];
for i = 1:N
    if idx0(i) == 1
        X1 = [X1 X(:,i)];
    elseif idx0(i) == 2
        X2 = [X2 X(:,i)];
    else
        X3 = [X3 X(:,i)];
    end
end
M0 = [];
N1 = size(X1,2);
mu1 = mean(X1')';
M0 = [M0 mu1];
W1 = X1 - mu1;
sig1 = (W1*W1')/N1;
S0(1) = {sig1}; 
N2 = size(X2,2);
mu2 = mean(X2')';
M0 = [M0 mu2];
W2 = X2 - mu2;
sig2 = (W2*W2')/N2;
S0(2) = {sig2}; 
N3 = size(X3,2);
mu3 = mean(X3')';
M0 = [M0 mu3];
W3 = X3 - mu3;
sig3 = (W3*W3')/N3;
S0(3) = {sig3}; 
% Perform E step and M step 
ds = sqrt(D);
Sold = S0; Mold = M0; cold = c0;
Nf = [];
er = 10^6;
Er = [];
while er >= t
   G = zeros(K,N);
   Nx = zeros(K,N);
   for k = 1:K
      muk = Mold(:,k);
      sigk = Sold{k};
      for n = 1:N
          xn = X(:,n);
          Nx(k,n) = mvnpdf(xn,muk,sigk);
      end
   end
   c = cold(:);
   for n = 1:N
       dn = c'*Nx(:,n);
       G(:,n) = (c.*Nx(:,n))/dn;
   end
   Nk = sum(G')';
   Mnew = [];
   for k = 1:K
       mnew = (X*G(k,:)')/Nk(k);
       Mnew = [Mnew mnew];
   end
   Snew = [];
   for k = 1:K
       mw = Mnew(:,k);
       Sw = zeros(D,D);
       for n = 1:N
           dw = X(:,n) - mw;
           Sw = Sw + G(k,n)*(dw*dw');
       end
       Snew{k} = Sw/Nk(k);
   end
   cnew = Nk/N;
   Nf = [Nf Nk];
   % Get the difference between old and new estimates.
   e1 = (norm(cnew - cold))/ds;
   e2 = 0; e3 = 0;
   for k = 1:K
       e2k = norm(Mnew(:,k)-Mold(:,k));
       e2 = e2 + e2k;
       e3k = norm(Snew{k} - Sold{k},'fro');
       e3 = e3 + e3k;
   end
   e2 = e2/(K*ds);
   e3 = e3/(K*D);
   er = e1 + e2 + e3;
   Er = [Er er];
   Sold = Snew;
   Mold = Mnew;
   cold = cnew;
end
% Output the results
Sf = Snew;
Mf = Mnew;
cf = cnew;
c = cf(:);
Gf = G;
Ct = [];
for i = 1:N
    gi = Gf(:,i);
    [~,ix] = max(gi);
    Ct = [Ct ix];
end
idc1 = find(Ct == 1);
C1 = X(:,idc1);
idc2 = find(Ct == 2);
C2 = X(:,idc2);
idc3 = find(Ct == 3);
C3 = X(:,idc3);
idc(1) = {id(idc1)};
idc(2) = {id(idc2)};
idc(3) = {id(idc3)};
C(1) = {C1};
C(2) = {C2};
C(3) = {C3};
M = length(Er);
% Prepare confusion matrix
c1 = idc{1}; c2 = idc{2}; c3 = idc{3};
t1 = zeros(1,length(c1)); 
t2 = zeros(1,length(c2));
t3 = zeros(1,length(c3));
for i = 1:length(c1)
    if c1(i) >= 1 & c1(i) <= 50
        t1(i) = 1;
    elseif c1(i) >= 51 & c1(i) <= 100
        t1(i) = 2;
    else t1(i) = 3;
    end
end
for i = 1:length(c2)
    if c2(i) >= 1 & c2(i) <= 50
        t2(i) = 1;
    elseif c2(i) >= 51 & c2(i) <= 100
        t2(i) = 2;
    else t2(i) = 3;
    end
end
for i = 1:length(c3)
    if c3(i) >= 1 & c3(i) <= 50
        t3(i) = 1;
    elseif c3(i) >= 51 & c3(i) <= 100
        t3(i) = 2;
    else t3(i) = 3;
    end
end
Mc = zeros(3,3);
Mc(1,1) = length(find(t1==1));
Mc(1,2) = length(find(t1==2));
Mc(1,3) = length(find(t1==3));
Mc(3,1) = length(find(t2==1));
Mc(3,2) = length(find(t2==2));
Mc(3,3) = length(find(t2==3));
Mc(2,1) = length(find(t3==1));
Mc(2,2) = length(find(t3==2));
Mc(2,3) = length(find(t3==3));
disp('Confusion matrix for EM-based clustering:')
Mc
disp('Accuracy in percent:')
tx = sum(sum(Mc));
tt = trace(Mc);
100*(tt/tx)
disp('Rate of dis-clustering:')
100*((tx-tt)/tx)