IMP logo
IMP Reference Guide  2.20.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.pmi.mmcif
14 import IMP.isd.gmm_tools
15 from math import sqrt
16 
17 
18 class GaussianEMRestraint(object):
19  """Fit Gaussian-decorated particles to an EM map
20  (also represented with a set of Gaussians)
21  @note This class wraps an isd::GaussianEMRestraint
22  """
23  def __init__(self, densities,
24  target_fn='',
25  target_ps=[],
26  cutoff_dist_model_model=0.0,
27  cutoff_dist_model_data=0.0,
28  target_mass_scale=1.0,
29  target_mass=None,
30  target_radii_scale=3.0,
31  model_radii_scale=1.0,
32  slope=0.0,
33  spherical_gaussians=False,
34  close_pair_container=None,
35  backbone_slope=False,
36  scale_target_to_mass=False,
37  weight=1.0,
38  target_is_rigid_body=False,
39  local=False):
40  """Constructor.
41  @param densities The Gaussian-decorated particles to be restrained
42  @param target_fn GMM file of the target density map
43  (alternatively, pass the ps)
44  @param target_ps List of Gaussians of the target map
45  (alternatively, pass the filename)
46  @param cutoff_dist_model_model Distance in model-model close
47  pair container
48  @param cutoff_dist_model_data Distance in model-data close pair
49  container. Usually can set to zero because we multiply the
50  target radii
51  @param target_mass_scale Scale up the target densities so that
52  the mass is accurate.
53  Needed if the GMM you generated was not already scaled.
54  To make it the same as model mass, set scale_to_target_mass=True
55  @param target_mass Sets the mass of the target density to the given
56  value. Default is None. This will override target_mass_scale
57  argument
58  @param target_radii_scale Scale the target density radii -
59  only used for the close pair container.
60  If you keep this at 3.0 or so you don't have to use cutoff dist.
61  @param model_radii_scale Scale the model density radii - only used
62  for the close pair container
63  @param slope Linear term added to help bring model into the density
64  @param spherical_gaussians Set to True for a speed bonus when
65  the model densities are spheres. (This means you don't have
66  to do a matrix multiplication if they rotate.)
67  @param close_pair_container Pass a close pair container for
68  the model if you already have one (e.g. for an excluded
69  volume restraint.) May give a speed bonus.
70  @param backbone_slope Only apply slope to backbone particles -
71  only matters for atomic
72  @param scale_target_to_mass Set True if you would need to scale
73  target to EXACTLY the model mass
74  @param weight The restraint weight
75  @param target_is_rigid_body Set True if you want to put the target
76  density particles into a rigid body that need to be sampled
77  (e.g.,when you need to fit one density
78  against another one). Default is False.
79  @param local Only consider density particles that are within the
80  specified model-density cutoff (experimental)
81  """
82 
83  # some parameters
84  self.label = "None"
85  self.sigmaissampled = False
86  self.sigmamaxtrans = 0.3
87  self.sigmamin = 1.0
88  self.sigmamax = 100.0
89  self.sigmainit = 1.0
90  self.label = "None"
91  self.densities = densities
92  self.em_root_hier = None
93 
94  # setup target GMM
95  self.m = self.densities[0].get_model()
96 
97  if scale_target_to_mass:
98  def hierarchy_mass(h):
99  leaves = IMP.atom.get_leaves(h)
100  return sum(IMP.atom.Mass(p).get_mass() for p in leaves)
101  target_mass = sum(hierarchy_mass(h) for h in densities)
102  print('will set target mass to', target_mass)
103  print('will scale target mass by', target_mass_scale)
104 
105  if target_fn != '':
106  self._set_dataset(target_fn)
107  self.target_ps = []
109  target_fn, self.target_ps, self.m,
110  radius_scale=target_radii_scale,
111  mass_scale=target_mass_scale)
112  elif target_ps != []:
113  self.target_ps = target_ps
114  else:
115  print('Gaussian EM restraint: must provide target density file '
116  'or properly set up target densities')
117  return
118 
119  if target_mass:
120  tmass = sum([IMP.atom.Mass(p).get_mass() for p in self.target_ps])
121  scale = target_mass/tmass
122  for p in self.target_ps:
123  ms = IMP.atom.Mass(p).get_mass()
124  IMP.atom.Mass(p).set_mass(ms*scale)
125 
126  for p, state in IMP.pmi.tools._all_protocol_outputs(densities[0]):
127  p.add_em3d_restraint(state, self.target_ps, self.densities,
128  self)
129 
130  # setup model GMM
131  self.model_ps = []
132  for h in self.densities:
133  self.model_ps += [k.get_particle() for k in IMP.atom.get_leaves(h)]
134  if model_radii_scale != 1.0:
135  for p in self.model_ps:
136  rmax = sqrt(max(IMP.core.Gaussian(p).get_variances())) * \
137  model_radii_scale
138  if not IMP.core.XYZR.get_is_setup(p):
140  else:
141  IMP.core.XYZR(p).set_radius(rmax)
142  # wrap target particles in rigid body if requested
143  if target_is_rigid_body:
144  self.rb = IMP.atom.create_rigid_body(self.target_ps)
145  else:
146  self.rb = None
147 
148  # sigma particle
149  self.sigmaglobal = IMP.pmi.tools.SetupNuisance(
150  self.m, self.sigmainit, self.sigmamin, self.sigmamax,
151  self.sigmaissampled).get_particle()
152 
153  # create restraint
154  print('target num particles', len(self.target_ps),
155  'total weight', sum(IMP.atom.Mass(p).get_mass()
156  for p in self.target_ps))
157  print('model num particles', len(self.model_ps),
158  'total weight', sum(IMP.atom.Mass(p).get_mass()
159  for p in self.model_ps))
160 
161  update_model = not spherical_gaussians
162  self.gaussianEM_restraint = IMP.isd.GaussianEMRestraint(
163  self.m,
164  IMP.get_indexes(self.model_ps),
165  IMP.get_indexes(self.target_ps),
166  self.sigmaglobal.get_particle().get_index(),
167  cutoff_dist_model_model,
168  cutoff_dist_model_data,
169  slope,
170  update_model, backbone_slope, local)
171  if target_fn != '':
172  self.gaussianEM_restraint.set_density_filename(target_fn)
173 
174  print('done EM setup')
175  self.rs = IMP.RestraintSet(self.m, 'GaussianEMRestraint')
176  self.rs.add_restraint(self.gaussianEM_restraint)
177  self.set_weight(weight)
178 
179  def _set_dataset(self, target_fn):
180  """Set the dataset to point to the input file"""
182  self.dataset = p.parse_file(target_fn)['dataset']
183 
184  def center_target_density_on_model(self):
185  target_com = IMP.algebra.Vector3D(0, 0, 0)
186  target_mass = 0.0
187  for p in self.target_ps:
188  mass = IMP.atom.Mass(p).get_mass()
189  pos = IMP.core.XYZ(p).get_coordinates()
190  target_com += pos * mass
191  target_mass += mass
192  target_com /= target_mass
193  print('target com', target_com)
194  model_com = IMP.algebra.Vector3D(0, 0, 0)
195  model_mass = 0.0
196  for p in self.model_ps:
197  mass = IMP.atom.Mass(p).get_mass()
198  pos = IMP.core.XYZ(p).get_coordinates()
199  model_com += pos * mass
200  model_mass += mass
201  model_com /= model_mass
202  print('model com', model_com)
203 
204  v = target_com - model_com
205  print('translating with', -v)
207  for p in self.target_ps:
208  IMP.core.transform(IMP.core.RigidBody(p), transformation)
209  # IMP.pmi.tools.translate_hierarchies(self.densities,v)
210 
211  def get_center_of_mass(self, target=True):
212  '''Returns the geometric center of the GMM particles
213  @param target = True - returns target map gmm COM
214  @param target = False - returns model gmm COM'''
215  com = IMP.algebra.Vector3D(0, 0, 0)
216  total_mass = 0.0
217  if target:
218  ps = self.target_ps
219  else:
220  ps = self.model_ps
221  for p in ps:
222  mass = IMP.atom.Mass(p).get_mass()
223  pos = IMP.core.XYZ(p).get_coordinates()
224  com += pos * mass
225  total_mass += mass
226  com /= total_mass
227  return com
228 
229  def center_target_density_on_origin(self):
230  target_com = self.get_center_of_mass()
231  print('target com', target_com)
232  model_com = IMP.algebra.Vector3D(0, 0, 0)
233  print('model com', model_com)
234  v = target_com - model_com
235  print('translating with', -v)
237  for p in self.target_ps:
238  IMP.core.transform(IMP.core.RigidBody(p), transformation)
239  # IMP.pmi.tools.translate_hierarchies(self.densities,v)
240 
241  def center_model_on_target_density(self, input_object):
242  hier = input_object.get_hierarchy()
243  target_com = self.get_center_of_mass()
244  print('target com', target_com)
245  model_com = self.get_center_of_mass(target=False)
246  print('model com', model_com)
247  v = target_com - model_com
248  print('translating with', v)
250 
251  rigid_bodies = set()
252  XYZRs = set()
253 
254  for p in IMP.atom.get_leaves(hier):
256  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
257  rigid_bodies.add(rb)
259  XYZRs.add(p)
260 
261  for rb in list(rigid_bodies):
262  IMP.core.transform(rb, transformation)
263 
264  for p in list(XYZRs):
265  IMP.core.transform(IMP.core.XYZ(p), transformation)
266 
267  def center_on_target_density(self):
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 self.model_ps:
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 set_weight(self, weight):
293  self.weight = weight
294  self.rs.set_weight(weight)
295 
296  def set_label(self, label):
297  self.label = label
298 
299  def add_to_model(self):
300  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs, add_to_rmf=True)
301 
302  def get_particles_to_sample(self):
303  ps = {}
304  if self.sigmaissampled:
305  ps["Nuisances_GaussianEMRestraint_sigma_" +
306  self.label] = ([self.sigmaglobal], self.sigmamaxtrans)
307  if self.rb:
308  ps["Rigid_Bodies_GaussianEMRestraint"] = (
309  [self.rb],
310  4,
311  0.03)
312  return ps
313 
314  def get_rigid_body(self):
315  if self.rb is None:
316  raise Exception("No rigid body created for GMM particles. Ensure "
317  "target_is_rigid_body is set to True")
318  return self.rb
319 
320  def get_density_as_hierarchy(self):
321  if self.em_root_hier is None:
322  self.em_root_hier = IMP.atom.Copy.setup_particle(
323  IMP.Particle(self.m), 0)
324  self.em_root_hier.set_name("GaussianEMRestraint_density_"
325  + self.label)
326  for p in self.target_ps:
327  self.em_root_hier.add_child(p)
328  return self.em_root_hier
329 
331  ''' Can add a target GMM to a Hierarchy.
332  For PMI2 a state object may also be passed'''
333  if type(inp) is IMP.pmi.topology.State:
334  inp.get_hierarchy().add_child(self.get_density_as_hierarchy())
335  elif type(inp) is IMP.atom.Hierarchy:
336  inp.add_child(self.get_density_as_hierarchy())
337  else:
338  raise Exception(
339  "Can only add a density to a PMI State object or "
340  "IMP.atom.Hierarchy. You passed a", type(inp))
341 
342  def get_restraint(self):
343  return self.rs
344 
345  def get_restraint_set(self):
346  return self.rs
347 
348  def get_output(self):
349  output = {}
350  score = self.weight * self.rs.unprotected_evaluate(None)
351  ccc = self.gaussianEM_restraint.get_cross_correlation_coefficient()
352 
353  output["_TotalScore"] = str(score)
354  output["GaussianEMRestraint_" +
355  self.label] = str(score)
356  output["GaussianEMRestraint_%s_CCC" % self.label] = ccc
357  output["GaussianEMRestraint_sigma_" +
358  self.label] = str(self.sigmaglobal.get_scale())
359  return output
360 
361  def evaluate(self):
362  return self.weight * self.rs.unprotected_evaluate(None)
363 
364  def write_target_gmm_to_mrc(self, fileout=None, voxel_size=5.0):
365  '''Writes target GMM file to MRC'''
366  if fileout is None:
367  fileout = "Gaussian_map_" + self.label + ".mrc"
368  IMP.isd.gmm_tools.write_gmm_to_map(self.target_ps, fileout, voxel_size)
369  return fileout
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:634
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: rigid_bodies.h:635
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:1724
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:41
ParticleIndexPairs get_indexes(const ParticlePairsTemp &ps)
Get the indexes from a list of particle pairs.
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:92
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:118
static Copy setup_particle(Model *m, ParticleIndex pi, Int number)
Create a decorator for the numberth copy.
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...
The general base class for IMP exceptions.
Definition: exception.h:48
VectorD< 3 > Vector3D
Definition: VectorD.h:425
IMP::core::RigidBody create_rigid_body(Hierarchy h)
Class to handle individual particles of a Model object.
Definition: Particle.h:43
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: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