#!/usr/bin/python

#=============================================================================================
# Test MBAR by performing statistical tests on a set of of 1D harmonic oscillators, for which
# the true free energy differences can be computed analytically.
#
# A number of replications of an experiment in which i.i.d. samples are drawn from a set of
# K harmonic oscillators are produced.  For each replicate, we estimate the dimensionless free
# energy differences and mean-square displacements (an observable), as well as their uncertainties.
#
# For a 1D harmonic oscillator, the potential is given by
#   V(x;K) = (K/2) * (x-x_0)**2
# where K denotes the spring constant.
#
# The equilibrium distribution is given analytically by
#   p(x;beta,K) = sqrt[(beta K) / (2 pi)] exp[-beta K (x-x_0)**2 / 2]
# The dimensionless free energy is therefore
#   f(beta,K) = - (1/2) * ln[ (2 pi) / (beta K) ]
#
#=============================================================================================

#=============================================================================================
# TODO
#=============================================================================================
# * Generate a plot after completion, similar to the plot from WHAM paper.
#=============================================================================================

#=============================================================================================
# VERSION CONTROL INFORMATION
#=============================================================================================
__version__ = "$Revision: 218 $ $Date: 2010-07-06 10:34:44 -0400 (Tue, 06 Jul 2010) $"
# $Date: 2010-07-06 10:34:44 -0400 (Tue, 06 Jul 2010) $
# $Revision: 218 $
# $LastChangedBy: mrshirts $
# $HeadURL: https://simtk.org/svn/pymbar/trunk/examples/harmonic-oscillators/harmonic-oscillators.py $
# $Id: harmonic-oscillators.py 218 2010-07-06 14:34:44Z mrshirts $

#=============================================================================================
# IMPORTS
#=============================================================================================
import pymbar
import numpy
import confidenceintervals
import matplotlib.pyplot as plt
import pdb
import sys
import entropy-utilities 
from entropy_utilities import covdiffs as covdiffs
from entropy_utilities import covdiffs3 as covdiffs3

#=============================================================================================
# PARAMETERS
#=============================================================================================

#if (len(sys.argv) < 2):
#  nall = 50 # default
#  generateplots = True
#else:
#  nall = int(sys.argv[1])

nall = 1000
generateplots = False

# Linear in K and linear in O
# K = Kscale*(0.1+0.9*lambda); lambda = 0, K=0.1*kscale; lambda=1, K = Kscale*1 
# O = Oscale*(lambda) 

Nstates = 11
K_limit = [0.1,1.0] 
O_limit = [0.0,1.0]
dK = (K_limit[1]-K_limit[0])/(Nstates-1.0)
dO = (O_limit[1]-O_limit[0])/(Nstates-1.0)

Kscale = 10.0
Oscale = 5
K_k = Kscale*numpy.arange(K_limit[0], K_limit[1]+0.5*dK,dK) # spring constants for each state
O_k = Oscale*numpy.arange(O_limit[0], O_limit[1]+0.5*dK,dO) # spring constants for each state
N_k = nall*numpy.ones(Nstates) # number of samples from each state (can be zero for some states)
beta = 1.0 # inverse temperature for all simulations
T = beta**(-1)

dfrac_samp = numpy.float(sys.argv[1])
dfrac_pert = numpy.float(sys.argv[2])

# two temperature differences that can be varied independently; 
# perturbation and actual samples

# how big is dT_samp?
dt_samp = T*dfrac_samp
Thigh_samp = T+dt_samp/2
Tlow_samp = T-dt_samp/2

#not the low beta, but the temperature equivalent to low beta
# betalow = 1.0/Tlow_samp
btlow_samp = (T-dt_samp/2)**(-1)
# betahigh = 1.0/Thigh_samp
bthigh_samp = (T+dt_samp/2)**(-1)

# how big is dT_pert?
dt_pert = T*dfrac_pert
Thigh_pert = T+dt_pert/2
Tlow_pert = T-dt_pert/2

#not the low beta, but the temperature equivalent to low beta
# betalow = 1.0/Tlow
btlow_pert = (T-dt_pert/2)**(-1)
# betahigh = 1.0/Thigh
bthigh_pert = (T+dt_pert/2)**(-1)

print "dfrac_samp:", dfrac_samp
print "Tlow_samp:", Tlow_samp
print "Thigh_samp:", Thigh_samp

print "dfrac_pert:", dfrac_pert
print "Tlow_pert:", Tlow_pert
print "Thigh_pert:", Thigh_pert

# Enthalpy differences
# a lot of the terms are the same
# dp[0] = dp[1] = -T/dt
# thus dp[1]f1 - dp[2]f2 = T*[f(T+dt/2) - f(T-dt/2)]/dt = u 

