% Simple 1D model system illustrating single large perturbation followed by many steps of Metropolis MC propagation in a harmonic potential.
% JDC 2011-08-12

% Define potential.
potential = @(x) (1/2)*(x.^2);

% Define run parameters.
nsteps_ncmc = 10000000; % number of NCMC steps
dx_ncmc = 1.0; % stddev of outer NCMC trial moves
nsteps_mmc = 1; % number of Metropolis MC steps per NCMC step
dx_mmc = 0.3; % stddev of inner MMC trial moves
nequil = round(0.1 * nsteps_ncmc); % discard 10% to equilibration
p_aK = 0.9; % probability of forward (alpha K) protocol
ignore_protocol_ratio = 0; % if 1, will ignore inclusion of P(\tilde{\Lambda})/P(\Lambda) in Eq. 11 acceptance computation.

% Jerome's overrides.
dx_ncmc=10; dx_mmc=1;

% Compute probability of reverse protocol.
p_Ka = 1.0 - p_aK;

% Initialize sampler.
x_current = 0.0; % sampler location
xhist = zeros(nsteps_ncmc,1); % xhist(i) is position of sampler at iteration i
naccepted_ncmc = 0; % number of accepted NCMC moves


% Sample.
for ncmc_iteration = 1:nsteps_ncmc
  if (rand < p_aK)
    % alpha K order

    % Make perturbation, accumulating work.
    x_proposed = x_current + dx_ncmc*randn();  
    work = potential(x_proposed) - potential(x_current);
    
    % Take propagation steps.
    naccepted_mmc = 0; % number of inner MMC steps accepted
    for mmc_step = 1:nsteps_mmc
      x_proposed_mmc = x_proposed + dx_mmc*randn();
      du = potential(x_proposed_mmc) - potential(x_proposed);
      if (du < 0.0) || (rand() < exp(-du))
	% Accept.
	x_proposed = x_proposed_mmc;
	naccepted_mmc = naccepted_mmc + 1;
      else
	% Reject.
      end
    end

    if ignore_protocol_ratio
      log_P_accept = - work;
    else
      % use Eq. 11
      log_P_accept = - work + log(p_Ka / p_aK);
    end
  else
    % K alpha order
    
    x_proposed = x_current;

    % Take propagation steps.
    naccepted_mmc = 0; % number of inner MMC steps accepted
    for mmc_step = 1:nsteps_mmc
      x_proposed_mmc = x_proposed + dx_mmc*randn();
      du = potential(x_proposed_mmc) - potential(x_proposed);
      if (du < 0.0) || (rand() < exp(-du))
	% Accept.
	x_proposed = x_proposed_mmc;
	naccepted_mmc = naccepted_mmc + 1;
      else
	% Reject.
      end
    end
    
    % Make perturbation, accumulating work.
    x_proposed_ncmc = x_proposed + dx_ncmc*randn();  
    work = potential(x_proposed_ncmc) - potential(x_proposed);
    x_proposed = x_proposed_ncmc;
    
    if ignore_protocol_ratio
      log_P_accept = - work;
    else
      % use eq 11
      log_P_accept = - work + log(p_aK / p_Ka);
    end
  end    
    
  % Accept or reject NCMC.
  if (log_P_accept > 0.0) || (rand() < exp(log_P_accept))
    % Accept.
    x_current = x_proposed;
    naccepted_ncmc = naccepted_ncmc + 1;
  else
    % Reject.
  end

  % Store.
  xhist(ncmc_iteration) = x_current;

  % Display.
  if mod(ncmc_iteration, 1000) == 0
    disp(sprintf('Iteration %d | work %8.3f | log(P_accept) %8.3f | accepted %8.1f%% of inner MMC steps', ncmc_iteration, work, log_P_accept, (naccepted_mmc/nsteps_mmc)*100.0));
  end
end

% Discard to equilibration.
xhist = xhist(nequil:end);

% Display.
disp(sprintf('Accepted %.3f%% of NCMC moves.', (naccepted_ncmc/nsteps_ncmc)*100.0));
clf;
nbins = 100;  % number of histogram bins
[N,x] = hist(xhist, nbins);
Px_obs = N / sum(N);
g = statistical_inefficiency(xhist); % compute statistical inefficiency
disp(sprintf('statistical inefficiency = %.1f iterations', g));
dPx_obs = sqrt(Px_obs .* (1-Px_obs)) ./ sqrt(N/g); % binomial
hold on;
Px_true = exp(-potential(x)); Px_true = Px_true / sum(Px_true);
%plot(x,N,'k-',x,Px,'r-')
errorbar(x,Px_obs,2*dPx_obs,2*dPx_obs,'k.'); % 95% confidence interval
plot(x,Px_true,'r-');
set(gca,'YTick', []);
xlabel('x');
ylabel('P(x)');
legend('observed', 'expected');

