IMP logo
IMP Reference Guide  develop.031dafb4d2,2024/05/21
The Integrative Modeling Platform
lib/IMP/pmi1/restraints/em.py
1 """@namespace IMP.pmi1.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.pmi1.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.pmi1.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(r'(.*\.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.pmi1.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.pmi1.tools.translate_hierarchies(self.densities,v)
260 
261  def center_model_on_target_density(self, input_object):
262  if type(input_object) is IMP.pmi1.representation.Representation:
263  hier = input_object.prot
264  else:
265  raise Exception("Input must be a Representation object")
266  target_com = self.get_center_of_mass()
267  print('target com', target_com)
268  model_com = self.get_center_of_mass(target=False)
269  print('model com', model_com)
270  v = target_com - model_com
271  print('translating with', v)
273 
274  rigid_bodies = set()
275  XYZRs = set()
276 
277  for p in IMP.atom.get_leaves(hier):
279  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
280  rigid_bodies.add(rb)
282  XYZRs.add(p)
283 
284  for rb in list(rigid_bodies):
285  IMP.core.transform(rb, transformation)
286 
287  for p in list(XYZRs):
288  IMP.core.transform(IMP.core.XYZ(p), transformation)
289 
290  def center_on_target_density(self):
291  target_com = self.get_center_of_mass()
292  print('target com', target_com)
293  model_com = self.get_center_of_mass(target=False)
294  print('model com', model_com)
295  v = target_com - model_com
296  print('translating with', v)
298 
299  rigid_bodies = set()
300  XYZRs = set()
301 
302  for p in self.model_ps:
304  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
305  rigid_bodies.add(rb)
307  XYZRs.add(p)
308 
309  for rb in list(rigid_bodies):
310  IMP.core.transform(rb, transformation)
311 
312  for p in list(XYZRs):
313  IMP.core.transform(IMP.core.XYZ(p), transformation)
314 
315 
316 
317 
318  def set_weight(self,weight):
319  self.weight = weight
320  self.rs.set_weight(weight)
321 
322  def set_label(self, label):
323  self.label = label
324 
325  def add_to_model(self):
326  IMP.pmi1.tools.add_restraint_to_model(self.m, self.rs, add_to_rmf=True)
327 
328  def get_particles_to_sample(self):
329  ps = {}
330  if self.sigmaissampled:
331  ps["Nuisances_GaussianEMRestraint_sigma_" +
332  self.label] = ([self.sigmaglobal], self.sigmamaxtrans)
333  if self.rb:
334  ps["Rigid_Bodies_GaussianEMRestraint"] = (
335  [self.rb],
336  4,
337  0.03)
338  return ps
339 
340  def get_rigid_body(self):
341  if self.rb is None:
342  raise Exception("No rigid body created for GMM particles. Ensure target_is_rigid_body is set to True")
343  return self.rb
344 
345  def get_density_as_hierarchy(self):
346  if self.em_root_hier is None:
347  self.em_root_hier = IMP.atom.Copy.setup_particle(IMP.Particle(self.m),0)
348  self.em_root_hier.set_name("GaussianEMRestraint_density_"+self.label)
349  for p in self.target_ps:
350  self.em_root_hier.add_child(p)
351  return self.em_root_hier
352 
354  ''' Can add a target GMM to a Hierarchy.
355  For PMI2 a state object may also be passed'''
356  if type(inp) is IMP.atom.Hierarchy:
357  inp.add_child(self.get_density_as_hierarchy())
358  else:
359  raise Exception("Can only add a density to an IMP.atom.Hierarchy. You passed a", type(inp))
360 
361  def get_restraint(self):
362  return self.rs
363 
364  def get_restraint_set(self):
365  return self.rs
366 
367  def get_output(self):
368  output = {}
369  score = self.weight * self.rs.unprotected_evaluate(None)
370  ccc = self.gaussianEM_restraint.get_cross_correlation_coefficient()
371 
372  output["_TotalScore"] = str(score)
373  output["GaussianEMRestraint_" +
374  self.label] = str(score)
375  output["GaussianEMRestraint_%s_CCC" % self.label] = ccc
376  output["GaussianEMRestraint_sigma_" +
377  self.label] = str(self.sigmaglobal.get_scale())
378  return output
379 
380  def evaluate(self):
381  return self.weight * self.rs.unprotected_evaluate(None)
382 
383  def write_target_gmm_to_mrc(self, fileout=None, voxel_size=5.0):
384  '''Writes target GMM file to MRC'''
385  if fileout is None:
386  fileout="Gaussian_map_" + self.label + ".mrc"
387  IMP.isd.gmm_tools.write_gmm_to_map(self.target_ps, fileout, voxel_size)
388  return fileout
389 
390 
391 #-------------------------------------------
392 
394  """Fit particles to an EM map. This creates a simulate density map and updates them every eval.
395  \note Wraps an em::FitRestraint
396  """
397  def __init__(self,
398  ps,
399  dmap,
400  resolution,
401  origin=None,
402  voxel_size=None,
403  weight=1.0,
404  label=""):
405  """Constructor
406  @param ps The particles to restrain. Currently these must be atomic particles.
407  @param map_fn The EM density map to fit to
408  @param resolution Map resolution
409  @param origin In case you need to tell IMP the correct origin
410  @param voxel_size In case you need to tell IMP the angstroms per pixel
411  @param weight The data weight
412  @param label Extra PMI label
413  """
414  print('FitRestraint: setup')
415  #print('\tmap_fn',map_fn)
416  print('\tresolution',resolution)
417  print('\tvoxel_size',voxel_size)
418  print('\torigin',origin)
419  print('\tweight',weight)
420 
421  # some parameters
422  self.mdl = ps[0].get_model()
423  self.label = label
424  self.dmap = dmap #IMP.em.read_map(map_fn,IMP.em.MRCReaderWriter())
425  #dh = self.dmap.get_header()
426  #dh.set_resolution(resolution)
427  if voxel_size:
428  self.dmap.update_voxel_size(voxel_size)
429  if origin is not None:
430  if type(origin)==IMP.algebra.Vector3D:
431  self.dmap.set_origin(origin)
432  elif type(origin)==list:
433  self.dmap.set_origin(*origin)
434  else:
435  print('FitRestraint did not recognize format of origin')
436  exit()
437  fr = IMP.em.FitRestraint(ps,self.dmap)
438  self.rs = IMP.RestraintSet(self.mdl,weight,"FitRestraint")
439  self.rs.add_restraint(fr)
440  self.set_weight(weight)
441 
442  def set_weight(self,weight):
443  self.weight = weight
444  self.rs.set_weight(weight)
445 
446  def set_label(self, label):
447  self.label = label
448 
449  def add_to_model(self):
451 
452  def get_restraint_set(self):
453  return self.rs
454 
455  def get_output(self):
456  output = {}
457  score = self.weight * self.rs.unprotected_evaluate(None)
458  output["_TotalScore"] = str(score)
459  output["EMRestraint_" + self.label] = str(score)
460  return output
461 
462  def evaluate(self):
463  return self.weight * self.rs.unprotected_evaluate(None)
464 
465 
466 #-------------------------------------------
467 
468 class ElectronMicroscopy2D(object):
469 
470  def __init__(
471  self,
472  representation,
473  images,
474  resolution=None):
475 
476  self.weight=1.0
477  self.m = representation.prot.get_model()
478  self.rs = IMP.RestraintSet(self.m, 'em2d')
479  self.label = "None"
480 
481  # IMP.atom.get_by_type
482  particles = IMP.pmi1.tools.select(
483  representation,
484  resolution=resolution)
485 
486  em2d = None
487  self.rs.add_restraint(em2d)
488 
489  def set_label(self, label):
490  self.label = label
491 
492  def add_to_model(self):
494 
495  def get_restraint(self):
496  return self.rs
497 
498  def set_weight(self,weight):
499  self.weight=weight
500  self.rs.set_weigth(self.weight)
501 
502  def get_output(self):
503  output = {}
504  score = self.weight*self.rs.unprotected_evaluate(None)
505  output["_TotalScore"] = str(score)
506  output["ElectronMicroscopy2D_" + self.label] = str(score)
507  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.
A member of a rigid body, it has internal (local) coordinates.
Definition: rigid_bodies.h:634
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: rigid_bodies.h:635
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.
def get_center_of_mass
Returns the geometric center of the GMM particles.
def write_target_gmm_to_mrc
Writes target GMM file to MRC.
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: XYZR.h:47
Object used to hold a set of restraints.
Definition: RestraintSet.h:41
Miscellaneous utilities.
Definition: /tools.py:1
ParticleIndexPairs get_indexes(const ParticlePairsTemp &ps)
Get the indexes from a list of particle pairs.
Set up the representation of all proteins and nucleic acid macromolecules.
The standard decorator for manipulating molecular structures.
Ints get_index(const ParticlesTemp &particles, const Subset &subset, const Subsets &excluded)
def add_target_density_to_hierarchy
Can add a target GMM to a Hierarchy.
void transform(XYZ a, const algebra::Transformation3D &tr)
Apply a transformation to the particle.
Restraint between two Gaussian Mixture Models, "model" and "density".
Fit Gaussian-decorated particles to an EM map (also represented with a set of Gaussians) ...
A decorator for a particle with x,y,z coordinates.
Definition: XYZ.h:30
def write_gmm_to_map
write density map from GMM.
Definition: gmm_tools.py:119
static Copy setup_particle(Model *m, ParticleIndex pi, Int number)
Create a decorator for the numberth copy.
Definition: Copy.h:42
def select
this function uses representation=SimplifiedModel it returns the corresponding selected particles rep...
Definition: /tools.py:727
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:48
VectorD< 3 > Vector3D
Definition: VectorD.h:408
IMP::core::RigidBody create_rigid_body(Hierarchy h)
Class to handle individual particles of a Model object.
Definition: Particle.h:43
A decorator for a rigid body.
Definition: rigid_bodies.h:82
def add_restraint_to_model
Add a PMI restraint to the model.
Definition: /tools.py:53
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