demux trajs

#!/usr/bin/env python import sys,os

import IMP import IMP.atom import IMP.container import IMP.isd from IMP.isd.shared_functions import sfo_common from IMP.isd.Statistics import Statistics from IMP.isd.Entry import Entry

from math import sqrt

kB= (1.381 * 6.02214) / 4184.0

class sfo(sfo_common): "shared functions object, published on all nodes"

def init_model(self, wd, seqfile, initpdb, noe, ff_temp=300.0): "loads pdb and restraints and creates particles and scales"

Create an IMP model and add a heavy atom-only protein from a PDB file

self.init_model_base(wd) m=self._m # print "loading protein and force field" prot, ff, rsb, rs = self.init_model_charmm_protein_and_ff(initpdb, "top.lib", "par.lib", IMP.atom.NonWaterPDBSelector(), IMP.isd.RepulsiveDistancePairScore(0,1), ff_temp=300.0) self._rs = {} self._rs['phys'] = rs self._rs['phys_bonded'] = rsb self._p={} self._p['prot'] = prot # print "NOE restraints" noe_rs, prior_rs, sigma, gamma = self.init_model_NOEs(prot, seqfile, noe, name='NOE', prior_rs=None, bounds_sigma=(1.0,0.1,100), bounds_gamma=(1.0,0.1,100)) self._rs['NOE'] = noe_rs self._rs['prior'] = prior_rs self._p['sigma'] = sigma self._p['gamma'] = gamma

def init_simulation(self,lambda_temp=1.0, tau=500.0): "prepare md and mc sims" self.inv_temp = lambda_temp self.md_tau = tau self._md = self._setup_md(thermostat = 'berendsen', temperature = 1/(kB*lambda_temp), coupling = tau) rs_scales=[self._rs['NOE'],self._rs['prior']] self._mc_sigma, self._nm_sigma = \ self.init_simulation_setup_scale_mc( self._p['sigma'], 1/(kB*lambda_temp), mc_restraints=rs_scales, nm_stepsize=0.1) self._mc_gamma, self._nm_gamma = \ self.init_simulation_setup_scale_mc( self._p['gamma'], 1/(kB*lambda_temp), mc_restraints=rs_scales, nm_stepsize=0.1)

def init_stats(self, prefix='p01',rate=1): #print statistics in files with prefix and at a certain rate (in gibbs #sampling steps). stat = Statistics(prefix,rate) self.stat = stat #general stuff to monitor: forcefield terms def get_ephys(): return self._rs['phys_bonded'].evaluate(False) \

self._rs['phys'].evaluate(False) stat.add_entry('global', entry=Entry('E_phys', '%10f', get_ephys)) stat.add_entry('global', entry=Entry('E_NOE', '%10f', self._rs['NOE'].evaluate, False)) stat.add_entry('global', entry=Entry('E_prior', '%10f', self._rs['prior'].evaluate, False)) #MD monitoring: target_temp, instant_temp, E_kinetic self.md_key = self.init_stats_add_md_category(stat) #MC monitoring: acceptance, stepsize, sigma self.sigma_mc_key = \ self.init_stats_add_mc_category(stat, name='mc_sigma', coord='sigma') #MC monitoring: acceptance, stepsize, gamma self.gamma_mc_key = \ self.init_stats_add_mc_category(stat, name='mc_gamma', coord='gamma')

def do_md(self,nsteps): "perform md simulation on protein for nsteps" for i in IMP.atom.get_leaves(self._p['prot']): IMP.core.XYZR(i).set_coordinates_are_optimized(True) self._p['sigma'].set_is_optimized(IMP.FloatKey("scale"),False) self._p['gamma'].set_is_optimized(IMP.FloatKey("scale"),False) self._md.optimize(nsteps) self.do_md_protein_statistics(self.md_key, nsteps, self._md, temperature=1/(kB*self.inv_temp), prot_coordinates=self.get_pdb(self._p['prot']))

