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