function [samples, energies, diagn] = hmc(f, x, options, gradf, varargin)
%HMC	Hybrid Monte Carlo sampling.
%
%	Description
%	SAMPLES = HMC(F, X, OPTIONS, GRADF) uses a  hybrid Monte Carlo
%	algorithm to sample from the distribution P ~ EXP(-F), where F is the
%	first argument to HMC. The Markov chain starts at the point X, and
%	the function GRADF is the gradient of the `energy' function F.
%
%	HMC(F, X, OPTIONS, GRADF, P1, P2, ...) allows additional arguments to
%	be passed to F() and GRADF().
%
%	[SAMPLES, ENERGIES, DIAGN] = HMC(F, X, OPTIONS, GRADF) also returns a
%	log of the energy values (i.e. negative log probabilities) for the
%	samples in ENERGIES and DIAGN, a structure containing diagnostic
%	information (position, momentum and acceptance threshold) for each
%	step of the chain in DIAGN.POS, DIAGN.MOM and DIAGN.ACC respectively.
%	All candidate states (including rejected ones) are stored in
%	DIAGN.POS.
%
%	[SAMPLES, ENERGIES, DIAGN] = HMC(F, X, OPTIONS, GRADF) also returns
%	the ENERGIES (i.e. negative log probabilities) corresponding to the
%	samples.  The DIAGN structure contains three fields:
%
%	POS the position vectors of the dynamic process.
%
%	MOM the momentum vectors of the dynamic process.
%
%	ACC the acceptance thresholds.
%
%	S = HMC('STATE') returns a state structure that contains the state of
%	the two random number generators RAND and RANDN and the momentum of
%	the dynamic process.  These are contained in fields  randstate,
%	randnstate and mom respectively.  The momentum state is only used for
%	a persistent momentum update.
%
%	HMC('STATE', S) resets the state to S.  If S is an integer, then it
%	is passed to RAND and RANDN and the momentum variable is randomised.
%	If S is a structure returned by HMC('STATE') then it resets the
%	generator to exactly the same state.
%
%	The optional parameters in the OPTIONS vector have the following
%	interpretations.
%
%	OPTIONS(1) is set to 1 to display the energy values and rejection
%	threshold at each step of the Markov chain. If the value is 2, then
%	the position vectors at each step are also displayed.
%
%	OPTIONS(5) is set to 1 if momentum persistence is used; default 0,
%	for complete replacement of momentum variables.
%
%	OPTIONS(7) defines the trajectory length (i.e. the number of leap-
%	frog steps at each iteration).  Minimum value 1.
%
%	OPTIONS(9) is set to 1 to check the user defined gradient function.
%
%	OPTIONS(14) is the number of samples retained from the Markov chain;
%	default 100.
%
%	OPTIONS(15) is the number of samples omitted from the start of the
%	chain; default 0.
%
%	OPTIONS(16) is the size of the acceptance window; default 1.
%
%	OPTIONS(17) defines the momentum used when a persistent update of
%	(leap-frog) momentum is used.  This is bounded to the interval [0,
%	1).
%
%	OPTIONS(18) is the step size used in leap-frogs; default 1/trajectory
%	length.
%
%	See also
%	METROP
%

%	Copyright (C) 1996-1998 Ian T Nabney, Christopher M Bishop
%	Copyright (C) 1999-2000 Aki Vehtari

%  The portion of the HMC algorithm involving "windows" is derived
%  from the C code for this function included in the Software for
%  Flexible Bayesian Modeling written by Radford Neal, Version of
%  1999-03-1 <http://www.cs.toronto.edu/~radford/fbm.software.html>. 
%  The translation to Matlab was performed by Aki Vehtari, who is
%  solely responsible for any errors introduced by this
%  translation. Radford Neal has given permission to distribute
%  this translation of his software under the terms of the GNU
%  Library General Public License version 2, or any essentially
%  similar licence, provided this acknowledgement is included.

% Global variable to store state of momentum variables: set by set_state
% Used to initialise variable if set
global HMC_MOM
if nargin <= 2
  if ~strcmp(f, 'state')
    error('Unknown argument to hmc');
  end
  switch nargin
    case 1
      samples = get_state(f);
      return;
    case 2
      set_state(f, x);
      return;
  end
end

display = options(1);
if (round(options(5) == 1))
  persistence = 1;
  % Set alpha to lie in [0, 1)
  alpha = max(0, options(17));
  alpha = min(1, alpha);
  salpha = sqrt(1-alpha*alpha);
else
  persistence = 0;
end
L = max(1, options(7)); % At least one step in leap-frogging
if options(14) > 0
  nsamples = options(14);
else
  nsamples = 100;	% Default
end
if options(15) >= 0
  nomit = options(15);
else
  nomit = 0;
end
if options(16) >= 1
  window = options(16);              % Window size.
else
  window = 1;                        % Default
end
if options(18) > 0
  step_size = options(18);	% Step size.
else
  step_size = 1/L;		% Default  
end
x = x(:)';		% Force x to be a row vector
nparams = length(x);

% Set up strings for evaluating potential function and its gradient.
f = fcnchk(f, length(varargin));
gradf = fcnchk(gradf, length(varargin));

% Check the gradient evaluation.
if (options(9))
  % Check gradients
  feval('gradchek', x, f, gradf, varargin{:});
end

samples = zeros(nsamples, nparams);	% Matrix of returned samples.
if nargout >= 2
  en_save = 1;
  energies = zeros(nsamples, 1);
else
  en_save = 0;
end
if nargout >= 3
  diagnostics = 1;
  diagn_pos = zeros(nsamples, nparams);
  diagn_mom = zeros(nsamples, nparams);
  diagn_acc = zeros(nsamples, 1);
else
  diagnostics = 0;
end

k = - nomit + 1;
Eold = feval(f, x, varargin{:});	% Evaluate starting energy.
nreject = 0;
if (~persistence | isempty(HMC_MOM))
  p = randn(1, nparams);		% Initialise momenta at random
else
  p = HMC_MOM;				% Initialise momenta from stored state
end
lambda = 1;
window_offset=0;

% Main loop.
while k <= nsamples

  xold = x;		    % Store starting position.
  pold = p;		    % Store starting momenta
  Hold = Eold + 0.5*(p*p'); % Recalculate Hamiltonian as momenta have changed

  if ~persistence
    % Choose a direction at random
    if (rand(1) < 0.5)
      lambda = -1;
    else
      lambda = 1;
    end
  end
  % Perturb step length.
  epsilon = lambda*step_size*(1.0 + 0.1*randn(1));

  % Decide on window offset.
  if window>1
    window_offset=fix(window*rand(1));
  end
    
  have_rej = 0;
  have_acc = 0;
  n = window_offset;
  dir = -1;
  while (dir==-1 | n~=L)

    if (dir==-1 & n==0)
      % Restore, next state should be original start state.
      if window_offset > 0
	x = xold;
	p = pold;
	n = window_offset;
      end
      H = Hold;
      dir = 1;
      stps = dir;
    else
      if (n<window | n>(L-window))
	% State in the accept and/or reject window.
	stps = dir;
      else
	% State not in the accept and/or reject window.
	stps = L-2*window+1;
      end
      % First half-step of leapfrog.
      p = p - dir*0.5*epsilon.*feval(gradf, x, varargin{:});
      x = x + dir*epsilon.*p;
      % Full leapfrog steps.
      for m = 1:(abs(stps)-1)
	p = p - dir*epsilon.*feval(gradf, x, varargin{:});
	x = x + dir*epsilon.*p;
      end
      % Final half-step of leapfrog.
      p = p - 0.5*epsilon.*feval(gradf, x, varargin{:});

      H = feval(f, x, varargin{:}) + 0.5*(p*p');
      
      n=n+stps;
    end


    if (window~=(L+1) & n<window)
      % Account for state in reject window.  Reject window can be
      % ignored if windows consist of the entire trajectory.
      if ~have_rej
	rej_free_energy = H;
      else
	rej_free_energy = -addlogs(-rej_free_energy, -H);
      end
      if (~have_rej | rand(1) < exp(rej_free_energy-H));
	x_rej=x;
	p_rej=p;
        have_rej = 1;
      end
    end

    if (n>(L-window))
      % Account for state in the accept window.
      if ~have_acc
	acc_free_energy = H;
      else
	acc_free_energy = -addlogs(-acc_free_energy, -H);
      end
      if (~have_acc | rand(1) < exp(acc_free_energy-H))
	x_acc=x;
	p_acc=p;
	have_acc = 1;
      end
    end
  end

  % Acceptance threshold.
  a = exp(rej_free_energy - acc_free_energy );

  if (diagnostics & k > 0)
    diagn_pos(k,:) = x_acc;
    diagn_mom(k,:) = p_acc;
    diagn_acc(k,:) = a;
  end
  if (display > 1)
    fprintf(1, 'New position is\n');
    disp(x);
  end

  % Take new state from the appropriate window.
  if a > rand(1)
    % Accept
    Eold=feval(f, x, varargin{:});      % Update energy
    x=x_acc;
    p=-p_acc;                           % Reverse momenta
    if (display > 0)
      fprintf(1, 'Finished step %4d  Threshold: %g\n', k, a);
    end
  else
    % Reject
    if k > 0
      nreject = nreject + 1;
    end
    x=x_rej;
    p=p_rej;
    if (display > 0)
      fprintf(1, '  Sample rejected %4d.  Threshold: %g\n', k, a);
    end
  end
  
  if k > 0
    samples(k,:) = x;			% Store sample.
    if en_save 
      energies(k) = Eold;		% Store energy.
    end
  end

  % Set momenta for next iteration
  if persistence
    p = -p;
    % Adjust momenta by a small random amount.
    p = alpha.*p + salpha.*randn(1, nparams);
  else
    p = randn(1, nparams);	% Replace all momenta.
  end

  k = k + 1;
end

if (display > 0)
  fprintf(1, '\nFraction of samples rejected:  %g\n', ...
    nreject/(nsamples));
end
if diagnostics
  diagn.pos = diagn_pos;
  diagn.mom = diagn_mom;
  diagn.acc = diagn_acc;
end
% Store final momentum value in global so that it can be retrieved later
HMC_MOM = p;
return

% Return complete state of sampler (including momentum)
function state = get_state(f)

global HMC_MOM
state.randstate = rand('state');
state.randnstate = randn('state');
state.mom = HMC_MOM;
return

% Set complete state of sampler (including momentum) or just set randn
% and rand with integer argument.
function set_state(f, x)

global HMC_MOM
if isnumeric(x)
  rand('state', x);
  randn('state', x);
  HMC_MOM = [];
else
  if ~isstruct(x)
    error('Second argument to hmc must be number or state structure');
  end
  if (~isfield(x, 'randstate') | ~isfield(x, 'randnstate') ...
      | ~isfield(x, 'mom'))
    error('Second argument to hmc must contain correct fields')
  end
  rand('state', x.randstate);
  randn('state', x.randnstate);
  HMC_MOM = x.mom;
end
return

% Add numbers represented by their logarithms.
% Computes log(exp(a)+exp(b)) in such a fashion that it works even
% when a and b have large magnitude.
function c=addlogs(a,b)
if a>b
  c = a + log(1+exp(b-a));
else
  c = b + log(1+exp(a-b));
end