IMP logo
IMP Reference Guide  2.12.0
The Integrative Modeling Platform
bayesianem/restraint.py
1 """@namespace IMP.bayesianem
2 Restraints for handling electron microscopy maps.
3 """
4 
5 from __future__ import print_function
6 import IMP
7 import IMP.core
8 import IMP.algebra
9 import IMP.atom
10 import IMP.container
11 import IMP.bayesianem
12 import IMP.isd
13 import IMP.pmi.tools
14 import IMP.pmi.mmcif
15 import IMP.isd.gmm_tools
16 import sys
17 from math import sqrt
18 
20  """Fit Gaussian-decorated particles to an EM map
21  (also represented with a set of Gaussians)
22  \note This class wraps an isd::GaussianEMRestraint
23  """
24  def __init__(self, densities,
25  target_fn='',
26  target_ps=[],
27  cutoff_dist_model_model=10.0,
28  cutoff_dist_model_data=10.0,
29  target_mass_scale=1.0,
30  target_mass=None,
31  target_radii_scale=3.0,
32  model_radii_scale=1.0,
33  slope=0.0,
34  spherical_gaussians=False,
35  close_pair_container=None,
36  backbone_slope=False,
37  scale_target_to_mass=False,
38  weight=1.0,
39  target_is_rigid_body=False):
40  """Constructor.
41  @param densities The Gaussian-decorated particles to be restrained
42  @param target_fn GMM file of the target density map
43  (alternatively, pass the ps)
44  @param target_ps List of Gaussians of the target map
45  (alternatively, pass the filename)
46  @param cutoff_dist_model_model Distance in model-model close
47  pair container
48  @param cutoff_dist_model_data Distance in model-data close pair
49  container. Usually can set to zero because we multiply the
50  target radii
51  @param target_mass_scale Scale up the target densities so that
52  the mass is accurate.
53  Needed if the GMM you generated was not already scaled.
54  To make it the same as model mass, set scale_to_target_mass=True
55  @param target_mass Sets the mass of the target density to the given value. Default is None. This
56  will override target_mass_scale argument
57  @param target_radii_scale Scale the target density radii -
58  only used for the close pair container.
59  If you keep this at 3.0 or so you don't have to use cutoff dist.
60  @param model_radii_scale Scale the model density radii - only used
61  for the close pair container
62  @param slope Linear term added to help bring model into the density
63  @param spherical_gaussians Set to True for a speed bonus when
64  the model densities are spheres. (This means you don't have
65  to do a matrix multiplication if they rotate.)
66  @param close_pair_container Pass a close pair container for
67  the model if you already have one (e.g. for an excluded
68  volume restraint.) May give a speed bonus.
69  @param backbone_slope Only apply slope to backbone particles -
70  only matters for atomic
71  @param scale_target_to_mass Set True if you would need to scale
72  target to EXACTLY the model mass
73  @param weight The restraint weight
74  @param target_is_rigid_body Set True if you want to put the target density particles
75  into a rigid body that need to be sampled (e.g.,when you need to fit one density
76  against another one). Default is False.
77  """
78  # some parameters
79  self.label = "None"
80  self.sigmaissampled = True
81  self.sigmamaxtrans = 0.3
82  self.sigmamin = 0.0001
83  self.sigmamax = 100.0
84  self.sigmainit = 1.0
85  self.label = "None"
86  self.densities = densities
87  self.em_root_hier = None
88 
89  # setup target GMM
90  self.m = self.densities[0].get_model()
91 
92  if scale_target_to_mass:
93  def hierarchy_mass(h):
94  leaves = IMP.atom.get_leaves(h)
95  return sum(IMP.atom.Mass(p).get_mass() for p in leaves)
96  target_mass = sum(hierarchy_mass(h) for h in densities)
97  print('will scale target mass by', target_mass_scale)
98 
99  if target_fn != '':
100  self._set_dataset(target_fn)
101  self.target_ps = []
103  target_fn,
104  self.target_ps,
105  self.m,
106  radius_scale=target_radii_scale,
107  mass_scale=target_mass_scale)
108  elif target_ps != []:
109  self.target_ps = target_ps
110  else:
111  print('Gaussian EM restraint: must provide target density file or properly set up target densities')
112  return
113 
114  if target_mass:
115  tmass=sum([IMP.atom.Mass(p).get_mass() for p in self.target_ps])
116  scale=target_mass/tmass
117  print('will set target mass to', target_mass,tmass,scale)
118  for p in self.target_ps:
119  ms=IMP.atom.Mass(p).get_mass()
120  IMP.atom.Mass(p).set_mass(ms*scale)
121 
122  for p, state in IMP.pmi.tools._all_protocol_outputs([], densities[0]):
123  p.add_em3d_restraint(state, self.target_ps, self.densities, self)
124 
125  # setup model GMM
126  self.model_ps = []
127  for h in self.densities:
128  self.model_ps += [ k.get_particle() for k in IMP.atom.get_leaves(h) ]
129  if model_radii_scale != 1.0:
130  for p in self.model_ps:
131  rmax = sqrt(max(IMP.core.Gaussian(p).get_variances())) * \
132  model_radii_scale
133  if not IMP.core.XYZR.get_is_setup(p):
135  else:
136  IMP.core.XYZR(p).set_radius(rmax)
137  #wrap target particles in rigid body if requested
138  if target_is_rigid_body:
139  #p = IMP.Particle(self.m)
140  #self.rb=IMP.core.RigidBody.setup_particle(p,self.target_ps)
141  self.rb=IMP.atom.create_rigid_body(self.target_ps)
142  else:
143  self.rb=None
144 
145  # sigma particle
146  self.sigmaglobal = IMP.pmi.tools.SetupNuisance(self.m, self.sigmainit,
147  self.sigmamin, self.sigmamax,
148  self.sigmaissampled).get_particle()
149 
150  # create restraint
151  print('target num particles', len(self.target_ps), \
152  'total weight', sum([IMP.atom.Mass(p).get_mass()
153  for p in self.target_ps]))
154  print('model num particles', len(self.model_ps), \
155  'total weight', sum([IMP.atom.Mass(p).get_mass()
156  for p in self.model_ps]))
157 
158  #Normalize
159  #
160  #mass_model=sum([IMP.atom.Mass(p).get_mass() for p in self.model_ps])
161  #for p in self.model_ps:
162  # pmass=IMP.atom.Mass(p).get_mass()
163  # newpmass=pmass/mass_model
164  # IMP.atom.Mass(p).set_mass(newpmass)
165 
166 
167  #mass_target=sum([IMP.atom.Mass(p).get_mass() for p in self.target_ps])
168  #for p in self.target_ps:
169  # pmass=IMP.atom.Mass(p).get_mass()
170  # newpmass=pmass/mass_target
171  # IMP.atom.Mass(p).set_mass(newpmass)
172 
173  update_model=not spherical_gaussians
174  log_score=False
175  self.gaussianEM_restraint = IMP.bayesianem.GaussianEMRestraint(
176  self.m,
177  IMP.get_indexes(self.model_ps),
178  IMP.get_indexes(self.target_ps),
179  self.sigmaglobal.get_particle().get_index(),
180  cutoff_dist_model_model,
181  cutoff_dist_model_data,
182  slope,
183  update_model, backbone_slope)
184 
185  print('done EM setup')
186  self.rs = IMP.RestraintSet(self.m, 'GaussianEMRestraint')
187  self.rs.add_restraint(self.gaussianEM_restraint)
188  self.set_weight(weight)
189 
190  def _set_dataset(self, target_fn):
191  """Set the dataset to point to the input file"""
193  self.dataset = p.parse_file(target_fn)['dataset']
194 
196  '''
197  aligns the center of mass of the target GMM on the center of mass of the model
198  '''
199  target_com = IMP.algebra.Vector3D(0, 0, 0)
200  target_mass = 0.0
201  for p in self.target_ps:
202  mass = IMP.atom.Mass(p).get_mass()
203  pos = IMP.core.XYZ(p).get_coordinates()
204  target_com += pos * mass
205  target_mass += mass
206  target_com /= target_mass
207  print('target com', target_com)
208  model_com = IMP.algebra.Vector3D(0, 0, 0)
209  model_mass = 0.0
210  for p in self.model_ps:
211  mass = IMP.atom.Mass(p).get_mass()
212  pos = IMP.core.XYZ(p).get_coordinates()
213  model_com += pos * mass
214  model_mass += mass
215  model_com /= model_mass
216  print('model com', model_com)
217 
218  v = target_com - model_com
219  print('translating with', -v)
221  for p in self.target_ps:
222  IMP.core.transform(IMP.core.RigidBody(p), transformation)
223  # IMP.pmi.tools.translate_hierarchies(self.densities,v)
224 
225  def get_center_of_mass(self, target=True):
226  '''Returns the geometric center of the GMM particles
227  @param target = True - returns target map gmm COM
228  @param target = False - returns model gmm COM'''
229  com = IMP.algebra.Vector3D(0, 0, 0)
230  total_mass = 0.0
231  if target:
232  ps = self.target_ps
233  else:
234  ps = self.model_ps
235  for p in ps:
236  mass = IMP.atom.Mass(p).get_mass()
237  pos = IMP.core.XYZ(p).get_coordinates()
238  com += pos * mass
239  total_mass += mass
240  com /= total_mass
241  return com
242 
244  '''
245  aligns the center of mass of the target GMM on the origin
246  '''
247  target_com = self.get_center_of_mass()
248  print('target com', target_com)
249  model_com = IMP.algebra.Vector3D(0, 0, 0)
250  print('model com', model_com)
251  v = target_com - model_com
252  print('translating with', -v)
254  for p in self.target_ps:
255  IMP.core.transform(IMP.core.RigidBody(p), transformation)
256  # IMP.pmi.tools.translate_hierarchies(self.densities,v)
257 
258  def center_model_on_target_density(self, input_object):
259  '''
260  aligns the model on the target density
261  @param input_objects IMP.pmi.representation.Representation or IMP.pmi.topology.State
262  '''
263  if type(input_object) is IMP.pmi.representation.Representation:
264  hier = input_object.prot
265  elif type(input_object) is IMP.pmi.topology.State:
266  hier = input_object.get_hierarchy()
267  else:
268  raise Exception("Input must be a Representation or topology.State object")
269  target_com = self.get_center_of_mass()
270  print('target com', target_com)
271  model_com = self.get_center_of_mass(target=False)
272  print('model com', model_com)
273  v = target_com - model_com
274  print('translating with', v)
276 
277  rigid_bodies = set()
278  XYZRs = set()
279 
280  for p in IMP.atom.get_leaves(hier):
282  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
283  rigid_bodies.add(rb)
285  XYZRs.add(p)
286 
287  for rb in list(rigid_bodies):
288  IMP.core.transform(rb, transformation)
289 
290  for p in list(XYZRs):
291  IMP.core.transform(IMP.core.XYZ(p), transformation)
292 
294  '''
295  align the model on target GMM
296  '''
297  target_com = self.get_center_of_mass()
298  print('target com', target_com)
299  model_com = self.get_center_of_mass(target=False)
300  print('model com', model_com)
301  v = target_com - model_com
302  print('translating with', v)
304 
305  rigid_bodies = set()
306  XYZRs = set()
307 
308  for p in self.model_ps:
310  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
311  rigid_bodies.add(rb)
313  XYZRs.add(p)
314 
315  for rb in list(rigid_bodies):
316  IMP.core.transform(rb, transformation)
317 
318  for p in list(XYZRs):
319  IMP.core.transform(IMP.core.XYZ(p), transformation)
320 
321 
322  def set_weight(self,weight):
323  '''
324  set the weight of the restraint
325  @param weight
326  '''
327  self.weight = weight
328  self.rs.set_weight(weight)
329 
330  def set_label(self, label):
331  self.label = label
332 
333  def add_to_model(self):
334  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)
335 
336  def get_particles_to_sample(self):
337  ps = {}
338  if self.sigmaissampled:
339  ps["Nuisances_GaussianEMRestraint_sigma_" +
340  self.label] = ([self.sigmaglobal], self.sigmamaxtrans)
341  if self.rb:
342  ps["Rigid_Bodies_GaussianEMRestraint_"+self.label] = (
343  [self.rb],
344  4,
345  0.25)
346  return ps
347 
348  def get_rigid_body(self):
349  if self.rb is None:
350  raise Exception("No rigid body created for GMM particles. Ensure target_is_rigid_body is set to True")
351  return self.rb
352 
354  '''
355  returns a hierarchy whose leaves are the gaussian particles of the target GMM
356  '''
357  if self.em_root_hier is None:
358  self.em_root_hier = IMP.atom.Copy.setup_particle(IMP.Particle(self.m),0)
359  self.em_root_hier.set_name("GaussianEMRestraint_density_"+self.label)
360  for p in self.target_ps:
361  self.em_root_hier.add_child(p)
362  return self.em_root_hier
363 
365  ''' Can add a target GMM to a Hierarchy.
366  For PMI2 a state object may also be passed'''
367  if type(inp) is IMP.pmi.topology.State:
368  inp.get_hierarchy().add_child(self.get_density_as_hierarchy())
369  elif type(inp) is IMP.atom.Hierarchy:
370  inp.add_child(self.get_density_as_hierarchy())
371  else:
372  raise Exception("Can only add a density to a PMI State object or IMP.atom.Hierarchy. You passed a", type(inp))
373 
374  def get_restraint_set(self):
375  return self.rs
376 
377  def get_output(self):
378  self.m.update()
379  output = {}
380  score = self.weight * self.rs.unprotected_evaluate(None)
381  output["_TotalScore"] = str(score)
382  output["GaussianEMRestraint_" +
383  self.label] = str(score)
384  output["GaussianEMRestraint_sigma_" +
385  self.label] = str(self.sigmaglobal.get_scale())
386  return output
387 
388  def evaluate(self):
389  return self.weight * self.rs.unprotected_evaluate(None)
390 
391  def write_target_gmm_to_mrc(self, fileout=None, voxel_size=5.0):
392  '''Writes target GMM file to MRC'''
393  if fileout is None:
394  fileout="Gaussian_map_" + self.label + ".mrc"
395  IMP.isd.gmm_tools.write_gmm_to_map(self.target_ps, fileout, voxel_size)
396  return fileout
397 
Tools for handling Gaussian Mixture Models.
Definition: gmm_tools.py:1
Add mass to a particle.
Definition: Mass.h:23
Simple 3D transformation class.
Support for the mmCIF file format.
Definition: mmcif.py:1
def get_center_of_mass
Returns the geometric center of the GMM particles.
Restraints for handling electron microscopy maps.
A member of a rigid body, it has internal (local) coordinates.
Definition: rigid_bodies.h:616
def center_model_on_target_density
aligns the model on the target density
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: rigid_bodies.h:617
Various classes to hold sets of particles.
static XYZR setup_particle(Model *m, ParticleIndex pi)
Definition: XYZR.h:48
def center_target_density_on_origin
aligns the center of mass of the target GMM on the origin
double get_mass(ResidueType c)
Get the mass from the residue type.
def add_target_density_to_hierarchy
Can add a target GMM to a Hierarchy.
Miscellaneous utilities.
Definition: tools.py:1
Creates a restraint between two Gaussian Mixture Models, "model" and "density".
Extract metadata from an EM density GMM file.
Definition: mmcif.py:1610
Fit Gaussian-decorated particles to an EM map (also represented with a set of Gaussians) ...
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: XYZR.h:47
Object used to hold a set of restraints.
Definition: RestraintSet.h:36
ParticleIndexPairs get_indexes(const ParticlePairsTemp &ps)
The standard decorator for manipulating molecular structures.
Ints get_index(const ParticlesTemp &particles, const Subset &subset, const Subsets &excluded)
void transform(XYZ a, const algebra::Transformation3D &tr)
Apply a transformation to the particle.
def add_restraint_to_model
Add a PMI restraint to the model.
Definition: tools.py:89
def center_on_target_density
align the model on target GMM
A decorator for a particle with x,y,z coordinates.
Definition: XYZ.h:30
def write_gmm_to_map
write density map from GMM.
Definition: gmm_tools.py:113
static Copy setup_particle(Model *m, ParticleIndex pi, Int number)
Definition: Copy.h:42
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...
def write_target_gmm_to_mrc
Writes target GMM file to MRC.
The general base class for IMP exceptions.
Definition: exception.h:49
VectorD< 3 > Vector3D
Definition: VectorD.h:421
IMP::core::RigidBody create_rigid_body(Hierarchy h)
Class to handle individual particles of a Model object.
Definition: Particle.h:41
Stores a list of Molecules all with the same State index.
A decorator for a rigid body.
Definition: rigid_bodies.h:82
def center_target_density_on_model
aligns the center of mass of the target GMM on the center of mass of the model
Functionality for loading, creating, manipulating and scoring atomic structures.
Hierarchies get_leaves(const Selection &h)
def decorate_gmm_from_text
read the output from write_gmm_to_text, decorate as Gaussian and Mass
Definition: gmm_tools.py:22
def get_density_as_hierarchy
returns a hierarchy whose leaves are the gaussian particles of the target GMM
Inferential scoring building on methods developed as part of the Inferential Structure Determination ...
A decorator for a particle with x,y,z coordinates and a radius.
Definition: XYZR.h:27