content: #!/usr/bin/python2.4
  
  import sys
  import numpy
  import pymbar # for MBAR analysis
  import timeseries # for timeseries analysis
  import os
  import os.path
  import pdb  # for debugging
  from optparse import OptionParser
  import MBAR_pmfQz
  import wham
  import MBAR_pmfQ
  import cPickle
  
  def parse_args():
      parser=OptionParser()
      #parser.add_option("-t", "--temprange", nargs=2, default=[300.0,450.0], type="float", dest="temprange", help="temperature range of replicas")
      parser.add_option("-r", "--replicas", default=24, type="int",dest="replicas", help="number of replicas (default: 24)")
      parser.add_option("-n", "--N_max", default=100000, type="int",dest="N_max", help="number of data points to read in (default: 100k)")
      parser.add_option("-s", "--skip", default=1, type="int",dest="skip", help="skip every n data points")
      parser.add_option("--direc", dest="direc", help="Qtraj_singleprot.txt file location")
      parser.add_option("--tfile", dest="tfile", default="/home/edz3fz/proteinmontecarlo/T.txt", help="file of temperatures (default: T.txt)")
      parser.add_option('--cpt', action="store_true", default=False, help="use checkpoint files, if they exist")
      (options,args) = parser.parse_args()
      return options
  
  def get_ukln(args,N_max,K,Z,beta_k,spring_constant,U_kn,z_kn,N_k):
      print 'Computing reduced potential energies...'
      u_kln = numpy.zeros([K,K,N_max], numpy.float32)
      for k in range(K):
     		for l in range(K):
     			#z_index = l/(len(T)) # z is outer dimension
     			#T_index = l%(len(T)) # T is inner dimension
     			dz = z_kn[k,0:N_k[k]] - Z[l]
     			u_kln[k,l,0:N_k[k]] = beta_k[l] * (U_kn[k,0:N_k[k]] + spring_constant[l]*(dz)**2)
      return u_kln
  
  def get_mbar(args, beta_k, Z, U_kn, N_k, u_kln):
      if args.cpt:
          if os.path.exists('%s/f_k_foldingcurve.npy' % args.direc):
              print 'Reading in free energies from %s/f_k.npy' % args.direc
              f_k = numpy.load('%s/f_k.npy' % args.direc)
              mbar = pymbar.MBAR(u_kln,N_k,initial_f_k = f_k, maximum_iterations=0,verbose=True,use_optimized=1)
              return mbar
      print 'Using WHAM to generate historgram-based initial guess of dimensionless free energies f_k...'
      #beta_k = numpy.array(beta_k.tolist()*len(Z)) 
      #f_k = wham.histogram_wham(beta_k, U_kn, N_k)
      print 'Initializing MBAR...'
      mbar = pymbar.MBAR(u_kln, N_k, #initial_f_k = f_k, 
                          use_optimized='', verbose=True)
      mbar_file = '%s/f_k_foldingcurve.npy' % args.direc
      print 'Saving free energies to %s' % mbar_file
      saving = True
      if saving:
      	numpy.save(mbar_file, mbar.f_k)
      return mbar
  
  def main():
      options = parse_args()
      kB = 0.00831447/4.184  #Boltzmann constant
      
      T = numpy.loadtxt(options.tfile)
      Z = numpy.arange(9,31.5,1.5)
      print 'Initial temperature states are', T
      print 'Distance states are', Z
      K = len(T)*len(Z)
      spring_constant = numpy.ones(K)
    
      # read in data
      U_kn, Q_kn, z_kn, N_max = MBAR_pmfQz.read_data(options, K, Z, T, spring_constant[0])
  
      # subsample the data
      U_kn, Q_kn, z_kn, N_k = MBAR_pmfQz.subsample(U_kn,Q_kn,z_kn,K,N_max)
   
      # insert unweighted states
      T_new = numpy.arange(200,410,10)
      T_new = numpy.array([200,210,220,230,235,240,245,250,255,260,265,270,275,280,285,290,295,300,305,310,315,320,325,330,335,340,345,350,375,400])
      Z_new = numpy.zeros(len(T_new))
      K_new = len(T_new)
      print 'inserting unweighted temperature states', T_new
      
      # update states
      print 'Inserting blank states'
      Z = Z.tolist()
      Z = [x for x in Z for _ in range(len(T))]
      Z = numpy.concatenate((numpy.array(Z),Z_new))
      T = numpy.array(T.tolist()*(K/len(T)))
      T = numpy.concatenate((T,T_new))
      K += K_new
      spring_constant = numpy.concatenate((spring_constant,numpy.zeros(K_new)))
      print 'all temperature states are ', T
      print 'all surface states are ', Z
      print 'there are a total of %i states' % K
      N_k = numpy.concatenate((N_k,numpy.zeros(K_new)))
      U_kn = numpy.concatenate((U_kn,numpy.zeros([K_new,N_max])))
      Q_kn = numpy.concatenate((Q_kn,numpy.zeros([K_new,N_max])))
      z_kn = numpy.concatenate((z_kn,numpy.zeros([K_new,N_max])))
  
      beta_k = 1/(kB*T)
  
      u_kln = get_ukln(options, N_max, K, Z, beta_k, spring_constant, U_kn, z_kn, N_k)
      
      print "Initializing MBAR..."
      # Use Adaptive Method (Both Newton-Raphson and Self-Consistent, testing which is better)
      mbar = get_mbar(options,beta_k,Z,U_kn,N_k,u_kln)
      
      print "Computing Expectations for E..."
      (E_expect, dE_expect) = mbar.computeExpectations(u_kln)*(beta_k)**(-1)
      
      print "Computing Expectations for E^2..."
      (E2_expect,dE2_expect) = mbar.computeExpectations(u_kln*u_kln)*(beta_k)**(-2)
      
      print "Computing Expectations for Q..."
      (Q,dQ) = mbar.computeExpectations(Q_kn)
      
      print "Computing Heat Capacity as ( <E^2> - <E>^2 ) / ( R*T^2 )..."
      Cv = numpy.zeros([K], numpy.float64)
      dCv = numpy.zeros([K], numpy.float64)
      for i in range(K):
             Cv[i] = (E2_expect[i] - (E_expect[i]*E_expect[i])) / ( kB * T[i] * T[i])
             dCv[i] = 2*dE_expect[i]**2 / (kB *T[i]*T[i])   # from propagation of error
  
      numpy.save(options.direc+'/foldingcurve_umbrella',numpy.array([T, Q, dQ]))
      numpy.save(options.direc+'/heatcap_umbrella',numpy.array([T, Cv, dCv]))
      
  #    pdb.set_trace()
  # 
  #    print 'Computing PMF(Q) at 325 K'
  #    nbins = 25
  #    target_temperature = 325
  #    target_beta = 1.0/(kB*target_temperature)
  #    nbins, bin_centers, bin_counts, bin_kn = get_bins(nbins,K,N_max,Q_kn)
  #    u_kn = target_beta*U_kn
  #    f_i, d2f_i = mbar.computePMF_states(u_kn, bin_kn, nbins)
  #    pmf_file = '%s/pmfQ_umbrella_%i.pkl' % (options.direc, target_temperature)
  #    f = file(pmf_file, 'wb')
  #    print 'Saving target temperature, bin centers, f_i, df_i to %s' % pmf_file
  #    cPickle.dump(target_temperature,f)
  #    cPickle.dump(bin_centers,f)
  #    cPickle.dump(f_i,f)
  #    cPickle.dump(d2f_i,f)
  #    f.close()
  #
  #    try:
  #        import matplotlib.pyplot as plt
  #        plt.figure(1)
  #        plt.plot(T,Q,'k')
  #        plt.errorbar(T, Q, yerr=dQ)
  #        plt.xlabel('Temperature (K)')
  #        plt.ylabel('Q fraction native contacts')
  #        plt.savefig(options.direc+'/foldingcurve_umbrella.png')
  #        plt.show()
  #    except:
  #        pass    
  #
  if __name__ == '__main__':
      main()
  