dpH_pert = [1,T/dt_pert,T/dt_pert]
dpH_samp = [1,T/dt_samp,T/dt_samp]

#dp[1] is (1/T)/*(dt/(T+dt/2))  = (T-dt/2)/(dt*T) = beta*(T+dt/2)/dt  
#dp[2] is (1/T)/*(dt/(T-dt/2))  = (T+dt/2)/(dt*T) = beta*(T-dt/2)/dt  
# thus dp[1]f1 - dp[2]f2 = beta*[(T+dt/2)f(T+dt/2) - beta*(T+dt/2)f(T+dt/2)]/dt
# Thus it will be beta*dF/dT = reduced entropy

dpS_pert = [1,beta/(dt_pert*btlow_pert),beta/(dt_pert*bthigh_pert)]
dpS_samp = dpS_pert = [1,beta/(dt_samp*btlow_samp),beta/(dt_samp*bthigh_samp)]

nreplicates = 200 # number of replicates of experiment for testing uncertainty estimate

def dhdl(Ks,Os,K,O,x):
  # assume O is always linear with zero intercept.
  lam = O/Os
  # For a 1D harmonic oscillator, the potential is given by
  #   V(x;K) = 0.5*Ks*(0.1+0.9*lam) * (x-Os*lambda))^2
  # where K denotes the spring constant, and O the offset
  # So dU/dlam = 0.5*0.9*Ks*(x-O(l))**2 - Ks*(0.1+0.9*lam)*(x-Os*lam)*Os. For now, assume linear scaling
  return 0.45*Ks*(x-O)**2 + K*Os*(x-Os*lam)

def setReplicate(replicate,type,df,ddf,du,ddu,ds,dds):

  df_err = df - Deltaf_ij_analytical
  du_err = du - Deltau_ij_analytical
  ds_err = ds - Deltas_ij_analytical

  # free energy
  str = 'df_' + type 
  replicate[str] = df.copy()
  str = 'ddf_' + type
  replicate[str] = ddf.copy()
  str = 'df_' + type + '_error'
  replicate[str] = df_err.copy()

  # average energy
  str = 'du_' + type 
  replicate[str] = du.copy()
  str = 'ddu_' + type
  replicate[str] = ddu.copy()
  str = 'du_' + type + '_error'
  replicate[str] = du_err.copy()

  #entropy
  str = 'ds_' + type 
  replicate[str] = ds.copy()
  str = 'dds_' + type
  replicate[str] = dds.copy()
  str = 'ds_' + type + '_error'
  replicate[str] = ds_err.copy()

def integrateij(dfdl,stddfdl,dsdl,stddsdl,dudl,stddudl,mode='trapezoid'):
  
  K = numpy.size(dfdl) # assume the others are the same length.
  # we are building a range of weights, for different integration lengths. 
  wk = numpy.zeros([K,K],float)
  dl = 1.0/(K-1)

  for Ki in range(1,K):
    if mode == 'trapezoid':
      wk[Ki,0] = 0.5*dl 
      wk[Ki,Ki] = 0.5*dl
      for k in range(1,Ki):
        wk[Ki,k] = dl
    elif (mode == 'simpsons'):
      if ((Ki%2)==0):
        # use trapezoid 
        wk[Ki,0] = 0.5*dl 
        wk[Ki,Ki] = 0.5*dl
        for k in range(1,Ki):
          wk[Ki,k] = dl
      else:
        wk[Ki,0] = dl/6 
        wk[Ki,Ki-1] = dl/6    
        for k in range(1,Ki-1):
          if (k%2):
            wk[Ki,k] = 4*dl/6
          else:
            wk[Ki,k] = 2*dl/6
    else:
      print "mode not implemented: garbage results!"

  df = numpy.zeros([K,K],float)
  ddf = numpy.zeros([K,K],float)
  du = numpy.zeros([K,K],float)
  ddu = numpy.zeros([K,K],float)
  ds = numpy.zeros([K,K],float)
  dds = numpy.zeros([K,K],float)

  for i in range(K):
    for j in range(i+1,K):
        l = j-i
        df[i,j] = numpy.dot(wk[l,0:l+1],dfdl[i:j+1])
        ddf[i,j] = numpy.sqrt(numpy.dot(wk[l,0:l+1]**2,stddfdl[i:j+1]**2))
        ds[i,j] = numpy.dot(wk[l,0:l+1],dsdl[i:j+1])
        dds[i,j] = numpy.sqrt(numpy.dot(wk[l,0:l+1]**2,stddsdl[i:j+1]**2))
        du[i,j] = numpy.dot(wk[l,0:l+1],dudl[i:j+1])
        ddu[i,j] = numpy.sqrt(numpy.dot(wk[l,0:l+1]**2,stddudl[i:j+1]**2))
  for i in range(K):
    for j in range(0,i):
       df[i,j] = df[j,i]
       ddf[i,j] = ddf[j,i]
       ds[i,j] = ds[j,i]
       dds[i,j] = dds[j,i]
       du[i,j] = du[j,i]
       ddu[i,j] = ddu[j,i]

  return (df,ddf,du,ddu,ds,dds)      
  
