custom_filter.py

## \example domino/custom_filter.py
# This example looks like the six particle optimization example except a
# filter is used instead of a restraint to remove the flip degree of
# freedom. The filter is written is Python, which makes for quick
# prototyping, but slower run times.

import IMP
import IMP.domino
import IMP.core
import IMP.container
import sys

IMP.setup_from_argv(sys.argv, "custom filter")

# set restraints


def create_scoring(m, ps):
    pairs = [[0, 1], [0, 2], [1, 3], [2, 3], [3, 4], [4, 5], [1, 5]]
    score = IMP.core.HarmonicDistancePairScore(1, 1)
    # the restraint will be broken apart during optimization
    pc = IMP.container.ListPairContainer(m,
                                         [(ps[p[0]], ps[p[1]]) for p in pairs],
                                         "Restrained pairs")
    pr = IMP.container.PairsRestraint(score, pc)
    pr.set_maximum_score(.01)
    d = IMP.core.DistanceToSingletonScore(IMP.core.HarmonicUpperBound(2, 1),
                                          IMP.algebra.Vector3D(2, 0, 0))
    return [pr]


def create_representation(m):
    ps = []
    for i in range(0, 6):
        p = m.add_particle("P%d" % i)
        IMP.core.XYZ.setup_particle(m, p, IMP.algebra.Vector3D(i, 0., 0.))
        ps.append(p)
    return ps


def create_discrete_states(m, ps):
    pst = IMP.domino.ParticleStatesTable()
    vs = [IMP.algebra.Vector3D(1, 0, 0),
          IMP.algebra.Vector3D(0, 1, 0),
          IMP.algebra.Vector3D(1, 1, 0),
          IMP.algebra.Vector3D(2, 1, 0),
          IMP.algebra.Vector3D(2, 0, 0)]
    vs = vs + [-v for v in vs]
    print(len(vs), "states for each particle")
    print(vs[1])
    states = IMP.domino.XYZStates(vs)
    # special case ps[0] to remove a sliding degree of freedom
    for p in ps[1:]:
        pst.set_particle_states(m.get_particle(p), states)
    return pst

# force particle p to be in state s


class MyFilterTable(IMP.domino.SubsetFilterTable):

    class MyFilter(IMP.domino.SubsetFilter):

        def __init__(self, pos, value):
            # print "Filtering with", pos, value
            IMP.domino.SubsetFilter.__init__(
                self, "MF" + str(pos) + " " + str(value))
            self.pos = pos
            self.value = value

        def get_next_state(self, pos, s):
            # suggest that the sampler try the correct state
            # this method is only called if the filter failed, so pos must be
            # self.pos
            if s[self.pos] > self.value:
                # a very large number
                return 2 ** 29
            else:
                return self.value

        def get_is_ok(self, state):
            # it is only OK if it has the required state
            ret = state[self.pos] == self.value
            return ret

    def get_strength(self, s, excluded):
        # return the maximum value since it dictates the state for the particle
        return 1

    def __init__(self, p, s):
        # set the name of the object to something vaguely useful
        IMP.domino.SubsetFilterTable.__init__(
            self, "MFT" + p.get_name() + " at " + str(s))
        self.p = p
        self.s = s

    def get_subset_filter(self, subset, excluded):
        # create a filter if self.p is in subset but not in excluded
        if self.p in subset and self.p not in sum([list(x) for x in excluded], []):
            # pass the position of self.p and the value that it must have
            return self.MyFilter(list(subset).index(self.p), self.s)
        else:
            return None


def create_sampler(m, ps, rs, pst):
    s = IMP.domino.DominoSampler(m, pst)
    s.set_restraints(rs)
    # s.set_log_level(IMP.VERBOSE)
    # the following lines recreate the defaults and so are optional
    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))
    filters[-1].set_log_level(IMP.SILENT)
    mf = MyFilterTable(m.get_particle(ps[1]), 0)
    # try with and without this line
    filters.append(mf)
    states = IMP.domino.BranchAndBoundAssignmentsTable(pst, filters)
    # states.set_log_level(IMP.SILENT);
    s.set_assignments_table(states)
    s.set_subset_filter_tables(filters)
    return s

IMP.set_log_level(IMP.TERSE)
m = IMP.Model()
m.set_log_level(IMP.SILENT)

print("creating representation")
ps = create_representation(m)
print("creating discrete states")
pst = create_discrete_states(m, ps)
print("creating score function")
rs = create_scoring(m, ps)

print("creating sampler")
s = create_sampler(m, ps, rs, pst)

# s.set_log_level(IMP.SILENT)
print("sampling")
cs = s.create_sample()

print("found ", cs.get_number_of_configurations(), "solutions")
sf = IMP.core.RestraintsScoringFunction(rs)
for i in range(cs.get_number_of_configurations()):
    cs.load_configuration(i)
    print("solution number:", i, " is:", sf.evaluate(False))
    for p in ps:
        print(IMP.core.XYZ(m, p).get_x(), IMP.core.XYZ(m, p).get_y())