IMP logo
IMP Reference Guide  develop.374319d7bb,2024/06/17
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 import IMP
6 import IMP.core
7 import IMP.algebra
8 import IMP.atom
9 import IMP.container
10 import IMP.isd
11 import IMP.pmi.tools
12 import IMP.pmi.mmcif
13 import IMP.isd.gmm_tools
14 from math import sqrt
15 
16 
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_mass=None,
29  target_radii_scale=3.0,
30  model_radii_scale=1.0,
31  slope=0.0,
32  spherical_gaussians=False,
33  close_pair_container=None,
34  backbone_slope=False,
35  scale_target_to_mass=False,
36  weight=1.0,
37  target_is_rigid_body=False,
38  local=False):
39  """Constructor.
40  @param densities The Gaussian-decorated particles to be restrained
41  @param target_fn GMM file of the target density map
42  (alternatively, pass the ps)
43  @param target_ps List of Gaussians of the target map
44  (alternatively, pass the filename)
45  @param cutoff_dist_model_model Distance in model-model close
46  pair container
47  @param cutoff_dist_model_data Distance in model-data close pair
48  container. Usually can set to zero because we multiply the
49  target radii
50  @param target_mass_scale Scale up the target densities so that
51  the mass is accurate.
52  Needed if the GMM you generated was not already scaled.
53  To make it the same as model mass, set scale_to_target_mass=True
54  @param target_mass Sets the mass of the target density to the given
55  value. Default is None. This will override target_mass_scale
56  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
75  density particles into a rigid body that need to be sampled
76  (e.g.,when you need to fit one density
77  against another one). Default is False.
78  @param local Only consider density particles that are within the
79  specified model-density cutoff (experimental)
80  """
81 
82  # some parameters
83  self.label = "None"
84  self.sigmaissampled = False
85  self.sigmamaxtrans = 0.3
86  self.sigmamin = 1.0
87  self.sigmamax = 100.0
88  self.sigmainit = 1.0
89  self.label = "None"
90  self.densities = densities
91  self.em_root_hier = None
92 
93  # setup target GMM
94  self.m = self.densities[0].get_model()
95 
96  if scale_target_to_mass:
97  def hierarchy_mass(h):
98  leaves = IMP.atom.get_leaves(h)
99  return sum(IMP.atom.Mass(p).get_mass() for p in leaves)
100  target_mass = sum(hierarchy_mass(h) for h in densities)
101  print('will set target mass to', target_mass)
102  print('will scale target mass by', target_mass_scale)
103 
104  if target_fn != '':
105  self._set_dataset(target_fn)
106  self.target_ps = []
108  target_fn, self.target_ps, self.m,
109  radius_scale=target_radii_scale,
110  mass_scale=target_mass_scale)
111  elif target_ps != []:
112  self.target_ps = target_ps
113  else:
114  print('Gaussian EM restraint: must provide target density file '
115  'or properly set up target densities')
116  return
117 
118  if target_mass:
119  tmass = sum([IMP.atom.Mass(p).get_mass() for p in self.target_ps])
120  scale = target_mass/tmass
121  for p in self.target_ps:
122  ms = IMP.atom.Mass(p).get_mass()
123  IMP.atom.Mass(p).set_mass(ms*scale)
124 
125  for p, state in IMP.pmi.tools._all_protocol_outputs(densities[0]):
126  p.add_em3d_restraint(state, self.target_ps, self.densities,
127  self)
128 
129  # setup model GMM
130  self.model_ps = []
131  for h in self.densities:
132  self.model_ps += [k.get_particle() for k in IMP.atom.get_leaves(h)]
133  if model_radii_scale != 1.0:
134  for p in self.model_ps:
135  rmax = sqrt(max(IMP.core.Gaussian(p).get_variances())) * \
136  model_radii_scale
137  if not IMP.core.XYZR.get_is_setup(p):
139  else:
140  IMP.core.XYZR(p).set_radius(rmax)
141  # wrap target particles in rigid body if requested
142  if target_is_rigid_body:
143  self.rb = IMP.atom.create_rigid_body(self.target_ps)
144  else:
145  self.rb = None
146 
147  # sigma particle
148  self.sigmaglobal = IMP.pmi.tools.SetupNuisance(
149  self.m, self.sigmainit, self.sigmamin, self.sigmamax,
150  self.sigmaissampled).get_particle()
151 
152  # create restraint
153  print('target num particles', len(self.target_ps),
154  'total weight', sum(IMP.atom.Mass(p).get_mass()
155  for p in self.target_ps))
156  print('model num particles', len(self.model_ps),
157  'total weight', sum(IMP.atom.Mass(p).get_mass()
158  for p in self.model_ps))
159 
160  update_model = not spherical_gaussians
161  self.gaussianEM_restraint = IMP.isd.GaussianEMRestraint(
162  self.m,
163  IMP.get_indexes(self.model_ps),
164  IMP.get_indexes(self.target_ps),
165  self.sigmaglobal.get_particle().get_index(),
166  cutoff_dist_model_model,
167  cutoff_dist_model_data,
168  slope,
169  update_model, backbone_slope, local)
170  if target_fn != '':
171  self.gaussianEM_restraint.set_density_filename(target_fn)
172 
173  print('done EM setup')
174  self.rs = IMP.RestraintSet(self.m, 'GaussianEMRestraint')
175  self.rs.add_restraint(self.gaussianEM_restraint)
176  self.set_weight(weight)
177 
178  def _set_dataset(self, target_fn):
179  """Set the dataset to point to the input file"""
181  self.dataset = p.parse_file(target_fn)['dataset']
182 
183  def center_target_density_on_model(self):
184  target_com = IMP.algebra.Vector3D(0, 0, 0)
185  target_mass = 0.0
186  for p in self.target_ps:
187  mass = IMP.atom.Mass(p).get_mass()
188  pos = IMP.core.XYZ(p).get_coordinates()
189  target_com += pos * mass
190  target_mass += mass
191  target_com /= target_mass
192  print('target com', target_com)
193  model_com = IMP.algebra.Vector3D(0, 0, 0)
194  model_mass = 0.0
195  for p in self.model_ps:
196  mass = IMP.atom.Mass(p).get_mass()
197  pos = IMP.core.XYZ(p).get_coordinates()
198  model_com += pos * mass
199  model_mass += mass
200  model_com /= model_mass
201  print('model com', model_com)
202 
203  v = target_com - model_com
204  print('translating with', -v)
206  for p in self.target_ps:
207  IMP.core.transform(IMP.core.RigidBody(p), transformation)
208  # IMP.pmi.tools.translate_hierarchies(self.densities,v)
209 
210  def get_center_of_mass(self, target=True):
211  '''Returns the geometric center of the GMM particles
212  @param target = True - returns target map gmm COM
213  @param target = False - returns model gmm COM'''
214  com = IMP.algebra.Vector3D(0, 0, 0)
215  total_mass = 0.0
216  if target:
217  ps = self.target_ps
218  else:
219  ps = self.model_ps
220  for p in ps:
221  mass = IMP.atom.Mass(p).get_mass()
222  pos = IMP.core.XYZ(p).get_coordinates()
223  com += pos * mass
224  total_mass += mass
225  com /= total_mass
226  return com
227 
228  def center_target_density_on_origin(self):
229  target_com = self.get_center_of_mass()
230  print('target com', target_com)
231  model_com = IMP.algebra.Vector3D(0, 0, 0)
232  print('model com', model_com)
233  v = target_com - model_com
234  print('translating with', -v)
236  for p in self.target_ps:
237  IMP.core.transform(IMP.core.RigidBody(p), transformation)
238  # IMP.pmi.tools.translate_hierarchies(self.densities,v)
239 
240  def center_model_on_target_density(self, input_object):
241  hier = input_object.get_hierarchy()
242  target_com = self.get_center_of_mass()
243  print('target com', target_com)
244  model_com = self.get_center_of_mass(target=False)
245  print('model com', model_com)
246  v = target_com - model_com
247  print('translating with', v)
249 
250  rigid_bodies = set()
251  XYZRs = set()
252 
253  for p in IMP.atom.get_leaves(hier):
255  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
256  rigid_bodies.add(rb)
258  XYZRs.add(p)
259 
260  for rb in list(rigid_bodies):
261  IMP.core.transform(rb, transformation)
262 
263  for p in list(XYZRs):
264  IMP.core.transform(IMP.core.XYZ(p), transformation)
265 
266  def center_on_target_density(self):
267  target_com = self.get_center_of_mass()
268  print('target com', target_com)
269  model_com = self.get_center_of_mass(target=False)
270  print('model com', model_com)
271  v = target_com - model_com
272  print('translating with', v)
274 
275  rigid_bodies = set()
276  XYZRs = set()
277 
278  for p in self.model_ps:
280  rb = IMP.core.RigidBodyMember(p).get_rigid_body()
281  rigid_bodies.add(rb)
283  XYZRs.add(p)
284 
285  for rb in list(rigid_bodies):
286  IMP.core.transform(rb, transformation)
287 
288  for p in list(XYZRs):
289  IMP.core.transform(IMP.core.XYZ(p), transformation)
290 
291  def set_weight(self, weight):
292  self.weight = weight
293  self.rs.set_weight(weight)
294 
295  def set_label(self, label):
296  self.label = label
297 
298  def add_to_model(self):
299  IMP.pmi.tools.add_restraint_to_model(self.m, self.rs, add_to_rmf=True)
300 
301  def get_particles_to_sample(self):
302  ps = {}
303  if self.sigmaissampled:
304  ps["Nuisances_GaussianEMRestraint_sigma_" +
305  self.label] = ([self.sigmaglobal], self.sigmamaxtrans)
306  if self.rb:
307  ps["Rigid_Bodies_GaussianEMRestraint"] = (
308  [self.rb],
309  4,
310  0.03)
311  return ps
312 
313  def get_rigid_body(self):
314  if self.rb is None:
315  raise Exception("No rigid body created for GMM particles. Ensure "
316  "target_is_rigid_body is set to True")
317  return self.rb
318 
319  def get_density_as_hierarchy(self):
320  if self.em_root_hier is None:
321  self.em_root_hier = IMP.atom.Copy.setup_particle(
322  IMP.Particle(self.m), 0)
323  self.em_root_hier.set_name("GaussianEMRestraint_density_"
324  + self.label)
325  for p in self.target_ps:
326  self.em_root_hier.add_child(p)
327  return self.em_root_hier
328 
330  ''' Can add a target GMM to a Hierarchy.
331  For PMI2 a state object may also be passed'''
332  if type(inp) is IMP.pmi.topology.State:
333  inp.get_hierarchy().add_child(self.get_density_as_hierarchy())
334  elif type(inp) is IMP.atom.Hierarchy:
335  inp.add_child(self.get_density_as_hierarchy())
336  else:
337  raise Exception(
338  "Can only add a density to a PMI State object or "
339  "IMP.atom.Hierarchy. You passed a", type(inp))
340 
341  def get_restraint(self):
342  return self.rs
343 
344  def get_restraint_set(self):
345  return self.rs
346 
347  def get_output(self):
348  output = {}
349  score = self.weight * self.rs.unprotected_evaluate(None)
350  ccc = self.gaussianEM_restraint.get_cross_correlation_coefficient()
351 
352  output["_TotalScore"] = str(score)
353  output["GaussianEMRestraint_" +
354  self.label] = str(score)
355  output["GaussianEMRestraint_%s_CCC" % self.label] = ccc
356  output["GaussianEMRestraint_sigma_" +
357  self.label] = str(self.sigmaglobal.get_scale())
358  return output
359 
360  def get_likelihood(self):
361  """Get the unweighted likelihood of the restraint"""
362  likelihood = self.gaussianEM_restraint.get_probability()
363  return likelihood
364 
365  def evaluate(self):
366  return self.weight * self.rs.unprotected_evaluate(None)
367 
368  def write_target_gmm_to_mrc(self, fileout=None, voxel_size=5.0):
369  '''Writes target GMM file to MRC'''
370  if fileout is None:
371  fileout = "Gaussian_map_" + self.label + ".mrc"
372  IMP.isd.gmm_tools.write_gmm_to_map(self.target_ps, fileout, voxel_size)
373  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:1715
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:84
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:408
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 get_likelihood
Get the unweighted likelihood of the restraint.
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