% Run model system simulation and analysis.

% PARAMETERS
xmin = -1.75;
xmax = +1.75;

x0 = -1; % initial point for simulation
k0 = 1; % initial temperature index for simulation
kT = 1.0; % thermal energy
beta = 1.0 / kT;

ntemps = 15;
%temperature_k = linspace(1.0, 10.0, ntemps);
temperature_k = 1.0*(10.0/1.0).^linspace(0,1,ntemps);
beta_k = temperature_k .^ (-1);

% Define potential
%K = 10;
%potential = @(x) K * (x-1).^2 .* (x+1).^2 + 0.1*x; % asymmetric double well

% Plot potential and stationary distribution at multiple temperatures.
npoints = 1000;
dx = (xmax - xmin) / (npoints-1);
xpoints = linspace(xmin, xmax, npoints);
Upoints = potential(xpoints);

figure(1);
clf;
subplot(2,1,1);
plot(xpoints, Upoints, 'k-');
axis([xmin xmax min(Upoints) max(Upoints)]);
xlabel('x');
ylabel('U(x)');
box on;

subplot(2,1,2);
hold on;
Pmax = 0.0;
for k = 1:ntemps
  beta = beta_k(k);
  Peq = exp(-beta*Upoints);
  Peq = Peq / sum(Peq) * dx;
  plot(xpoints, Peq);
  Pmax = max(Pmax, max(Peq));
end
axis([xmin xmax 0 Pmax]);
xlabel('x');
ylabel('\pi(x)');
set(gca,'XTick',[]);
box on;

filename = 'model-expanded-ensemble-potential.eps';
exportfig(gcf, filename, 'width', 7.5, 'height', 10, 'color', 'cmyk');
system(sprintf('epstopdf %s', filename));

% Compute free energies of each temperature.
Z_k = zeros(1,ntemps);
f_k = zeros(1,ntemps);
for k = 1:ntemps
  beta = beta_k(k);
  Peq = exp(-beta*Upoints);
  Z_k(k) = sum(Peq) * dx;
  f_k(k) = - log(Z_k(k));
end
f_k

% Set log weights for expanded ensembles simulation.
g_k = f_k;

% Run expanded ensemble simulation.
niterations = 100000;
nshowiterations = 5000;
x = x0;
k = k0;
U = potential(x);
nsteps_per_iteration = 100;
dx = 0.1;
nxattempted = 0;
nxaccepted = 0;
temperature_update_scheme = 'neighbor';
%temperature_update_scheme = 'gibbs   ';
k_t = zeros(niterations+1,1);
x_t = zeros(niterations+1,1);
k_t(1) = k0;
x_t(1) = x0;
for i = 1:niterations
%  if (i % 100 == 0)
%    disp(i);
%  end

  % Update temperature.
  if (temperature_update_scheme == 'neighbor')
    if (rand < 0.5)
      k_proposed = k-1;
    else
      k_proposed = k+1;
    end
    if (k_proposed >= 1) && (k_proposed <= ntemps)
      log_P_accept = -beta_k(k_proposed)*U + g_k(k_proposed) + beta_k(k) - g_k(k);
      if (rand < exp(log_P_accept))
	k = k_proposed;
      end
    end
  else
    % Gibbs update of temperature index.
    P_k = exp(-beta_k * U + g_k);
    P_k = P_k / sum(P_k);
    k = randsample(1:ntemps,1,true,P_k);
  end

  % Update position.
  for j = 1:nsteps_per_iteration
    x_proposed = x + dx * randn;
    U_proposed = potential(x_proposed);
    log_P_accept = -beta_k(k) * U_proposed + beta_k(k) * U;
    nxattempted = nxattempted + 1;
    if (rand < exp(log_P_accept))
      % Accept.
      nxaccepted = nxaccepted + 1;
      x = x_proposed;
      U = U_proposed;
    end  
  end

  k_t(i+1) = k;
  x_t(i+1) = x;
end

% Compute correlation times.
g_k = statistical_inefficiency(k_t);
g_x = statistical_inefficiency(x_t);

disp(sprintf('Accepted %d / %d (%.1f%%) position moves.', nxaccepted, nxattempts, nxaccepted/nxattempted*100));
figure(2);
clf;
subplot(2,1,1);
plot(k_t, 'r-');
axis([0 nshowiterations 0 ntemps+1]);
title(sprintf('g_k = %.1f', g_k));
xlabel('iteration');
ylabel('temperature index');
box on;

subplot(2,1,2);
plot(x_t, 'k.');
axis([0 nshowiterations xmin xmax]);
title(sprintf('g_x = %.1f', g_x));
xlabel('iteration');
ylabel('x');
box on;

filename = 'model-expanded-ensemble-neighbor.eps';
%filename = 'model-expanded-ensemble-gibbs.eps';
exportfig(gcf, filename, 'width', 7.5, 'height', 10, 'color', 'cmyk');
system(sprintf('epstopdf %s', filename));

