Squashed 'modules/pmi/' changes from 67640f9..d2651da

d2651da Don't run slow examples in debug mode.
3e770ee Fix print output in Python 2.
ba22e63 Remove unused imports and add docstrings
485f81b Use new restraint base, relates salilab/pmi#156
f7fc45f Add tests for parameter restraints
bb9c499 Add mix-in for nuisance setup in restraints, Relates salilab/pmi#156
dd500d5 Allow weight to be modified at any time
af9a472 updates to restraint base; new handling of weights, labels, restraint name; helper method for restraint creation
514acda comply with PEP8, fixed docstrings, reordered methods
b419fc6 use new restraint base
103a943 use new restraint base
9af9741 use new restraint base
ffc1013 added helper method for getting label suffix
7f38ff0 fix output key lookup
c456568 comply with PEP8, format and add docstrings
99165d2 use new restraint base
5ff5c8b fixed reference to undeclared variable
9e14ea7 comply to PEP8, format and add docstrings

git-subtree-dir: modules/pmi
git-subtree-split: d2651dac6b8a91b6541839818ed29ad9a3f45571

modules/pmi/examples/atomistic.py

@@ -3,6 +3,7 @@
 with a secondary structure elastic network to speed things up.
 """
 
+from __future__ import print_function
 import IMP
 import RMF
 import IMP.atom
@@ -15,6 +16,10 @@
 import sys
 
 IMP.setup_from_argv(sys.argv, "Simulation of an atomic system")
+if IMP.get_is_quick_test():
+    print("This example is too slow to test in debug mode - run without")
+    print("internal tests enabled, or without the --run-quick-test flag")
+    sys.exit(0)
 
 # Setup System and add a State
 mdl = IMP.Model()

modules/pmi/examples/automatic.py

@@ -2,6 +2,7 @@
 """This script shows how to use the BuildSystem macro to construct large systems with minimal code
 """
 
+from __future__ import print_function
 import IMP
 import RMF
 import IMP.atom
@@ -16,6 +17,10 @@
 import sys
 
 IMP.setup_from_argv(sys.argv, "Automatic setup of a large system")
+if IMP.get_is_quick_test():
+    print("This example is too slow to test in debug mode - run without")
+    print("internal tests enabled, or without the --run-quick-test flag")
+    sys.exit(0)
 
 # This is the topology table format.
 #  It allows you to create many components in a simple way

modules/pmi/examples/multiscale.py

@@ -2,6 +2,7 @@
 """This script shows how to represent a system at multiple scales and do basic sampling.
 """
 
+from __future__ import print_function
 import IMP
 import RMF
 import IMP.atom
@@ -15,6 +16,10 @@
 import sys
 
 IMP.setup_from_argv(sys.argv, "Representation at multiple scales")
+if IMP.get_is_quick_test():
+    print("This example is too slow to test in debug mode - run without")
+    print("internal tests enabled, or without the --run-quick-test flag")
+    sys.exit(0)
 
 ###################### SYSTEM SETUP #####################
 # Read sequences etc

modules/pmi/pyext/src/restraints/__init__.py

@@ -16,42 +16,71 @@ class RestraintBase(object):
 
     """Base class for PMI restraints, which wrap `IMP.Restraint`(s)."""
 
-    def __init__(self, m, rname=None):
+    def __init__(self, m, name=None, label=None, weight=1.):
         """Constructor.
         @param m The model object
-        @param rname The name of the primary restraint set that is wrapped.
+        @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.
         """
         self.m = m
-        self.label = None
-        self.weight = 1.
-        if rname is None:
-            rname = self.__class__.__name__
-        self.rs = IMP.RestraintSet(self.m, rname)
-        self.restraint_sets = [self.rs]
-
-    def set_weight(self, weight):
-        """Set the weight of the restraint.
-        @param weight Restraint weight
-        """
+        self.restraint_sets = []
+        self._label_is_set = False
         self.weight = weight
-        self.rs.set_weight(weight)
+        self._label = None
+        self._label_suffix = ""
+        self.set_label(label)
+
+        if not name:
+            self.name = self.__class__.__name__
+        else:
+            self.name = str(name)
+
+        self.rs = self._create_restraint_set(name=None)
 
     def set_label(self, label):
-        """Set the label used in outputs.
+        """Set the unique label used in outputs and particle/restraint names.
         @param label Label
         """