def do_mc(self,nsteps): """perform mc on scales for nsteps, updating stepsizes to target 50% acceptance. Don't make nsteps too small (say << 50). """ prot = self._p['prot'] sigma = self._p['sigma'] gamma = self._p['gamma'] for i in IMP.atom.get_leaves(prot): IMP.core.XYZR(i).set_coordinates_are_optimized(False) sigma.set_is_optimized(IMP.FloatKey("scale"),True) gamma.set_is_optimized(IMP.FloatKey("scale"),False) self._mc_and_update_nm(nsteps, self._mc_sigma, self._nm_sigma, self.sigma_mc_key, adjust_stepsize=True) self.stat.update(self.sigma_mc_key, 'sigma', sigma.get_scale()) sigma.set_is_optimized(IMP.FloatKey("scale"),False) gamma.set_is_optimized(IMP.FloatKey("scale"),True) self._mc_and_update_nm(nsteps,self._mc_gamma,self._nm_gamma, self.gamma_mc_key) self.stat.update(self.gamma_mc_key, 'gamma', gamma.get_scale())

def write_stats(self): #increment global counter self.stat.increment_counter('global', 1) #write statistics if necessary self.stat.write_stats()

def set_inv_temp(self, inv_temp): "sets inverse temperature of mc and md sims (used in replica exchange)" #MD: temperature and rescale velocities self._md.set_thermostat(2, 1.0/(kB*inv_temp), self.md_tau) self._md.rescale_velocities(sqrt(self.inv_temp/inv_temp)) #MC: temperature self._mc_sigma.set_temperature(1.0/inv_temp) self._mc_gamma.set_temperature(1.0/inv_temp) #store it locally as well self.inv_temp = inv_temp

def get_inv_temp(self): return self.inv_temp

def get_state(self): """get_state and set_state work in tandem. They should only respect two rules:

get_state should return a dict that contains a key called inv_temp which is the inverse temperature used by replica-exchange
set_state should be able to set the state of this node when given the output of a get_state call. """ state={} state['inv_temp']=self.get_inv_temp() state['gamma_mc_stepsize']=self._nm_gamma.get_sigma() state['sigma_mc_stepsize']=self._nm_sigma.get_sigma() return state

def set_state(self,state): """set the state of this node according to the passed dict. See get_state. """ self.set_inv_temp(state['inv_temp']) self._nm_gamma.set_sigma(state['gamma_mc_stepsize']) self._nm_sigma.set_sigma(state['sigma_mc_stepsize'])

if name == 'main': sfo=sfo() sfo.init_model('./','sequence.dat', 'generated2.pdb', 'NOE_HN-full_7A_sparse100.tbl') sfo.init_simulation() sfo.init_stats() for i in xrange(100): print "\r%d" % i, sys.stdout.flush() sfo.do_md(10) sfo.do_mc(10) sfo.write_stats()

#!/usr/bin/env python
if __name__ == '__main__':
    #retrieve number of replicas and replicanums list
    replicanums = []
    stepno = []
    fl=open('replicanums.txt')
    tokens=(fl.readline()).split()
    nreps = len(tokens)-1
    replicanums.append([int(i)-1 for i in tokens[1:]])
    stepno.append(int(tokens[0]))
    for line in fl:
        tokens=line.split()
        replicanums.append([int(i)-1 for i in tokens[1:]])
        stepno.append(int(tokens[0]))
    #demux p??_stats.txt files
    infiles = [open('r%02d_stats.txt' % i) for i in xrange(1,nreps+1)]
    outfiles = [open('p%02d_stats.txt' % i, 'w') for i in xrange(1,nreps+1)]
    for inf,outf in zip(infiles,outfiles):
        outf.write(inf.readline())
    lines = [fl.readline() for fl in infiles]
    for (no,rn) in zip(stepno,replicanums):
        while (lines[0] != '' and int(lines[0].split()[0]) == no):
            for i in xrange(nreps):
                outfiles[i].write(lines[rn[i]])
            lines = [fl.readline() for fl in infiles]

harmonic

#!/usr/bin/env python import sys,os

import IMP import IMP.atom import IMP.container import IMP.isd from IMP.isd.shared_functions import sfo_common from IMP.isd.Statistics import Statistics from IMP.isd.Entry import Entry

from math import sqrt

kB= (1.381 * 6.02214) / 4184.0

class sfo(sfo_common): "shared functions object, published on all nodes"

