restraint_cache.py

## \example domino/restraint_cache.py
# Caching restraint scores so that restraints are not evaluated
# repeatedly for the same configuration is an important part of
# domino. Without caching, sub assignments that are shared between
# subsets will be rescored. The IMP::domino::RestraintCache provides
# a centralized place for this. To use it, one creates one and then
# adds the restraints you want to use for filtering and scoring to
# it. You can then pass the cache to the
# IMP::domino::RestraintScoreFilterTable and it will filter based on
# those restraints. You can also extract scores from the table
# directly, using it to manage the loading of particle states.

from __future__ import print_function
import IMP.domino
import IMP.algebra
import IMP.container
import IMP.atom
import IMP.rmf
import RMF
import sys

IMP.setup_from_argv(sys.argv, "restraint cache")

resolution = .5

with IMP.SetLogState(IMP.SILENT):
    m = IMP.Model()

    # create some particles and a restraint that scores based on the
    # length of the chain
    ps = [IMP.core.XYZR.setup_particle(IMP.Particle(m))
          for i in range(0, int(2.0 / resolution))]
    eps = [IMP.core.XYZR.setup_particle(IMP.Particle(m))
           for i in range(0, 2)]
    eps[0].set_coordinates(IMP.algebra.Vector3D(-1, 0, 0))
    eps[1].set_coordinates(IMP.algebra.Vector3D(1, 0, 0))
    for p in ps:
        p.set_radius(resolution / 2.0)
    for p in eps:
        p.set_radius(.4 * resolution)
    order = [eps[0]] + ps + [eps[1]]
    # create a hierarchy with the particles to aid writing them to RMF
    h = IMP.atom.Hierarchy.setup_particle(IMP.Particle(m))
    for i, p in enumerate(order):
        IMP.atom.Mass.setup_particle(p, 1)
        h.add_child(IMP.atom.Hierarchy.setup_particle(p))
        color = IMP.display.get_jet_color(float(i) / (len(order) - 1))
        IMP.display.Colored.setup_particle(p, color)

    r = IMP.container.PairsRestraint(
        IMP.core.HarmonicUpperBoundSphereDistancePairScore(0, 1),
        IMP.container.ConsecutivePairContainer(m, order))
    # allow a maximum distance of 2
    r.set_maximum_score(0.01)

    # dumb way to generate points on a gride
    g = IMP.algebra.DenseFloatGrid3D(
        resolution, IMP.algebra.get_unit_bounding_box_3d())
    # create some random sites
    vs = [g.get_center(i) for i in g.get_all_indexes()]
    print("States are", vs)
    pst = IMP.domino.ParticleStatesTable()
    states = IMP.domino.XYZStates(vs)
    for p in ps:
        pst.set_particle_states(p, states)

    # now create a restraint cache
    # cache the most recently used one million scores
    rc = IMP.domino.RestraintCache(pst, 1000000)
    r.set_log_level(IMP.SILENT)
    rc.add_restraints([r])

    s = IMP.domino.DominoSampler(m, pst)
    s.set_restraints([r])
    s.set_check_level(IMP.NONE)
    ef = IMP.domino.ExclusionSubsetFilterTable(pst)
    # pass the cache to the restraint score based filter
    # it will use all the restraints in the cache
    rssft = IMP.domino.RestraintScoreSubsetFilterTable(rc)
    s.set_subset_filter_tables([ef, rssft])

    # create a subset with all the particles
    # Subsets keep the particles in sorted order and so are different than
    # a simple list of particles
    alls = IMP.domino.Subset(ps)

    # get all the assignments that fit
    sts = s.get_sample_assignments(alls)

    # create a temporary file to write things to
    rmf = RMF.create_rmf_file(IMP.create_temporary_file_name("cache", ".rmf"))
    print("saving configurations to", rmf.get_name())

    IMP.rmf.add_hierarchy(rmf, h)

    # the restraint cache processes the restraints to make them more efficient
    # for use with domino. So we want to get back the processed restraints
    # before looking at them further
    domino_restraints = rc.get_restraints()

    for r in domino_restraints:
        IMP.show_restraint_hierarchy(r)

    IMP.rmf.add_restraints(rmf, domino_restraints)
    # for each assignment load it, and add the configuration to an rmf file
    print("found assignments", sts)
    # we don't care about messages during saving
    for i, s in enumerate(sts):
        IMP.domino.load_particle_states(alls, s, pst)
        # go through the restraints and get the score
        # here, we only care about its side effect of setting the last score
        for r in domino_restraints:
            rc.load_last_score(r, alls, s)
        # save the configuration and scores to the rmf
        IMP.rmf.save_frame(rmf)
    print("done")