-        self.label = label
+        if self._label_is_set:
+            raise ValueError("Label has already been set.")
+        if not label:
+            self._label = ""
+            self._label_suffix = ""
+        else:
+            self._label = str(label)
+            self._label_suffix = "_" + self._label
+            self._label_is_set = True
+
+    @property
+    def label(self):
+        return self._label
+
+    def set_weight(self, weight):
+        """Set the weight to apply to all internal restraints.
+        @param weight Weight
+        """
+        self.weight = weight
+        for rs in self.restraint_sets:
+            rs.set_weight(self.weight)
 
     def add_to_model(self):
         """Add the restraint to the model."""
-        IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)
+        self._label_is_set = True
+        for rs in self.restraint_sets:
+            IMP.pmi.tools.add_restraint_to_model(self.m, rs)
 
     def evaluate(self):
         """Evaluate the score of the restraint."""
+        self._label_is_set = True
         return self.weight * self.rs.unprotected_evaluate(None)
 
     def get_restraint_set(self):
         """Get the primary restraint set."""
+        self._label_is_set = True
         return self.rs
 
     def get_restraint(self):
@@ -60,39 +89,98 @@ def get_restraint(self):
 
     def get_restraint_for_rmf(self):
         """Get the restraint for visualization in an RMF file."""
+        self._label_is_set = True
         return self.rs
 
     def get_particles_to_sample(self):
         """Get any created particles which should be sampled."""
+        self._label_is_set = True
         return {}
 
     def get_output(self):
         """Get outputs to write to stat files."""
-        self.m.update()
         output = {}
+        self.m.update()
         score = self.evaluate()
         output["_TotalScore"] = str(score)
 
-        suffix = "_Score"
-        if self.label is not None:
-            suffix += "_" + str(self.label)
-
+        suffix = "_Score" + self._label_suffix
         for rs in self.restraint_sets:
             out_name = rs.get_name() + suffix
             output[out_name] = str(
                 self.weight * rs.unprotected_evaluate(None))
+        return output
+
+    def _create_restraint_set(self, name=None):
+        """Create ``IMP.RestraintSet``."""
+        if not name:
+            name = self.name
+        else:
+            name = self.name + "_" + str(name)
+        rs = IMP.RestraintSet(self.m, name)
+        rs.set_weight(self.weight)
+        self.restraint_sets.append(rs)
+        return rs
+
+
+class _RestraintNuisanceMixin(object):
 