def init_model(self, wd, seqfile, initpdb, noe, ff_temp=300.0): "loads pdb and restraints and creates particles and scales"

Create an IMP model and add a heavy atom-only protein from a PDB file

self._rs['phys'].evaluate(False) stat.add_entry('global', entry=Entry('E_phys', '%10f', get_ephys)) stat.add_entry('global', entry=Entry('E_NOE', '%10f', self._rs['NOE'].evaluate, False)) stat.add_entry('global', entry=Entry('E_prior', '%10f', self._rs['prior'].evaluate, False)) #MD monitoring: target_temp, instant_temp, E_kinetic self.md_key = self.init_stats_add_md_category(stat) #MC monitoring: acceptance, stepsize, sigma self.sigma_mc_key = \ self.init_stats_add_mc_category(stat, name='mc_sigma', coord='sigma') #MC monitoring: acceptance, stepsize, gamma self.gamma_mc_key = \ self.init_stats_add_mc_category(stat, name='mc_gamma', coord='gamma')

def write_stats(self): #increment global counter self.stat.increment_counter('global', 1) #write statistics if necessary self.stat.write_stats()

def get_inv_temp(self): return self.inv_temp

def get_state(self): """get_state and set_state work in tandem. They should only respect two rules:

get_state should return a dict that contains a key called inv_temp which is the inverse temperature used by replica-exchange
set_state should be able to set the state of this node when given the output of a get_state call. """ state={} state['inv_temp']=self.get_inv_temp() state['gamma_mc_stepsize']=self._nm_gamma.get_sigma() state['sigma_mc_stepsize']=self._nm_sigma.get_sigma() return state

#!/usr/bin/env python
import sys,os,errno
import atexit
import random
try:
    from IMP.isd.PyroGrid import PyroGrid as Grid
#   from IMP.isd.FileBasedGrid import FileBasedGrid as Grid
except ImportError:
    print >> sys.stderr, "This example needs the Python Pyro module"
    sys.exit(0)
from IMP.isd.hosts import create_host_list
import IMP.container
import IMP.isd
import shared_functions as sf
from IMP.isd.Replica import ReplicaTracker
IMP.set_log_level(IMP.NONE)
###simulation settings
#where to output files
outfolder=os.path.join(os.getcwd(), 'results')
#temp dir
tmpdir=os.path.join(os.getcwd(), 'tmp')
#number of replicas / hosts
nreps = 10
#lambda scaling distribution
kB= (1.381 * 6.02214) / 4184.0
lambda_1 = 1.0/(kB*100)
lambda_N = 1.0/(kB*1000)
lambdas=[lambda_N*(lambda_1/lambda_N)**((float(nreps)-k)/(nreps-1))
        for k in xrange(1,nreps+1)]
#thermostat coupling constant (berendsen, in fs)
tau=[500.0]*nreps
#stat_rate is the rate at which to print out traj statistics
#(in units of gibbs sampling steps)
stat_rate=[1]*nreps
#list of files relative to the current dir to copy over to all nodes
filelist=['shared_functions.py'] #add whatever you want
#prefix of output files
nums=[os.path.join(outfolder,'r%02d' % (i+1)) for i in xrange(nreps)]
#number of gibbs sampling steps
n_gibbs = 1000
#number of mc steps
n_mc = 1000
#where to run sims
hostlist = ['localhost']*nreps
#replica exchange scheme
rex_scheme='standard'
#replica exchange exchange method
rex_xchg='gromacs'
rexlog = os.path.join(outfolder,'replicanums.txt')
#misc
imppy = os.path.abspath(
        os.path.join(os.getenv('IMP_ISD_DATA'),'../../tools/imppy.sh'))
src_path = os.path.abspath(
        os.path.join(os.getenv('IMP_ISD_DATA'),'../lib/IMP/isd'))
showX11 = False
grid_debug = False
grid_verbose = False
X11_delay = 1.0
window_size = '80x25'
#pyroGrid
shared_temp_path = True
terminate_during_publish = False
nshost = None
def mkdir_p(path):
    "mkdir -p, taken from stackoverflow"
    try:
        os.makedirs(path)
    except OSError, exc:
        if exc.errno == errno.EEXIST:
            pass
        else: raise
