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