%% set up data
trainpat = load('mnist_train.csv','-ascii'); %training patterns; 60000 patterns x [category label, 784 pixel values]
trainpat(:,2:end) = trainpat(:,2:end)/255; %convert image intensity values from 0:255 to unit interval [0,1]
testpat = load('mnist_test.csv','-ascii'); %test patterns; 10000 patterns x [category label, 784 pixel values]
testpat(:,2:end) = testpat(:,2:end)/255; %convert image intensity values from 0:255 to unit interval [0,1]
%% train network
ntrain = 100; %number of training patterns to use (up to 60k)
nepoch = 1000; %number of update cycles
e = .0001; %learning rate
momentum = .9; %momentum learning: exponential filtering of past gradient estimates (following Hinton, 2002, sec 8)
nhid = 100; %number of hidden units
[W,bvis,bhid] = RBMtrain(trainpat(1:ntrain,2:end),nhid,nepoch,e,momentum); %train and return learned weights
%% plot the latent features, 50 per figure
for j=1:nhid
    if mod(j,50)==1
        figure(1+ceil(j/50)),clf
        colormap(gray)
    end
    subplot(5,10,mod(j-1,50)+1)
    surf(reshape(W(:,j),28,28)) %receptive field for this feature
    axis([0 28 0 28])
    view(90,90)
end
%% evaluate reconstruction error on training patterns
rmseTrain = 0; %tracking reconstruction error
figure(2+ceil(nhid/50)),clf %example plots for first 25 training stimuli
colormap(gray)
for i=1:ntrain
    vpos = trainpat(i,2:end)'; %input pattern
    hpos = rand(nhid,1) < 1./(1+exp(-W'*vpos-bhid)); %sample values for hidden units
    vneg = 1./(1+exp(-W*hpos-bvis)); %mean reconstructed values for visible units 
    rmseTrain = rmseTrain + mean((vpos-vneg).^2); %add mean squared reconstruction error for this pattern
    if i<=25 %make example plots of original stimulus and reconstruction
        subplot(5,10,2*i-1) %plot original stimulus
        surf(reshape(trainpat(i,2:end),28,28))
        axis([0 28 0 28])
        view(90,90)
        subplot(5,10,2*i) %plot reconstructed stimulus
        surf(reshape(vneg,28,28))
        axis([0 28 0 28])
        view(90,90)
    end
end
rmseTrain = sqrt(rmseTrain/ntrain); %convert to root mean squared error
%% evaluate reconstruction error on test patterns
ntest = 1000; %number of test patterns to use (up to 10k)
figure(3+ceil(nhid/50)),clf %example plots for first 25 test stimuli
colormap(gray)
rmseTest = 0; %tracking reconstruction error
for i=1:ntest
    vpos = testpat(i,2:end)'; %input pattern
    hpos = rand(nhid,1) < 1./(1+exp(-W'*vpos-bhid)); %sample values for hidden units
    vneg = 1./(1+exp(-W*hpos-bvis)); %mean reconstructed values for visible units 
    rmseTest = rmseTest + mean((vpos-vneg).^2); %add mean squared reconstruction error for this pattern
    if i<=25 %make example plots of original stimulus and reconstruction
        subplot(5,10,2*i-1) %plot original stimulus
        surf(reshape(testpat(i,2:end),28,28))
        axis([0 28 0 28])
        view(90,90)
        subplot(5,10,2*i) %plot reconstructed stimulus
        surf(reshape(vneg,28,28))
        axis([0 28 0 28])
        view(90,90)
    end
end
rmseTest = sqrt(rmseTest/ntest); %convert to root mean squared error
disp(['Train rmse: ' num2str(rmseTrain) ', Test rmse: ' num2str(rmseTest)])
