Python MBAR.getFreeEnergyDifferences方法代码示例

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
    def update_logZ_with_mbar(self):
        """
        Use MBAR to update logZ estimates.
        """
        if not self.ncfile:
            raise Exception("Cannot update logZ using MBAR since no NetCDF file is storing history.")

        if not self.sampler.update_scheme == 'global-jump':
            raise Exception("Only global jump is implemented right now.")

        if not self.ncfile:
            raise Exception("Must have a storage file attached to use MBAR updates")

        # Extract relative energies.
        if self.verbose:
            print('Updating logZ estimate with MBAR...')
        initial_time = time.time()
        from pymbar import MBAR
        #first = int(self.iteration / 2)
        first = 0
        u_kn = np.array(self.ncfile.variables['u_k'][first:,:]).T
        [N_k, bins] = np.histogram(self.ncfile.variables['state_index'][first:], bins=(np.arange(self.sampler.nstates+1) - 0.5))
        mbar = MBAR(u_kn, N_k)
        Deltaf_ij, dDeltaf_ij, Theta_ij = mbar.getFreeEnergyDifferences(compute_uncertainty=True, uncertainty_method='approximate')
        self.logZ[:] = -mbar.f_k[:]
        self.logZ -= self.logZ[0]
        final_time = time.time()
        elapsed_time = final_time - initial_time
        self._timing['MBAR time'] = elapsed_time
        if self.verbose:
            print('MBAR time    %8.3f s' % elapsed_time)

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
def test_exponential_mbar_free_energies():
    """Exponential Distribution Test: can MBAR calculate correct free energy differences?"""
    test = exponential_distributions.ExponentialTestCase(rates)
    x_kn, u_kln, N_k_output = test.sample(N_k, mode='u_kln')
    eq(N_k, N_k_output)

    mbar = MBAR(u_kln, N_k)
    fe, fe_sigma = mbar.getFreeEnergyDifferences()
    fe, fe_sigma = fe[0,1:], fe_sigma[0,1:]

    fe0 = test.analytical_free_energies()
    fe0 = fe0[1:] - fe0[0]

    z = (fe - fe0) / fe_sigma
    eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
def test_harmonic_oscillators_mbar_free_energies():
    """Harmonic Oscillators Test: can MBAR calculate correct free energy differences?"""
    test = harmonic_oscillators.HarmonicOscillatorsTestCase(O_k, k_k)
    x_kn, u_kln, N_k_output = test.sample(N_k)

    eq(N_k, N_k_output)

    mbar = MBAR(u_kln, N_k)
    fe, fe_sigma = mbar.getFreeEnergyDifferences()
    fe, fe_sigma = fe[0,1:], fe_sigma[0,1:]

    fe0 = test.analytical_free_energies()
    fe0 = fe0[1:] - fe0[0]

    z = (fe - fe0) / fe_sigma
    eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
def test_harmonic_oscillators_mbar_free_energies():
    """Harmonic Oscillators Test: can MBAR calculate correct free energy differences?"""
    test = harmonic_oscillators.HarmonicOscillatorsTestCase(O_k, k_k)
    x_n, u_kn, origin = test.sample(N_k)
    u_ijn, N_k_output = convert_ukn_to_uijn(u_kn)
    
    eq(N_k, N_k_output.values)

    mbar = MBAR(u_ijn.values, N_k)
    fe, fe_sigma = mbar.getFreeEnergyDifferences()
    fe, fe_sigma = fe[0], fe_sigma[0]

    fe0 = test.analytical_free_energies()

    z = (fe - fe0) / fe_sigma
    z = z[1:]  # First component is undetermined.
    eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
