IMP logo
IMP Reference Guide  2.8.0
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.metadata
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  if representation:
129  for p, state in representation._protocol_output:
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 
178  print('done EM setup')
179  self.rs = IMP.RestraintSet(self.m, 'GaussianEMRestraint')
180  self.rs.add_restraint(self.gaussianEM_restraint)
181  self.set_weight(weight)
182 
183  def _set_dataset(self, target_fn, representation):
184  """Set the dataset to point to the input file"""
185  if representation:
186  self.dataset = representation.get_file_dataset(target_fn)
187  if self.dataset:
188  return
189  l = IMP.pmi.metadata.FileLocation(target_fn,
190  details="Electron microscopy density map, "
191  "represented as a Gaussian Mixture "
192  "Model (GMM)")
193  self.dataset = IMP.pmi.metadata.EMDensityDataset(l)
194  # If the GMM was derived from an MRC file that exists, add that too
195  m = re.match('(.*\.mrc)\..*\.txt$', target_fn)
196  if m and os.path.exists(m.group(1)):
197  l = IMP.pmi.metadata.FileLocation(path=m.group(1),
198  details='Original MRC file from which the GMM was derived')
199  self.dataset.add_parent(IMP.pmi.metadata.EMDensityDataset(l))
200 
201  def center_target_density_on_model(self):
202  target_com = IMP.algebra.Vector3D(0, 0, 0)
203  target_mass = 0.0
204  for p in self.target_ps:
205  mass = IMP.atom.Mass(p).get_mass()
206  pos = IMP.core.XYZ(p).get_coordinates()
207  target_com += pos * mass
208  target_mass += mass
209  target_com /= target_mass
210  print('target com', target_com)
211  model_com = IMP.algebra.Vector3D(0, 0, 0)
212  model_mass = 0.0
213  for p in self.model_ps:
214  mass = IMP.atom.Mass(p).get_mass()
215  pos = IMP.core.XYZ(p).get_coordinates()
216  model_com += pos * mass
217  model_mass += mass
218  model_com /= model_mass
219  print('model com', model_com)
220 
221  v = target_com - model_com
222  print('translating with', -v)
224  for p in self.target_ps:
225  IMP.core.transform(IMP.core.RigidBody(p), transformation)
226  # IMP.pmi.tools.translate_hierarchies(self.densities,v)
227 
228  def get_center_of_mass(self, target=True):
229  '''Returns the geometric center of the GMM particles
230  @param target = True - returns target map gmm COM
231  @param target = False - returns model gmm COM'''
232  com = IMP.algebra.Vector3D(0, 0, 0)
233  total_mass = 0.0
234  if target:
235  ps = self.target_ps
236  else:
237  ps = self.model_ps
238  for p in ps:
239  mass = IMP.atom.Mass(p).get_mass()
240  pos = IMP.core.XYZ(p).get_coordinates()
241  com += pos * mass
242  total_mass += mass
243  com /= total_mass
244  return com
245 
246  def center_target_density_on_origin(self):
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  if type(input_object) is IMP.pmi.representation.Representation:
260  hier = input_object.prot
261  elif type(input_object) is IMP.pmi.topology.State:
262  hier = input_object.get_hierarchy()
263  else:
264  raise Exception("Input must be a Representation or topology.State object")
265  target_com = self.get_center_of_mass()
266  print('target com', target_com)
267  model_com = self.get_center_of_mass(target=False)
268  print('model com', model_com)
269  v = target_com - model_com
270  print('translating with', v)
272 
273  rigid_bodies = set()
274  XYZRs = set()
275 
276  for p in IMP.atom.get_leaves(hier):
278  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
279  rigid_bodies.add(rb)
281  XYZRs.add(p)
282 
283  for rb in list(rigid_bodies):
284  IMP.core.transform(rb, transformation)
285 
286  for p in list(XYZRs):
287  IMP.core.transform(IMP.core.XYZ(p), transformation)
288 
289  def center_on_target_density(self):
290  target_com = self.get_center_of_mass()
291  print('target com', target_com)
292  model_com = self.get_center_of_mass(target=False)
293  print('model com', model_com)
294  v = target_com - model_com
295  print('translating with', v)
297 
298  rigid_bodies = set()
299  XYZRs = set()
300 
301  for p in self.model_ps:
303  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
304  rigid_bodies.add(rb)
306  XYZRs.add(p)
307 
308  for rb in list(rigid_bodies):
309  IMP.core.transform(rb, transformation)
310 
311  for p in list(XYZRs):
312  IMP.core.transform(IMP.core.XYZ(p), transformation)
313 
314 
315 
316 
317  def set_weight(self,weight):
318  self.weight = weight
319  self.rs.set_weight(weight)
320 
321  def set_label(self, label):
322  self.label = label
323 
324  def add_to_model(self):
325  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)
326 
327  def get_particles_to_sample(self):
328  ps = {}
329  if self.sigmaissampled:
330  ps["Nuisances_GaussianEMRestraint_sigma_" +
331  self.label] = ([self.sigmaglobal], self.sigmamaxtrans)
332  if self.rb:
333  ps["Rigid_Bodies_GaussianEMRestraint"] = (
334  [self.rb],
335  4,
336  0.03)
337  return ps
338 
339  def get_rigid_body(self):
340  if self.rb is None:
341  raise Exception("No rigid body created for GMM particles. Ensure target_is_rigid_body is set to True")
342  return self.rb
343 
344  def get_density_as_hierarchy(self):
345  if self.em_root_hier is None:
346  self.em_root_hier = IMP.atom.Copy.setup_particle(IMP.Particle(self.m),0)
347  self.em_root_hier.set_name("GaussianEMRestraint_density_"+self.label)
348  for p in self.target_ps:
349  self.em_root_hier.add_child(p)
350  return self.em_root_hier
351 
353  ''' Can add a target GMM to a Hierarchy.
354  For PMI2 a state object may also be passed'''
355  if type(inp) is IMP.pmi.topology.State:
356  inp.get_hierarchy().add_child(self.get_density_as_hierarchy())
357  elif type(inp) is IMP.atom.Hierarchy:
358  inp.add_child(self.get_density_as_hierarchy())
359  else:
360  raise Exception("Can only add a density to a PMI State object or IMP.atom.Hierarchy. You passed a", type(inp))
361 
362  def get_restraint_set(self):
363  return self.rs
364 
365  def get_output(self):
366  self.m.update()
367  output = {}
368  score = self.weight * self.rs.unprotected_evaluate(None)
369  ccc = self.gaussianEM_restraint.get_cross_correlation_coefficient()
370 
371  output["_TotalScore"] = str(score)
372  output["GaussianEMRestraint_" +
373  self.label] = str(score)
374  output["GaussianEMRestraint_%s_CCC" % self.label] = ccc
375  output["GaussianEMRestraint_sigma_" +
376  self.label] = str(self.sigmaglobal.get_scale())
377  return output
378 
379  def evaluate(self):
380  return self.weight * self.rs.unprotected_evaluate(None)
381 
382  def write_target_gmm_to_mrc(self, fileout=None, voxel_size=5.0):
383  '''Writes target GMM file to MRC'''
384  if fileout is None:
385  fileout="Gaussian_map_" + self.label + ".mrc"
386  IMP.isd.gmm_tools.write_gmm_to_map(self.target_ps, fileout, voxel_size)
387  return fileout
388 
389 
390 #-------------------------------------------
391 
393  """Fit particles to an EM map. This creates a simulate density map and updates them every eval.
394  \note Wraps an em::FitRestraint
395  """
396  def __init__(self,
397  ps,
398  dmap,
399  resolution,
400  origin=None,
401  voxel_size=None,
402  weight=1.0,
403  label=""):
404  """Constructor
405  @param ps The particles to restrain. Currently these must be atomic particles.
406  @param map_fn The EM density map to fit to
407  @param resolution Map resolution
408  @param origin In case you need to tell IMP the correct origin
409  @param voxel_size In case you need to tell IMP the angstroms per pixel
410  @param weight The data weight
411  @param label Extra PMI label
412  """
413  print('FitRestraint: setup')
414  #print('\tmap_fn',map_fn)
415  print('\tresolution',resolution)
416  print('\tvoxel_size',voxel_size)
417  print('\torigin',origin)
418  print('\tweight',weight)
419 
420  # some parameters
421  self.mdl = ps[0].get_model()
422  self.label = label
423  self.dmap = dmap #IMP.em.read_map(map_fn,IMP.em.MRCReaderWriter())
424  #dh = self.dmap.get_header()
425  #dh.set_resolution(resolution)
426  if voxel_size:
427  self.dmap.update_voxel_size(voxel_size)
428  if origin is not None:
429  if type(origin)==IMP.algebra.Vector3D:
430  self.dmap.set_origin(origin)
431  elif type(origin)==list:
432  self.dmap.set_origin(*origin)
433  else:
434  print('FitRestraint did not recognize format of origin')
435  exit()
436  fr = IMP.em.FitRestraint(ps,self.dmap)
437  self.rs = IMP.RestraintSet(self.mdl,weight,"FitRestraint")
438  self.rs.add_restraint(fr)
439  self.set_weight(weight)
440 
441  def set_weight(self,weight):
442  self.weight = weight
443  self.rs.set_weight(weight)
444 
445  def set_label(self, label):
446  self.label = label
447 
448  def add_to_model(self):
449  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)
450 
451  def get_restraint_set(self):
452  return self.rs
453 
454  def get_output(self):
455  self.mdl.update()
456  output = {}
457  score = self.weight * self.rs.unprotected_evaluate(None)
458  output["_TotalScore"] = str(score)
459  output["EMRestraint_" + self.label] = str(score)
460  return output
461 
462  def evaluate(self):
463  return self.weight * self.rs.unprotected_evaluate(None)
464 
465 
466 #-------------------------------------------
467 
468 class ElectronMicroscopy2D(object):
469 
470  def __init__(
471  self,
472  representation,
473  images,
474  resolution=None):
475 
476  self.weight=1.0
477  self.m = representation.prot.get_model()
478  self.rs = IMP.RestraintSet(self.m, 'em2d')
479  self.label = "None"
480 
481  # IMP.atom.get_by_type
482  particles = IMP.pmi.tools.select(
483  representation,
484  resolution=resolution)
485 
486  em2d = None
487  self.rs.add_restraint(em2d)
488 
489  def set_label(self, label):
490  self.label = label
491 
492  def add_to_model(self):
493  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)
494 
495  def get_restraint(self):
496  return self.rs
497 
498  def set_weight(self,weight):
499  self.weight=weight
500  self.rs.set_weigth(self.weight)
501 
502  def get_output(self):
503  self.m.update()
504  output = {}
505  score = self.weight*self.rs.unprotected_evaluate(None)
506  output["_TotalScore"] = str(score)
507  output["ElectronMicroscopy2D_" + self.label] = str(score)
508  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.
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:508
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: rigid_bodies.h:509
Classes for attaching metadata to PMI objects.
Definition: metadata.py:1
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
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:38
A 3D electron microscopy dataset.
Definition: metadata.py:92
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:111
static Copy setup_particle(Model *m, ParticleIndex pi, Int number)
Definition: Copy.h:42
An individual file or directory.
Definition: metadata.py:167
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:700
A decorator for a rigid body.
Definition: rigid_bodies.h:75
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