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