IMP logo
IMP Reference Guide  2.14.0
The Integrative Modeling Platform
pmi/restraints/__init__.py
1 """@namespace IMP.pmi.restraints
2  Classes to handle different kinds of restraints.
3 
4 PMI restraints generally wrap IMP restraints. Typical features in PMI restraints are:
5  - Easy setup: for example, you can usually create one with a PMI [Molecule](@ref IMP::pmi::topology::Molecule) or a slice from one.
6  - Fast setup from data files. For example you can set up the [CrossLinkingMassSpectrometryRestraint](@ref IMP::pmi::restraints::crosslinking::CrossLinkingMassSpectrometryRestraint) by reading in a cross-link file into a [database](@ref IMP::pmi::io::crosslink::CrossLinkDataBase).
7  - Useful output: reporting functions which are put into log files when running [ReplicaExchange](@ref IMP::pmi::macros::ReplicaExchange0).
8 """
9 
10 import IMP
11 import IMP.pmi
12 import IMP.pmi.tools
13 
14 
15 class RestraintBase(object):
16  _include_in_rmf = False
17 
18  """Base class for PMI restraints, which wrap `IMP.Restraint`(s)."""
19 
20  def __init__(self, m, name=None, label=None, weight=1.,
21  restraint_set_class=IMP.RestraintSet):
22  """Constructor.
23  @param m The model object
24  @param name The name of the primary restraint set that is wrapped.
25  This is used for outputs and particle/restraint names
26  and should be set by the child class.
27  @param label A unique label to be used in outputs and
28  particle/restraint names.
29  @param weight The weight to apply to all internal restraints.
30  @param restraint_set_class The class to use for the restraint set
31  """
32  self.model = m
33  self.restraint_sets = []
34  self._label_is_set = False
35  self.weight = weight
36  self._label = None
37  self._label_suffix = ""
38  self.set_label(label)
39 
40  if not name:
41  self.name = self.__class__.__name__
42  else:
43  self.name = str(name)
44 
45  self.rs = self._create_restraint_set(name=None, cls=restraint_set_class)
46 
47  def set_label(self, label):
48  """Set the unique label used in outputs and particle/restraint names.
49  @param label Label
50  """
51  if self._label_is_set:
52  raise ValueError("Label has already been set.")
53  if not label:
54  self._label = ""
55  self._label_suffix = ""
56  else:
57  self._label = str(label)
58  self._label_suffix = "_" + self._label
59  self._label_is_set = True
60 
61  @property
62  def label(self):
63  return self._label
64 
65  def set_weight(self, weight):
66  """Set the weight to apply to all internal restraints.
67  @param weight Weight
68  """
69  self.weight = weight
70  for rs in self.restraint_sets:
71  rs.set_weight(self.weight)
72 
73  def add_to_model(self):
74  """Add the restraint to the model."""
75  self._label_is_set = True
76  for rs in self.restraint_sets:
78  add_to_rmf=self._include_in_rmf)
79 
80  def evaluate(self):
81  """Evaluate the score of the restraint."""
82  self._label_is_set = True
83  return self.weight * self.rs.unprotected_evaluate(None)
84 
85  def get_restraint_set(self):
86  """Get the primary restraint set."""
87  self._label_is_set = True
88  return self.rs
89 
90  def get_restraint(self):
91  """Get the primary restraint set. Identical to `get_restraint_set`."""
92  return self.get_restraint_set()
93 
94  def get_restraint_for_rmf(self):
95  """Get the restraint for visualization in an RMF file."""
96  self._label_is_set = True
97  return self.rs
98 
99  def get_particles_to_sample(self):
100  """Get any created particles which should be sampled."""
101  self._label_is_set = True
102  return {}
103 
104  def get_output(self):
105  """Get outputs to write to stat files."""
106  output = {}
107  score = self.evaluate()
108  output["_TotalScore"] = str(score)
109 
110  suffix = "_Score" + self._label_suffix
111  for rs in self.restraint_sets:
112  out_name = rs.get_name() + suffix
113  output[out_name] = str(
114  self.weight * rs.unprotected_evaluate(None))
115  return output
116 
117  def _create_restraint_set(self, name=None, cls=IMP.RestraintSet):
118  """Create ``IMP.RestraintSet``."""
119  if not name:
120  name = self.name
121  else:
122  name = self.name + "_" + str(name)
123  rs = cls(self.model, name)
124  rs.set_weight(self.weight)
125  self.restraint_sets.append(rs)
126  rs.set_was_used(True)
127  return rs
128 
129 
130 class _RestraintNuisanceMixin(object):
131 
132  """Mix-in to add nuisance particle creation functionality to restraint.
133 
134  This class must only be inherited if also inheriting
135  IMP.pmi.restraints.RestraintBase.
136  """
137 
138  def __init__(self, *args, **kwargs):
139  super(_RestraintNuisanceMixin, self).__init__(*args, **kwargs)
140  self.sampled_nuisances = {}
141  self.nuisances = {}
142 
143  def _create_nuisance(self, init_val, min_val, max_val, max_trans, name,
144  is_sampled=False):
145  """Create nuisance particle.
146  @param init_val Initial value of nuisance
147  @param min_val Minimum value of nuisance
148  @param max_val Maximum value of nuisance
149  @param max_trans Maximum move to apply to nuisance
150  @param name Name of particle
151  @param is_sampled Nuisance is a sampled particle
152  \see IMP.pmi.tools.SetupNuisance
153  """
154  nuis = IMP.pmi.tools.SetupNuisance(
155  self.model, init_val, min_val, max_val,
156  isoptimized=is_sampled).get_particle()
157  nuis_name = self.name + "_" + name
158  nuis.set_name(nuis_name)
159  self.nuisances[nuis_name] = nuis
160  if is_sampled:
161  self.sampled_nuisances[nuis_name] = (nuis, max_trans)
162  return nuis
163 
164  def get_particles_to_sample(self):
165  """Get any created particles which should be sampled."""
166  ps = super(_RestraintNuisanceMixin, self).get_particles_to_sample()
167  for name, (nuis, max_trans) in self.sampled_nuisances.items():
168  ps["Nuisances_" + name + self._label_suffix] = ([nuis], max_trans)
169  return ps
170 
171  def get_output(self):
172  """Get outputs to write to stat files."""
173  output = super(_RestraintNuisanceMixin, self).get_output()
174  for nuis_name, nuis in self.nuisances.items():
175  output[nuis_name + self._label_suffix] = str(nuis.get_scale())
176  return output
177 
178 
179 class _NuisancesBase(object):
180 
181  """This base class is used to provide nuisance setup and interface
182  for the ISD cross-link restraints"""
183 
184  sigma_dictionary = {}
185  psi_dictionary = {}
186 
187  def create_length(self):
188  """Create a nuisance on the length of the cross-link."""
189  lengthinit = 10.0
190  self.lengthissampled = True
191  lengthminnuis = 0.0000001
192  lengthmaxnuis = 1000.0
193  lengthmin = 6.0
194  lengthmax = 30.0
195  lengthtrans = 0.2
196  length = IMP.pmi.tools.SetupNuisance(self.m, lengthinit,
197  lengthminnuis, lengthmaxnuis,
198  self.lengthissampled
199  ).get_particle()
200  self.rslen.add_restraint(
202  self.m,
203  length,
204  1000000000.0,
205  lengthmax,
206  lengthmin))
207 
208  def create_sigma(self, resolution):
209  """Create a nuisance on the structural uncertainty."""
210  if isinstance(resolution, str):
211  sigmainit = 2.0
212  else:
213  sigmainit = resolution + 2.0
214  self.sigmaissampled = True
215  sigmaminnuis = 0.0000001
216  sigmamaxnuis = 1000.0
217  sigmamin = 0.01
218  sigmamax = 100.0
219  sigmatrans = 0.5
220  sigma = IMP.pmi.tools.SetupNuisance(self.m, sigmainit, sigmaminnuis,
221  sigmamaxnuis, self.sigmaissampled
222  ).get_particle()
223  self.sigma_dictionary[resolution] = (
224  sigma,
225  sigmatrans,
226  self.sigmaissampled)
227  self.rssig.add_restraint(
229  self.m,
230  sigma,
231  1000000000.0,
232  sigmamax,
233  sigmamin))
234  # self.rssig.add_restraint(IMP.isd.JeffreysRestraint(self.sigma))
235 
236  def get_sigma(self, resolution):
237  """Get the nuisance on structural uncertainty."""
238  if resolution not in self.sigma_dictionary:
239  self.create_sigma(resolution)
240  return self.sigma_dictionary[resolution]
241 
242  def create_psi(self, value):
243  """Create a nuisance on the inconsistency."""
244  if isinstance(value, str):
245  psiinit = 0.5
246  else:
247  psiinit = value
248  self.psiissampled = True
249  psiminnuis = 0.0000001
250  psimaxnuis = 0.4999999
251  psimin = 0.01
252  psimax = 0.49
253  psitrans = 0.1
254  psi = IMP.pmi.tools.SetupNuisance(self.m, psiinit,
255  psiminnuis, psimaxnuis,
256  self.psiissampled).get_particle()
257  self.psi_dictionary[value] = (
258  psi,
259  psitrans,
260  self.psiissampled)
261  self.rspsi.add_restraint(
263  self.m,
264  psi,
265  1000000000.0,
266  psimax,
267  psimin))
268  self.rspsi.add_restraint(IMP.isd.JeffreysRestraint(self.m, psi))
269 
270  def get_psi(self, value):
271  """Get the nuisance on the inconsistency."""
272  if value not in self.psi_dictionary:
273  self.create_psi(value)
274  return self.psi_dictionary[value]
Miscellaneous utilities.
Definition: tools.py:1
Uniform distribution with harmonic boundaries.
Definition: UniformPrior.h:20
Object used to hold a set of restraints.
Definition: RestraintSet.h:36
def add_restraint_to_model
Add a PMI restraint to the model.
Definition: tools.py:84
Python classes to represent, score, sample and analyze models.