IMP logo
IMP Reference Guide  2.10.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  output = {}
373  score = self.weight * self.rs.unprotected_evaluate(None)
374  ccc = self.gaussianEM_restraint.get_cross_correlation_coefficient()
375 
376  output["_TotalScore"] = str(score)
377  output["GaussianEMRestraint_" +
378  self.label] = str(score)
379  output["GaussianEMRestraint_%s_CCC" % self.label] = ccc
380  output["GaussianEMRestraint_sigma_" +
381  self.label] = str(self.sigmaglobal.get_scale())
382  return output
383 
384  def evaluate(self):
385  return self.weight * self.rs.unprotected_evaluate(None)
386 
387  def write_target_gmm_to_mrc(self, fileout=None, voxel_size=5.0):
388  '''Writes target GMM file to MRC'''
389  if fileout is None:
390  fileout="Gaussian_map_" + self.label + ".mrc"
391  IMP.isd.gmm_tools.write_gmm_to_map(self.target_ps, fileout, voxel_size)
392  return fileout
393 
394 
395 #-------------------------------------------
396 
398  """Fit particles to an EM map. This creates a simulate density map and updates them every eval.
399  @note Wraps an em::FitRestraint
400  """
401  def __init__(self,
402  ps,
403  dmap,
404  resolution,
405  origin=None,
406  voxel_size=None,
407  weight=1.0,
408  label=""):
409  """Constructor
410  @param ps The particles to restrain. Currently these must be atomic particles.
411  @param map_fn The EM density map to fit to
412  @param resolution Map resolution
413  @param origin In case you need to tell IMP the correct origin
414  @param voxel_size In case you need to tell IMP the angstroms per pixel
415  @param weight The data weight
416  @param label Extra PMI label
417  """
418  print('FitRestraint: setup')
419  #print('\tmap_fn',map_fn)
420  print('\tresolution',resolution)
421  print('\tvoxel_size',voxel_size)
422  print('\torigin',origin)
423  print('\tweight',weight)
424 
425  # some parameters
426  self.mdl = ps[0].get_model()
427  self.label = label
428  self.dmap = dmap #IMP.em.read_map(map_fn,IMP.em.MRCReaderWriter())
429  #dh = self.dmap.get_header()
430  #dh.set_resolution(resolution)
431  if voxel_size:
432  self.dmap.update_voxel_size(voxel_size)
433  if origin is not None:
434  if type(origin)==IMP.algebra.Vector3D:
435  self.dmap.set_origin(origin)
436  elif type(origin)==list:
437  self.dmap.set_origin(*origin)
438  else:
439  print('FitRestraint did not recognize format of origin')
440  exit()
441  fr = IMP.em.FitRestraint(ps,self.dmap)
442  self.rs = IMP.RestraintSet(self.mdl,weight,"FitRestraint")
443  self.rs.add_restraint(fr)
444  self.set_weight(weight)
445 
446  def set_weight(self,weight):
447  self.weight = weight
448  self.rs.set_weight(weight)
449 
450  def set_label(self, label):
451  self.label = label
452 
453  def add_to_model(self):
454  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)
455 
456  def get_restraint_set(self):
457  return self.rs
458 
459  def get_output(self):
460  output = {}
461  score = self.weight * self.rs.unprotected_evaluate(None)
462  output["_TotalScore"] = str(score)
463  output["EMRestraint_" + self.label] = str(score)
464  return output
465 
466  def evaluate(self):
467  return self.weight * self.rs.unprotected_evaluate(None)
468 
469 
470 #-------------------------------------------
471 
472 class ElectronMicroscopy2D(object):
473 
474  def __init__(
475  self,
476  representation,
477  images,
478  resolution=None):
479 
480  self.weight=1.0
481  self.m = representation.prot.get_model()
482  self.rs = IMP.RestraintSet(self.m, 'em2d')
483  self.label = "None"
484 
485  # IMP.atom.get_by_type
486  particles = IMP.pmi.tools.select(
487  representation,
488  resolution=resolution)
489 
490  em2d = None
491  self.rs.add_restraint(em2d)
492 
493  def set_label(self, label):
494  self.label = label
495 
496  def add_to_model(self):
497  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)
498 
499  def get_restraint(self):
500  return self.rs
501 
502  def set_weight(self,weight):
503  self.weight=weight
504  self.rs.set_weigth(self.weight)
505 
506  def get_output(self):
507  output = {}
508  score = self.weight*self.rs.unprotected_evaluate(None)
509  output["_TotalScore"] = str(score)
510  output["ElectronMicroscopy2D_" + self.label] = str(score)
511  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: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
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:69
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:726
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