IMP logo
IMP Reference Guide  2.11.1
The Integrative Modeling Platform
lib/IMP/pmi/restraints/em.py
1 """@namespace IMP.pmi.restraints.em
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.isd
12 import IMP.pmi.tools
13 import IMP.pmi.mmcif
14 import IMP.isd.gmm_tools
15 import sys
16 import re
17 import os
18 from math import sqrt
19 
20 class GaussianEMRestraint(object):
21  """Fit Gaussian-decorated particles to an EM map
22  (also represented with a set of Gaussians)
23  @note This class wraps an isd::GaussianEMRestraint
24  """
25  def __init__(self, densities,
26  target_fn='',
27  target_ps=[],
28  cutoff_dist_model_model=0.0,
29  cutoff_dist_model_data=0.0,
30  target_mass_scale=1.0,
31  target_mass=None,
32  target_radii_scale=3.0,
33  model_radii_scale=1.0,
34  slope=0.0,
35  spherical_gaussians=False,
36  close_pair_container=None,
37  backbone_slope=False,
38  scale_target_to_mass=False,
39  weight=1.0,
40  target_is_rigid_body=False,
41  local=False,
42  representation=None):
43  """Constructor.
44  @param densities The Gaussian-decorated particles to be restrained
45  @param target_fn GMM file of the target density map
46  (alternatively, pass the ps)
47  @param target_ps List of Gaussians of the target map
48  (alternatively, pass the filename)
49  @param cutoff_dist_model_model Distance in model-model close
50  pair container
51  @param cutoff_dist_model_data Distance in model-data close pair
52  container. Usually can set to zero because we multiply the
53  target radii
54  @param target_mass_scale Scale up the target densities so that
55  the mass is accurate.
56  Needed if the GMM you generated was not already scaled.
57  To make it the same as model mass, set scale_to_target_mass=True
58  @param target_mass Sets the mass of the target density to the given value. Default is None. This
59  will override target_mass_scale argument
60  @param target_radii_scale Scale the target density radii -
61  only used for the close pair container.
62  If you keep this at 3.0 or so you don't have to use cutoff dist.
63  @param model_radii_scale Scale the model density radii - only used
64  for the close pair container
65  @param slope Linear term added to help bring model into the density
66  @param spherical_gaussians Set to True for a speed bonus when
67  the model densities are spheres. (This means you don't have
68  to do a matrix multiplication if they rotate.)
69  @param close_pair_container Pass a close pair container for
70  the model if you already have one (e.g. for an excluded
71  volume restraint.) May give a speed bonus.
72  @param backbone_slope Only apply slope to backbone particles -
73  only matters for atomic
74  @param scale_target_to_mass Set True if you would need to scale
75  target to EXACTLY the model mass
76  @param weight The restraint weight
77  @param target_is_rigid_body Set True if you want to put the target density particles
78  into a rigid body that need to be sampled (e.g.,when you need to fit one density
79  against another one). Default is False.
80  @param local Only consider density particles that are within the
81  specified model-density cutoff (experimental)
82  """
83 
84  # some parameters
85  self.label = "None"
86  self.sigmaissampled = False
87  self.sigmamaxtrans = 0.3
88  self.sigmamin = 1.0
89  self.sigmamax = 100.0
90  self.sigmainit = 1.0
91  self.label = "None"
92  self.densities = densities
93  self.em_root_hier = None
94 
95  # setup target GMM
96  self.m = self.densities[0].get_model()
97 
98  if scale_target_to_mass:
99  def hierarchy_mass(h):
100  leaves = IMP.atom.get_leaves(h)
101  return sum(IMP.atom.Mass(p).get_mass() for p in leaves)
102  target_mass = sum(hierarchy_mass(h) for h in densities)
103  print('will set target mass to', target_mass)
104  print('will scale target mass by', target_mass_scale)
105 
106  if target_fn != '':
107  self._set_dataset(target_fn, representation)
108  self.target_ps = []
110  target_fn,
111  self.target_ps,
112  self.m,
113  radius_scale=target_radii_scale,
114  mass_scale=target_mass_scale)
115  elif target_ps != []:
116  self.target_ps = target_ps
117  else:
118  print('Gaussian EM restraint: must provide target density file or properly set up target densities')
119  return
120 
121  if target_mass:
122  tmass=sum([IMP.atom.Mass(p).get_mass() for p in self.target_ps])
123  scale=target_mass/tmass
124  for p in self.target_ps:
125  ms=IMP.atom.Mass(p).get_mass()
126  IMP.atom.Mass(p).set_mass(ms*scale)
127 
128  for p, state in IMP.pmi.tools._all_protocol_outputs([representation],
129  densities[0]):
130  p.add_em3d_restraint(state, self.target_ps, self.densities,
131  self)
132 
133  # setup model GMM
134  self.model_ps = []
135  for h in self.densities:
136  self.model_ps += [ k.get_particle() for k in IMP.atom.get_leaves(h) ]
137  if model_radii_scale != 1.0:
138  for p in self.model_ps:
139  rmax = sqrt(max(IMP.core.Gaussian(p).get_variances())) * \
140  model_radii_scale
141  if not IMP.core.XYZR.get_is_setup(p):
143  else:
144  IMP.core.XYZR(p).set_radius(rmax)
145  #wrap target particles in rigid body if requested
146  if target_is_rigid_body:
147  #p = IMP.Particle(self.m)
148  #self.rb=IMP.core.RigidBody.setup_particle(p,self.target_ps)
149  self.rb=IMP.atom.create_rigid_body(self.target_ps)
150  else:
151  self.rb=None
152 
153  # sigma particle
154  self.sigmaglobal = IMP.pmi.tools.SetupNuisance(self.m, self.sigmainit,
155  self.sigmamin, self.sigmamax,
156  self.sigmaissampled).get_particle()
157 
158  # create restraint
159  print('target num particles', len(self.target_ps), \
160  'total weight', sum([IMP.atom.Mass(p).get_mass()
161  for p in self.target_ps]))
162  print('model num particles', len(self.model_ps), \
163  'total weight', sum([IMP.atom.Mass(p).get_mass()
164  for p in self.model_ps]))
165 
166  update_model=not spherical_gaussians
167  log_score=False
168  self.gaussianEM_restraint = IMP.isd.GaussianEMRestraint(
169  self.m,
170  IMP.get_indexes(self.model_ps),
171  IMP.get_indexes(self.target_ps),
172  self.sigmaglobal.get_particle().get_index(),
173  cutoff_dist_model_model,
174  cutoff_dist_model_data,
175  slope,
176  update_model, backbone_slope, local)
177  if target_fn != '':
178  self.gaussianEM_restraint.set_density_filename(target_fn)
179 
180  print('done EM setup')
181  self.rs = IMP.RestraintSet(self.m, 'GaussianEMRestraint')
182  self.rs.add_restraint(self.gaussianEM_restraint)
183  self.set_weight(weight)
184 
185  def _set_dataset(self, target_fn, representation):
186  """Set the dataset to point to the input file"""
188  self.dataset = p.parse_file(target_fn)['dataset']
189 
190  def center_target_density_on_model(self):
191  target_com = IMP.algebra.Vector3D(0, 0, 0)
192  target_mass = 0.0
193  for p in self.target_ps:
194  mass = IMP.atom.Mass(p).get_mass()
195  pos = IMP.core.XYZ(p).get_coordinates()
196  target_com += pos * mass
197  target_mass += mass
198  target_com /= target_mass
199  print('target com', target_com)
200  model_com = IMP.algebra.Vector3D(0, 0, 0)
201  model_mass = 0.0
202  for p in self.model_ps:
203  mass = IMP.atom.Mass(p).get_mass()
204  pos = IMP.core.XYZ(p).get_coordinates()
205  model_com += pos * mass
206  model_mass += mass
207  model_com /= model_mass
208  print('model com', model_com)
209 
210  v = target_com - model_com
211  print('translating with', -v)
213  for p in self.target_ps:
214  IMP.core.transform(IMP.core.RigidBody(p), transformation)
215  # IMP.pmi.tools.translate_hierarchies(self.densities,v)
216 
217  def get_center_of_mass(self, target=True):
218  '''Returns the geometric center of the GMM particles
219  @param target = True - returns target map gmm COM
220  @param target = False - returns model gmm COM'''
221  com = IMP.algebra.Vector3D(0, 0, 0)
222  total_mass = 0.0
223  if target:
224  ps = self.target_ps
225  else:
226  ps = self.model_ps
227  for p in ps:
228  mass = IMP.atom.Mass(p).get_mass()
229  pos = IMP.core.XYZ(p).get_coordinates()
230  com += pos * mass
231  total_mass += mass
232  com /= total_mass
233  return com
234 
235  def center_target_density_on_origin(self):
236  target_com = self.get_center_of_mass()
237  print('target com', target_com)
238  model_com = IMP.algebra.Vector3D(0, 0, 0)
239  print('model com', model_com)
240  v = target_com - model_com
241  print('translating with', -v)
243  for p in self.target_ps:
244  IMP.core.transform(IMP.core.RigidBody(p), transformation)
245  # IMP.pmi.tools.translate_hierarchies(self.densities,v)
246 
247  def center_model_on_target_density(self, input_object):
248  if type(input_object) is IMP.pmi.representation.Representation:
249  hier = input_object.prot
250  elif type(input_object) is IMP.pmi.topology.State:
251  hier = input_object.get_hierarchy()
252  else:
253  raise Exception("Input must be a Representation or topology.State object")
254  target_com = self.get_center_of_mass()
255  print('target com', target_com)
256  model_com = self.get_center_of_mass(target=False)
257  print('model com', model_com)
258  v = target_com - model_com
259  print('translating with', v)
261 
262  rigid_bodies = set()
263  XYZRs = set()
264 
265  for p in IMP.atom.get_leaves(hier):
267  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
268  rigid_bodies.add(rb)
270  XYZRs.add(p)
271 
272  for rb in list(rigid_bodies):
273  IMP.core.transform(rb, transformation)
274 
275  for p in list(XYZRs):
276  IMP.core.transform(IMP.core.XYZ(p), transformation)
277 
278  def center_on_target_density(self):
279  target_com = self.get_center_of_mass()
280  print('target com', target_com)
281  model_com = self.get_center_of_mass(target=False)
282  print('model com', model_com)
283  v = target_com - model_com
284  print('translating with', v)
286 
287  rigid_bodies = set()
288  XYZRs = set()
289 
290  for p in self.model_ps:
292  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
293  rigid_bodies.add(rb)
295  XYZRs.add(p)
296 
297  for rb in list(rigid_bodies):
298  IMP.core.transform(rb, transformation)
299 
300  for p in list(XYZRs):
301  IMP.core.transform(IMP.core.XYZ(p), transformation)
302 
303 
304 
305 
306  def set_weight(self,weight):
307  self.weight = weight
308  self.rs.set_weight(weight)
309 
310  def set_label(self, label):
311  self.label = label
312 
313  def add_to_model(self):
314  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs, add_to_rmf=True)
315 
316  def get_particles_to_sample(self):
317  ps = {}
318  if self.sigmaissampled:
319  ps["Nuisances_GaussianEMRestraint_sigma_" +
320  self.label] = ([self.sigmaglobal], self.sigmamaxtrans)
321  if self.rb:
322  ps["Rigid_Bodies_GaussianEMRestraint"] = (
323  [self.rb],
324  4,
325  0.03)
326  return ps
327 
328  def get_rigid_body(self):
329  if self.rb is None:
330  raise Exception("No rigid body created for GMM particles. Ensure target_is_rigid_body is set to True")
331  return self.rb
332 
333  def get_density_as_hierarchy(self):
334  if self.em_root_hier is None:
335  self.em_root_hier = IMP.atom.Copy.setup_particle(IMP.Particle(self.m),0)
336  self.em_root_hier.set_name("GaussianEMRestraint_density_"+self.label)
337  for p in self.target_ps:
338  self.em_root_hier.add_child(p)
339  return self.em_root_hier
340 
342  ''' Can add a target GMM to a Hierarchy.
343  For PMI2 a state object may also be passed'''
344  if type(inp) is IMP.pmi.topology.State:
345  inp.get_hierarchy().add_child(self.get_density_as_hierarchy())
346  elif type(inp) is IMP.atom.Hierarchy:
347  inp.add_child(self.get_density_as_hierarchy())
348  else:
349  raise Exception("Can only add a density to a PMI State object or IMP.atom.Hierarchy. You passed a", type(inp))
350 
351  def get_restraint(self):
352  return self.rs
353 
354  def get_restraint_set(self):
355  return self.rs
356 
357  def get_output(self):
358  output = {}
359  score = self.weight * self.rs.unprotected_evaluate(None)
360  ccc = self.gaussianEM_restraint.get_cross_correlation_coefficient()
361 
362  output["_TotalScore"] = str(score)
363  output["GaussianEMRestraint_" +
364  self.label] = str(score)
365  output["GaussianEMRestraint_%s_CCC" % self.label] = ccc
366  output["GaussianEMRestraint_sigma_" +
367  self.label] = str(self.sigmaglobal.get_scale())
368  return output
369 
370  def evaluate(self):
371  return self.weight * self.rs.unprotected_evaluate(None)
372 
373  def write_target_gmm_to_mrc(self, fileout=None, voxel_size=5.0):
374  '''Writes target GMM file to MRC'''
375  if fileout is None:
376  fileout="Gaussian_map_" + self.label + ".mrc"
377  IMP.isd.gmm_tools.write_gmm_to_map(self.target_ps, fileout, voxel_size)
378  return fileout
379 
380 
381 #-------------------------------------------
382 
384  """Fit particles to an EM map. This creates a simulate density map and updates them every eval.
385  @note Wraps an em::FitRestraint
386  """
387  def __init__(self,
388  ps,
389  dmap,
390  resolution,
391  origin=None,
392  voxel_size=None,
393  weight=1.0,
394  label=""):
395  """Constructor
396  @param ps The particles to restrain. Currently these must be atomic particles.
397  @param map_fn The EM density map to fit to
398  @param resolution Map resolution
399  @param origin In case you need to tell IMP the correct origin
400  @param voxel_size In case you need to tell IMP the angstroms per pixel
401  @param weight The data weight
402  @param label Extra PMI label
403  """
404  print('FitRestraint: setup')
405  #print('\tmap_fn',map_fn)
406  print('\tresolution',resolution)
407  print('\tvoxel_size',voxel_size)
408  print('\torigin',origin)
409  print('\tweight',weight)
410 
411  # some parameters
412  self.mdl = ps[0].get_model()
413  self.label = label
414  self.dmap = dmap #IMP.em.read_map(map_fn,IMP.em.MRCReaderWriter())
415  #dh = self.dmap.get_header()
416  #dh.set_resolution(resolution)
417  if voxel_size:
418  self.dmap.update_voxel_size(voxel_size)
419  if origin is not None:
420  if isinstance(origin, IMP.algebra.Vector3D):
421  self.dmap.set_origin(origin)
422  elif isinstance(origin, list):
423  self.dmap.set_origin(*origin)
424  else:
425  print('FitRestraint did not recognize format of origin')
426  exit()
427  fr = IMP.em.FitRestraint(ps,self.dmap)
428  self.rs = IMP.RestraintSet(self.mdl,weight,"FitRestraint")
429  self.rs.add_restraint(fr)
430  self.set_weight(weight)
431 
432  def set_weight(self,weight):
433  self.weight = weight
434  self.rs.set_weight(weight)
435 
436  def set_label(self, label):
437  self.label = label
438 
439  def add_to_model(self):
440  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)
441 
442  def get_restraint_set(self):
443  return self.rs
444 
445  def get_output(self):
446  output = {}
447  score = self.weight * self.rs.unprotected_evaluate(None)
448  output["_TotalScore"] = str(score)
449  output["EMRestraint_" + self.label] = str(score)
450  return output
451 
452  def evaluate(self):
453  return self.weight * self.rs.unprotected_evaluate(None)
454 
455 
456 #-------------------------------------------
457 
458 class ElectronMicroscopy2D(object):
459 
460  def __init__(
461  self,
462  representation,
463  images,
464  resolution=None):
465 
466  self.weight=1.0
467  self.m = representation.prot.get_model()
468  self.rs = IMP.RestraintSet(self.m, 'em2d')
469  self.label = "None"
470 
471  # IMP.atom.get_by_type
472  particles = IMP.pmi.tools.select(
473  representation,
474  resolution=resolution)
475 
476  em2d = None
477  self.rs.add_restraint(em2d)
478 
479  def set_label(self, label):
480  self.label = label
481 
482  def add_to_model(self):
483  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)
484 
485  def get_restraint(self):
486  return self.rs
487 
488  def set_weight(self,weight):
489  self.weight=weight
490  self.rs.set_weigth(self.weight)
491 
492  def get_output(self):
493  output = {}
494  score = self.weight*self.rs.unprotected_evaluate(None)
495  output["_TotalScore"] = str(score)
496  output["ElectronMicroscopy2D_" + self.label] = str(score)
497  return output
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
Fit Gaussian-decorated particles to an EM map (also represented with a set of Gaussians) ...
A member of a rigid body, it has internal (local) coordinates.
Definition: rigid_bodies.h:575
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: rigid_bodies.h:576
def write_target_gmm_to_mrc
Writes target GMM file to MRC.
Various classes to hold sets of particles.
static XYZR setup_particle(Model *m, ParticleIndex pi)
Definition: XYZR.h:48
double get_mass(ResidueType c)
Get the mass from the residue type.
Miscellaneous utilities.
Definition: tools.py:1
Extract metadata from an EM density GMM file.
Definition: mmcif.py:1633
def add_target_density_to_hierarchy
Can add a target GMM to a Hierarchy.
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.
Restraint between two Gaussian Mixture Models, "model" and "density".
def add_restraint_to_model
Add a PMI restraint to the model.
Definition: tools.py:86
A decorator for a particle with x,y,z coordinates.
Definition: XYZ.h:30
def get_center_of_mass
Returns the geometric center of the GMM particles.
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...
Calculate score based on fit to EM map.
Definition: FitRestraint.h:39
The general base class for IMP exceptions.
Definition: exception.h:49
VectorD< 3 > Vector3D
Definition: VectorD.h:395
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.
def select
this function uses representation=SimplifiedModel it returns the corresponding selected particles rep...
Definition: tools.py:545
A decorator for a rigid body.
Definition: rigid_bodies.h:82
Set up the representation of all proteins and nucleic acid macromolecules.
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
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