% c10p10.m
%
% demonstration of the Helmholtz machine using the wake-sleep
% algorithm. 

rand('state',105);

ndim=4;        % dimension of the square of bars
ndimsq=ndim^2;
nu=400;        % number of examples
pbar=0.4;      % probability that a given bar is on
pout=0.9;      % extra noise added at the end

ii=ones(1,ndim);
oo=zeros(1,ndim);

conv=zeros(ndimsq,ndim);
for i=1:ndim
  conv((i-1)*ndim+(1:ndim),i)=ones(ndim,1);
end
	       
bars=conv*(rand(ndim,nu)<pbar);  % which bars are on in which examples
u=bars.*(rand(ndimsq,nu)<pout) + (1-bars).*(rand(ndimsq,nu)>pout);
                                 % generate output noise

nhid=ndim+1;   % number of hidden units
W=zeros(nhid,ndimsq);
G=zeros(ndimsq,nhid);
g=zeros(nhid,1);    % generative biases for v
h=zeros(ndimsq,1);  % generative biases for u
w=zeros(nhid,1);    % recognition biases for v

sig=inline('1./(1+exp(-x))');  % sigmoid function
				 
mini=20;            % mini-batches of 20 input cases at a time
nits=100000;
eta=0.1/mini;       % learning rate
for it=1:nits
  %wake phase
  rp=randperm(nu);  
  au=u(:,rp(1:mini));   % choose a random set of mini patterns
  v = rand(nhid,mini) < sig(W*au + repmat(w,1,mini));    % propagate u->v and sample
  gv = sig(g);          % top down prediction of v
  g = g + eta*(v - repmat(gv,1,mini))*ones(mini,1); 
                        % delta rule for g
  gpu = sig(G*v + repmat(h,1,mini));   % top down prediction of u
  h = h + eta*(au - gpu)*ones(mini,1); % delta rule for h
  G = G + eta*(au - gpu)*v';           % delta rule for G
  sv = v;                              % for plotting
  %sleep phase
  gv = sig(g);          % top down prediction of v
  v = rand(nhid,mini)<repmat(gv,1,mini); % top down sample of v
  gu = rand(ndimsq,mini) < sig(G*v + repmat(h,1,mini)); % top down sample of u
  rv = sig(W*gu + repmat(w,1,mini));
  w = w + eta*(v - rv)*ones(mini,1);   % delta rule for w
  W = W + eta*(v - rv)*gu';            % delta rule for W
  if (mod(it,200)==0)
    disp(sprintf('%d',it));
    clf;
    colormap(gray);
    for i=1:nhid
      subplot(4,nhid,i);
      imagesc(reshape(G(:,i),ndim,ndim),[min(G(:)) max(G(:))]);
      set(gca,'xtick',[],'ytick',[]);
      title(num2str(sig(g(i))));       % title is gen prob
      colorbar;
      subplot(4,nhid,i+nhid);
      imagesc(reshape(W(i,:),ndim,ndim),[min(W(:)) max(W(:))]);
      set(gca,'xtick',[],'ytick',[]);
      colorbar;
      subplot(4,nhid,i+2*nhid);
      imagesc(reshape(u(:,rp(i)),ndim,ndim),[0 1]);
      set(gca,'xtick',[],'ytick',[]);
      title(rp(i));                    % title is pattern number
      colorbar;
      subplot(4,nhid,i+3*nhid);
      imagesc(reshape(gpu(:,i),ndim,ndim),[0 1]);
      set(gca,'xtick',[],'ytick',[]);  
      title(sprintf('%1d',sv(:,i)));   % title is hidden unit state
      colorbar;
    end
    drawnow;
  end
end

      








