Python MBAR._computeUnnormalizedLogWeights方法代码示例

# 需要导入模块: from pymbar import MBAR [as 别名]
# 或者: from pymbar.MBAR import _computeUnnormalizedLogWeights [as 别名]
def write_pdb_resampled(ncfile, atoms, output_pdb_filename, reference_state=0, ndiscard=0, nsamples=1000):
    """
    Resample configurations with probability in specified state, concatenating into PDB file.

    ARGUMENTS
       ncfile (NetCDF) - input YANK netcdf file
       atoms (list) - atom records from reference PDB file
       output_pdb_filename (string) - PDB file of resampled configurations to write

    OPTIONAL ARGUMENTS
       reference_state (int) - state to reweight to
       ndiscard (int) - number of iterations to discard to equilibration
       nsamples (int) - number of sampls to generate

    """

    import numpy

    # 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 to collect by thermodynamic state.
    print "Deconvoluting replicas..."
    u_kln = numpy.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."

    # Discard initial data to equilibration.
    u_kln = u_kln[:,:,ndiscard:]
    [K,L,N] = u_kln.shape
    N_k = N * numpy.ones([K], numpy.int32)

    # Compute snapshot energies.
    print "Computing snapshot energies..."
    u_kn = numpy.zeros([K,N], numpy.float32)
    # Get temperature.
    temperature = ncfile.groups['thermodynamic_states'].variables['temperatures'][reference_state] * units.kelvin
    kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA # Boltzmann constant
    kT = kB * temperature # thermal energy
    # Deserialize system.
    system = openmm.XmlSerializer.deserialize(str(ncfile.groups['thermodynamic_states'].variables['systems'][reference_state]))
    # Create Context.
    integrator = openmm.VerletIntegrator(1.0 * units.femtoseconds)
    context = openmm.Context(system, integrator)
    # Turn off restraints.
    context.setParameter('restraint_lambda', 0.0)
    for n in range(N):
        iteration = ndiscard + n
        for replica_index in range(K):        
            # Recompute energies.
            state_index = ncfile.variables['states'][iteration,replica_index]
            positions = ncfile.variables['positions'][iteration,replica_index]
            context.setPositions(positions)
            openmm_state = context.getState(getEnergy=True)
            u_kn[state_index,n] = openmm_state.getPotentialEnergy() / kT

    #===================================================================================================
    # Initialize MBAR.
    #===================================================================================================   
   
    # Initialize MBAR (computing free energy estimates, which may take a while)
    print "Initializing MBAR (warning: using all correlated data)..."
    mbar = MBAR(u_kln, N_k, verbose=False) 

    # Get snapshot weights.
    #u_kn = numpy.squeeze(u_kln[:,state,:])
    log_w_kn = mbar._computeUnnormalizedLogWeights(u_kn)
    f = _logsum(log_w_kn[mbar.indices])
    w_kn = numpy.exp(log_w_kn - f)
    p_kn = w_kn / w_kn[mbar.indices].sum()

    # Form linear list of potential snapshots to choose.
    [K, N] = p_kn.shape
    snapshots = list()
    probabilities = list()
    for replica in range(K):
        for sample in range(N):
            state = ncfile.variables['states'][ndiscard+sample, replica]
            snapshots.append( (replica, ndiscard+sample) )
            probabilities.append( p_kn[state,sample] )
    probabilities = numpy.array(probabilities)

    # Draw samples.
    state_hist = numpy.zeros([K], numpy.int32)
    import numpy.random
    sample_indices = numpy.random.choice(range(len(snapshots)), size=[nsamples], p=probabilities)
    title = 'generated by resampling'
    outfile = open(output_pdb_filename, 'w')
#.........这里部分代码省略.........