% Demonstrate (stochastic) gradient descent in 1D datasets with:
% - MLP regression: f(x) = W2.s(W1.x + w1) + w2 (one hidden sigmoidal layer)
%   trained with least-squares error.

% Dataset for 1D regression
rng(1778);
N = 10; Nv = 50; D = 1; Dy = 1;
a = -0.6; b = 0.4; f = @(x) a*x + b; s = 0.1;		% Linear fcn & noise
%a = 3; b = -0.5; f = @(x) sin(a*x) + b; s = 0.05;	% Nonlinear fcn & noise
X = 3*rand(N,D); Y = f(X) + s*randn(N,Dy);		% Training set
Xv = 3*rand(Nv,D); Yv = f(Xv) + s*randn(Nv,Dy);		% Validation set
ax = [-1 4 -2 1];					% Range: y-x plot

% Error of the true function (> optimal error because of the noise added)
e = Y - f(X); ef = e(:)'*e(:)/2; ef = ef/N;		% Optimal error

% MLP hyperparameters
H = 2;							% # hidden units
l = 0;							% Weight decay

% Initial weights
g0.W1 = randn(H,D+1)/10; g0.W2 = randn(Dy,H+1)/10; g0.type='mlp';
maxit = 200;						% Max # iterations
I = ceil(maxit/100);				% Save every Ith iteration

% MLP, iteratively with gradient descent
eta = 3e-2;						% Step size
[GDg,GDe] = mlpgd(X,Y,Xv,Yv,l,eta,maxit,g0,I);
GDe = diag([N Nv].^(-1))*GDe;

% MLP, iteratively with stochastic gradient descent
Eta = 1e-1;						% Step size
B = 1;							% Minibatch size
[SGDg,SGDe] = mlpsgd(X,Y,Xv,Yv,l,eta,maxit,B,g0,I);
SGDe = diag([N Nv].^(-1))*SGDe;


% ---------------------------------- Plots ----------------------------------

% Dataset and regression function estimates
x = linspace(ax(1),ax(2),100)'; fx = f(x);
GDx = mlp(x,GDg{1}); SGDx = mlp(x,SGDg{1});
figure(1); clf; hold on;
set(gcf,'Name','1D dataset; MLP regression by least-squares error');
h = plot(X,Y,'bo',Xv,Yv,'g+',x,fx,'k-');
h1 = plot(x,GDx,'r-','LineWidth',2); h2 = plot(x,SGDx,'m-','LineWidth',2);
legend([h;h1;h2],'training','validation','true f','GD f','SGD f');
axis(ax); daspect([1 1 1]); box on; xlabel('x'); ylabel('y');
title(['H=' num2str(H) ' hidden units, weight decay \lambda=' num2str(l)...
       ', iteration (epoch): 0']);

% Learning curves: error over iterations
figure(2); clf;
set(gcf,'Name','1D dataset; MLP regression by least-squares error');
ii=80; z = [GDe(:,ii:end);SGDe(:,ii:end)];	% zoom in on long-term iterates
for i=1:2
  subplot(2,1,i);
  plot([0 maxit],[ef ef],'k-',...
       0:I:maxit,GDe(1,:),'r-',0:I:maxit,SGDe(1,:),'m-',...
       0:I:maxit,GDe(2,:),'r--',0:I:maxit,SGDe(2,:),'m--');
  hold on; box on; xlabel('iteration (epoch)');
  h3{i} = plot(0,GDe(1,1),'ro',0,SGDe(1,1),'mo');
  if i==1
    axis([0 maxit 0 max(GDe(1,1),SGDe(1,1))]); ylabel('error');
    title(['H=' num2str(H) ' hidden units, weight decay \lambda=' num2str(l)...
           '; \eta_{GD} = ' num2str(eta) ', \eta_{SGD} = ' num2str(Eta) ...
           ', B = ' num2str(B)]);
    legend('true f','GD f (training)','SGD f (training)',...
           'GD f (validation)','SGD f (validation)');
  else
    axis([0 maxit min(z(:)) max(z(:))]); ylabel('error (zoom in)');
  end
end

% Plots over iterates
for i=2:length(GDg)
  pause;
  GDx = mlp(x,GDg{i}); SGDx = mlp(x,SGDg{i});
  set(0,'CurrentFigure',1); delete(h1); delete(h2);
  h1 = plot(x,GDx,'r-','LineWidth',2); h2 = plot(x,SGDx,'m-','LineWidth',2);
  title(['H=' num2str(H) ' hidden units, weight decay \lambda=' num2str(l)...
         ', iteration (epoch): ' num2str(I*(i-1))]);
  set(0,'CurrentFigure',2); delete(h3{1}); delete(h3{2});
  subplot(2,1,1); h3{1} = plot((i-1)*I,GDe(1,i),'ro',(i-1)*I,SGDe(1,i),'mo');
  subplot(2,1,2); h3{2} = plot((i-1)*I,GDe(1,i),'ro',(i-1)*I,SGDe(1,i),'mo');
  drawnow;