def test_mbar_free_energies():

    """Can MBAR calculate moderately correct free energy differences?"""

    for system_generator in system_generators:
        name, test = system_generator()
        x_n, u_kn, N_k_output, s_n = test.sample(N_k, mode='u_kn')
        eq(N_k, N_k_output)
        mbar = MBAR(u_kn, N_k)

        fe, fe_sigma, Theta_ij = mbar.getFreeEnergyDifferences()
        fe, fe_sigma = fe[0,1:], fe_sigma[0,1:]

        fe0 = test.analytical_free_energies()
        fe0 = fe0[1:] - fe0[0]

        z = (fe - fe0) / fe_sigma
        eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
    def run_mbar(self, ndiscard=0, nuse=None):
        """Estimate free energies of all alchemical states.

        Parameters
        ----------
        ndiscard : int, optinoal, default=0
            number of iterations to discard to equilibration
        nuse : int, optional, default=None
            maximum number of iterations to use (after discarding)

        Returns
        -------
        
        Deltaf_ij : np.ndarray, shape=(n_states, n_states)
            The statewise free energy differences

        dDeltaf_ij : np.ndarray, shape=(n_states, n_states)
            The statewise free energy difference uncertainties

        """    
        
        u_kln_replica, u_kln, u_n = self.get_u_kln()

        u_kln_replica, u_kln, u_n, N_k, N = self.equilibrate_and_subsample(u_kln_replica, u_kln, u_n, ndiscard=ndiscard, nuse=nuse)

        logger.info("Initialing MBAR and computing free energy differences...")
        mbar = MBAR(u_kln, N_k, verbose = False, method = 'self-consistent-iteration', maximum_iterations = 50000) # use slow self-consistent-iteration (the default)

        # Get matrix of dimensionless free energy differences and uncertainty estimate.
        logger.info("Computing covariance matrix...")
        (Deltaf_ij, dDeltaf_ij) = mbar.getFreeEnergyDifferences(uncertainty_method='svd-ew')
       
        logger.info("\n%-24s %16s\n%s" % ("Deltaf_ij", "current state", pd.DataFrame(Deltaf_ij).to_string()))
        logger.info("\n%-24s %16s\n%s" % ("Deltaf_ij", "current state", pd.DataFrame(dDeltaf_ij).to_string()))        
                
        return (Deltaf_ij, dDeltaf_ij)

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
def compute_hydration_energy(entry, parameters, platform_name="CPU"):
    """
    Compute hydration energy of a single molecule given a GBSA parameter set.

    ARGUMENTS

    molecule (OEMol) - molecule with GBSA atom types
    parameters (dict) - parameters for GBSA atom types

    RETURNS

    energy (float) - hydration energy in kcal/mol

    """


    platform = openmm.Platform.getPlatformByName('CPU')

    from pymbar import MBAR

    timestep = 2 * units.femtoseconds

    molecule = entry['molecule']
    iupac_name = entry['iupac']
    cid = molecule.GetData('cid')

    # Retrieve OpenMM System.
    vacuum_system = entry['system']
    solvent_system = copy.deepcopy(entry['solvated_system'])

    # Get nonbonded force.
    forces = { solvent_system.getForce(index).__class__.__name__ : solvent_system.getForce(index) for index in range(solvent_system.getNumForces()) }
    nonbonded_force = forces['NonbondedForce']
    gbsa_force = forces['CustomGBForce']

    # Build indexable list of atoms.
    atoms = [atom for atom in molecule.GetAtoms()]
    natoms = len(atoms)


    # Create context for solvent system.
    timestep = 2.0 * units.femtosecond
    solvent_integrator = openmm.VerletIntegrator(timestep)


    # Create context for vacuum system.
    vacuum_integrator = openmm.VerletIntegrator(timestep)

    # Assign GBSA parameters.
    for (atom_index, atom) in enumerate(atoms):
        [charge, sigma, epsilon] = nonbonded_force.getParticleParameters(atom_index)
        atomtype = atom.GetStringData("gbsa_type") # GBSA atomtype
        radius = parameters['%s_%s' % (atomtype, 'radius')] * units.angstroms
        scalingFactor = parameters['%s_%s' % (atomtype, 'scalingFactor')]
        gbsa_force.setParticleParameters(atom_index, [charge, radius, scalingFactor])

    solvent_context = openmm.Context(solvent_system, solvent_integrator,platform)
    vacuum_context = openmm.Context(vacuum_system, vacuum_integrator, platform)

    # Compute energy differences.
    temperature = entry['temperature']
    kT = kB * temperature
    beta = 1.0 / kT

    initial_time = time.time()
    x_n = entry['x_n']
    u_n = entry['u_n']
    nsamples = len(u_n)
    nstates = 3 # number of thermodynamic states
    u_kln = np.zeros([3,3,nsamples], np.float64)
    for sample in range(nsamples):
        positions = units.Quantity(x_n[sample,:,:], units.nanometers)

        u_kln[0,0,sample] = u_n[sample]

        vacuum_context.setPositions(positions)
        vacuum_state = vacuum_context.getState(getEnergy=True)
        u_kln[0,1,sample] = beta * vacuum_state.getPotentialEnergy()

        solvent_context.setPositions(positions)
        solvent_state = solvent_context.getState(getEnergy=True)
        u_kln[0,2,sample] = beta * solvent_state.getPotentialEnergy()

    N_k = np.zeros([nstates], np.int32)
    N_k[0] = nsamples


    mbar = MBAR(u_kln, N_k)
    try:
        df_ij, ddf_ij, _ = mbar.getFreeEnergyDifferences()
    except linalg.LinAlgError:
        return np.inf

    DeltaG_in_kT = df_ij[1,2]
    dDeltaG_in_kT = ddf_ij[1,2]

    final_time = time.time()
    elapsed_time = final_time - initial_time
    #print "%48s | %48s | reweighting took %.3f s" % (cid, iupac_name, elapsed_time)

