IMP logo

IMP.domino examples

four_proteins_optimization.py

Look for the positions of 4 proteins that optimize a scoring function.

import IMP
import IMP.domino
import IMP.core
import IMP.helper
import random,math
NUM_STATES=5

def create_random_transformation(angle_step=30.,translation_step=4.):
    zero_vec=IMP.algebra.Vector3D(0.,0.,0.)
    z_vec=IMP.algebra.Vector3D(0.,0.,1.)
    rand_t=IMP.algebra.Transformation3D(
        IMP.algebra.get_rotation_in_radians_about_axis(
        z_vec,
        random.uniform(-math.radians(angle_step),math.radians(angle_step))),
        IMP.algebra.get_random_vector_in(IMP.algebra.Sphere3D(zero_vec,translation_step)))
    return rand_t

#randomly sample NUM_STATES states for each particle
def setup_discrete_sampling_space(ps):
    mdl=ps[0].get_model()
    atts=IMP.core.XYZ.get_xyz_keys()
    discrete_set = IMP.domino.MappedDiscreteSet(ps,atts)
    for p in ps:
        for i in range(NUM_STATES):
            #create a new state
            rand_t=create_random_transformation()
            new_p=IMP.Particle(mdl)
            IMP.domino.Transformation.setup_particle(new_p,rand_t)
            discrete_set.add_state(new_p)
            discrete_set.add_mapped_state(p,new_p)
    return discrete_set

#set restraints
def setup_restraints(mhs):
    mdl=mhs[0].get_particle().get_model()
    rsrs=[]
    pairs=[[0,1],[0,2],[0,3],[1,3]]
    simple_rsrs=[]
    for [i,j] in pairs:
        my_rbs=IMP.core.RigidBodies()
        my_ps=IMP.Particles()
        my_ps.append(mhs[i].get_particle())
        my_ps.append(mhs[j].get_particle())
        my_rbs.append(IMP.core.RigidBody(my_ps[0]))
        my_rbs.append(IMP.core.RigidBody(my_ps[0]))
        rsrs.append([IMP.helper.create_simple_distance(my_ps),
        #rsrs.append([IMP.helper.create_simple_excluded_volume_on_rigid_bodies(my_rbs),
                     my_ps])
    return rsrs

def fast_enumerate(sampler,rsrs,mdl,num_sols):
    """
    Enumerate the discrete sampling space to find the best combination
    """
    jt_filename = IMP.domino.get_example_path("1z5s.jt")
    jt = IMP.domino.JunctionTree()
    IMP.domino.read_junction_tree(jt_filename,jt)
    re = IMP.domino.RestraintEvaluator(sampler)
    d_opt = IMP.domino.DominoOptimizer(jt,mdl,re)
    #add restraints to be evaluates by the optimizer
    for r in rsrs:
        d_opt.add_restraint(r[0].get_restraint(),r[1])
    #set the sampling space
    d_opt.set_sampling_space(sampler)
    #to see the graph, uncomment this line
    #d_opt.show_restraint_graph()
    d_opt.set_number_of_solutions(num_sols)
    d_opt.optimize(num_sols)
    return d_opt

IMP.set_log_level(IMP.SILENT)
mdl=IMP.Model()
sel=IMP.atom.CAlphaPDBSelector()
#1. set up the particles
mhs_ps=IMP.Particles()
mhs=IMP.atom.Hierarchies()
for prot_name,chain_id in [["1z5s_A.pdb","A"],["1z5s_B.pdb","B"],["1z5s_C.pdb","C"],["1z5s_D.pdb","D"]]:
    mhs.append(IMP.atom.read_pdb(IMP.domino.get_example_path(prot_name),
                                 mdl,sel))
    mhs[-1].get_particle().add_attribute(IMP.domino.node_name_key(),chain_id)
    IMP.atom.add_radii(mhs[-1])
    IMP.atom.setup_as_rigid_body(mhs[-1])
    mhs_ps.append(mhs[-1].get_particle())

#2. set up the discrete set of states
discrete_set=setup_discrete_sampling_space(mhs_ps)

#3. add restraints (defining the scoring function)
restraints=setup_restraints(mhs)

#4. set the sampler
sampler = IMP.domino.CartesianProductSampler(discrete_set,mhs_ps)
#here alternative samplers can be PermutationSampler.

