IMP logo
IMP Reference Guide  2.21.0
The Integrative Modeling Platform
restraints/saxs.py
1 """@namespace IMP.pmi.restraints.saxs
2 Restraints for handling small angle x-ray (SAXS) data.
3 """
4 
5 from __future__ import print_function
6 import math
7 import IMP
8 import IMP.core
9 import IMP.algebra
10 import IMP.atom
11 import IMP.container
12 import IMP.pmi.tools
13 import IMP.pmi.restraints
14 import IMP.saxs
15 import IMP.isd
16 try:
17  import IMP.isd2
18 except ImportError:
19  pass
20 import warnings
21 
22 
23 class SAXSRestraint(IMP.pmi.restraints.RestraintBase):
24  """Basic SAXS restraint."""
25 
26  _include_in_rmf = True
27 
28  def __init__(self, input_objects, saxs_datafile, weight=1.0,
29  ff_type=IMP.saxs.HEAVY_ATOMS, label=None, maxq="standard"):
30  """Builds the restraint.
31  @param input_objects A list of hierarchies or PMI objects that the
32  SAXS restraint will be applied to. This hierarchy must be
33  atomic unless ff_type=IMP.saxs.RESIDUES is used.
34  @param saxs_datafile the SAXS .dat file.
35  @param weight Restraint score coefficient
36  @param ff_type the form factor to use, of the following types:
37  - IMP.saxs.HEAVY_ATOMS: use form factors with implicit
38  hydrogens
39  - IMP.saxs.ALL_ATOMS: use individual form factors for all
40  atoms. Does not build missing hydrogens.
41  - IMP.saxs.CA_ATOMS: use residue based form factors
42  centered at CA atoms
43  - IMP.saxs.RESIDUES: use residue based form factors
44  using per-residue beads
45  @param label Label for the restraint in outputs
46  @param maxq Maximum q value that the restraint will be evaluated at.
47  If set to 'standard' (the default), the following values will
48  be used (these values were eyeballed by comparing ALL_ATOM
49  calculated SAXS profiles to those calculated with the reduced
50  representations, so could be improved):
51  - For ff_type = ALL_ATOMS: 0.5
52  - HEAVY_ATOMS: 0.4
53  - CA_ATOMS and RESIDUES: 0.15
54  """
55  # Get all hierarchies.
56  hiers = IMP.pmi.tools.input_adaptor(input_objects,
57  flatten=True)
58  model = list(hiers)[0].get_model()
59  super(SAXSRestraint, self).__init__(model, label=label, weight=weight)
60 
61  # Determine maxq to compare computed and experimental profiles
62  if maxq == "standard":
63  if ff_type == IMP.saxs.CA_ATOMS or ff_type == IMP.saxs.RESIDUES:
64  maxq = 0.15
65  elif ff_type == IMP.saxs.HEAVY_ATOMS:
66  maxq = 0.4
67  else:
68  maxq = 0.5
69  elif isinstance(maxq, float):
70  if maxq < 0.01 or maxq > 4.0:
71  raise Exception(
72  "SAXSRestraint: maxq must be set between 0.01 and 4.0")
73  if (ff_type == IMP.saxs.CA_ATOMS or ff_type == IMP.saxs.RESIDUES) \
74  and maxq > 0.15:
75  warnings.warn("SAXSRestraint: for residue-resolved form "
76  "factors, a maxq > 0.15 is not recommended!",
78  else:
79  raise Exception(
80  "SAXSRestraint: maxq must be set to a number between 0.01 "
81  "and 4.0")
82 
83  # Depending on the type of FF used, get the correct particles
84  # from the hierarchies list and create an IMP::saxs::Profile()
85  # at the appropriate maxq.
86  if ff_type == IMP.saxs.RESIDUES:
87  self.particles = IMP.atom.Selection(
88  hiers, resolution=1).get_selected_particles()
89  self.profile = IMP.saxs.Profile(saxs_datafile, False, maxq)
90  elif ff_type == IMP.saxs.CA_ATOMS:
91  self.particles = IMP.atom.Selection(
92  hiers, atom_type=IMP.atom.AT_CA).get_selected_particles()
93  self.profile = IMP.saxs.Profile(saxs_datafile, False, maxq)
94  elif ff_type == IMP.saxs.HEAVY_ATOMS or ff_type == IMP.saxs.ALL_ATOMS:
95  self.particles = IMP.atom.Selection(
96  hiers, resolution=0).get_selected_particles()
97  self.profile = IMP.saxs.Profile(saxs_datafile, False, maxq)
98  else:
99  raise Exception(
100  "SAXSRestraint: Must provide an IMP.saxs atom type: "
101  "RESIDUES, CA_ATOMS, HEAVY_ATOMS or ALL_ATOMS")
102 
103  if len(self.particles) == 0:
104  raise Exception("SAXSRestraint: There are no selected particles")
105 
106  self.restraint = IMP.saxs.Restraint(self.particles, self.profile,
107  ff_type)
108  self.rs.add_restraint(self.restraint)
109 
110 
111 class SAXSISDRestraint(IMP.pmi.restraints.RestraintBase):
112 
113  """Basic SAXS restraint using ISD."""
114 
115  def __init__(self, representation, profile, resolution=0, weight=1,
116  ff_type=IMP.saxs.HEAVY_ATOMS, label=None):
117 
118  if not hasattr(IMP, 'isd2'):
119  raise ImportError("Module isd2 not installed. "
120  "Cannot use SAXSISDRestraint")
121 
122  model = representation.prot.get_model()
123  super(SAXSISDRestraint, self).__init__(model, label=label,
124  weight=weight)
125 
126  self.taumaxtrans = 0.05
127  self.prof = IMP.saxs.Profile(profile)
128 
129  self.atoms = IMP.pmi.tools.select(
130  representation,
131  resolution=resolution)
132 
133  # gamma nuisance
134  self.gamma = IMP.pmi.tools.SetupNuisance(
135  self.model, 1., 0., None, False).get_particle()
136 
137  # sigma nuisance
138  self.sigma = IMP.pmi.tools.SetupNuisance(
139  self.model, 10.0, 0., None, False).get_particle()
140 
141  # tau nuisance, optimized
142  self.tau = IMP.pmi.tools.SetupNuisance(self.model, 1., 0., None, False,
143  ).get_particle()
144 
145  # c1 and c2, optimized
146  self.c1 = IMP.pmi.tools.SetupNuisance(self.model, 1.0, 0.95, 1.05,
147  True).get_particle()
148  self.c2 = IMP.pmi.tools.SetupNuisance(self.model, 0.0, -2., 4.,
149  True).get_particle()
150 
151  # weight, optimized
152  self.w = IMP.pmi.tools.SetupWeight(self.model).get_particle()
153  IMP.isd.Weight(self.w).set_weights_are_optimized(False)
154 
155  # take identity covariance matrix for the start
156  self.cov = [[1 if i == j else 0 for j in range(self.prof.size())]
157  for i in range(self.prof.size())]
158 
159  print("create saxs restraint")
160  self.saxs = IMP.isd2.SAXSRestraint(self.prof, self.sigma, self.tau,
161  self.gamma, self.w, self.c1,
162  self.c2)
163  self.saxs.add_scatterer(self.atoms, self.cov, ff_type)
164 
165  self.rs.add_restraint(self.saxs)
166 
167  # self.saxs_stuff={'nuis':(sigma,gamma),'cov':cov,
168  # 'exp':prof,'th':tmp}
169 
170  self.rs2 = self._create_restraint_set('Prior')
171  # jeffreys restraints for nuisances
172  j1 = IMP.isd.JeffreysRestraint(self.model, self.sigma)
173  self.rs2.add_restraint(j1)
174  j2 = IMP.isd.JeffreysRestraint(self.model, self.tau)
175  self.rs2.add_restraint(j2)
176  j3 = IMP.isd.JeffreysRestraint(self.model, self.gamma)
177  self.rs2.add_restraint(j3)
178 
179  pw = IMP.isd.Weight(self.w)
180  pw.set_weights(pw.get_unit_simplex().get_barycenter())
181  pw.set_weights_are_optimized(True)
182 
183  def optimize_sigma(self):
184  """Set sigma to the value that maximizes its conditional likelihood"""
185  self.model.update()
186  sigma2hat = self.saxs.get_sigmasq_scale_parameter() \
187  / (self.saxs.get_sigmasq_shape_parameter() + 1)
188  IMP.isd.Scale(self.sigma).set_scale(math.sqrt(sigma2hat))
189 
190  def optimize_gamma(self):
191  """Set gamma to the value that maximizes its conditional likelihood"""
192  self.model.update()
193  gammahat = math.exp(self.saxs.get_loggamma_variance_parameter() *
194  self.saxs.get_loggamma_jOg_parameter())
195  IMP.isd.Scale(self.gamma).set_scale(gammahat)
196 
197  def optimize_tau(self, ltaumin=-2, ltaumax=3, npoints=100):
198  values = []
199  self.model.update()
200  IMP.atom.write_pdb(self.atoms, 'tauvals.pdb')
201  fl = open('tauvals.txt', 'w')
202  for tauval in self._logspace(ltaumin, ltaumax, npoints):
203  IMP.isd.Scale(self.tau).set_scale(tauval)
204  try:
205  values.append((self.model.evaluate(False), tauval))
206  except: # noqa: E722
207  pass
208  fl.write('%G %G\n' % (values[-1][1], values[-1][0]))
209  values.sort()
210  ltcenter = math.log(values[0][1]) / math.log(10)
211  spacing = (ltaumax - ltaumin) / float(npoints)
212  values = []
213  for tauval in self._logspace(
214  ltcenter - 2 * spacing, ltcenter + 2 * spacing,
215  npoints):
216  IMP.isd.Scale(self.tau).set_scale(tauval)
217  values.append((self.model.evaluate(False), tauval))
218  fl.write('%G %G\n' % (values[-1][1], values[-1][0]))
219  values.sort()
220  IMP.isd.Scale(self.tau).set_scale(values[0][1])
221 
222  def get_gamma_value(self):
223  """Get value of gamma."""
224  return self.gamma.get_scale()
225 
226  def set_taumaxtrans(self, taumaxtrans):
227  self.taumaxtrans = taumaxtrans
228 
229  def draw_sigma(self):
230  """Draw 1/sigma2 from gamma distribution."""
231  self.model.update()
232  self.saxs.draw_sigma()
233 
234  def draw_gamma(self):
235  """Draw gamma from lognormal distribution."""
236  self.model.update()
237  self.saxs.draw_gamma()
238 
239  def update_covariance_matrix(self):
240  c1 = IMP.isd.Nuisance(self.c1).get_nuisance()
241  c2 = IMP.isd.Nuisance(self.c2).get_nuisance()
242  # tau = IMP.isd.Nuisance(self.tau).get_nuisance()
243  tau = 1.0
244  self.cov = IMP.isd2.compute_relative_covariance(self.atoms, c1, c2,
245  tau, self.prof)
246  # for i in xrange(len(self.cov)):
247  # for j in xrange(len(self.cov)):
248  # self.cov[i][j] = self.cov[i][j]/tau**2
249  self.saxs.set_cov(0, self.cov)
250 
251  def write_covariance_matrix(self, fname):
252  fl = open(fname, 'w')
253  for line in self.cov:
254  for i in line:
255  fl.write('%G ' % i)
256  fl.write('\n')
257 
258  def get_output(self):
259  output = super(SAXSISDRestraint, self).get_output()
260  suffix = self._get_label_suffix()
261  output["SAXSISDRestraint_Sigma" +
262  suffix] = str(self.sigma.get_scale())
263  output["SAXSISDRestraint_Tau" + suffix] = str(self.tau.get_scale())
264  output["SAXSISDRestraint_Gamma" +
265  suffix] = str(self.gamma.get_scale())
266  return output
267 
268  @staticmethod
269  def _logspace(a, b, num=100):
270  """Mimic numpy's logspace function"""
271  for i in range(num):
272  val = a + float(b - a) / float(num - 1) * i
273  yield 10 ** val
Add weights to a particle.
Definition: Weight.h:29
Various classes to hold sets of particles.
def optimize_gamma
Set gamma to the value that maximizes its conditional likelihood.
Calculate score based on fit to SAXS profile.
Miscellaneous utilities.
Definition: tools.py:1
void write_pdb(const Selection &mhd, TextOutput out, unsigned int model=1)
Add scale parameter to particle.
Definition: Scale.h:24
Classes to handle different kinds of restraints.
def input_adaptor
Adapt things for PMI (degrees of freedom, restraints, ...) Returns list of list of hierarchies...
Definition: tools.py:888
def optimize_sigma
Set sigma to the value that maximizes its conditional likelihood.
def draw_gamma
Draw gamma from lognormal distribution.
Add nuisance parameter to particle.
Definition: Nuisance.h:27
Basic SAXS restraint using ISD.
def __init__
Builds the restraint.
Basic functionality that is expected to be used by a wide variety of IMP users.
General purpose algebraic and geometric methods that are expected to be used by a wide variety of IMP...
The general base class for IMP exceptions.
Definition: exception.h:48
def draw_sigma
Draw 1/sigma2 from gamma distribution.
Functionality for loading, creating, manipulating and scoring atomic structures.
Select hierarchy particles identified by the biological name.
Definition: Selection.h:70
Support for small angle X-ray scattering (SAXS) data.
Warning for probably incorrect input parameters.
Inferential scoring building on methods developed as part of the Inferential Structure Determination ...