def launch_grid():
    hosts = create_host_list(hostlist,tmpdir)
    for host in hosts:
        #ugly hack
        host.init_cmd = imppy + ' !'
    #pyro grid
    grid = Grid(hosts, src_path, showX11, X11_delay, grid_debug,
            grid_verbose, shared_temp_path, nshost, terminate_during_publish)
    #uncomment the following lines and comments the two previous ones to use file based grid
    #file based grid
    #grid = Grid(hosts, src_path, showX11, X11_delay, grid_debug,
    #        grid_verbose)
    #grid.shared_temp_path = shared_temp_path
    if showX11:
        grid.window_size = window_size
    grid.copy_files('./', filelist)
    #grid.copy_files(src_path,["shared_functions.py"])
    #start grid on all nodes
    grid.start()
    return grid
def main():
    # launch grid
    print "launching grid"
    grid = launch_grid()
    #publish the shared object
    print "publishing sfo"
    sfo = sf.sfo()
    sfo_id = grid.publish(sfo)
    print "communication test"
    #get a worker to do a task
    #proxy = grid.acquire_service(sfo_id)
    #print proxy.hello().get()
    #grid.release_service(proxy)
    print "broadcast test"
    print grid.gather(grid.broadcast(sfo_id, 'hello'))
    #call init on all nodes
    print "initializing model"
    mkdir_p(tmpdir)
    mkdir_p(outfolder)
    requests = grid.broadcast(sfo_id, 'init_model', tmpdir)
    #wait til init is done
    results = grid.gather(requests)
    #turn off verbose noise (works because IMP.NONE is picklable, being an int.
    grid.gather(grid.broadcast(sfo_id, 'set_checklevel', IMP.NONE))
    grid.gather(grid.broadcast(sfo_id, 'set_loglevel', IMP.NONE))
    # evaluate the score of the whole system (without derivatives, False flag)
    print "initial energy"
    grid.gather(grid.broadcast(sfo_id, 'm', 'evaluate', False))
    #berendsen 300K tau=0.5ps
    #perform two independent MC moves for sigma and gamma
    print "initializing simulation and statistics"
    grid.gather(grid.scatter(sfo_id, 'init_simulation',
        zip(lambdas, tau)))
    grid.gather(grid.scatter(sfo_id, 'init_stats', zip(nums, stat_rate)))
    replica = ReplicaTracker(nreps, lambdas, grid, sfo_id,
            rexlog=rexlog,
            scheme=rex_scheme, xchg=rex_xchg,
            tune_temps=False)
    print "start gibbs sampling loop"
    for i in range(n_gibbs):
        print "\rgibbs step %d" % i,
        sys.stdout.flush()
        #print " mc"
        grid.gather(grid.broadcast(sfo_id, 'do_mc_and_update_stepsize', n_mc))
        grid.gather(grid.broadcast(sfo_id, 'write_stats'))
        #print " swaps"
        replica.replica_exchange()
        #print " stats"
        replica.write_rex_stats()
    print "terminating grid"
    grid.terminate()
    print "done."
if __name__ == '__main__':
    main()

shared functions

#!/usr/bin/env python import sys,os

import IMP import IMP.atom import IMP.container import IMP.isd from IMP.isd.shared_functions import sfo_common from IMP.isd.Statistics import Statistics from IMP.isd.Entry import Entry

from math import sqrt

kB= (1.381 * 6.02214) / 4184.0

class sfo(sfo_common): "shared functions object, published on all nodes"

def init_model(self, wd, seqfile, initpdb, noe, ff_temp=300.0): "loads pdb and restraints and creates particles and scales"

Create an IMP model and add a heavy atom-only protein from a PDB file

self._rs['phys'].evaluate(False) stat.add_entry('global', entry=Entry('E_phys', '%10f', get_ephys)) stat.add_entry('global', entry=Entry('E_NOE', '%10f', self._rs['NOE'].evaluate, False)) stat.add_entry('global', entry=Entry('E_prior', '%10f', self._rs['prior'].evaluate, False)) #MD monitoring: target_temp, instant_temp, E_kinetic self.md_key = self.init_stats_add_md_category(stat) #MC monitoring: acceptance, stepsize, sigma self.sigma_mc_key = \ self.init_stats_add_mc_category(stat, name='mc_sigma', coord='sigma') #MC monitoring: acceptance, stepsize, gamma self.gamma_mc_key = \ self.init_stats_add_mc_category(stat, name='mc_gamma', coord='gamma')