#.........这里部分代码省略.........

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
  # Generate independent data samples from K one-dimensional harmonic oscillators centered at q = 0.
  #=============================================================================================
  
  randomsample = testsystems.harmonic_oscillators.HarmonicOscillatorsTestCase(O_k=O_k, K_k=K_k, beta=beta)
  [x_kn,u_kln,N_k] = randomsample.sample(N_k,mode='u_kln')

  # get the unreduced energies
  U_kln = u_kln/beta

  #=============================================================================================
  # Estimate free energies and expectations.
  #=============================================================================================

  # Initialize the MBAR class, determining the free energies.
  mbar = MBAR(u_kln, N_k, method = 'adaptive',relative_tolerance=1.0e-10,verbose=False) # use fast Newton-Raphson solver
  (Deltaf_ij_estimated, dDeltaf_ij_estimated) = mbar.getFreeEnergyDifferences()
  
  # Compute error from analytical free energy differences.
  Deltaf_ij_error = Deltaf_ij_estimated - Deltaf_ij_analytical

  # Estimate the expectation of the mean-squared displacement at each condition.
  if observe == 'RMS displacement':
    A_kn = numpy.zeros([K,K,N_max], dtype = numpy.float64);
    for k in range(0,K):
      for l in range(0,K):
        A_kn[k,l,0:N_k[k]] = (x_kn[k,0:N_k[k]] - O_k[l])**2 # observable is the squared displacement

  # observable is the potential energy, a 3D array since the potential energy is a function of 
  # thermodynamic state
  elif observe == 'potential energy':
    A_kn = U_kln

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
      print 'E_%d evaluated in model_%d'%(k,l), u_kln[k,l,n]


# Initialize MBAR with reduced energies u_kln and number of uncorrelated configurations from each state N_k.
# 
# u_kln[k,l,n] is the reduced potential energy beta*U_l(x_kn), where U_l(x) is the potential energy function for state l,
# beta is the inverse temperature, and and x_kn denotes uncorrelated configuration n from state k.
#
# N_k[k] is the number of configurations from state k stored in u_knm
# 
# Note that this step may take some time, as the relative dimensionless free energies f_k are determined at this point.
mbar = MBAR(u_kln, N_k)

