text: # Copyright 2016 James Hensman, alexggmatthews
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  # http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  
  from __future__ import division, print_function
  import numpy as np
  
  
  def sample_HMC(f, num_samples, Lmin, Lmax, epsilon, x0, verbose=False,
                 thin=1, burn=0, RNG=np.random.RandomState(0),
                 return_logprobs=False):
      """
      A straight-forward HMC implementation. The mass matrix is assumed to be the
      identity.
  
      f is a python function that returns the energy and its gradient
  
        f(x) = E(x), dE(x)/dx
  
      we then generate samples from the distribution
  
        pi(x) = exp(-E(x))/Z
  
      - num_samples is the number of samples to generate.
      - Lmin, Lmax, epsilon are parameters of the HMC procedure to be tuned.
      - x0 is the starting position for the procedure.
      - verbose is a flag which turns on the display of the running accept ratio.
      - thin is an integer which specifies the thinning interval
      - burn is an integer which specifies how many initial samples to discard.
      - RNG is a random number generator
      - return_logprobs is a boolean indicating whether to return the log densities alongside the samples.
  
      The total number of iterations is given by
  
        burn + thin * num_samples
  
      The return shape is always num_samples x D.
  
      The leafrog (Verlet) integrator works by picking a random number of steps
      uniformly between Lmin and Lmax, and taking steps of length epsilon.
      """
  
      # an array to store the logprobs in (even if the user doesn't want them)
      logprob_track = np.empty(num_samples)
  
      # burn some samples if needed.
      if burn > 0:
          if verbose:
              print("burn-in sampling started")
          samples = sample_HMC(f, num_samples=burn, Lmin=Lmin, Lmax=Lmax,
                               epsilon=epsilon, x0=x0,
                               verbose=verbose, thin=1, burn=0, RNG=RNG)
          if verbose:
              print("burn-in sampling ended")
          x0 = samples[-1]
  
      D = x0.size
      samples = np.zeros((num_samples, D))
      samples[0] = x0.copy()
      x = x0.copy()
      logprob, grad = f(x0)
      logprob, grad = -logprob, -grad
      logprob_track[0] = logprob
  
      accept_count_batch = 0
  
      for t in range(1, num_samples * thin):
  
          # Output acceptance rate every 100 iterations
          if ((t+1) % 100) == 0:
              if verbose:
                  print("Iteration: ", t+1,
                        "\t Acc Rate: ", 1. * accept_count_batch, "%")
              accept_count_batch = 0
  
          # make a copy of the old state.
          x_old, logprob_old, grad_old = x.copy(), logprob, grad.copy()
          p_old = RNG.randn(D)
  
          # Standard HMC - begin leapfrogging
          premature_reject = False
          p = p_old + 0.5 * epsilon * grad
          for l in range(RNG.randint(Lmin, Lmax)):
              x += epsilon * p
              logprob, grad = f(x)
              logprob, grad = -logprob, -grad
              if np.any(np.isnan(grad)):  # pragma: no cover
                  premature_reject = True
                  break
              p += epsilon * grad
          p -= 0.5*epsilon * grad
          # leapfrogging done
  
          # reject the proposal if there are numerical errors.
          if premature_reject:  # pragma: no cover
              print("warning: numerical instability.\
                    Rejecting this proposal prematurely")
              x, logprob, grad = x_old, logprob_old, grad_old
              if t % thin == 0:
                  samples[t // thin] = x_old
                  logprob_track[t // thin] = logprob_old
              continue
  
          # work out whether to accept the proposal
          log_accept_ratio = logprob - 0.5 * p.dot(p) -\
              logprob_old + 0.5 * p_old.dot(p_old)
          logu = np.log(RNG.rand())
  
          if logu < log_accept_ratio:  # accept
              if t % thin == 0:
                  samples[t // thin] = x
                  logprob_track[t // thin] = logprob
              accept_count_batch += 1
          else:  # reject
              if t % thin == 0:
                  samples[t // thin] = x_old
                  logprob_track[t // thin] = logprob_old
              x, logprob, grad = x_old, logprob_old, grad_old
      if return_logprobs:
          return samples, logprob_track
      else:
          return samples
  