def write_stats(self): #increment global counter self.stat.increment_counter('global', 1) #write statistics if necessary self.stat.write_stats()

def get_inv_temp(self): return self.inv_temp

def get_state(self): """get_state and set_state work in tandem. They should only respect two rules:

get_state should return a dict that contains a key called inv_temp which is the inverse temperature used by replica-exchange
set_state should be able to set the state of this node when given the output of a get_state call. """ state={} state['inv_temp']=self.get_inv_temp() state['gamma_mc_stepsize']=self._nm_gamma.get_sigma() state['sigma_mc_stepsize']=self._nm_sigma.get_sigma() return state

#!/usr/bin/env python
import sys,os
import IMP
import IMP.atom
import IMP.container
import IMP.isd
from IMP.isd.shared_functions import sfo_common
from IMP.isd.Statistics import Statistics
from IMP.isd.Entry import Entry
from math import sqrt
kB= (1.381 * 6.02214) / 4184.0
class sfo(sfo_common):
    "shared functions object, published on all nodes"
    def init_model(self, wd, seqfile, initpdb, noe, ff_temp=300.0):
        "loads pdb and restraints and creates particles and scales"
        # Create an IMP model and add a heavy atom-only protein from a PDB file
        self.init_model_base(wd)
        m=self._m
        #
        print "loading protein and force field"
        prot, ff, rsb, rs = self.init_model_charmm_protein_and_ff(initpdb,
                                "top.lib", "par.lib",
                                IMP.atom.NonWaterPDBSelector(),
                                IMP.isd.RepulsiveDistancePairScore(0,1),
                                ff_temp=300.0)
        self._rs = {}
        self._rs['phys'] = rs
        self._rs['phys_bonded'] = rsb
        self._p={}
        self._p['prot'] = prot
        #
        print "NOE restraints"
        noe_rs, prior_rs, sigma, gamma = self.init_model_NOEs(prot,
                        seqfile, noe, name='NOE', prior_rs=None,
                        bounds_sigma=(1.0,0.1,100),
                        bounds_gamma=(1.0,0.1,100))
        self._rs['NOE'] = noe_rs
        self._rs['prior'] = prior_rs
        self._p['sigma'] = sigma
        self._p['gamma'] = gamma
    def init_simulation(self,lambda_temp=1.0, tau=500.0):
        "prepare md and mc sims"
        self.inv_temp = lambda_temp
        self.md_tau = tau
        self._md = self._setup_md(thermostat = 'berendsen',
                         temperature = 1/(kB*lambda_temp),
                         coupling = tau)
        rs_scales=[self._rs['NOE'],self._rs['prior']]
        self._mc_sigma, self._nm_sigma = \
                self.init_simulation_setup_scale_mc(
                        self._p['sigma'], 1/(kB*lambda_temp),
                            mc_restraints=rs_scales, nm_stepsize=0.1)
        self._mc_gamma, self._nm_gamma = \
                self.init_simulation_setup_scale_mc(
                        self._p['gamma'], 1/(kB*lambda_temp),
                            mc_restraints=rs_scales, nm_stepsize=0.1)
    def init_stats(self, prefix='p01',rate=1):
        #print statistics in files with prefix and at a certain rate (in gibbs
        #sampling steps).
        stat = Statistics(prefix,rate)
        self.stat = stat
        #general stuff to monitor: forcefield terms
        def get_ephys():
            return self._rs['phys_bonded'].evaluate(False) \
                    + self._rs['phys'].evaluate(False)
        stat.add_entry('global', entry=Entry('E_phys', '%10f', get_ephys))
        stat.add_entry('global', entry=Entry('E_NOE', '%10f',
                                        self._rs['NOE'].evaluate, False))
        stat.add_entry('global', entry=Entry('E_prior', '%10f',
                                        self._rs['prior'].evaluate, False))
        #MD monitoring: target_temp, instant_temp, E_kinetic
        self.md_key = self.init_stats_add_md_category(stat)
        #MC monitoring: acceptance, stepsize, sigma
        self.sigma_mc_key = \
                self.init_stats_add_mc_category(stat, name='mc_sigma',
                                                        coord='sigma')
        #MC monitoring: acceptance, stepsize, gamma
        self.gamma_mc_key = \
                self.init_stats_add_mc_category(stat, name='mc_gamma',
                                                        coord='gamma')
    def do_md(self,nsteps):
        "perform md simulation on protein for nsteps"
        for i in IMP.atom.get_leaves(self._p['prot']):
            IMP.core.XYZR(i).set_coordinates_are_optimized(True)
        self._p['sigma'].set_is_optimized(IMP.FloatKey("scale"),False)
        self._p['gamma'].set_is_optimized(IMP.FloatKey("scale"),False)
        self._md.optimize(nsteps)
        self.do_md_protein_statistics(self.md_key, nsteps, self._md,
                temperature=1/(kB*self.inv_temp),
                prot_coordinates=self.get_pdb(self._p['prot']))
    def do_mc(self,nsteps):
        """perform mc on scales for nsteps, updating stepsizes to target
        50% acceptance. Don't make nsteps too small (say << 50).
        """
        prot = self._p['prot']
        sigma = self._p['sigma']
        gamma = self._p['gamma']
        for i in IMP.atom.get_leaves(prot):
            IMP.core.XYZR(i).set_coordinates_are_optimized(False)
        sigma.set_is_optimized(IMP.FloatKey("scale"),True)
        gamma.set_is_optimized(IMP.FloatKey("scale"),False)
        self._mc_and_update_nm(nsteps, self._mc_sigma, self._nm_sigma,
                self.sigma_mc_key, adjust_stepsize=True)
        self.stat.update(self.sigma_mc_key, 'sigma', sigma.get_scale())
        sigma.set_is_optimized(IMP.FloatKey("scale"),False)
        gamma.set_is_optimized(IMP.FloatKey("scale"),True)
        self._mc_and_update_nm(nsteps,self._mc_gamma,self._nm_gamma,
                self.gamma_mc_key)
        self.stat.update(self.gamma_mc_key, 'gamma', gamma.get_scale())
    def write_stats(self):
        #increment global counter
        self.stat.increment_counter('global', 1)
        #write statistics if necessary
        self.stat.write_stats()
    def set_inv_temp(self, inv_temp):
        "sets inverse temperature of mc and md sims (used in replica exchange)"
        #MD: temperature and rescale velocities
        self._md.set_thermostat(2, 1.0/(kB*inv_temp), self.md_tau)
        self._md.rescale_velocities(sqrt(self.inv_temp/inv_temp))
        #MC: temperature
        self._mc_sigma.set_temperature(1.0/inv_temp)
        self._mc_gamma.set_temperature(1.0/inv_temp)
        #store it locally as well
        self.inv_temp = inv_temp
    def get_inv_temp(self):
        return self.inv_temp
    def get_state(self):
        """get_state and set_state work in tandem. They should only respect two
        rules:
            - get_state should return a dict that contains a key called inv_temp
              which is the inverse temperature used by replica-exchange
            - set_state should be able to set the state of this node when given
              the output of a get_state call.
        """
        state={}
        state['inv_temp']=self.get_inv_temp()
        state['gamma_mc_stepsize']=self._nm_gamma.get_sigma()
        state['sigma_mc_stepsize']=self._nm_sigma.get_sigma()
        return state
    def set_state(self,state):
        """set the state of this node according to the passed dict.
        See get_state.
        """
        self.set_inv_temp(state['inv_temp'])
        self._nm_gamma.set_sigma(state['gamma_mc_stepsize'])
        self._nm_sigma.set_sigma(state['sigma_mc_stepsize'])
if __name__ == '__main__':
    sfo=sfo()
    sfo.init_model('./','sequence.dat', 'generated2.pdb',
            'NOE_HN-full_7A_sparse100.tbl')
    sfo.init_simulation()
    sfo.init_stats()
    for i in xrange(100):
        print "\r%d" % i,
        sys.stdout.flush()
        sfo.do_md(10)
        sfo.do_mc(10)
        sfo.write_stats()