rigid_body_excluded_volume.py

## \example domino/rigid_body_excluded_volume.py
# This example shows using two rigid bodies and doing excluded volume with
# them.

import IMP
import IMP.core
import IMP.algebra
import IMP.domino
import IMP.atom
import math

# create a rigid body per helix


def create_representation():
    m = IMP.kernel.Model()
    h0 = IMP.atom.read_pdb(IMP.domino.get_example_path(
        'helix_0.pdb'), m, IMP.atom.CAlphaPDBSelector())
    h1 = IMP.atom.read_pdb(IMP.domino.get_example_path(
        'helix_1.pdb'), m, IMP.atom.CAlphaPDBSelector())
    for h in [h0, h1]:
        IMP.atom.create_rigid_body(h)
    return (m, [h0, h1])


def create_excluded_volume(m, helices):
    # this is the interesting function:
    # it uses a KClosePairsPair score to generate the list of close atoms on
    # the fly
    all = []
    for h in helices:
        all.extend(IMP.atom.get_by_type(h, IMP.atom.ATOM_TYPE))
    lsc = IMP.container.ListSingletonContainer(all)
    evr = IMP.core.ExcludedVolumeRestraint(lsc, 1)
    evr.set_maximum_score(.01)
    return [evr]


# creating the discrete states for domino
def create_discrete_states(m, helices):
    pst = IMP.domino.ParticleStatesTable()
    rot00 = IMP.algebra.get_rotation_about_axis(
        IMP.algebra.Vector3D(0, 1, 0), math.pi / 2.0)
    trs = []
    zv = IMP.algebra.Vector3D(0, 0, 0)
    trs.append(IMP.algebra.ReferenceFrame3D(
        IMP.algebra.Transformation3D(rot00, zv)))
    for dx in range(0, 15):
        tr = IMP.algebra.Vector3D(1.0 * dx, 0, 0)
        trs.append(IMP.algebra.ReferenceFrame3D(
            IMP.algebra.Transformation3D(rot00, tr)))
    pstate = IMP.domino.RigidBodyStates(trs)
    for h in helices:
        pst.set_particle_states(
            IMP.core.RigidMember(h).get_rigid_body(), pstate)
    return pst

# setting up domino (and filters)


def create_sampler(m, rs, pst):
    s = IMP.domino.DominoSampler(m, pst)
    filters = []
    # do not allow particles with the same ParticleStates object
    # to have the same state index
    filters.append(IMP.domino.ExclusionSubsetFilterTable(pst))
    rc = IMP.domino.RestraintCache(pst)
    rc.add_restraints(rs)
    # filter states that score worse than the cutoffs in the Model
    filters.append(IMP.domino.RestraintScoreSubsetFilterTable(rc))
    states = IMP.domino.BranchAndBoundAssignmentsTable(pst, filters)
    s.set_assignments_table(states)
    s.set_subset_filter_tables(filters)
    return s


def display(m, helices, name):
    m.update()
    w = IMP.display.PymolWriter(name)
    for i, h in enumerate(helices):
        g = IMP.atom.HierarchyGeometry(h)
        g.set_color(IMP.display.get_display_color(i))
        w.add_geometry(g)

IMP.base.set_log_level(IMP.base.SILENT)
print("creating representation")
(m, helices) = create_representation()

print("creating score function")
rs = create_excluded_volume(m, helices)

print("creating discrete states")
pst = create_discrete_states(m, helices)

print("creating sampler")
s = create_sampler(m, rs, pst)
m.set_log_level(IMP.base.SILENT)
IMP.base.set_log_level(IMP.base.VERBOSE)
print("sampling")
cs = s.create_sample()

print("found ", cs.get_number_of_configurations(), "solutions")
score = []
for i in range(cs.get_number_of_configurations()):
    cs.load_configuration(i)
    ss = m.evaluate(False)
    score.append(ss)
    print("** solution number:", i, " is:", ss)
    display(m, helices, "sol_" + str(i) + ".pym")