+    """Mix-in to add nuisance particle creation functionality to restraint.
+
+    This class must only be inherited if also inheriting
+    IMP.pmi.restraints.RestraintBase.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(_RestraintNuisanceMixin, self).__init__(*args, **kwargs)
+        self.sampled_nuisances = {}
+        self.nuisances = {}
+
+    def _create_nuisance(self, init_val, min_val, max_val, max_trans, name,
+                         is_sampled=False):
+        """Create nuisance particle.
+        @param init_val Initial value of nuisance
+        @param min_val Minimum value of nuisance
+        @param max_val Maximum value of nuisance
+        @param max_trans Maximum move to apply to nuisance
+        @param is_sampled Nuisance is a sampled particle
+        \see IMP.pmi.tools.SetupNuisance
+        """
+        nuis = IMP.pmi.tools.SetupNuisance(
+            self.m, init_val, min_val, max_val,
+            isoptimized=is_sampled).get_particle()
+        nuis_name = self.name + "_" + name
+        nuis.set_name(nuis_name)
+        self.nuisances[nuis_name] = nuis
+        if is_sampled:
+            self.sampled_nuisances[nuis_name] = (nuis, max_trans)
+        return nuis
+
+    def get_particles_to_sample(self):
+        """Get any created particles which should be sampled."""
+        ps = super(_RestraintNuisanceMixin, self).get_particles_to_sample()
+        for name, (nuis, max_trans) in self.sampled_nuisances.iteritems():
+            ps["Nuisances_" + name + self._label_suffix] = ([nuis], max_trans)
+        return ps
+
+    def get_output(self):
+        """Get outputs to write to stat files."""
+        output = super(_RestraintNuisanceMixin, self).get_output()
+        for nuis_name, nuis in self.nuisances.iteritems():
+            output[nuis_name + self._label_suffix] = str(nuis.get_scale())
         return output
 
 
 class _NuisancesBase(object):
-    ''' This base class is used to provide nuisance setup and interface
-    for the ISD cross-link restraints '''
-    sigma_dictionary={}
-    psi_dictionary={}
+
+    """This base class is used to provide nuisance setup and interface
+    for the ISD cross-link restraints"""
+
+    sigma_dictionary = {}
+    psi_dictionary = {}
 
     def create_length(self):
-        ''' a nuisance on the length of the cross-link '''
+        """Create a nuisance on the length of the cross-link."""
         lengthinit = 10.0
         self.lengthissampled = True
         lengthminnuis = 0.0000001
@@ -101,9 +189,9 @@ def create_length(self):
         lengthmax = 30.0
         lengthtrans = 0.2
         length = IMP.pmi.tools.SetupNuisance(self.m, lengthinit,
-                                                    lengthminnuis,
-                                                    lengthmaxnuis,
-                                                    lengthissampled).get_particle()
+                                             lengthminnuis, lengthmaxnuis,
+                                             self.lengthissampled
+                                             ).get_particle()
         self.rslen.add_restraint(
             IMP.isd.UniformPrior(
                 self.m,
@@ -113,8 +201,8 @@ def create_length(self):
                 lengthmin))
 
     def create_sigma(self, resolution):
-        ''' a nuisance on the structural uncertainty '''
-        if isinstance(resolution,str):
+        """Create a nuisance on the structural uncertainty."""
+        if isinstance(resolution, str):
             sigmainit = 2.0
         else:
             sigmainit = resolution + 2.0
@@ -124,8 +212,9 @@ def create_sigma(self, resolution):
         sigmamin = 0.01
         sigmamax = 100.0
         sigmatrans = 0.5
-        sigma = IMP.pmi.tools.SetupNuisance(self.m, sigmainit,
-                                                 sigmaminnuis, sigmamaxnuis, self.sigmaissampled).get_particle()
+        sigma = IMP.pmi.tools.SetupNuisance(self.m, sigmainit, sigmaminnuis,
+                                            sigmamaxnuis, self.sigmaissampled
+                                            ).get_particle()
         self.sigma_dictionary[resolution] = (
             sigma,
             sigmatrans,
@@ -140,13 +229,14 @@ def create_sigma(self, resolution):
         # self.rssig.add_restraint(IMP.isd.JeffreysRestraint(self.sigma))
 
     def get_sigma(self, resolution):
-        if not resolution in self.sigma_dictionary:
+        """Get the nuisance on structural uncertainty."""
+        if resolution not in self.sigma_dictionary:
             self.create_sigma(resolution)
         return self.sigma_dictionary[resolution]
 
     def create_psi(self, value):
-        ''' a nuisance on the inconsistency '''
-        if isinstance(value,str):
+        """Create a nuisance on the inconsistency."""
+        if isinstance(value, str):
             psiinit = 0.5
         else:
             psiinit = value
@@ -157,8 +247,8 @@ def create_psi(self, value):
         psimax = 0.49
         psitrans = 0.1
         psi = IMP.pmi.tools.SetupNuisance(self.m, psiinit,
-                                               psiminnuis, psimaxnuis,
-                                               self.psiissampled).get_particle()
+                                          psiminnuis, psimaxnuis,
+                                          self.psiissampled).get_particle()
         self.psi_dictionary[value] = (
             psi,
             psitrans,
@@ -173,6 +263,7 @@ def create_psi(self, value):
         self.rspsi.add_restraint(IMP.isd.JeffreysRestraint(self.m, psi))
 
     def get_psi(self, value):
-        if not value in self.psi_dictionary:
+        """Get the nuisance on the inconsistency."""
+        if value not in self.psi_dictionary:
             self.create_psi(value)
         return self.psi_dictionary[value]