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