# Extract dimensionless free energy differences and their statistical uncertainties.
(Deltaf_ij, dDeltaf_ij) = mbar.getFreeEnergyDifferences()
print 'Deltaf_ij', Deltaf_ij
print 'dDeltaf_ij', dDeltaf_ij
beta = 1.0 # keep in units kT
print 'Unit-bearing (units kT) free energy difference f_1K = f_K - f_1: %f +- %f' % ( (1./beta) * Deltaf_ij[0,K-1], (1./beta) * dDeltaf_ij[0,K-1])
f_df = np.zeros( (nlambda, 2) )  # first column is Deltaf_ij[0,:], second column is dDeltaf_ij[0,:]
f_df[:,0] = Deltaf_ij[0,:]
f_df[:,1] = dDeltaf_ij[0,:]
print 'Writing %s...'%args.bayesfactorfile
savetxt(args.bayesfactorfile, f_df)
print '...Done.'


# Compute the expectation of some observable A(x) at each state i, and associated uncertainty matrix.
# Here, A_kn[k,n] = A(x_{kn})
#(A_k, dA_k) = mbar.computeExpectations(A_kn)

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
    def MBAR_analysis(self, debug = False):
	"""MBAR analysis for populations and BICePs score"""
	# load necessary data first
	self.load_data()

	# Suppose the energies sampled from each simulation are u_kln, where u_kln[k,l,n] is the reduced potential energy
	#   of snapshot n \in 1,...,N_k of simulation k \in 1,...,K evaluated at reduced potential for state l.
	self.K = self.nlambda   # number of thermodynamic ensembles
	# N_k[k] will denote the number of correlated snapshots from state k
	N_k = np.array( [len(self.traj[i]['trajectory']) for i in range(self.nlambda)] )
	nsnaps = N_k.max()
	u_kln = np.zeros( (self.K, self.K, nsnaps) )
	nstates = int(self.states)
	print 'nstates', nstates
	states_kn = np.zeros( (self.K, nsnaps) )

	# Get snapshot energies rescored in the different ensembles
	"""['step', 'E', 'accept', 'state', 'sigma_noe', 'sigma_J', 'sigma_cs', 'sigma_pf''gamma']
	[int(step), float(self.E), int(accept), int(self.state), int(self.sigma_noe_index), int(self.sigma_J_index), int(self.sigma_cs_H_index), int(self.sigma_cs_Ha_index), int(self.sigma_cs_N_index), int(self.sigma_cs_Ca_index), int(self.sigma_pf_index), int(self.gamma_index)]               	"""

	for n in range(nsnaps):

  		for k in range(self.K):
    			for l in range(self.K):
				if debug:
      					print 'step', self.traj[k]['trajectory'][n][0],
      				if k==l:
          				print 'E%d evaluated in model %d'%(k,k), self.traj[k]['trajectory'][n][1],
          				u_kln[k,k,n] = self.traj[k]['trajectory'][n][1]
          			state, sigma_noe_index, sigma_J_index, sigma_cs_H_index, sigma_cs_Ha_index, sigma_cs_N_index, sigma_cs_Ca_index, sigma_pf_index, gamma_index = self.traj[k]['trajectory'][n][3:] 	# IMPORTANT: make sure the order of these parameters is the same as the way they are saved in PosteriorSampler
          			print 'state, sigma_noe_index, sigma_J_index, sigma_cs_H_index, sigma_cs_Ha_index, sigma_cs_N_index, sigma_cs_Ca_index, sigma_pf_index, gamma_index', state, sigma_noe_index, sigma_J_index, sigma_cs_H_index, sigma_cs_Ha_index, sigma_cs_N_index, sigma_cs_Ca_index, sigma_pf_index, gamma_index 
          			states_kn[k,n] = state
          			sigma_noe = self.traj[k]['allowed_sigma_noe'][sigma_noe_index]
          			sigma_J = self.traj[k]['allowed_sigma_J'][sigma_J_index]
          			sigma_cs_H = self.traj[k]['allowed_sigma_cs_H'][sigma_cs_H_index]
          			sigma_cs_Ha = self.traj[k]['allowed_sigma_cs_Ha'][sigma_cs_Ha_index]
          			sigma_cs_N = self.traj[k]['allowed_sigma_cs_N'][sigma_cs_N_index]
          			sigma_cs_Ca = self.traj[k]['allowed_sigma_cs_Ca'][sigma_cs_Ca_index]
          			sigma_pf = self.traj[k]['allowed_sigma_pf'][sigma_pf_index]
          			u_kln[k,l,n] = self.sampler[l].neglogP(0, state, sigma_noe, sigma_J, sigma_cs_H, sigma_cs_Ha, sigma_cs_N, sigma_cs_Ca, sigma_pf, gamma_index)
				if debug:
					print 'E_%d evaluated in model_%d'%(k,l), u_kln[k,l,n]


	# Initialize MBAR with reduced energies u_kln and number of uncorrelated configurations from each state N_k.
 	# u_kln[k,l,n] is the reduced potential energy beta*U_l(x_kn), where U_l(x) is the potential energy function for state l,
	# beta is the inverse temperature, and and x_kn denotes uncorrelated configuration n from state k.
	# N_k[k] is the number of configurations from state k stored in u_knm
	# Note that this step may take some time, as the relative dimensionless free energies f_k are determined at this point.
	mbar = MBAR(u_kln, N_k)

	# Extract dimensionless free energy differences and their statistical uncertainties.