end

% Training & validation error over number of hidden units H
% We use GD with a large maxit to ensure with are close to a minimum
eta = 5e-3; maxit = 1e5;
HH = [1 2 5 10 30 50]; l = 0;
eH = zeros(2,length(HH)); gH = cell(1,length(HH));
gHx = zeros(length(x),length(HH));
for i=1:length(HH)
  g0.W1 = randn(HH(i),D+1)/10; g0.W2 = randn(Dy,HH(i)+1)/10; g0.type='mlp';
  [G,E] = mlpgd(X,Y,Xv,Yv,l,eta,maxit,g0,maxit);
  gH{i} = G{end}; eH(:,i) = diag([N Nv].^(-1))*E(:,end);
  gHx(:,i) = mlp(x,gH{i});
end
figure(3); clf;
set(gcf,'Name','1D dataset; MLP regression by least-squares error');
plot(HH,eH(1,:),'bo-',HH,eH(2,:),'ro-',HH([1 end]),[ef ef],'k-');
legend('training','validation','true'); xlabel('H'); ylabel('error');
title(['Weight decay \lambda=' num2str(l)]);
figure(4); clf;
set(gcf,'Name','1D dataset; MLP regression by least-squares error');
h = plot(X,Y,'bo',Xv,Yv,'g+',x,fx,'k-');
hold on; h1 = plot(x,gHx,'LineWidth',2); hold off;
legend([h(end);h1],['true f';repmat('H = ',length(HH),1) num2str(HH')]);
axis(ax); daspect([1 1 1]); box on; xlabel('x'); ylabel('y');
title(['Weight decay \lambda=' num2str(l)]);

% Training & validation error over weight decay parameter l
H = 10; ll = [10.^([-15 -5 -2 -0.5])];
el = zeros(2,length(ll)); gl = cell(1,length(ll));
glx = zeros(length(x),length(ll)); gw = zeros(H*(D+1)+Dy*(H+1),length(ll));
g0.W1 = randn(H,D+1)/10; g0.W2 = randn(Dy,H+1)/10; g0.type='mlp';
for i=1:length(ll)
  [G,E] = mlpgd(X,Y,Xv,Yv,ll(i),eta,maxit,g0,maxit);
  gl{i} = G{end}; el(:,i) = diag([N Nv].^(-1))*E(:,end);
  glx(:,i) = mlp(x,gl{i}); gw(:,i) = [gl{i}.W1(:);gl{i}.W2(:)];
end
figure(5); clf;
set(gcf,'Name','1D dataset; MLP regression by least-squares error');
subplot(2,1,1);
semilogx(ll,el(1,:),'bo-',ll,el(2,:),'ro-',ll([1 end]),[ef ef],'k-');
legend('training','validation','true'); xlabel('\lambda'); ylabel('error');
title(['H=' num2str(H) ' hidden units']);
subplot(2,1,2); semilogx(ll,mean(abs(gw),1),'bo-',ll,gw,'k.');
xlabel('\lambda'); ylabel({'weights: distribution','and mean magnitude'});
figure(6); clf;
set(gcf,'Name','1D dataset; MLP regression by least-squares error');
h = plot(X,Y,'bo',Xv,Yv,'g+',x,fx,'k-');
hold on; h1 = plot(x,glx,'LineWidth',2); hold off;
legend([h(end);h1],['true f         ';...
                    repmat('\lambda = ',length(ll),1) num2str(ll','%4.0e')]);
axis(ax); daspect([1 1 1]); box on; xlabel('x'); ylabel('y');
title(['H=' num2str(H) ' hidden units']);


% Suggestions of things to try:
% - Note: the runtime can be long even for small datasets.
% - SGD is much faster than GD if N is large enough: for the linear function:
%     N=100; H=2; eta=3e-3; Eta=1e-1; maxit=200;
% - For the nonlinear fcn:
%   . This gives a reasonably well estimated function:
%       H=10; eta=1e-2; Eta=0.3; maxit=2e4; I=100;
%     which gets better if increasing maxit (but takes longer).
%     What happens if taking the step size too big? Try with GD:
%       H=10; eta=2e-2; Eta=0.3; maxit=2e4; I=100;
%   . This shows overfitting (run with GD only because SGD will take too long):
%       H=10; eta=1e-2; maxit=3e5; I=3000;
% - Number of hidden units H = 1 2 5 10 100.
% - Weight decay l = 0 1e-5 1e-4 1e-3 1e-2 1e-1.
% - GD: eta = ...
% - SGD: Eta = ...
%   shuffling or not shuffling.
% - Vary the initial weights: U[-0.1,0.1] works, U[-10,10] gets stuck.
% - Dataset: vary noise level s, #points N, true function.
% - Local optima.

