Creating a PMI Restraint

Implementing a Restraint in PMI

PMI is able to utilize IMP Restraints by wrapping them with python classes so they can be passed to sampling macros.

All PMI restraints must inherit IMP.pmi.restraints.RestraintBase. A simple restraint can be created like

class DistanceRestraint(IMP.pmi.restraints.RestraintBase):

    """Simple harmonic distance restraint between two XYZ particles."""

    def __init__(self, p1, p2, d, k, name="DistanceRestraint", label=None, weight=1.):
        """Setup distance restraint.
        @param p1 First XYZ particle
        @param p2 Second XYZ particle
        @param d Expected distance
        @param k Stiffness constant
        @param name The name of the primary restraint set that is wrapped.
                    This is used for outputs and particle/restraint names
                    and should be set by the child class.
        @param label A unique label to be used in outputs and
                     particle/restraint names.
        @param weight The weight to apply to all internal restraints.
        """
        model = p1.get_model()
        super(DistanceRestraint, self).__init__(model, name=name, label=label,
                                                weight=weight)
        f = IMP.core.Harmonic(d, k)
        s = IMP.core.DistancePairScore(f)
        r = IMP.core.PairRestraint(self.model, s, (p1, p2))
        self.rs.add_restraint(r)

While name is the general name of the restraint, label is used to differentiate restraints of the same name in outputs. Note that the two main ingredients of PMI restraint creation are instantiation of the base class and adding a restraint to self.rs. self.rs should contain all wrapped IMP Restraints that contribute to the scoring function. Another restraint set that is written to outputs can be created with self._create_restraint_set, but this set will not contribute to the total score. The below example demonstrates this.

It is not uncommon for a PMI Restraint to create an IMP.isd.Nuisance particle (e.g. for sigma in a Gaussian distribution). The value of the nuisance must be saved to outputs, and IMP.pmi.dof.DOF must be passed the nuisance for sampling. All of this functionality may be added to a restraint by also inheriting the IMP.pmi.restraints._RestraintNuisanceMixin:

class GaussianRestraint(IMP.pmi.restraints._RestraintNuisanceMixin,
                        IMP.pmi.restraints.RestraintBase):

    """Restraint value of nuisance particle with a Gaussian."""

    def __init__(self, p, mean_val, name="GaussianRestraint", label=None,
                 weight=1.):
        """Setup Gaussian restraint.
        @param p Nuisance particle
        @param mean_val Mean of the Gaussian distribution
        @param name The name of the primary restraint set that is wrapped.
                    This is used for outputs and particle/restraint names
                    and should be set by the child class.
        @param label A unique label to be used in outputs and
                     particle/restraint names.
        @param weight The weight to apply to all internal restraints.
        """
        model = p.get_model()
        super(GaussianRestraint, self).__init__(model, name=name, label=label,
                                                weight=weight)

        # Create tracked Nuisances
        self.mu = self._create_nuisance(mean_val, None, None, None,
                                        "Mu", is_sampled=False)
        self.sigma = self._create_nuisance(1., 0., None, .1, "Sigma",
                                           is_sampled=True)

        # Create wrapped restraints
        r = IMP.isd.GaussianRestraint(p, self.mu, self.sigma)
        self.rs.add_restraint(r)
        r = IMP.isd.JeffreysRestraint(self.model, self.sigma)
        self.rs.add_restraint(r)

        # Store Jeffreys prior in outputs
        self.rs_jeffreys = self._create_restraint_set(
            "Sigma_JeffreysPrior")
        self.rs_jeffreys.add_restraint(r)

Note that IMP.pmi.restraints._RestraintNuisanceMixin is listed first in inheritance; this is necessary for proper set-up. In the above example, the total score consists of the Gaussian likelihood and Jeffreys prior on the sigma particle, but the Jeffreys score will also appear in the outputs, as will the values of the Mu and Sigma particles. When the restraint is passed to IMP.pmi.dof.DOF.get_nuisances_from_restraint, sigma will be sampled, while mu will remain constant.