#	(Deltaf_ij, dDeltaf_ij) = mbar.getFreeEnergyDifferences()
	#(Deltaf_ij, dDeltaf_ij, Theta_ij) = mbar.getFreeEnergyDifferences(uncertainty_method='svd-ew')
	(Deltaf_ij, dDeltaf_ij, Theta_ij) = mbar.getFreeEnergyDifferences(uncertainty_method='approximate')
	#print 'Deltaf_ij', Deltaf_ij
	#print 'dDeltaf_ij', dDeltaf_ij
	beta = 1.0 # keep in units kT
	#print 'Unit-bearing (units kT) free energy difference f_1K = f_K - f_1: %f +- %f' % ( (1./beta) * Deltaf_ij[0,K-1], (1./beta) * dDeltaf_ij[0,K-1])
	self.f_df = np.zeros( (self.nlambda, 2) )  # first column is Deltaf_ij[0,:], second column is dDeltaf_ij[0,:]
	self.f_df[:,0] = Deltaf_ij[0,:]
	self.f_df[:,1] = dDeltaf_ij[0,:]

	# Compute the expectation of some observable A(x) at each state i, and associated uncertainty matrix.
        # Here, A_kn[k,n] = A(x_{kn})
        #(A_k, dA_k) = mbar.computeExpectations(A_kn)
        self.P_dP = np.zeros( (nstates, 2*self.K) )  # left columns are P, right columns are dP
	if debug:
        	print 'state\tP\tdP'
    	for i in range(nstates):
        	A_kn = np.where(states_kn==i,1,0)
        	(p_i, dp_i) = mbar.computeExpectations(A_kn, uncertainty_method='approximate')
        	self.P_dP[i,0:self.K] = p_i
        	self.P_dP[i,self.K:2*self.K] = dp_i
		print i
        	for p in p_i: print p,
        	for dp in dp_i: print dp,
		print 
	pops, dpops = self.P_dP[:,0:self.K], self.P_dP[:,self.K:2*self.K]

	# save results
	self.save_MBAR()

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
#      print 'E_%d evaluated in model_%d'%(k,l), u_kln[k,l,n]


# Initialize MBAR with reduced energies u_kln and number of uncorrelated configurations from each state N_k.
# 
# u_kln[k,l,n] is the reduced potential energy beta*U_l(x_kn), where U_l(x) is the potential energy function for state l,
# beta is the inverse temperature, and and x_kn denotes uncorrelated configuration n from state k.
#
# N_k[k] is the number of configurations from state k stored in u_knm
# 
# Note that this step may take some time, as the relative dimensionless free energies f_k are determined at this point.
mbar = MBAR(u_kln, N_k)

