% function mh(iters, burnin, psig, pausevery)
%
% Metropolis-Hastings sampler demonstration
%
% iters: how many iterations total
% burnin: how many iterations for burnin
% psig: std dev of proposal step (e.g. .1, .25)
% pausevery: stop and display progress this often (0 for never)
%
% For example,
%   mh(5000, 1000, .25, 0)
%   mh(1000, 300, .1, 20)

function mh(iters, burnin, psig, pausevery)

parallel = 1;				% number of parallel
                                        % trajectories to display
                                        % at once

[xs, ys] = meshgrid(-1:.05:1);

% plot function to integrate
figure(1);
zs = foo(xs,ys);
surfl(xs,ys,zs);
shading interp;
xlabel X
ylabel Y
zlabel f(X,Y)
set(gca,'CameraPosition', [13.017 -4.10426 29.1589])
% print -depsc foo.eps

% get true integral x^2*foo(x,y) by fine discretization
zs = zs./sum(zs(:));
fprintf('True E(x^2) %f\n', sum(sum(zs.*xs.*xs)));

% contour plot of function to integrate, on top of which we will
% plot progress of MH
figure(2);
contourf(xs,ys,foo(xs,ys));

% initial point for MH
x = zeros(parallel,1);
y = x;
v = foo(x,y);

% record the samples and acceptance decisions
samples = zeros(parallel, iters);
accept = zeros(parallel, iters);

% the MH sampler
for i = 1:iters

  % sample from proposal distribution (and possibly plot)
  nx = x + randn(size(x))*psig;
  ny = y + randn(size(y))*psig;
  nv = foo(nx, ny);
  if (pausevery > 0 && mod(i, pausevery)==0)
    h = line(nx, ny, 'Marker', 'o', 'MarkerSize', 15, 'LineWidth', 2, ...
	     'Color', [1 .5 .5], 'LineStyle', 'none');
  end

  % compute acceptance probability and decide which proposals to
  % accept; note Q is symmetric, so we don't need its ratio
  accprob = nv ./ v;
  mask = rand(size(x))>accprob;
  nx(mask) = x(mask);
  ny(mask) = y(mask);

  % add this step to the plot
  if (pausevery > 0)
    line([x nx]', [y ny]');
  else
    line(nx, ny, 'Marker', '.');
  end
  
  % do the update and remember the sample
  x = nx;
  y = ny;
  samples(:,i) = x.^2;
  accept(:,i) = ~mask;

  % pause for a plot if desired
  if (pausevery > 0 && mod(i, pausevery)==0)
    fprintf('acceptance rate %f\n', mean(mean(accept(:,1:i))));
    if (i > burnin)
      fprintf('current estimate %f\n', mean(mean(samples(:,burnin:i))));
    end
    pause
    delete(h)
  end
end

fprintf('final acceptance rate %f\n', mean(mean(accept)));
fprintf('final estimate %f\n', mean(mean(samples)));

% print -depsc mh-path.eps

return;


function zs = foo(xs, ys)

zs = exp(sin(2*pi*(xs+ys))-xs-2*(xs-.3).^2-2*ys.^2)+2*exp(-(xs-ys).^2);
zs(xs<-1) = 0;
zs(xs>1) = 0;
zs(ys<-1) = 0;
zs(ys>1) = 0;

return;