makeplots = False
# Uncomment the following line to seed the random number generated to produce reproducible output.
numpy.random.seed(1)

#=============================================================================================
# MAIN
#=============================================================================================

# Determine number of simulations.
K = numpy.size(N_k)
if numpy.shape(K_k) != numpy.shape(N_k): raise "K_k and N_k must have same dimensions."

# Determine maximum number of samples to be drawn for any state.
N_max = numpy.max(N_k)

# Compute widths of sampled distributions.
# For a harmonic oscillator with spring constant K,
# x ~ Normal(x_0, sigma^2), where sigma = 1/sqrt(beta K)
sigma_k = (beta * K_k)**-0.5
print "Gaussian widths:"
print sigma_k

# Compute the absolute dimensionless free energies of each oscillator analytically.
# f = - ln(sqrt((2 pi)/(beta K)) )
print 'Computing dimensionless free energies analytically...'
f_k_analytical = - numpy.log(numpy.sqrt(2 * numpy.pi) * sigma_k )
u_k_analytical = (1/2)*numpy.ones([K],float)  # By equipartition
s_k_analytical = u_k_analytical - f_k_analytical 

# Compute true free energy differences.
Deltaf_ij_analytical = numpy.zeros([K,K], dtype = numpy.float64)
Deltau_ij_analytical = numpy.zeros([K,K], dtype = numpy.float64)
Deltas_ij_analytical = numpy.zeros([K,K], dtype = numpy.float64)

for i in range(0,K):
  for j in range(0,K):
    Deltaf_ij_analytical[i,j] = f_k_analytical[j] - f_k_analytical[i]
    Deltau_ij_analytical[i,j] = u_k_analytical[j] - u_k_analytical[i]
    Deltas_ij_analytical[i,j] = s_k_analytical[j] - s_k_analytical[i]

print Deltaf_ij_analytical[0,:]
# DEBUG info
print "This script will perform %d replicates of an experiment where samples are drawn from %d harmonic oscillators." % (nreplicates, K)
print "The harmonic oscillators have equilibrium positions"
print O_k
print "and spring constants"
print K_k
print "and the following number of samples will be drawn from each (can be zero if no samples drawn):"
print N_k
print ""

# Conduct a number of replicates of the same experiment
replicates = [] # storage for one hash for each replicate
# Initialize storage for samples.
x_kn = numpy.zeros([K,N_max], dtype = numpy.float64) # x_kn[k,n] is the coordinate x of independent snapshot n of simulation k
U_kln = numpy.zeros([K,K,N_max], dtype = numpy.float64) # U_kmt(k,m,t) is the value of the energy for snapshot t from simulation k at potential m
dhdl_kln = numpy.zeros([K,K,N_max], dtype = numpy.float64) # dhdl_kmt(k,m,t) is the value of the derivative of the energy for snapshot t from simulation k at potential m

# alternate storage variables.
xt_kn = numpy.zeros([K,N_max], dtype = numpy.float64) # x_kn[k,n] is the coordinate x of independent snapshot n of simulation k
Ut_kln = numpy.zeros([K,K,N_max], dtype = numpy.float64) # U_kmt(k,m,t) is the value of the energy for snapshot t from simulation k at potential m