# Extract dimensionless free energy differences and their statistical uncertainties.
#(Deltaf_ij, dDeltaf_ij) = mbar.getFreeEnergyDifferences()
(Deltaf_ij, dDeltaf_ij, Theta_ij) = mbar.getFreeEnergyDifferences(uncertainty_method='approximate')
#print 'Deltaf_ij', Deltaf_ij
#print 'dDeltaf_ij', dDeltaf_ij
beta = 1.0 # keep in units kT
#print 'Unit-bearing (units kT) free energy difference f_1K = f_K - f_1: %f +- %f' % ( (1./beta) * Deltaf_ij[0,K-1], (1./beta) * dDeltaf_ij[0,K-1])
f_df = np.zeros( (nlambda, 2) )  # first column is Deltaf_ij[0,:], second column is dDeltaf_ij[0,:]
f_df[:,0] = Deltaf_ij[0,:]
f_df[:,1] = dDeltaf_ij[0,:]
print 'Writing %s...'%args.bayesfactorfile
savetxt(args.bayesfactorfile, f_df)
print '...Done.'


# Compute the expectation of some observable A(x) at each state i, and associated uncertainty matrix.
# Here, A_kn[k,n] = A(x_{kn})
#(A_k, dA_k) = mbar.computeExpectations(A_kn)

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
kT = unit.AVOGADRO_CONSTANT_NA * unit.BOLTZMANN_CONSTANT_kB * integrator.getTemperature()
for k in range(nstates):
    for iteration in range(niterations):
        print('state %5d iteration %5d / %5d' % (k, iteration, niterations))
        # Set alchemical state
        context.setParameter('lambda', lambdas[k])
        # Run some dynamics
        integrator.step(nsteps)
        # Compute energies at all alchemical states
        for l in range(nstates):
            context.setParameter('lambda', lambdas[l])
            u_kln[k,l,iteration] = context.getState(getEnergy=True).getPotentialEnergy() / kT

# Estimate free energy of Lennard-Jones particle insertion
from pymbar import MBAR, timeseries
# Subsample data to extract uncorrelated equilibrium timeseries
N_k = np.zeros([nstates], np.int32) # number of uncorrelated samples
for k in range(nstates):
    [nequil, g, Neff_max] = timeseries.detectEquilibration(u_kln[k,k,:])
    indices = timeseries.subsampleCorrelatedData(u_kln[k,k,:], g=g)
    N_k[k] = len(indices)
    u_kln[k,:,0:N_k[k]] = u_kln[k,:,indices].T
# Compute free energy differences and statistical uncertainties
mbar = MBAR(u_kln, N_k)
[DeltaF_ij, dDeltaF_ij, Theta_ij] = mbar.getFreeEnergyDifferences()

