IMP logo
IMP Reference Guide  2.12.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  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  hier = input_object.get_hierarchy()
249  target_com = self.get_center_of_mass()
250  print('target com', target_com)
251  model_com = self.get_center_of_mass(target=False)
252  print('model com', model_com)
253  v = target_com - model_com
254  print('translating with', v)
256 
257  rigid_bodies = set()
258  XYZRs = set()
259 
260  for p in IMP.atom.get_leaves(hier):
262  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
263  rigid_bodies.add(rb)
265  XYZRs.add(p)
266 
267  for rb in list(rigid_bodies):
268  IMP.core.transform(rb, transformation)
269 
270  for p in list(XYZRs):
271  IMP.core.transform(IMP.core.XYZ(p), transformation)
272 
273  def center_on_target_density(self):
274  target_com = self.get_center_of_mass()
275  print('target com', target_com)
276  model_com = self.get_center_of_mass(target=False)
277  print('model com', model_com)
278  v = target_com - model_com
279  print('translating with', v)
281 
282  rigid_bodies = set()
283  XYZRs = set()
284 
285  for p in self.model_ps:
287  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
288  rigid_bodies.add(rb)
290  XYZRs.add(p)
291 
292  for rb in list(rigid_bodies):
293  IMP.core.transform(rb, transformation)
294 
295  for p in list(XYZRs):
296  IMP.core.transform(IMP.core.XYZ(p), transformation)
297 
298 
299 
300 
301  def set_weight(self,weight):
302  self.weight = weight
303  self.rs.set_weight(weight)
304 
305  def set_label(self, label):
306  self.label = label
307 
308  def add_to_model(self):
309  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs, add_to_rmf=True)
310 
311  def get_particles_to_sample(self):
312  ps = {}
313  if self.sigmaissampled:
314  ps["Nuisances_GaussianEMRestraint_sigma_" +
315  self.label] = ([self.sigmaglobal], self.sigmamaxtrans)
316  if self.rb:
317  ps["Rigid_Bodies_GaussianEMRestraint"] = (
318  [self.rb],
319  4,
320  0.03)
321  return ps
322 
323  def get_rigid_body(self):
324  if self.rb is None:
325  raise Exception("No rigid body created for GMM particles. Ensure target_is_rigid_body is set to True")
326  return self.rb
327 
328  def get_density_as_hierarchy(self):
329  if self.em_root_hier is None:
330  self.em_root_hier = IMP.atom.Copy.setup_particle(IMP.Particle(self.m),0)
331  self.em_root_hier.set_name("GaussianEMRestraint_density_"+self.label)
332  for p in self.target_ps:
333  self.em_root_hier.add_child(p)
334  return self.em_root_hier
335 
337  ''' Can add a target GMM to a Hierarchy.
338  For PMI2 a state object may also be passed'''
339  if type(inp) is IMP.pmi.topology.State:
340  inp.get_hierarchy().add_child(self.get_density_as_hierarchy())
341  elif type(inp) is IMP.atom.Hierarchy:
342  inp.add_child(self.get_density_as_hierarchy())
343  else:
344  raise Exception("Can only add a density to a PMI State object or IMP.atom.Hierarchy. You passed a", type(inp))
345 
346  def get_restraint(self):
347  return self.rs
348 
349  def get_restraint_set(self):
350  return self.rs
351 
352  def get_output(self):
353  output = {}
354  score = self.weight * self.rs.unprotected_evaluate(None)
355  ccc = self.gaussianEM_restraint.get_cross_correlation_coefficient()
356 
357  output["_TotalScore"] = str(score)
358  output["GaussianEMRestraint_" +
359  self.label] = str(score)
360  output["GaussianEMRestraint_%s_CCC" % self.label] = ccc
361  output["GaussianEMRestraint_sigma_" +
362  self.label] = str(self.sigmaglobal.get_scale())
363  return output
364 
365  def evaluate(self):
366  return self.weight * self.rs.unprotected_evaluate(None)
367 
368  def write_target_gmm_to_mrc(self, fileout=None, voxel_size=5.0):
369  '''Writes target GMM file to MRC'''
370  if fileout is None:
371  fileout="Gaussian_map_" + self.label + ".mrc"
372  IMP.isd.gmm_tools.write_gmm_to_map(self.target_ps, fileout, voxel_size)
373  return fileout
374 
375 
376 #-------------------------------------------
377 
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 class ElectronMicroscopy2D(object):
454 
455  def __init__(
456  self,
457  representation,
458  images,
459  resolution=None):
460 
461  self.weight=1.0
462  self.m = representation.prot.get_model()
463  self.rs = IMP.RestraintSet(self.m, 'em2d')
464  self.label = "None"
465 
466  # IMP.atom.get_by_type
467  particles = IMP.pmi.tools.select(
468  representation,
469  resolution=resolution)
470 
471  em2d = None
472  self.rs.add_restraint(em2d)
473 
474  def set_label(self, label):
475  self.label = label
476 
477  def add_to_model(self):
478  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)
479 
480  def get_restraint(self):
481  return self.rs
482 
483  def set_weight(self,weight):
484  self.weight=weight
485  self.rs.set_weigth(self.weight)
486 
487  def get_output(self):
488  output = {}
489  score = self.weight*self.rs.unprotected_evaluate(None)
490  output["_TotalScore"] = str(score)
491  output["ElectronMicroscopy2D_" + self.label] = str(score)
492  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:1610
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:89
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: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: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