#5. optimize
num_sols=5
d_opt=fast_enumerate(sampler,restraints,mdl,num_sols)

#6. print solutions
rg = d_opt = d_opt.get_graph()
for i in range(num_sols):
    comb = rg.get_opt_combination(i)
    print "solution number:",i," consist of the configuration: ",comb.key(), " with score ",comb.get_total_score()

six_particles_optimization.py

Look for the positions of 6 particles that optimize a scoring function.

import IMP
import IMP.domino
import IMP.core
NUM_PARTICLES=6
NUM_STATES=40

#randomly sample NUM_STATES states for each particle
def setup_discrete_sampling_space(ps,bb):
    mdl=ps[0].get_model()
    atts=IMP.core.XYZ.get_xyz_keys()
    discrete_set = IMP.domino.MappedDiscreteSet(ps,atts)
    for p in ps:
        for i in range(NUM_STATES):
            #create a new state
            new_p=IMP.Particle(mdl)
            IMP.core.XYZ.setup_particle(new_p,IMP.algebra.get_random_vector_in(bb))
            discrete_set.add_state(new_p)
            discrete_set.add_mapped_state(p,new_p)
    return discrete_set

#set restraints
def setup_restraints(ps):
    mdl=ps[0].get_model()
    rsrs=[]
    pairs=[[0,1],[0,2],[1,2],[2,3],[3,4],[4,5],[3,5]]
    sf = IMP.core.Harmonic(1.0, 0.5)
    for pair in pairs:
        r=IMP.core.DistanceRestraint(sf, ps[pair[0]], ps[pair[1]])
        r.set_was_used(True)
        rsrs.append(r)
        mdl.add_restraint(r)#different than doing mdl.add_restraint(rsrs) - Daniel ?
    return rsrs
def fast_enumerate(sampler,rsrs,mdl,num_sols):
    """
    Enumerate the discrete sampling space to find the best combination
    """
    jt_filename = IMP.domino.get_example_path("six_particles.jt")
    jt = IMP.domino.JunctionTree()
    IMP.domino.read_junction_tree(jt_filename,jt)
    re = IMP.domino.RestraintEvaluator(sampler)
    d_opt = IMP.domino.DominoOptimizer(jt,mdl,re)
    #add restraints to be evaluates by the optimizer
    for r in rsrs:
        d_opt.add_restraint(r)
    #set the sampling space
    d_opt.set_sampling_space(sampler)
    #if you wish to see how the graph looks like remove this
    #line from comment
    #d_opt.show_restraint_graph()
    d_opt.set_number_of_solutions(num_sols)
    d_opt.optimize(num_sols)
    return d_opt

#### Example main code ####
IMP.set_log_level(IMP.SILENT)
mdl=IMP.Model()

#1. set up the particles
print "setting up particles"
ps=IMP.Particles()
for i in range(NUM_PARTICLES):
    p=IMP.Particle(mdl)
    IMP.core.XYZR.setup_particle(p,IMP.algebra.Sphere3D(IMP.algebra.Vector3D(0.,0.,0.),1.))
    ps.append(p)
    p.add_attribute(IMP.domino.node_name_key(),str(i))

#2. set up the discrete set of states
print "setting up a discrete set of states"
ub=IMP.algebra.Vector3D(-10.,-10.,-10.)
lb=IMP.algebra.Vector3D(10.,10.,10.)
bb=IMP.algebra.BoundingBox3D(ub,lb)
discrete_set=setup_discrete_sampling_space(ps,bb)

#3. add restraints (defining the scoring function)
print "setting up restraints"
restraints=setup_restraints(ps)

#4. set the sampler
print "setting up sampling"
sampler = IMP.domino.CartesianProductSampler(discrete_set,ps)
#alternative samplers can be PermutationSampler.

#5. optimize
print "optimizing"
num_sols=5
d_opt=fast_enumerate(sampler,restraints,mdl,num_sols)

#6. print results
rg = d_opt.get_graph()
for i in range(num_sols):
    sol = rg.get_opt_combination(i)
    print "solution number:",i," is: ", sol.key() ," with score:", sol.get_total_score()

Generated on Mon Mar 8 23:08:33 2010 for IMP by doxygen 1.5.8