print('DeltaF_ij (kT):')
print(DeltaF_ij)
print('dDeltaF_ij (kT):')
print(dDeltaF_ij)

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
def estimate_free_energies(ncfile, ndiscard=0, nuse=None, g=None):
    """
    Estimate free energies of all alchemical states.

    Parameters
    ----------
    ncfile : NetCDF
       Input YANK netcdf file
    ndiscard : int, optional, default=0
       Number of iterations to discard to equilibration
    nuse : int, optional, default=None
       Maximum number of iterations to use (after discarding)
    g : int, optional, default=None
       Statistical inefficiency to use if desired; if None, will be computed.

    TODO
    ----
    * Automatically determine 'ndiscard'.

    """

    # Get current dimensions.
    niterations = ncfile.variables['energies'].shape[0]
    nstates = ncfile.variables['energies'].shape[1]
    natoms = ncfile.variables['energies'].shape[2]

    # Extract energies.
    logger.info("Reading energies...")
    energies = ncfile.variables['energies']
    u_kln_replica = np.zeros([nstates, nstates, niterations], np.float64)
    for n in range(niterations):
        u_kln_replica[:,:,n] = energies[n,:,:]
    logger.info("Done.")

    # Deconvolute replicas
    logger.info("Deconvoluting replicas...")
    u_kln = np.zeros([nstates, nstates, niterations], np.float64)
    for iteration in range(niterations):
        state_indices = ncfile.variables['states'][iteration,:]
        u_kln[state_indices,:,iteration] = energies[iteration,:,:]
    logger.info("Done.")

    # Compute total negative log probability over all iterations.
    u_n = np.zeros([niterations], np.float64)
    for iteration in range(niterations):
        u_n[iteration] = np.sum(np.diagonal(u_kln[:,:,iteration]))
    #logger.info(u_n

    # DEBUG
    outfile = open('u_n.out', 'w')
    for iteration in range(niterations):
        outfile.write("%8d %24.3f\n" % (iteration, u_n[iteration]))
    outfile.close()

    # Discard initial data to equilibration.
    u_kln_replica = u_kln_replica[:,:,ndiscard:]
    u_kln = u_kln[:,:,ndiscard:]
    u_n = u_n[ndiscard:]

    # Truncate to number of specified conforamtions to use
    if (nuse):
        u_kln_replica = u_kln_replica[:,:,0:nuse]
        u_kln = u_kln[:,:,0:nuse]
        u_n = u_n[0:nuse]

    # Subsample data to obtain uncorrelated samples
    N_k = np.zeros(nstates, np.int32)
    indices = timeseries.subsampleCorrelatedData(u_n, g=g) # indices of uncorrelated samples
    #print u_n # DEBUG
    #indices = range(0,u_n.size) # DEBUG - assume samples are uncorrelated
    N = len(indices) # number of uncorrelated samples
    N_k[:] = N
    u_kln[:,:,0:N] = u_kln[:,:,indices]
    logger.info("number of uncorrelated samples:")
    logger.info(N_k)
    logger.info("")

    #===================================================================================================
    # Estimate free energy difference with MBAR.
    #===================================================================================================

    # Initialize MBAR (computing free energy estimates, which may take a while)
    logger.info("Computing free energy differences...")
    mbar = MBAR(u_kln, N_k)

    # Get matrix of dimensionless free energy differences and uncertainty estimate.
    logger.info("Computing covariance matrix...")

    try:
        # pymbar 2
        (Deltaf_ij, dDeltaf_ij) = mbar.getFreeEnergyDifferences()
    except ValueError:
        # pymbar 3
        (Deltaf_ij, dDeltaf_ij, theta_ij) = mbar.getFreeEnergyDifferences()

#    # Matrix of free energy differences
    logger.info("Deltaf_ij:")
    for i in range(nstates):
        str_row = ""
        for j in range(nstates):
#.........这里部分代码省略.........

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
print("======================================")
print("      Initializing MBAR               ")
print("======================================")

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

# Initialize the MBAR class, determining the free energies.
mbar = MBAR(u_kln, N_k, relative_tolerance=1.0e-10, verbose=True)
# Get matrix of dimensionless free energy differences and uncertainty estimate.

print("=============================================")
print("      Testing getFreeEnergyDifferences       ")
print("=============================================")

results = mbar.getFreeEnergyDifferences()
Delta_f_ij_estimated = results['Delta_f']
dDelta_f_ij_estimated = results['dDelta_f']

# Compute error from analytical free energy differences.
Delta_f_ij_error = Delta_f_ij_estimated - Delta_f_ij_analytical

print("Error in free energies is:")
print(Delta_f_ij_error)
print("Uncertainty in free energies is:")
print(dDelta_f_ij_estimated)

