IMP  2.4.0
The Integrative Modeling Platform
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.base
9 import IMP.algebra
10 import IMP.atom
11 import IMP.container
12 import IMP.isd
13 
14 
15 class GaussianEMRestraint(object):
16 
17  def __init__(self, densities,
18  target_fn='',
19  target_ps=[],
20  cutoff_dist_model_model=0.0,
21  cutoff_dist_model_data=0.0,
22  overlap_threshold=0.0,
23  target_mass_scale=1.0,
24  target_radii_scale=3.0,
25  model_radii_scale=1.0,
26  slope=0.0,
27  spherical_gaussians=False,
28  pointwise_restraint=False,
29  local_mm=False,
30  close_pair_container=None,
31  backbone_slope=False,
32  scale_target_to_mass=False,
33  weight=1.0):
34  global sys, tools
35  import sys
36  import IMP.pmi.tools as tools
37  import IMP.isd.gmm_tools
38  from math import sqrt
39 
40 
41  # some parameters
42  self.label = "None"
43  self.sigmaissampled = False
44  self.sigmamaxtrans = 0.3
45  self.sigmamin = 1.0
46  self.sigmamax = 100.0
47  self.sigmainit = 2.0
48  self.label = "None"
49  self.densities = densities
50 
51  # setup target GMM
52  self.m = self.densities[0].get_model()
53  if scale_target_to_mass:
54  target_mass_scale = sum((IMP.atom.Mass(p).get_mass() for h in densities for p in IMP.atom.get_leaves(h)))
55  print('will scale target mass by', target_mass_scale)
56  if target_fn != '':
57  self.target_ps = []
58  IMP.isd.gmm_tools.decorate_gmm_from_text(
59  target_fn,
60  self.target_ps,
61  self.m,
62  radius_scale=target_radii_scale,
63  mass_scale=target_mass_scale)
64  elif target_ps != []:
65  self.target_ps = target_ps
66  else:
67  print('Gaussian EM restraint: must provide target density file or properly set up target densities')
68  return
69 
70  # setup model GMM
71  self.model_ps = []
72  for h in self.densities:
73  self.model_ps += IMP.atom.get_leaves(h)
74  if model_radii_scale != 1.0:
75  for p in self.model_ps:
76  rmax = sqrt(max(IMP.core.Gaussian(p).get_variances())) * \
77  model_radii_scale
78  if not IMP.core.XYZR.get_is_setup(p):
79  IMP.core.XYZR.setup_particle(p, rmax)
80  else:
81  IMP.core.XYZR(p).set_radius(rmax)
82 
83  # sigma particle
84  self.sigmaglobal = tools.SetupNuisance(self.m, self.sigmainit,
85  self.sigmamin, self.sigmamax,
86  self.sigmaissampled).get_particle()
87 
88  # create restraint
89  print('target num particles', len(self.target_ps), \
90  'total weight', sum([IMP.atom.Mass(p).get_mass()
91  for p in self.target_ps]))
92  print('model num particles', len(self.model_ps), \
93  'total weight', sum([IMP.atom.Mass(p).get_mass()
94  for p in self.model_ps]))
95 
96  update_model=not spherical_gaussians
97  log_score=False
98  if not pointwise_restraint:
99  self.gaussianEM_restraint = \
100  IMP.isd.GaussianEMRestraint(self.m,
101  IMP.get_indexes(self.model_ps),
102  IMP.get_indexes(self.target_ps),
103  self.sigmaglobal.get_particle().get_index(),
104  cutoff_dist_model_model,
105  cutoff_dist_model_data,
106  slope,
107  update_model, backbone_slope)
108  else:
109  print('USING POINTWISE RESTRAINT - requires isd_emxl')
110  import IMP.isd_emxl
111  print('update model?',update_model)
112  self.gaussianEM_restraint = \
113  IMP.isd_emxl.PointwiseGaussianEMRestraint(self.m,
114  IMP.get_indexes(self.model_ps),
115  IMP.get_indexes(self.target_ps),
116  self.sigmaglobal.get_particle().get_index(),
117  cutoff_dist_model_model,
118  cutoff_dist_model_data,
119  slope,
120  update_model,
121  log_score,
122  close_pair_container)
123 
124  print('done EM setup')
125  self.rs = IMP.RestraintSet(self.m, 'GaussianEMRestraint')
126  self.rs.add_restraint(self.gaussianEM_restraint)
127  self.set_weight(weight)
128 
129  def center_target_density_on_model(self):
130  target_com = IMP.algebra.Vector3D(0, 0, 0)
131  target_mass = 0.0
132  for p in self.target_ps:
133  mass = IMP.atom.Mass(p).get_mass()
134  pos = IMP.core.XYZ(p).get_coordinates()
135  target_com += pos * mass
136  target_mass += mass
137  target_com /= target_mass
138  print('target com', target_com)
139  model_com = IMP.algebra.Vector3D(0, 0, 0)
140  model_mass = 0.0
141  for p in self.model_ps:
142  mass = IMP.atom.Mass(p).get_mass()
143  pos = IMP.core.XYZ(p).get_coordinates()
144  model_com += pos * mass
145  model_mass += mass
146  model_com /= model_mass
147  print('model com', model_com)
148 
149  v = target_com - model_com
150  print('translating with', -v)
151  transformation = IMP.algebra.Transformation3D(IMP.algebra.Vector3D(-v))
152  for p in self.target_ps:
153  IMP.core.transform(IMP.core.RigidBody(p), transformation)
154  # IMP.pmi.tools.translate_hierarchies(self.densities,v)
155 
156  def center_model_on_target_density(self, representation):
157  target_com = IMP.algebra.Vector3D(0, 0, 0)
158  target_mass = 0.0
159  for p in self.target_ps:
160  mass = IMP.atom.Mass(p).get_mass()
161  pos = IMP.core.XYZ(p).get_coordinates()
162  target_com += pos * mass
163  target_mass += mass
164  target_com /= target_mass
165  print('target com', target_com)
166  model_com = IMP.algebra.Vector3D(0, 0, 0)
167  model_mass = 0.0
168  for p in self.model_ps:
169  mass = IMP.atom.Mass(p).get_mass()
170  pos = IMP.core.XYZ(p).get_coordinates()
171  model_com += pos * mass
172  model_mass += mass
173  model_com /= model_mass
174  print('model com', model_com)
175 
176  v = target_com - model_com
177  print('translating with', v)
178  transformation = IMP.algebra.Transformation3D(IMP.algebra.Vector3D(v))
179 
180  rigid_bodies = set()
181  XYZRs = set()
182 
183  for p in IMP.atom.get_leaves(representation.prot):
184  if IMP.core.RigidMember.get_is_setup(p):
185  rb = IMP.core.RigidMember(p).get_rigid_body()
186  rigid_bodies.add(rb)
187  elif IMP.core.NonRigidMember.get_is_setup(p):
188  rb = IMP.core.RigidMember(p).get_rigid_body()
189  rigid_bodies.add(rb)
190  elif IMP.core.XYZR.get_is_setup(p):
191  XYZRs.add(p)
192 
193  for rb in list(rigid_bodies):
194  IMP.core.transform(rb, transformation)
195 
196  for p in list(XYZRs):
197  IMP.core.transform(IMP.core.XYZ(p), transformation)
198 
199 
200  def set_weight(self,weight):
201  self.weight = weight
202  self.rs.set_weight(weight)
203 
204  def set_label(self, label):
205  self.label = label
206 
207  def add_to_model(self):
208  self.m.add_restraint(self.rs)
209 
210  def get_particles_to_sample(self):
211  ps = {}
212  if self.sigmaissampled:
213  ps["Nuisances_GaussianEMRestraint_sigma_" +
214  self.label] = ([self.sigmaglobal], self.sigmamaxtrans)
215  return ps
216 
217  def get_hierarchy(self):
218  return self.prot
219 
220  def get_density_as_hierarchy(self):
221  f=IMP.atom.Fragment().setup_particle(IMP.Particle(self.m))
222  f.set_name("GaussianEMRestraint_density_"+self.label)
223  for p in self.target_ps:
224  f.add_child(p)
225  return f
226 
227  def get_restraint_set(self):
228  return self.rs
229 
230  def get_output(self):
231  self.m.update()
232  output = {}
233  score = self.weight * self.rs.unprotected_evaluate(None)
234  output["_TotalScore"] = str(score)
235  output["GaussianEMRestraint_" +
236  self.label] = str(score)
237  output["GaussianEMRestraint_sigma_" +
238  self.label] = str(self.sigmaglobal.get_scale())
239  return output
240 
241  def evaluate(self):
242  return self.weight * self.rs.unprotected_evaluate(None)
243 
244 
245 #-------------------------------------------
246 
247 class ElectronMicroscopy2D(object):
248 
249  def __init__(
250  self,
251  representation,
252  images,
253  resolution=None):
254 
255  self.weight=1.0
256  self.m = representation.prot.get_model()
257  self.rs = IMP.RestraintSet(self.m, 'em2d')
258  self.label = "None"
259 
260  # IMP.atom.get_by_type
261  particles = IMP.pmi.tools.select(
262  representation,
263  resolution=resolution)
264 
265  em2d = None
266  self.rs.add_restraint(em2d)
267 
268  def set_label(self, label):
269  self.label = label
270 
271  def add_to_model(self):
272  self.m.add_restraint(self.rs)
273 
274  def get_restraint(self):
275  return self.rs
276 
277  def set_weight(self,weight):
278  self.weight=weight
279  self.rs.set_weigth(self.weight)
280 
281  def get_output(self):
282  self.m.update()
283  output = {}
284  score = self.weight*self.rs.unprotected_evaluate(None)
285  output["_TotalScore"] = str(score)
286  output["ElectronMicroscopy2D_" + self.label] = str(score)
287  return output
IMP::isd.gmm_tools
Tools for handling Gaussian Mixture Models.
Definition: gmm_tools.py:1
IMP::atom::Mass
Add mass to a particle.
Definition: Mass.h:23
IMP::algebra::Transformation3D
Simple 3D transformation class.
Definition: Transformation3D.h:30
IMP::atom::Fragment
A decorator to associate a particle with a part of a protein/DNA/RNA.
Definition: Fragment.h:20
IMP::domino::get_index
Ints get_index(const kernel::ParticlesTemp &particles, const Subset &subset, const Subsets &excluded)
IMP::container
Various classes to hold sets of particles.
IMP::atom::get_mass
double get_mass(ResidueType c)
Get the mass from the residue type.
IMP::pmi.tools
Miscellaneous utilities.
Definition: tools.py:1
IMP::kernel::RestraintSet
Object used to hold a set of restraints.
Definition: kernel/RestraintSet.h:36
IMP::base
Low level functionality (logging, error handling, profiling, command line flags etc) that is used by ...
IMP::core::NonRigidMember::get_is_setup
static bool get_is_setup(const IMP::kernel::ParticleAdaptor &p)
Definition: rigid_bodies.h:500
IMP::core::XYZR::get_is_setup
static bool get_is_setup(const IMP::kernel::ParticleAdaptor &p)
Definition: XYZR.h:47
IMP::core::XYZR::setup_particle
static XYZR setup_particle(kernel::Model *m, ParticleIndex pi)
Definition: XYZR.h:48
IMP::core::transform
void transform(XYZ a, const algebra::Transformation3D &tr)
Apply a transformation to the particle.
IMP::isd::GaussianEMRestraint
Creates a restraint between two Gaussian Mixture Models, "model" and "density".
Definition: GaussianEMRestraint.h:77
IMP::core::XYZ
A decorator for a particle with x,y,z coordinates.
Definition: XYZ.h:30
IMP::kernel::Particle
Class to handle individual model particles.
Definition: kernel/Particle.h:37
IMP::core
Basic functionality that is expected to be used by a wide variety of IMP users.
IMP::algebra
General purpose algebraic and geometric methods that are expected to be used by a wide variety of IMP...
IMP::core::RigidMember::get_is_setup
static bool get_is_setup(const IMP::kernel::ParticleAdaptor &p)
Definition: rigid_bodies.h:479
IMP::algebra::Vector3D
VectorD< 3 > Vector3D
Definition: VectorD.h:395
IMP::pmi.tools.select
def select
this function uses representation=SimplifiedModel it returns the corresponding selected particles rep...
Definition: tools.py:633
IMP::core::RigidBody
A decorator for a rigid body.
Definition: rigid_bodies.h:75
IMP::atom
Functionality for loading, creating, manipulating and scoring atomic structures.
IMP::atom::get_leaves
Hierarchies get_leaves(const Selection &h)
IMP::isd.gmm_tools.decorate_gmm_from_text
def decorate_gmm_from_text
read the output from write_gmm_to_text, decorate as Gaussian and Mass
Definition: gmm_tools.py:22
IMP::isd
Inferential scoring building on methods developed as part of the Inferential Structure Determination ...
IMP::core::XYZR
A decorator for a particle with x,y,z coordinates and a radius.
Definition: XYZR.h:27