IMP logo
IMP Reference Guide  2.20.2
The Integrative Modeling Platform
bayesianem/restraint.py
1 """@namespace IMP.bayesianem
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.bayesianem
12 import IMP.isd
13 import IMP.pmi.tools
14 import IMP.pmi.mmcif
15 import IMP.isd.gmm_tools
16 import sys
17 from math import sqrt
18 
20  """Fit Gaussian-decorated particles to an EM map
21  (also represented with a set of Gaussians)
22  \note This class wraps an IMP::bayesianem::GaussianEMRestraint
23  """
24  def __init__(self, densities,
25  target_fn='',
26  target_ps=[],
27  cutoff_dist_model_model=10.0,
28  cutoff_dist_model_data=10.0,
29  target_mass_scale=1.0,
30  target_mass=None,
31  target_radii_scale=3.0,
32  model_radii_scale=1.0,
33  slope=0.0,
34  spherical_gaussians=False,
35  close_pair_container=None,
36  backbone_slope=False,
37  scale_target_to_mass=False,
38  weight=1.0,
39  target_is_rigid_body=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 value. Default is None. This
56  will override target_mass_scale argument
57  @param target_radii_scale Scale the target density radii -
58  only used for the close pair container.
59  If you keep this at 3.0 or so you don't have to use cutoff dist.
60  @param model_radii_scale Scale the model density radii - only used
61  for the close pair container
62  @param slope Linear term added to help bring model into the density
63  @param spherical_gaussians Set to True for a speed bonus when
64  the model densities are spheres. (This means you don't have
65  to do a matrix multiplication if they rotate.)
66  @param close_pair_container Pass a close pair container for
67  the model if you already have one (e.g. for an excluded
68  volume restraint.) May give a speed bonus.
69  @param backbone_slope Only apply slope to backbone particles -
70  only matters for atomic
71  @param scale_target_to_mass Set True if you would need to scale
72  target to EXACTLY the model mass
73  @param weight The restraint weight
74  @param target_is_rigid_body Set True if you want to put the target density particles
75  into a rigid body that need to be sampled (e.g.,when you need to fit one density
76  against another one). Default is False.
77  """
78  # some parameters
79  self.label = "None"
80  self.sigmaissampled = True
81  self.sigmamaxtrans = 0.3
82  self.sigmamin = 0.0001
83  self.sigmamax = 100.0
84  self.sigmainit = 1.0
85  self.label = "None"
86  self.densities = densities
87  self.em_root_hier = None
88 
89  # setup target GMM
90  self.m = self.densities[0].get_model()
91 
92  if scale_target_to_mass:
93  def hierarchy_mass(h):
94  leaves = IMP.atom.get_leaves(h)
95  return sum(IMP.atom.Mass(p).get_mass() for p in leaves)
96  target_mass = sum(hierarchy_mass(h) for h in densities)
97  print('will scale target mass by', target_mass_scale)
98 
99  if target_fn != '':
100  self._set_dataset(target_fn)
101  self.target_ps = []
103  target_fn,
104  self.target_ps,
105  self.m,
106  radius_scale=target_radii_scale,
107  mass_scale=target_mass_scale)
108  elif target_ps != []:
109  self.target_ps = target_ps
110  else:
111  print('Gaussian EM restraint: must provide target density file or properly set up target densities')
112  return
113 
114  if target_mass:
115  tmass=sum([IMP.atom.Mass(p).get_mass() for p in self.target_ps])
116  scale=target_mass/tmass
117  print('will set target mass to', target_mass,tmass,scale)
118  for p in self.target_ps:
119  ms=IMP.atom.Mass(p).get_mass()
120  IMP.atom.Mass(p).set_mass(ms*scale)
121 
122  for p, state in IMP.pmi.tools._all_protocol_outputs(densities[0]):
123  p.add_em3d_restraint(state, self.target_ps, self.densities, self)
124 
125  # setup model GMM
126  self.model_ps = []
127  for h in self.densities:
128  self.model_ps += [ k.get_particle() for k in IMP.atom.get_leaves(h) ]
129  if model_radii_scale != 1.0:
130  for p in self.model_ps:
131  rmax = sqrt(max(IMP.core.Gaussian(p).get_variances())) * \
132  model_radii_scale
133  if not IMP.core.XYZR.get_is_setup(p):
135  else:
136  IMP.core.XYZR(p).set_radius(rmax)
137  #wrap target particles in rigid body if requested
138  if target_is_rigid_body:
139  #p = IMP.Particle(self.m)
140  #self.rb=IMP.core.RigidBody.setup_particle(p,self.target_ps)
141  self.rb=IMP.atom.create_rigid_body(self.target_ps)
142  else:
143  self.rb=None
144 
145  # sigma particle
146  self.sigmaglobal = IMP.pmi.tools.SetupNuisance(self.m, self.sigmainit,
147  self.sigmamin, self.sigmamax,
148  self.sigmaissampled).get_particle()
149 
150  # create restraint
151  print('target num particles', len(self.target_ps), \
152  'total weight', sum([IMP.atom.Mass(p).get_mass()
153  for p in self.target_ps]))
154  print('model num particles', len(self.model_ps), \
155  'total weight', sum([IMP.atom.Mass(p).get_mass()
156  for p in self.model_ps]))
157 
158  #Normalize
159  #
160  #mass_model=sum([IMP.atom.Mass(p).get_mass() for p in self.model_ps])
161  #for p in self.model_ps:
162  # pmass=IMP.atom.Mass(p).get_mass()
163  # newpmass=pmass/mass_model
164  # IMP.atom.Mass(p).set_mass(newpmass)
165 
166 
167  #mass_target=sum([IMP.atom.Mass(p).get_mass() for p in self.target_ps])
168  #for p in self.target_ps:
169  # pmass=IMP.atom.Mass(p).get_mass()
170  # newpmass=pmass/mass_target
171  # IMP.atom.Mass(p).set_mass(newpmass)
172 
173  update_model=not spherical_gaussians
174  log_score=False
175  self.gaussianEM_restraint = IMP.bayesianem.GaussianEMRestraint(
176  self.m,
177  IMP.get_indexes(self.model_ps),
178  IMP.get_indexes(self.target_ps),
179  self.sigmaglobal.get_particle().get_index(),
180  cutoff_dist_model_model,
181  cutoff_dist_model_data,
182  slope,
183  update_model, backbone_slope)
184  if target_fn != '':
185  self.gaussianEM_restraint.set_density_filename(target_fn)
186 
187  print('done EM setup')
188  self.rs = IMP.RestraintSet(self.m, 'GaussianEMRestraint')
189  self.rs.add_restraint(self.gaussianEM_restraint)
190  self.set_weight(weight)
191 
192  def _set_dataset(self, target_fn):
193  """Set the dataset to point to the input file"""
195  self.dataset = p.parse_file(target_fn)['dataset']
196 
198  '''
199  aligns the center of mass of the target GMM on the center of mass of the model
200  '''
201  target_com = IMP.algebra.Vector3D(0, 0, 0)
202  target_mass = 0.0
203  for p in self.target_ps:
204  mass = IMP.atom.Mass(p).get_mass()
205  pos = IMP.core.XYZ(p).get_coordinates()
206  target_com += pos * mass
207  target_mass += mass
208  target_com /= target_mass
209  print('target com', target_com)
210  model_com = IMP.algebra.Vector3D(0, 0, 0)
211  model_mass = 0.0
212  for p in self.model_ps:
213  mass = IMP.atom.Mass(p).get_mass()
214  pos = IMP.core.XYZ(p).get_coordinates()
215  model_com += pos * mass
216  model_mass += mass
217  model_com /= model_mass
218  print('model com', model_com)
219 
220  v = target_com - model_com
221  print('translating with', -v)
223  for p in self.target_ps:
224  IMP.core.transform(IMP.core.RigidBody(p), transformation)
225  # IMP.pmi.tools.translate_hierarchies(self.densities,v)
226 
227  def get_center_of_mass(self, target=True):
228  '''Returns the geometric center of the GMM particles
229  @param target = True - returns target map gmm COM
230  @param target = False - returns model gmm COM'''
231  com = IMP.algebra.Vector3D(0, 0, 0)
232  total_mass = 0.0
233  if target:
234  ps = self.target_ps
235  else:
236  ps = self.model_ps
237  for p in ps:
238  mass = IMP.atom.Mass(p).get_mass()
239  pos = IMP.core.XYZ(p).get_coordinates()
240  com += pos * mass
241  total_mass += mass
242  com /= total_mass
243  return com
244 
246  '''
247  aligns the center of mass of the target GMM on the origin
248  '''
249  target_com = self.get_center_of_mass()
250  print('target com', target_com)
251  model_com = IMP.algebra.Vector3D(0, 0, 0)
252  print('model com', model_com)
253  v = target_com - model_com
254  print('translating with', -v)
256  for p in self.target_ps:
257  IMP.core.transform(IMP.core.RigidBody(p), transformation)
258  # IMP.pmi.tools.translate_hierarchies(self.densities,v)
259 
260  def center_model_on_target_density(self, input_object):
261  '''
262  aligns the model on the target density
263  @param input_objects IMP.pmi.representation.Representation or IMP.pmi.topology.State
264  '''
265  if type(input_object) is IMP.pmi.representation.Representation:
266  hier = input_object.prot
267  elif type(input_object) is IMP.pmi.topology.State:
268  hier = input_object.get_hierarchy()
269  else:
270  raise Exception("Input must be a Representation or topology.State object")
271  target_com = self.get_center_of_mass()
272  print('target com', target_com)
273  model_com = self.get_center_of_mass(target=False)
274  print('model com', model_com)
275  v = target_com - model_com
276  print('translating with', v)
278 
279  rigid_bodies = set()
280  XYZRs = set()
281 
282  for p in IMP.atom.get_leaves(hier):
284  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
285  rigid_bodies.add(rb)
287  XYZRs.add(p)
288 
289  for rb in list(rigid_bodies):
290  IMP.core.transform(rb, transformation)
291 
292  for p in list(XYZRs):
293  IMP.core.transform(IMP.core.XYZ(p), transformation)
294 
296  '''
297  align the model on target GMM
298  '''
299  target_com = self.get_center_of_mass()
300  print('target com', target_com)
301  model_com = self.get_center_of_mass(target=False)
302  print('model com', model_com)
303  v = target_com - model_com
304  print('translating with', v)
306 
307  rigid_bodies = set()
308  XYZRs = set()
309 
310  for p in self.model_ps:
312  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
313  rigid_bodies.add(rb)
315  XYZRs.add(p)
316 
317  for rb in list(rigid_bodies):
318  IMP.core.transform(rb, transformation)
319 
320  for p in list(XYZRs):
321  IMP.core.transform(IMP.core.XYZ(p), transformation)
322 
323 
324  def set_weight(self,weight):
325  '''
326  set the weight of the restraint
327  @param weight
328  '''
329  self.weight = weight
330  self.rs.set_weight(weight)
331 
332  def set_label(self, label):
333  self.label = label
334 
335  def add_to_model(self):
336  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)
337 
338  def get_particles_to_sample(self):
339  ps = {}
340  if self.sigmaissampled:
341  ps["Nuisances_GaussianEMRestraint_sigma_" +
342  self.label] = ([self.sigmaglobal], self.sigmamaxtrans)
343  if self.rb:
344  ps["Rigid_Bodies_GaussianEMRestraint_"+self.label] = (
345  [self.rb],
346  4,
347  0.25)
348  return ps
349 
350  def get_rigid_body(self):
351  if self.rb is None:
352  raise Exception("No rigid body created for GMM particles. Ensure target_is_rigid_body is set to True")
353  return self.rb
354 
356  '''
357  returns a hierarchy whose leaves are the gaussian particles of the target GMM
358  '''
359  if self.em_root_hier is None:
360  self.em_root_hier = IMP.atom.Copy.setup_particle(IMP.Particle(self.m),0)
361  self.em_root_hier.set_name("GaussianEMRestraint_density_"+self.label)
362  for p in self.target_ps:
363  self.em_root_hier.add_child(p)
364  return self.em_root_hier
365 
367  ''' Can add a target GMM to a Hierarchy.
368  For PMI2 a state object may also be passed'''
369  if type(inp) is IMP.pmi.topology.State:
370  inp.get_hierarchy().add_child(self.get_density_as_hierarchy())
371  elif type(inp) is IMP.atom.Hierarchy:
372  inp.add_child(self.get_density_as_hierarchy())
373  else:
374  raise Exception("Can only add a density to a PMI State object or IMP.atom.Hierarchy. You passed a", type(inp))
375 
376  def get_restraint_set(self):
377  return self.rs
378 
379  def get_output(self):
380  self.m.update()
381  output = {}
382  score = self.weight * self.rs.unprotected_evaluate(None)
383  output["_TotalScore"] = str(score)
384  output["GaussianEMRestraint_" +
385  self.label] = str(score)
386  output["GaussianEMRestraint_sigma_" +
387  self.label] = str(self.sigmaglobal.get_scale())
388  return output
389 
390  def evaluate(self):
391  return self.weight * self.rs.unprotected_evaluate(None)
392 
393  def write_target_gmm_to_mrc(self, fileout=None, voxel_size=5.0):
394  '''Writes target GMM file to MRC'''
395  if fileout is None:
396  fileout="Gaussian_map_" + self.label + ".mrc"
397  IMP.isd.gmm_tools.write_gmm_to_map(self.target_ps, fileout, voxel_size)
398  return fileout
399 
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
def get_center_of_mass
Returns the geometric center of the GMM particles.
Restraints for handling electron microscopy maps.
A member of a rigid body, it has internal (local) coordinates.
Definition: rigid_bodies.h:634
def center_model_on_target_density
aligns the model on the target density
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
def center_target_density_on_origin
aligns the center of mass of the target GMM on the origin
double get_mass(ResidueType c)
Get the mass from the residue type.
def add_target_density_to_hierarchy
Can add a target GMM to a Hierarchy.
Miscellaneous utilities.
Definition: tools.py:1
Creates a restraint between two Gaussian Mixture Models, "model" and "density".
Extract metadata from an EM density GMM file.
Definition: mmcif.py:1724
Fit Gaussian-decorated particles to an EM map (also represented with a set of Gaussians) ...
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.
def add_restraint_to_model
Add a PMI restraint to the model.
Definition: tools.py:92
def center_on_target_density
align the model on target GMM
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: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...
def write_target_gmm_to_mrc
Writes target GMM file to MRC.
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
def center_target_density_on_model
aligns the center of mass of the target GMM on the center of mass of the model
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
def get_density_as_hierarchy
returns a hierarchy whose leaves are the gaussian particles of the target GMM
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