% Lloyd's algorithm - 02/02/12
% Jean-Marc Berthommé
%
% - 04/13/11:
%   . 1st version: reproduce Lloyd's algorithm on a uniform pdf
% - 04/15/11:
%   . center of mass calculation
% - 05/02/11:
%   . global update in relation to ccvt.m
% - 05/12/11:
%   . energy and capacity error displayed
%   . code optimization: for vs. repmat
% - 02/02/12:
%   . entropy interpretation added

function lloyd
n = 2^5;                    % nb of sites:  30 - 1024
r = 2^7;                    % resolution : 100 -  200
m = r^2;                    % nb of points
k = 0;                      % iteration counter
l = 0;                      % final repetition counter
meth = 'flo';               % sites generation method: rnd/uni/flo
dopt = 'col';               % display option: col/bw
stop = false;               % main loop flag
fig1 = false;               % fig1 flag
fig2 = false;               % fig2 flag
x = 1:r; y = 1:r;           % x,y values
[X Y] = meshgrid(x,y);      % image grid
scr = get(0,'ScreenSize');  % screen size

[sx, sy, n] = generate_sites(n, r, meth);
lab = label(sx, sy, X, Y);
[e,dc,h] = calc(n, m, sx, sy, lab, 'Random Init°');
fig1 = disp_label(scr, fig1, k, x, y, lab, sx, sy, dopt);
[fig2, K,E,Dc,H] = disp_graph(scr,fig2, [],[],[],[], k,e,dc,h);

while ~stop
    [sx, sy] = centroidal_move(n, sx, sy, lab, X, Y);
    newlab = label(sx, sy, X, Y); k=k+1;

    [e,dc,h] = calc(n, m, sx, sy, newlab, '');
    disp_label(scr, fig1, k, x, y, newlab, sx, sy, dopt);
    [~,K,E,Dc,H] = disp_graph(scr,fig2, K,E,Dc,H, k,e,dc,h);

    if E(end-1)-E(end) > eps, lab = newlab; % energy criterion
    else                        l=l+1; if l>5, stop = true; end;
    end
end
calc(n, m, sx, sy, lab, 'Final Lloyd');
pause; clear all; close all;

%*************************************************************************%
function [sx, sy, n] = generate_sites(n, r, meth)
if     strcmp(meth,'uni')    % uniform
    no = floor(sqrt(n)); n = no^2;
    l = linspace(1,r,no+2);
    l = 7+l(2:no+1);
    sx = reshape(ones(no,1)*l, 1, n);
    sy = reshape(l'*ones(1,no), 1, n);
elseif strcmp(meth,'rnd')    % random
    sx = r*rand(1,n);
    sy = r*rand(1,n);
elseif strcmp(meth,'flo')    % flower
    sx = 3*r/8 + r/4*rand(1,n);
    sy = 3*r/8 + r/4*rand(1,n);
else
    error('Wrong point generation method');
end

function D2 = sq_dist(X,Y)   % square distances
nx = size(X,1); ny = size(Y,1);
D2 = sum(X.^2,2)*ones(1,ny) + ones(nx,1)*sum(Y.^2,2)' - 2*(X*Y');

function lab = label(sx, sy, X, Y)
D2 = sq_dist([sx',sy'],[X(:),Y(:)]); % TODO: can be optimized! -> kNN!
[~, lab] = min(D2, [], 1);
lab = reshape(lab, size(X));

function [sx, sy] = centroidal_move(n, sx, sy, A, X, Y)
% LAB = repmat(lab,[1,1,n]);
% [~,~,V] = ndgrid(1:r,1:r,1:n);
% TEST = reshape(LAB == V, r*r, n, 1);
% XX = repmat(X(:),1,n); YY = repmat(Y(:),1,n);
% Xv = zeros(r*r, n);    Yv = zeros(r*r, n);
% Xv(TEST) = XX(TEST); Yv(TEST) = YY(TEST);
% sx = sum(Xv)./sum(TEST); sy = sum(Yv)./sum(TEST);

for i=1:n  % "for" is faster than the code above with a "repmat"!
    sx(i) = mean(X(A == i));
    sy(i) = mean(Y(A == i));
end

function Etot = energy(n, m, sx, sy, A)
r = sqrt(m);
[X Y] = meshgrid(1:r,1:r);
D2 = sq_dist([sx',sy'],[X(:),Y(:)]);
Id = ones(n,1)*(A(:)');
[~, Yd] = meshgrid(1:m,1:n);
Idx = (Yd == Id);
E = D2(Idx);                 % energy vector
Etot = sum(E);               % global energy

function H = entropy(P)
% Entropy of a state of n micro-states
% P: micro-states probabilities vector (whatever in line or in column)
if abs(sum(P)-1) > 1000*eps
    error('Sum of probabilities not equal to 1: %0.4f', sum(P));
end
Ok = (P > eps) & (P < 1-eps);    % avoids NaNs as p.log(p) -> 0 in 0+
H = sum(P(Ok).*log2(1./P(Ok)));  % normalization in bits

function [E, dc, H] = calc(n, m, sx, sy, A, str)
% calculate energy, capacity error and entropy
E = energy(n, m, sx, sy, A); % energy

C = histc(A(:),1:n);         % capacities
H = entropy(C/m);            % entropy [bits]

cs = m/n;                    % uniform capacity c "star"
dc = 1/n*sum((C/cs - 1).^2); % capacity error

fprintf('-> E = %0.0f - dc = %0.4f - H = %0.4f bits %s\n', E, dc, H, str);

function fig1 = disp_label(scr, fig1, k, x, y, lab, sx, sy, dopt)
f1 = figure(1);
if ~fig1, fig1 = true; set(gcf,'Color',[0.2,0.2,0.2]);
   set(f1,'Position',[2*scr(3)/3, scr(4)/2, scr(3)/3, scr(4)/2-50*2]); end;

tit = sprintf('-> k = %d <-', k);
if strcmp(dopt,'col')        % color
    imagesc(x,y,lab); axis image;
    hold on; voronoi(sx, sy, 'w'); hold off;
elseif strcmp(dopt, 'bw')    % black & white
    plot(sx, sy, 'k.'); axis image;
    axis([min(x) max(x) min(y) max(y)]);
else
    error('Wrong display option');
end
title(tit, 'Color', 'w');

function [fig2, K,E,Dc,H] = disp_graph(scr,fig2, K,E,Dc,H, k,e,dc,h)
f2 = figure(2);
if ~fig2, fig2 = true; set(gcf,'Color',[0.2,0.2,0.2]);
          set(f2,'Position',[2*scr(3)/3, 0, scr(3)/3, scr(4)/2-45*2]); end;

K = cat(1,K,k); E = cat(1,E,e); Dc = cat(1,Dc,dc); H = cat(1,H,h);
subplot(3,1,1); plot(K, E, 'r-'); title('global energy','color','w');
subplot(3,1,2); plot(K, Dc, 'b-'); title('capacity error','color','w');
subplot(3,1,3); plot(K, H, 'm-'); title('entropy','color','w');

-> E = 21774920 - dc = 2.6842 - H = 3.5352 bits Random Init°
-> E = 4078270 - dc = 1.2385 - H = 4.1530 bits 
-> E = 2952635 - dc = 0.8488 - H = 4.3814 bits 
-> ...
-> E = 1401286 - dc = 0.0054 - H = 4.9962 bits 
-> E = 1401286 - dc = 0.0054 - H = 4.9962 bits 
-> E = 1401286 - dc = 0.0054 - H = 4.9962 bits Final Lloyd