for replicate_index in range(0,nreplicates):
  print "Performing replicate %d / %d" % (replicate_index+1, nreplicates)

  replicate = {}
  #=============================================================================================
  # Generate independent data samples from K one-dimensional harmonic oscillators centered at q = 0.
  #=============================================================================================
  
  print 'generating samples...'

  # Generate samples.
  for k in range(0,K):
    # generate N_k[k] independent samples with spring constant K_k[k]
    x_kn[k,0:N_k[k]] = numpy.random.normal(O_k[k], sigma_k[k], [N_k[k]])
    
    # compute potential energy and derivative of all samples in all potentials
    for l in range(0,K):      
      U_kln[k,l,:] = (K_k[l]/2.0) * (x_kn[k,:]-O_k[l])**2
      dhdl_kln[k,l,:] = dhdl(Kscale,Oscale,K_k[l],O_k[l],x_kn[k,:])
    
  #=============================================================================================
  # Estimate free energies and expectations.
  #=============================================================================================

  # Estimate free energies from simulation using MBAR.
  print "Estimating relative free energies from simulation (this may take a while)..."

  # Compute reduced potential energies.
  u_kln = beta * U_kln

  # Initialize the MBAR class, determining the free energies.
  mbar = pymbar.MBAR(u_kln, N_k, method = 'adaptive') # use fast Newton-Raphson solver

  #=============================================================================================
  # Compute entropies and enthalpies with MBAR.  This uses s - f - <u>, from MBAR.
  #=============================================================================================

  # Get matrix of dimensionless free energy differences and uncertainty estimate.
  #print "Computing covariance matrix..."
  #(Deltaf_ij_estimated, dDeltaf_ij_estimated) = mbar.getFreeEnergyDifferences()

  (Deltaf_ij_estimated, dDeltaf_ij_estimated,Deltau_ij_estimated, dDeltau_ij_estimated,Deltas_ij_estimated, dDeltas_ij_estimated) = mbar.computeEntropyAndEnthalpy() 
  
  setReplicate(replicate,'mbar',Deltaf_ij_estimated,dDeltaf_ij_estimated,Deltau_ij_estimated,dDeltau_ij_estimated,Deltas_ij_estimated,dDeltas_ij_estimated)

  #=============================================================================================
  # Compute entropies and enthalpies with endpoints, but full df
  #=============================================================================================

  # endpoint estimate of energies.
  u_k_estimated = numpy.zeros([K],numpy.float64) 
  du_k_estimated = numpy.zeros([K],numpy.float64) 
  
  # 'standard' expectation averages
  for k in range(K):
    u_k_estimated[k] = numpy.average(u_kln[k,k,0:N_k[k]])
    du_k_estimated[k]  = numpy.sqrt(numpy.var(u_kln[k,k,0:N_k[k]])/(N_k[k]-1))

  for i in range(K):
    for j in range(K):
      Deltau_ij_estimated[i,j] = u_k_estimated[i] - u_k_estimated[j] 
      dDeltau_ij_estimated[i,j] = numpy.sqrt(du_k_estimated[i]**2 + du_k_estimated[j]**2)
      # these uncertainties aren't quite right.
      Deltas_ij_estimated[i,j] = Deltau_ij_estimated[i,j] - replicate['df_mbar'][i,j]
      dDeltas_ij_estimated[i,j] = numpy.sqrt(dDeltau_ij_estimated[i,j]**2 + replicate['ddf_mbar'][i,j]**2)

  setReplicate(replicate,'endpoint',Deltaf_ij_estimated,dDeltaf_ij_estimated,Deltau_ij_estimated,dDeltau_ij_estimated,Deltas_ij_estimated,dDeltas_ij_estimated)

  #=============================================================================================
  # Compute free energies, entropies, and enthalpies with TI (single state averages)
  #=============================================================================================
  # assume trapezoidal rule

  dfdl_k_estimated = numpy.zeros([K],numpy.float64) 
  ddfdl_k_estimated = numpy.zeros([K],numpy.float64) 
  dudl_k_estimated = numpy.zeros([K],numpy.float64) 
  ddudl_k_estimated = numpy.zeros([K],numpy.float64) 
  dsdl_k_estimated = numpy.zeros([K],numpy.float64) 
  ddsdl_k_estimated = numpy.zeros([K],numpy.float64) 

  for k in range(K):
    dfdl_k_estimated[k] = numpy.average(dhdl_kln[k,k,0:N_k[k]])
    ddfdl_k_estimated[k]  = numpy.sqrt(numpy.var(dhdl_kln[k,k,0:N_k[k]])/(N_k[k]-1))
    dsdl_k_estimated[k] =  -1*numpy.cov(dhdl_kln[k,k,0:N_k[k]],u_kln[k,k,0:N_k[k]])[0,1]
    ddsdl_k_estimated[k] =  numpy.sqrt(numpy.var(u_kln[k,k,0:N_k[k]]*dhdl_kln[k,k,0:N_k[k]])/(N_k[k]-1))
    dudl_k_estimated[k] =  dfdl_k_estimated[k] + dsdl_k_estimated[k] 
    du_varterm = dhdl_kln[k,k,0:N_k[k]] - (u_kln[k,k,0:N_k[k]]*dhdl_kln[k,k,0:N_k[k]]-(dfdl_k_estimated[k]*u_k_estimated[k]))
    ddudl_k_estimated[k] =  numpy.sqrt(numpy.var(du_varterm)/(N_k[k]-1))

  (Deltaf_ij_estimated, dDeltaf_ij_estimated,Deltau_ij_estimated, dDeltau_ij_estimated,Deltas_ij_estimated, dDeltas_ij_estimated) = integrateij(dfdl_k_estimated,ddfdl_k_estimated,dsdl_k_estimated,ddsdl_k_estimated,dudl_k_estimated,ddudl_k_estimated)

  setReplicate(replicate,'ti',Deltaf_ij_estimated,dDeltaf_ij_estimated,Deltau_ij_estimated,dDeltau_ij_estimated,Deltas_ij_estimated,dDeltas_ij_estimated)

  #=============================================================================================
  # Compute free energies, entropies, and enthalpies with TI (reweighted - STILL TODO)
  #=============================================================================================
  # assume trapezoidal rule

  dfdl_k_estimated = numpy.zeros([K],numpy.float64) 
  ddfdl_k_estimated = numpy.zeros([K],numpy.float64) 
  dudl_k_estimated = numpy.zeros([K],numpy.float64) 
  ddudl_k_estimated = numpy.zeros([K],numpy.float64) 
  dsdl_k_estimated = numpy.zeros([K],numpy.float64) 
  ddsdl_k_estimated = numpy.zeros([K],numpy.float64) 

  for k in range(K):
    dfdl_k_estimated[k] = numpy.average(dhdl_kln[k,k,0:N_k[k]])
    ddfdl_k_estimated[k]  = numpy.sqrt(numpy.var(dhdl_kln[k,k,0:N_k[k]])/(N_k[k]-1))
    dsdl_k_estimated[k] =  -1*numpy.cov(dhdl_kln[k,k,0:N_k[k]],u_kln[k,k,0:N_k[k]])[0,1]
    ddsdl_k_estimated[k] =  numpy.sqrt(numpy.var(u_kln[k,k,0:N_k[k]]*dhdl_kln[k,k,0:N_k[k]])/(N_k[k]-1))
    dudl_k_estimated[k] =  dfdl_k_estimated[k] + dsdl_k_estimated[k] 
    du_varterm = dhdl_kln[k,k,0:N_k[k]] - (u_kln[k,k,0:N_k[k]]*dhdl_kln[k,k,0:N_k[k]]-(dfdl_k_estimated[k]*u_k_estimated[k]))
    ddudl_k_estimated[k] =  numpy.sqrt(numpy.var(du_varterm)/(N_k[k]-1))

  (Deltaf_ij_estimated, dDeltaf_ij_estimated,Deltau_ij_estimated, dDeltau_ij_estimated,Deltas_ij_estimated, dDeltas_ij_estimated) = integrateij(dfdl_k_estimated,ddfdl_k_estimated,dsdl_k_estimated,ddsdl_k_estimated,dudl_k_estimated,ddudl_k_estimated)

  #setReplicate(replicate,'ti-reweighted',Deltaf_ij_estimated,dDeltaf_ij_estimated,Deltau_ij_estimated,dDeltau_ij_estimated,Deltas_ij_estimated,dDeltas_ij_estimated)

  #=============================================================================================
  # Compute entropies with temperature perturbed finite difference derivatives from FEP
  #=============================================================================================

  # allocate space for perturbed reduced energies
  upert = numpy.zeros([K,3*K,N_max])
  # how big is dT?

  upert[:,0:K,:] = beta * U_kln  
  upert[:,K:2*K,:] = btlow_pert * U_kln
  upert[:,2*K:3*K,:] = bthigh_pert * U_kln

  (Deltaf_ij_pert,dDeltaf_ij_pert) = mbar.computePerturbedFreeEnergies(upert)
  var_ij = dDeltaf_ij_pert**2

  # now compute the free energy differences:
  Deltaf_ij_estimated = replicate['df_mbar']
  dDeltaf_ij_estimated = replicate['ddf_mbar']
  Deltas_ij_estimated = dpS_pert[1]*Deltaf_ij_pert[K:2*K,K:2*K] - dpS_pert[2]*Deltaf_ij_pert[2*K:3*K,2*K:3*K]
  Deltau_ij_estimated = replicate['df_mbar'] + Deltas_ij_estimated

  dDeltas_ij_estimated = covdiffs(dDeltaf_ij_pert[1*K:3*K,1*K:3*K],K,dpS_pert)
  dDeltau_ij_estimated = covdiffs3(dDeltaf_ij_pert,K,dpS_pert)

  setReplicate(replicate,'finitefepF',Deltaf_ij_estimated,dDeltaf_ij_estimated,Deltau_ij_estimated,dDeltau_ij_estimated,Deltas_ij_estimated,dDeltas_ij_estimated)

  #=============================================================================================
  # Compute ENTHALPIES with temperature finite difference derivatives from FEP
  #=============================================================================================

  # now compute the free energy differences:

  Deltaf_ij_estimated = replicate['df_mbar']
  dDeltaf_ij_estimated = replicate['ddf_mbar']
  Deltau_ij_estimated = dpH_pert[1]*Deltaf_ij_pert[K:2*K,K:2*K] - dpH_pert[2]*Deltaf_ij_pert[2*K:3*K,2*K:3*K]
  Deltas_ij_estimated = Deltau_ij_estimated - replicate['df_mbar']

  dDeltas_ij_estimated = numpy.zeros([K,K],dtype=numpy.float64)
  dDeltau_ij_estimated = numpy.zeros([K,K],dtype=numpy.float64)

  setReplicate(replicate,'finitefepf',Deltaf_ij_estimated,dDeltaf_ij_estimated,Deltau_ij_estimated,dDeltau_ij_estimated,Deltas_ij_estimated,dDeltas_ij_estimated)

  #=============================================================================================
  # Compute entropies with temperature finite difference derivatives from direct sampling
  #=============================================================================================

  print 'generating high t samples...'

  sigmat_k = (btlow_samp * K_k)**-0.5

  # Generate samples.
  for k in range(0,K):
    # generate N_k[k] independent samples with spring constant K_k[k]
    xt_kn[k,0:N_k[k]] = numpy.random.normal(O_k[k], sigmat_k[k], [N_k[k]])
    
    # compute potential energy and derivative of all samples in all potentials
    for l in range(0,K):      
      Ut_kln[k,l,:] = (K_k[l]/2.0) * (xt_kn[k,:]-O_k[l])**2

  ulow_kln = btlow_samp * Ut_kln
  # Initialize the MBAR class, determining the free energies.
  mbart = pymbar.MBAR(ulow_kln, N_k, method = 'adaptive') # use fast Newton-Raphson solver
  (Deltaf_ij_low, dDeltaf_ij_low) = mbart.getFreeEnergyDifferences()

  print 'generating samples low t samples...'
  sigmat_k = (bthigh_samp * K_k)**-0.5

  # Generate samples.
  for k in range(0,K):
    # generate N_k[k] independent samples with spring constant K_k[k]
    xt_kn[k,0:N_k[k]] = numpy.random.normal(O_k[k], sigmat_k[k], [N_k[k]])
    
    # compute potential energy and derivative of all samples in all potentials
    for l in range(0,K):      
      Ut_kln[k,l,:] = (K_k[l]/2.0) * (xt_kn[k,:]-O_k[l])**2

  uhigh_kln = bthigh_samp * Ut_kln

  # Initialize the MBAR class, determining the free energies.
  mbar = pymbar.MBAR(uhigh_kln, N_k, method = 'adaptive') # use fast Newton-Raphson solver
  (Deltaf_ij_high, dDeltaf_ij_high) = mbar.getFreeEnergyDifferences()

  dp = [1,beta/(dt_pert*btlow_pert),beta/(dt_pert*bthigh_pert)]
  Deltas_ij_estimated = dpS_samp[1]*Deltaf_ij_high - dpS_samp[2]*Deltaf_ij_low
  dDeltas_ij_estimated = numpy.sqrt((dpS_samp[1]*dDeltaf_ij_low)**2 + (dpS_samp[2]*dDeltaf_ij_high)**2)
  
  Deltaf_ij_estimated = replicate['df_mbar']
  dDeltaf_ij_estimated = replicate['ddf_mbar']

  Deltau_ij_estimated = Deltaf_ij_estimated + Deltas_ij_estimated
  dDeltau_ij_estimated = numpy.sqrt(dDeltaf_ij_estimated**2 + dDeltas_ij_estimated**2)

  setReplicate(replicate,'finitedF',Deltaf_ij_estimated,dDeltaf_ij_estimated,Deltau_ij_estimated,dDeltau_ij_estimated,Deltas_ij_estimated,dDeltas_ij_estimated)

  #=============================================================================================
  # Compute ENTHALPIES with temperature finite difference derivatives from direct sampling
  #=============================================================================================

  Deltau_ij_estimated = dpH_samp[1]*Deltaf_ij_high - dpH_samp[2]*Deltaf_ij_low
  dDeltau_ij_estimated = numpy.sqrt((dpH_samp[1]**2)*dDeltaf_ij_low**2 + (dpH_samp[2]**2)*dDeltaf_ij_high**2)

  Deltaf_ij_estimated = replicate['df_mbar']
  dDeltaf_ij_estimated = replicate['ddf_mbar']

  Deltas_ij_estimated = Deltau_ij_estimated - Deltaf_ij_estimated
  dDeltas_ij_estimated = numpy.sqrt(dDeltau_ij_estimated**2 + dDeltaf_ij_estimated**2)

  setReplicate(replicate,'finitedf',Deltaf_ij_estimated,dDeltaf_ij_estimated,Deltau_ij_estimated,dDeltau_ij_estimated,Deltas_ij_estimated,dDeltas_ij_estimated)

  #=============================================================================================
  # Compute entropies with temperature finite difference derivatives using MBAR on two energies
  # with variable dt.
  #=============================================================================================

  # allocate space for perturbed reduced energies
  upert = numpy.zeros([2*K,4*K,N_max])

  Npert = numpy.zeros([4*K],int)
  # only first 2*K are sampled

  Npert[0*K:1*K] = N_k[:]
  Npert[1*K:2*K] = N_k[:]
  upert[0:K,0*K:1*K,:] = ulow_kln  
  upert[0:K,1*K:2*K,:] = (bthigh_samp/btlow_samp)*ulow_kln  
  upert[0:K,2*K:3*K,:] = (btlow_pert/btlow_samp)*ulow_kln  
  upert[0:K,3*K:4*K,:] = (bthigh_pert/btlow_samp)*ulow_kln  

  upert[K:2*K,0*K:1*K,:] = (btlow_samp/bthigh_samp)*ulow_kln  
  upert[K:2*K,1*K:2*K,:] = uhigh_kln  
  upert[K:2*K,2*K:3*K,:] = (btlow_pert/bthigh_samp)*ulow_kln  
  upert[K:2*K,3*K:4*K,:] = (bthigh_pert/bthigh_samp)*ulow_kln  

  # Initialize the MBAR class, determining the free energies, then taking derivatives at an arbitrary dt

  mbar = pymbar.MBAR(upert, Npert, method = 'adaptive') # use fast Newton-Raphson solver
  (Deltaf_ij, dDeltaf_ij) = mbar.getFreeEnergyDifferences()

  Deltaf_ij_estimated = replicate['df_mbar']
  dDeltaf_ij_estimated = replicate['ddf_mbar']

  Deltas_ij_estimated = dpS_pert[1]*Deltaf_ij[3*K:4*K,3*K:4*K] - dpS_pert[2]*Deltaf_ij[2*K:3*K,2*K:3*K]
  dDeltas_ij_estimated = covdiffs(Deltaf_ij[2*K:4*K,2*K:4*K],K,dpS_pert)
  Deltau_ij_estimated = Deltaf_ij_estimated + Deltas_ij_estimated
  dDeltau_ij_estimated = numpy.sqrt(Deltaf_ij_estimated**2 + Deltas_ij_estimated**2)
                                 
  setReplicate(replicate,'finitedF_mbar2',Deltaf_ij_estimated,dDeltaf_ij_estimated,Deltau_ij_estimated,dDeltau_ij_estimated,Deltas_ij_estimated,dDeltas_ij_estimated)

  #=============================================================================================
  # Compute enthalpies with temperature finite difference derivatives using MBAR on two energies
  # with variable dt.
  #=============================================================================================

  Deltau_ij_estimated = dpH_pert[1]*Deltaf_ij[3*K:4*K,3*K:4*K] - dpH_pert[2]*Deltaf_ij[2*K:3*K,2*K:3*K]
  dDeltau_ij_estimated = covdiffs(Deltaf_ij[2*K:4*K,2*K:4*K],K,dpH_pert)

  Deltas_ij_estimated = Deltau_ij_estimated - Deltaf_ij_estimated
  dDeltas_ij_estimated = numpy.sqrt(Deltaf_ij_estimated**2 + Deltau_ij_estimated**2)

  setReplicate(replicate,'finitedf_mbar2',Deltaf_ij_estimated,dDeltaf_ij_estimated,Deltau_ij_estimated,dDeltau_ij_estimated,Deltas_ij_estimated,dDeltas_ij_estimated)

  #=============================================================================================
  # Compute entropies with temperature finite difference derivatives using MBAR on three energies
  # with variable dt.
  #=============================================================================================

  # allocate space for perturbed reduced energies
  upert = numpy.zeros([3*K,5*K,N_max])

  Npert = numpy.zeros(5*K,int)
  Npert[0*K:1*K] = N_k[:]
  Npert[1*K:2*K] = N_k[:]
  Npert[2*K:3*K] = N_k[:]

  upert[0*K:1*K,0*K:1*K,:] = u_kln
  upert[0*K:1*K,1*K:2*K,:] = (btlow_samp/beta)*u_kln
  upert[0*K:1*K,2*K:3*K,:] = (bthigh_samp/beta)*u_kln
  upert[0*K:1*K,3*K:4*K,:] = (btlow_pert/beta)*u_kln
  upert[0*K:1*K,4*K:5*K,:] = (bthigh_pert/beta)*u_kln

  upert[1*K:2*K,0*K:1*K,:] = (beta/btlow_samp)*ulow_kln
  upert[1*K:2*K,1*K:2*K,:] = ulow_kln
  upert[1*K:2*K,2*K:3*K,:] = (bthigh_samp/btlow_samp)*ulow_kln
  upert[1*K:2*K,3*K:4*K,:] = (btlow_pert/btlow_samp)*ulow_kln
  upert[1*K:2*K,4*K:5*K,:] = (bthigh_pert/btlow_samp)*ulow_kln

  upert[2*K:3*K,0*K:1*K,:] = (beta/bthigh_samp)*uhigh_kln
  upert[2*K:3*K,1*K:2*K,:] = (btlow_samp/bthigh_samp)*uhigh_kln
  upert[2*K:3*K,2*K:3*K,:] = uhigh_kln
  upert[2*K:3*K,3*K:4*K,:] = (btlow_pert/bthigh_samp)*uhigh_kln
  upert[2*K:3*K,4*K:5*K,:] = (bthigh_pert/bthigh_samp)*uhigh_kln

  # Initialize the MBAR class, determining the free energies, then taking derivatives at an arbitrary dt

  mbar = pymbar.MBAR(upert, Npert, method = 'adaptive') # use fast Newton-Raphson solver
  (Deltaf_ij, dDeltaf_ij) = mbar.getFreeEnergyDifferences()

  Deltaf_ij_estimated = Deltaf_ij[0:K,0:K]
  dDeltaf_ij_estimated = dDeltaf_ij[0:K,0:K]

  Deltas_ij_estimated = dpS_pert[1]*Deltaf_ij[4*K:5*K,4*K:5*K] - dpS_pert[2]*Deltaf_ij[3*K:4*K,3*K:4*K]
  dDeltas_ij_estimated = covdiffs(Deltaf_ij[3*K:5*K,3*K:5*K],K,dpS_pert)
  Deltau_ij_estimated = Deltaf_ij_estimated + Deltas_ij_estimated;
  dDeltau_ij_estimated = numpy.sqrt(Deltaf_ij_estimated**2 + Deltas_ij_estimated**2)
                                 
  setReplicate(replicate,'finitedF_mbar3',Deltaf_ij_estimated,dDeltaf_ij_estimated,Deltau_ij_estimated,dDeltau_ij_estimated,Deltas_ij_estimated,dDeltas_ij_estimated)

  #=============================================================================================
  # Compute enthalpies with temperature finite difference derivatives using MBAR on three energies
  # with variable dt.
  #=============================================================================================

  Deltaf_ij_estimated = Deltaf_ij[0:K,0:K]
  dDeltaf_ij_estimated = dDeltaf_ij[0:K,0:K]

  Deltau_ij_estimated = dpH_pert[1]*Deltaf_ij[4*K:5*K,4*K:5*K] - dpH_pert[2]*Deltaf_ij[3*K:4*K,3*K:4*K]
  dDeltau_ij_estimated = covdiffs(Deltaf_ij[3*K:5*K,3*K:5*K],K,dpH_pert)

  ddf_part = numpy.zeros([3*K,3*K])
  ddf_part[0:K,0:K]     = dDeltaf_ij[0:K,0:K]
  ddf_part[0:K,K:3*K]   = dDeltaf_ij[0:K,3*K:5*K]
  ddf_part[K:3*K,0:K]   = dDeltaf_ij[3*K:5*K,0:K]
  ddf_part[K:3*K,K:3*K] = dDeltaf_ij[3*K:5*K,3*K:5*K]
  
  # f = u-s, s = u-f
  Deltas_ij_estimated = Deltau_ij_estimated - Deltaf_ij_estimated
  dDeltas_ij_estimated = covdiffs3(ddf_part,K,dpH_pert)

  setReplicate(replicate,'finitedf_mbar3',Deltaf_ij_estimated,dDeltaf_ij_estimated,Deltau_ij_estimated,dDeltau_ij_estimated,Deltas_ij_estimated,dDeltas_ij_estimated)

  #=============================================================================================
  # Store data for this replicate.
  #=============================================================================================  
  replicates.append(replicate)

strs = ('mbar','endpoint','ti','finitedF','finitedf','finitefepF','finitefepf','finitedF_mbar2','finitedf_mbar2','finitedF_mbar3','finitedf_mbar3')
vals = numpy.zeros(nreplicates,float)
stds = numpy.zeros(nreplicates,float)
errs = numpy.zeros(nreplicates,float)
das  = ['df','ds','du']
print "               sample_mean sample_std analytical"
for istr in strs:
   for da in das:  
      for index,replicate in enumerate(replicates):
         str = da + '_' + istr
         vals[index] = replicate[str][0,K-1]
         str = 'd'+ da + '_' + istr
         stds[index] = replicate[str][0,K-1]
         str = da + '_' + istr + '_error'
         errs[index] = replicate[str][0,K-1]
      
      sample_std = numpy.std(vals) 
      sample_mean = numpy.average(vals)
      analytical_average = numpy.average(stds)
      bias_mean = numpy.average(errs)
      print "%4s%15s%10.5f%10.5f%10.5f%10.5f" % (da,istr,sample_mean,sample_std,analytical_average,bias_mean)