% [C,L,e,code] = kmeans1(X,K[,init,maxit,tol,p,I])
% K-means clustering, with weights and partially fixed codebook
%
% This is like kmeans.m but it minimises the following:
%   min_{Z,C}{ sum_n{ sum_k{ pk.znk.|xn - ck|² } } }
% where C is the means matrix, Z the binary assignment variables, pk>0 is the
% weight for mean ck, and ck is fixed (not optimised over) for k in 1:I (where
% I>=0 is an integer). In regular kmeans, pk=1 and I=0. This version of kmeans
% is useful to have a (partially) predetermined codebook (say, the means can
% be one of -1, 0 or +1) and to enforce sparsity (by having pk<<1 for the 0
% mean).
%
% kmeans1 requires as initialisation an explicit means matrix "init" of KxD,
% where init(1:I,:) contains the fixed means (which won't be optimised over).
% To initialise means I+1 to K, you can run the usual kmeans like this (using
% kmeans++ as example):
%   init(I+1:end,:) = kmeans(X,K-I,'kmeans++',0);
%
% See kmeans.m for other input/output arguments.
% In:
%   init: a KxD matrix containing the initial K cluster means.
%   p: Kx1 vector of nonnegative weights for the K means. Default: ones.
%   I: integer in [0,K], means 1 to I won't be changed. Default: 0.
%
% Any non-mandatory argument can be given the value [] to force it to take
% its default value.

% Copyright (c) 2017 by Miguel A. Carreira-Perpinan

function [C,L,e,code] = kmeans1(X,K,init,maxit,tol,p,I)

% ---------- Argument defaults ----------
if ~exist('init','var') || isempty(init) init = 'kmeans++'; end;
if ~exist('maxit','var') || isempty(maxit) maxit = Inf; end;
if ~exist('tol','var') || isempty(tol) tol = 0; end;
if ~exist('p','var') p = []; end;
if ~exist('I','var') || isempty(I) I = 0; end;
% ---------- End of "argument defaults" ----------

% Given the initial means C, we do an assignment step
C = init; [L,e] = km_le(X,C,p);
if maxit == 0, code = 1; return; end
if I==K, code = 0; return; end

% Iteration
code = -1; i = 0;
while code < 0
  oldL = L; oldC = C;			% Keep if back up needed
  C = km_m(X,L,C,I);			% New means,...
  [L,e(end+1)] = km_le(X,C,p);		% ...labels and error value
  i = i + 1;
  % Stopping condition
  if e(end) >= (1-tol)*e(end-1), code = 0;
  elseif i >= maxit, code = 1;
  end
end

% Back up if the last iteration increased the error function. This can
% happen due to numerical inaccuracy if two points Xn are very close.
if e(end) > e(end-1), i = i - 1; e = e(1:end-1); C = oldC; L = oldL; end


% [L,e] = km_le(X,C,p) Labels and error function given means
%
% Input/output arguments as above.
% Out:
%   e: value of the error function.

function [L,e] = km_le(X,C,p)

if isempty(p) && (size(X,2)==1)
  % Fast algorithm for 1D data using binary search to assign point X(n) to its
  % closest codebook entry (on the sorted codebook). It avoids constructing an
  % intermediate matrix of NxK and takes O(N.logK) runtime.
  % This can differ slightly from the general computation if there are ties,
  % which may be broken differently.
  [S,J] = sort(C);
  L = J(discretize(X,[-Inf;(S(2:end)+S(1:end-1))/2;Inf])); % needs Matlab2015a
  e = sum((X - C(L)).^2);
elseif isempty(p)
  % This takes O(N.K.D) runtime and creates an intermediate matrix of O(N.K)
  % containing the point-centroid distances.
  [e,L] = min(sqdist(X,C),[],2); e = sum(e);
else
  [e,L] = min(bsxfun(@times,sqdist(X,C),p'),[],2); e = sum(e);
end


% C = km_m(X,L,C,I) Means given labels
%
% Input/output arguments as above.

function C = km_m(X,L,C,I)

% $$$ % This takes runtime O(N.(K-I)+N.D) but is easier to vectorise.
% $$$ for k=I+1:size(C,1), C(k,:) = mean(X(L==k,:),1); end

% This takes runtime O(N.D) but is harder to vectorise.
% Yet, it is much faster with Matlab 2015b...
if I < size(C,1)
  C(I+1:end,:) = 0; c = zeros(size(C,1),1);
  for n=find(L>I)', C(L(n),:) = C(L(n),:) + X(n,:); c(L(n)) = c(L(n)) + 1; end
  C(I+1:end,:) = bsxfun(@rdivide,C(I+1:end,:),c(I+1:end));

  % Dead means have no associated data points. Set each to a random data point.
  z = find(isnan(C(:,1))); C(z,:) = X(randperm(size(X,1),length(z)),:);
end

