%
% cwm.m
% Neil Gershenfeld  (c) 12/6/96
% d-dimensional cluster-weighted modeling
% local quadratic models, covariances, momentum
% Matlab version 5 version
%
clear all
clf
max_plot = 6000;% maximum number of points to plot
min_var = 1e-4;	% constant to add to variances (resolution)
%
% load data
%
%ntest = 10000;
%nstart = 1000;
%load vio4a;
%x(:,4) = wavedata((nstart-12):(nstart+ntest-1-12))';
%x(:,3) = wavedata((nstart-6):(nstart+ntest-1-6))';
%x(:,2) = wavedata((nstart-3):(nstart+ntest-1-3))';
%x(:,1) = wavedata((nstart-1):(nstart+ntest-1-1))';
%y(:,1) = wavedata((nstart-0):(nstart+ntest-1-0))';
%load vio4p;
%x(:,6) = wavedata((nstart-10):(nstart+ntest-1-10))';
%x(:,5) = wavedata((nstart-0):(nstart+ntest-1-0))';
%load vio4v;
%x(:,8) = wavedata((nstart-10):(nstart+ntest-1-10))';
%x(:,7) = wavedata((nstart-0):(nstart+ntest-1-0))';
%return
%load a.dat;
%load a.cont;
a = sin(1:500)';
%a = rand(2000,1);
%load lorgenbt001;
%a=x';
%x=[];
%a = a(end:-1:1);
%ntest = length(a);
ntest = 500;
%x(:,10) = a(1:(ntest-50));
%x(:,9) = a(2:(ntest-49));
%x(:,8) = a(3:(ntest-48));
%x(:,7) = a(4:(ntest-47));
%x(:,6) = a(5:(ntest-46));
%x(:,5) = a(6:(ntest-45));
%x(:,4) = a(7:(ntest-44));
%x(:,3) = a(9:(ntest-43));
%x(:,2) = a(10:(ntest-42));
%x(:,1) = a(11:(ntest-41));
%y(:,1) = a(12:(ntest-40));
load /tmp/army.dat
x(:,1) = army(:,1);
x(:,2) = army(:,2);
y(:,1) = army(:,3);
%xh = .1;
%yh = .1;
%ntest = 10000;
%for point = 1:ntest
%   xnew = yh + 1.0 - 1.4 * xh * xh;
%   yh = 0.3 * xh;
%   xh = xnew;
%   data(point,1) = xh;
%   end
%data = data(ntest:-1:1);
%x(:,4) = data(1:(ntest-20));
%x(:,3) = data(2:(ntest-19));
%x(:,2) = data(3:(ntest-18));
%x(:,1) = data(4:(ntest-17));
%y(:,1) = data(11:(ntest-10));
%load neil.lag;
%npts = length(neil);
%y = sum(neil((1:(npts-9)),(1:10))')';
%ylag = neil(:,2:5);
%ylag = sum(neil((1:(npts-9)),(11:21))')';
%y = abs(y) .* (2*(randn(size(y)) > 0) - 1);
%ylag = abs(ylag) .* (2*(randn(size(ylag)) > 0) - 1);
%x = [ylag neil((10:npts),[22])];
%x = [ylag neil(:,[22 24])];
%y = neil(:,1);
%x = [neil(:,2:5) neil(:,[22])];
%x = [neil(:,2:21)];
%load bs2.mat;
%npts = length(data);
%y = data(1:npts,1);
%x = data(1:npts,2:4);
%
% make data zero mean unit variance and plot
%
[npts dim] = size(x);
fprintf('read %d points in %iD\n',npts,dim);
ymean = mean(y);
ystd = std(y);
%ymean = 0;
%ystd = 1;
x = x - kron(ones(npts,1),mean(x));
x = x ./ kron(ones(npts,1),std(x));
y = y - ymean;
y = y/ystd;
subplot(1,1,1)
if (npts < max_plot)
   if (dim > 1)
      plot3(x(:,1),x(:,2),y,'r*');
      grid off
      drawnow
      rotate3d on
   else
      plot(x,y,'r*');
      grid off
      drawnow
      end
   end
%
% get input info and separate out cross-validation set
%
nclusters = input('Number of clusters?  ');
niterations = input('Number of iterations?  ');
momentum = input('Momentum (0-1)?  ');
%cross_frac = input('Fraction of data to use for cross-validation (0-1)?  ');
%momentum = 0.0;
cross_frac = 0.1;
input_type = input('Variance clusters (0) or covariance cluster (1)?  ');
VARIANCE = 0;
COVARIANCE = 1;
output_type = input('Linear models (0) or quadratic models (1)?  ');
LINEAR = 0;
QUADRATIC = 1;
fprintf('momentum = %f, cross-validation fraction = %f\n',momentum,cross_frac);
index = rand(npts,1) > cross_frac;
%index = (1:npts) > (npts*cross_frac);
xtemp = x(find(index),:);
ytemp = y(find(index));
x_c = x(find(~index),:);
y_c = y(find(~index));
x = xtemp;
y = ytemp;
npts = length(y);
ncross = length(y_c);
fprintf('using %d points in-sample and %d points for cross validation\n\n',...
   npts,ncross);
fprintf('predicted error (in-sample out-of-sample), true error, data likelihood:\n');
%
% generate random starting clusters
%
mu = .8*rand(nclusters,dim).*kron(ones(nclusters,1),max(x)-min(x)) ...
     + kron(ones(nclusters,1),min(x));
dmu = zeros(nclusters,dim);
varx = ones(nclusters,dim);
for cluster = 1:nclusters
   covarx(:,:,cluster) = diag(ones(1,dim));
   end
vary = ones(1,nclusters);
varyest = ones(1,nclusters);
pc = ones(1,nclusters)/nclusters;
f = zeros(npts,nclusters);
f_c = zeros(ncross,nclusters);
fi = ones(npts,1);
fi_c = ones(ncross,1);
fi(:,2:(dim+1)) = x;
fi_c(:,(2:(dim+1))) = x_c;
if (output_type == QUADRATIC)
   for j = 1:dim
      fi(:,((2+dim*j):(1+dim*(j+1)))) = x .* kron(x(:,j),ones(1,dim));
      end
   for j = 1:dim
      fi_c(:,((2+dim*j):(1+dim*(j+1)))) = x_c .* kron(x_c(:,j),ones(1,dim));
      end
   end
%
% E-M iteration
%
step = 1;
while (step <= niterations)
   %
   % plot
   %
   subplot(1,1,1)
   if (dim > 1)
      plot(mu(:,1),mu(:,2),'ro')
      hold on
      for cluster = 1:nclusters
         if (input_type == VARIANCE)
            plot([mu(cluster,1)-sqrt(varx(cluster,1)) ...
               mu(cluster,1)+sqrt(varx(cluster,1))], ...
               [mu(cluster,2) mu(cluster,2)]);
            plot([mu(cluster,1) mu(cluster,1)], ...
               [mu(cluster,2)-sqrt(varx(cluster,2)) ...
               mu(cluster,2)+sqrt(varx(cluster,2))]);
	  else
            [vec,val] = eig(covarx(:,:,cluster));
            for d = 1:dim
               plot([mu(cluster,1)-sqrt(val(d,d))*vec(1,d) ...
                  mu(cluster,1)+sqrt(val(d,d))*vec(1,d)], ...
                 [mu(cluster,2)-sqrt(val(d,d))*vec(2,d) ...
                 mu(cluster,2)+sqrt(val(d,d))*vec(2,d)]);
                 end
              end
	   end
      grid off
      drawnow
      hold off
   else
      plot(mu,a,'ro')
      grid off
      drawnow
      end
   %
   % calculate forward probability
   %
   for cluster = 1:nclusters
      if (input_type == COVARIANCE)
         cinv(:,:,cluster) = pinv(covarx(:,:,cluster));
         px(:,cluster) = (sqrt(det(cinv(:,:,cluster)))/((2*pi)^(dim/2))) ...
            * exp(-dot((x-kron(ones(npts,1),mu(cluster,:)))', ...
           (cinv(:,:,cluster)*(x-kron(ones(npts,1),mu(cluster,:)))'))'/2);
         px_c(:,cluster) = (sqrt(det(cinv(:,:,cluster)))/((2*pi)^(dim/2))) ...
            * exp(-dot((x_c-kron(ones(ncross,1),mu(cluster,:)))', ...
            (cinv(:,:,cluster)*(x_c-kron(ones(ncross,1),mu(cluster,:)))'))'/2);
      else
         px(:,cluster) = prod((exp(-(x-kron(ones(npts,1),mu(cluster,:))).^2 ...
            ./ (2*kron(ones(npts,1),varx(cluster,:)))) ...
            ./ kron(ones(npts,1),sqrt(2*pi*varx(cluster,:)))),2);
         px_c(:,cluster) = prod((exp(-(x_c-kron(ones(ncross,1),mu(cluster,:))).^2 ...
            ./ (2*kron(ones(ncross,1),varx(cluster,:)))) ...
            ./ kron(ones(ncross,1),sqrt(2*pi*varx(cluster,:)))),2);
         end
      py(:,cluster) = exp(-(y-f(:,cluster)).^2 / (2*vary(cluster))) ...
          / sqrt(2*pi*vary(cluster));
      py_c(:,cluster) = exp(-(y_c-f_c(:,cluster)).^2 / (2*vary(cluster))) ...
          / sqrt(2*pi*vary(cluster));
      end
   %
   % calculate posterior and total probability
   %
   p = sum((py .* px .* kron(ones(npts,1),pc)),2);
   p_c = sum((py_c .* px_c .* kron(ones(ncross,1),pc)),2);
   pp = (py .* px .* kron(ones(npts,1),pc)) ./ ...
       kron(p,ones(1,nclusters));
   pp_c = (py_c .* px_c .* kron(ones(ncross,1),pc)) ./ ...
       kron(p_c,ones(1,nclusters));
   pc = sum(pp) / npts;
   step_prob(step) = mean(p);
   step_prob_c(step) = mean(p_c);
   %
   % update clusters
   %
   d2 = 0;
   for cluster = 1:nclusters
      munew = momentum*dmu(cluster,:) + ...
         sum(x .* kron(pp(:,cluster),ones(1,dim))) / (npts*pc(cluster));
      dmu(cluster,:) = munew - mu(cluster,:);
      mu(cluster,:) = munew;
      varx(cluster,:) = min_var + ...
         sum((x - kron(ones(npts,1),mu(cluster,:))).^2 .* ...
            kron(pp(:,cluster),ones(1,dim))) / (npts*pc(cluster));
      if (input_type == COVARIANCE)
         for d = 1:dim
            covarx(d,:,cluster) =  sum( ...
              (x - kron(ones(npts,1),mu(cluster,:))) .* ...
                 (kron(x(:,d),ones(1,dim)) - ones(npts,dim)*mu(cluster,d)) ...
                 .* kron(pp(:,cluster),ones(1,dim))) / (npts * pc(cluster));
           end
         covarx(:,:,cluster) = covarx(:,:,cluster) + min_var*diag(ones(1,dim));
         end
      %
      % regress and evaluate model
      %
      if (output_type == QUADRATIC)
         A = (sum(kron(y,ones(1,(1+dim+dim*dim))) .* fi .* ...
            kron(pp(:,cluster),ones(1,(1+dim+dim*dim)))) / (npts*pc(cluster)))';
         for j = 1:(1+dim+dim*dim)
            B(j,:) = sum(kron(fi(:,j),ones(1,(1+dim+dim*dim))) .* fi .* ...
               kron(pp(:,cluster),ones(1,(1+dim+dim*dim)))) / (npts*pc(cluster));
            end
      else
          A = (sum(kron(y,ones(1,(1+dim))) .* fi .* ...
            kron(pp(:,cluster),ones(1,(1+dim)))) / (npts*pc(cluster)))';
         for j = 1:(1+dim)
            B(j,:) = sum(kron(fi(:,j),ones(1,(1+dim))) .* fi .* ...
               kron(pp(:,cluster),ones(1,(1+dim)))) / (npts*pc(cluster));
            end
	 end
     
      beta(cluster,:) = (pinv(B) * A)';
      f(:,cluster) = sum(kron(ones(npts,1),beta(cluster,:)) .* fi,2);
      f_c(:,cluster) = sum(kron(ones(ncross,1),beta(cluster,:)) .* fi_c,2);
%      evals = svd(covarx(:,:,cluster));
%      r = 0.1 * sqrt(max(evals));
%      d2_m = sum(sqrt(2./(pi*evals)) .* exp(-r^2./(2*evals)) * r ...
%           ./ erf(r./sqrt(2*evals)));
%      d2 = d2 + d2_m * pc(cluster);
      vary(1,cluster) = min_var + sum((y - f(:,cluster)).^2 .* pp(:,cluster)) ...
         / (npts*pc(cluster));
      end % cluster update loop
   %
   % find in-sample and out-of-sample predictions
   %
   ypred = sum((kron(ones(npts,1),pc) .* f .* px),2) ./ ...
      sum((kron(ones(npts,1),pc) .* px),2);
   ypred_c = sum((kron(ones(ncross,1),pc) .* f_c .* px_c),2) ./ ...
      sum((kron(ones(ncross,1),pc) .* px_c),2);
   varyest = sum(kron(((y-ypred).^2),ones(1,nclusters)) ...
      .* pp) ./ (npts*pc);
   yvarpred = sum((kron(ones(npts,1),pc) .* ...
      (f.^2+kron(ones(npts,1),varyest)) ...
      .* px),2) ./ sum((kron(ones(npts,1),pc) .* px),2) - ypred.^2;
   yvarpred_c = sum((kron(ones(ncross,1),pc) .* ...
      (f_c.^2+kron(ones(ncross,1),varyest)) ...
      .* px_c),2) ./ sum((kron(ones(ncross,1),pc) .* px_c),2) - ypred_c.^2;
   step_pred_error(step) = sqrt(mean(yvarpred));
   step_pred_error_c(step) = sqrt(mean(yvarpred_c));
   step_error(step) = std(ypred-y);
   step_error_c(step) = std(ypred_c-y_c);
   fprintf('%d: pred %7.3g %7.3g true %7.3g %7.3g prob %7.3g %7.3g min %7.3g dim %7.3g\n', ...
      step,step_pred_error(step),step_pred_error_c(step), ...
      step_error(step),step_error_c(step), ...
      step_prob(step),step_prob_c(step),npts*min(pc),d2);
   if (step == niterations)
      input('CR to plot convergence ');
      subplot(2,2,1)
      plot(step_error,'b-')
      hold on
      plot(step_pred_error,'b--')
      plot(step_error_c,'r-')
      plot(step_pred_error_c,'r--')
      hold off
      title('Error')
      subplot(2,2,2)
      plot(step_prob,'b-')
      hold on
      plot(step_prob_c,'r-')
      hold off
      title('Probability')
      subplot(2,2,3)
      hist((y-ypred)./sqrt(yvarpred),20);
      title('In-sample normalized error');
      subplot(2,2,4)
      hist((y_c-ypred_c)./sqrt(yvarpred_c),20);
      title('Out-of-sample normalized error');
      fprintf('variances: ');
      fprintf('%.2f ',sum(varx .* kron(pc',ones(1,dim))));
      fprintf('\n');
      answer = input('Total number of iterations? (CR to end) ');
      if ~isempty(answer)
         niterations = answer;
         end
      end
   step = step + 1;
   end % iterate step loop
%
% plot model errors vs first two axes
%
if (npts < max_plot)
   subplot(1,1,1)
   if (dim > 1)
      plot3(mu(:,1),mu(:,2),beta(:,1),'k+','MarkerSize',10,'LineWidth',4);
      grid off
      hold on
      plot3([x(:,1) x(:,1)]',[x(:,2) x(:,2)]',[y ypred]','b');
      plot3(x(:,1),x(:,2),y,'r+');
      drawnow
      hold off
   else
      plot(mu,a,'go');
      grid off
      hold on
      plot([x x]',[y ypred]','b');
      plot(x,y,'r+');
      drawnow
      hold off
      end
   end
%
% plot prediction surface if in 2D
%
if (dim == 2)
   input('CR to plot 2D prediction surface and input density ');
   subplot(1,1,1)
   x1min = min(x(:,1));
   x1max = max(x(:,1));
   x2min = min(x(:,2));
   x2max = max(x(:,2));
   gridsize = 25;
   dx1 = (x1max-x1min)/gridsize;
   dx2 = (x2max-x2min)/gridsize;
   [x1grid,x2grid] = meshgrid(x1min:dx1:x1max, x2min:dx2:x2max);
   [npts1,npts2] = size(x1grid);
   ngrid = npts1*npts2;
   xgrid = [reshape(x1grid,ngrid,1) reshape(x2grid,ngrid,1)];
   px = [];
   f = [];
   figrid = [];
   fidata = [];
   figrid(:,1) = ones(ngrid,1);
   fidata(:,1) = ones(npts,1);
   figrid(:,2:(dim+1)) = xgrid;
   fidata(:,2:(dim+1)) = x;
   if (output_type == QUADRATIC)
      for j = 1:dim
         figrid(:,((2+dim*j):(1+dim*(j+1)))) = xgrid .* kron(xgrid(:,j),ones(1,dim));
         fidata(:,((2+dim*j):(1+dim*(j+1)))) = x .* kron(x(:,j),ones(1,dim));
         end
      end
   for cluster = 1:nclusters
      if (input_type == COVARIANCE)
         cinv = pinv(covarx(:,:,cluster));
         pxgridc(:,cluster) = (sqrt(det(cinv))/((2*pi)^(dim/2))) ...
            * exp(-dot((xgrid-kron(ones(ngrid,1),mu(cluster,:)))', ...
            (cinv*(xgrid-kron(ones(ngrid,1),mu(cluster,:)))'))'/2);
         pxdatac(:,cluster) = (sqrt(det(cinv))/((2*pi)^(dim/2))) ...
            * exp(-dot((x-kron(ones(npts,1),mu(cluster,:)))', ...
            (cinv*(x-kron(ones(npts,1),mu(cluster,:)))'))'/2);
      else
         pxgridc(:,cluster) = prod((exp(-(xgrid-kron(ones(ngrid,1),mu(cluster,:))).^2 ...
            ./ (2*kron(ones(ngrid,1),varx(cluster,:)))) ...
            ./ kron(ones(ngrid,1),sqrt(2*pi*varx(cluster,:)))),2);
         pxdatac(:,cluster) = prod((exp(-(x-kron(ones(npts,1),mu(cluster,:))).^2 ...
            ./ (2*kron(ones(npts,1),varx(cluster,:)))) ...
            ./ kron(ones(npts,1),sqrt(2*pi*varx(cluster,:)))),2);
	 end
      fgrid(:,cluster) = sum(kron(ones(ngrid,1),beta(cluster,:)) .* figrid,2);
      fdata(:,cluster) = sum(kron(ones(npts,1),beta(cluster,:)) .* fidata,2);
      end
   pxgrid = sum((kron(ones(ngrid,1),pc) .* pxgridc),2);
   pxdata = sum((kron(ones(npts,1),pc) .* pxdatac),2);
   ysurfpred = sum((kron(ones(ngrid,1),pc) .* fgrid .* pxgridc),2) ./ pxgrid;
   ysurfvarypred = sum((kron(ones(ngrid,1),pc) .* ...
      (fgrid.^2+kron(ones(ngrid,1),varyest.^2)) ...
      .* pxgridc),2) ./ pxgrid - ysurfpred.^2;
   yvarygrid = reshape(ysurfvarypred,npts1,npts2);
   ygrid = reshape(ysurfpred,npts1,npts2);
   densegrid = reshape(pxgrid,npts1,npts2);
   surf(x1grid,x2grid,ygrid,densegrid);
   colormap('default');
   hold on
   plot3(x(:,1),x(:,2),y,'r+')
   hold off
   drawnow
return   
   input('CR to plot 2D prediction surface and estimated error ');
   for point = 1:npts
      for cluster = 1:nclusters
         pygridc(:,cluster) = exp(-(y(point)-fgrid(:,cluster)).^2 / (2*vary(cluster))) ...
            / sqrt(2*pi*vary(cluster));
         pydatac(:,cluster) = exp(-(y(point)-fdata(:,cluster)).^2 / (2*vary(cluster))) ...
            / sqrt(2*pi*vary(cluster));
	 end
      pyxgrid(point,:) = sum((pygridc .* pxgridc .* kron(ones(ngrid,1),pc)),2)';
      pyxdata = sum((pydatac .* pxdatac .* kron(ones(npts,1),pc)),2);
      pydata(point,1) = sum(pyxdata ./ pxdata) / npts;
      end
   arg = kron(y,ones(1,ngrid)) .* pyxgrid ./ kron(pydata,ones(1,ngrid));
   varest = (sum(arg.^2)/npts - (sum(arg)/npts).^2)' ./ (npts * pxgrid.^2);
   vargrid = reshape(varest,npts1,npts2);
   colormap('default');
   surf(x1grid,x2grid,ygrid,vargrid);
   hold on
   plot3(x(:,1),x(:,2),y,'r+')
   hold off
   drawnow
   end % 2D surface plot
return

%
% calculate out-of-sample iterated predictions for embedded data
%
ncalc = input('Number of iterated predictions to to forecast? (CR to skip) ');
if ~isempty(ncalc)
   px = [];
   f = [];
   yplot = [];
   yerrplot = [];
   xvec = [y(end) x(end,1:(dim-1))];
   for point = 1:ncalc
      for cluster = 1:nclusters
         px(cluster) = prod(exp(-(xvec-mu(cluster,:)).^2 ...
            ./ (2*varx(cluster,:))) ./ sqrt(2*pi*varx(cluster,:)));
         f(cluster) = a(cluster) + sum(b(cluster,:) ...
            .* (xvec - mu(cluster,:)));
         end
      ynew = sum(pc .* f .* px) / sum(pc .* px);
      yerror = sum(pc .* (f.^2 + varyest.^2) .* px) / sum(pc .* px) ...
         - ynew^2;
      yplot = [yplot ynew];
      yerrplot = [yerrplot sqrt(yerror)];
      yerr = [yerr ynew-ycont(point)];
      xvec = [ynew xvec(1:(dim-1))];
      end
   yplot = yplot*ystd+ymean;
   ycontplot = ycont(1:ncalc)*ystd+ymean;
   yerrplot = yerrplot*ystd;
   subplot(1,1,1)
   plot(1:ncalc,yplot,'g')
   hold on
   plot(1:ncalc,ycontplot,'r+')
   plot([(1:ncalc)' (1:ncalc)']',[ycontplot yplot']','b');
   plot(1:ncalc,yerrplot,'r--');
   grid off
   hold off
   fprintf('iterated prediction error: %8.3g\n',std(yerr));
   end
return