print("Standard deviations away is:")
# mathematical manipulation to avoid dividing by zero errors; we don't care
# about the diagnonals, since they are identically zero.
df_ij_mod = dDelta_f_ij_estimated + numpy.identity(K)

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import getFreeEnergyDifferences [as 别名]
def estimate_free_energies(ncfile, ndiscard = 0, nuse = None):
    """Estimate free energies of all alchemical states.

    ARGUMENTS
       ncfile (NetCDF) - input YANK netcdf file

    OPTIONAL ARGUMENTS
       ndiscard (int) - number of iterations to discard to equilibration
       nuse (int) - maximum number of iterations to use (after discarding)

    TODO: Automatically determine 'ndiscard'.
    """

    # Get current dimensions.
    niterations = ncfile.variables['energies'].shape[0]
    nstates = ncfile.variables['energies'].shape[1]
    natoms = ncfile.variables['energies'].shape[2]

    # Extract energies.
    print "Reading energies..."
    energies = ncfile.variables['energies']
    u_kln_replica = zeros([nstates, nstates, niterations], float64)
    for n in range(niterations):
        u_kln_replica[:,:,n] = energies[n,:,:]
    print "Done."

    # Deconvolute replicas
    print "Deconvoluting replicas..."
    u_kln = zeros([nstates, nstates, niterations], float64)
    for iteration in range(niterations):
        state_indices = ncfile.variables['states'][iteration,:]
        u_kln[state_indices,:,iteration] = energies[iteration,:,:]
    print "Done."

    # Compute total negative log probability over all iterations.
    u_n = zeros([niterations], float64)
    for iteration in range(niterations):
        u_n[iteration] = sum(diagonal(u_kln[:,:,iteration]))
    #print u_n

    # DEBUG
    outfile = open('u_n.out', 'w')
    for iteration in range(niterations):
        outfile.write("%8d %24.3f\n" % (iteration, u_n[iteration]))
    outfile.close()

    # Discard initial data to equilibration.
    u_kln_replica = u_kln_replica[:,:,ndiscard:]
    u_kln = u_kln[:,:,ndiscard:]
    u_n = u_n[ndiscard:]

    # Truncate to number of specified conforamtions to use
    if (nuse):
        u_kln_replica = u_kln_replica[:,:,0:nuse]
        u_kln = u_kln[:,:,0:nuse]
        u_n = u_n[0:nuse]
    
    # Subsample data to obtain uncorrelated samples
    N_k = zeros(nstates, int32)
    indices = timeseries.subsampleCorrelatedData(u_n) # indices of uncorrelated samples
    #indices = range(0,u_n.size) # DEBUG - assume samples are uncorrelated
    N = len(indices) # number of uncorrelated samples
    N_k[:] = N      
    u_kln[:,:,0:N] = u_kln[:,:,indices]
    print "number of uncorrelated samples:"
    print N_k
    print ""

    #===================================================================================================
    # Estimate free energy difference with MBAR.
    #===================================================================================================   
   
    # Initialize MBAR (computing free energy estimates, which may take a while)
    print "Computing free energy differences..."
    mbar = MBAR(u_kln, N_k, verbose = False, method = 'adaptive', maximum_iterations = 50000) # use slow self-consistent-iteration (the default)
    #mbar = MBAR(u_kln, N_k, verbose = False, method = 'self-consistent-iteration', maximum_iterations = 50000) # use slow self-consistent-iteration (the default)
    #mbar = MBAR(u_kln, N_k, verbose = True, method = 'Newton-Raphson') # use faster Newton-Raphson solver

    # Get matrix of dimensionless free energy differences and uncertainty estimate.
    print "Computing covariance matrix..."
    (Deltaf_ij, dDeltaf_ij) = mbar.getFreeEnergyDifferences(uncertainty_method='svd-ew')
   
#    # Matrix of free energy differences
    print "Deltaf_ij:"
    for i in range(nstates):
        for j in range(nstates):
            print "%8.3f" % Deltaf_ij[i,j],
        print ""        
    
#    print Deltaf_ij
#    # Matrix of uncertainties in free energy difference (expectations standard deviations of the estimator about the true free energy)
    print "dDeltaf_ij:"
    for i in range(nstates):
        for j in range(nstates):
            print "%8.3f" % dDeltaf_ij[i,j],
        print ""        

    # Return free energy differences and an estimate of the covariance.
    return (Deltaf_ij, dDeltaf_ij)