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