IMP  2.4.0
The Integrative Modeling Platform
new/em.py
1 #!/usr/bin/env python
2 from __future__ import print_function
3 import IMP
4 import IMP.atom
5 import IMP.em
6 import IMP.isd_emxl
7 import IMP.isd_emxl.gmm_tools
8 import IMP.pmi.tools as tools
9 import IMP.pmi.sampling_tools as sampling_tools
10 from math import sqrt
11 
12 class EMRestraint(object):
13  def __init__(self,
14  root,
15  map_fn,
16  resolution,
17  origin=None,
18  voxel_size=None,
19  weight=1.0,
20  label="",
21  selection_dict=None):
22  """ create a FitRestraint. can provide rigid bodies instead of individual particles """
23 
24  print('FitRestraint: setup')
25  print('\tmap_fn',map_fn)
26  print('\tresolution',resolution)
27  print('\tvoxel_size',voxel_size)
28  print('\torigin',origin)
29  print('\tweight',weight)
30 
31  # some parameters
32  self.mdl = root.get_model()
33  self.label = label
34  self.dmap = IMP.em.read_map(map_fn,IMP.em.MRCReaderWriter())
35  dh = self.dmap.get_header()
36  dh.set_resolution(resolution)
37  if voxel_size:
38  self.dmap.update_voxel_size(voxel_size)
39  if type(origin)==IMP.algebra.Vector3D:
40  self.dmap.set_origin(origin)
41  elif type(origin)==list:
42  self.dmap.set_origin(*origin)
43  else:
44  print('FitRestraint did not recognize format of origin')
45  exit()
46  if selection_dict:
47  ps=IMP.atom.Selection(root,**selection_dict).get_selected_particles()
48  else:
49  ps=IMP.atom.get_leaves(root)
50  fr = IMP.em.FitRestraint(ps,self.dmap)
51  self.rs = IMP.RestraintSet(self.mdl,weight,"FitRestraint")
52  self.rs.add_restraint(fr)
53  self.set_weight(weight)
54 
55  def set_weight(self,weight):
56  self.weight = weight
57  self.rs.set_weight(weight)
58 
59  def set_label(self, label):
60  self.label = label
61 
62  def add_to_model(self):
63  self.mdl.add_restraint(self.rs)
64 
65  def get_restraint_set(self):
66  return self.rs
67 
68  def get_output(self):
69  self.mdl.update()
70  output = {}
71  score = self.weight * self.rs.unprotected_evaluate(None)
72  output["_TotalScore"] = str(score)
73  output["EMRestraint_" + self.label] = str(score)
74  return output
75 
76  def evaluate(self):
77  return self.weight * self.rs.unprotected_evaluate(None)
78 
79 class GaussianEMRestraint(object):
80  def __init__(self,
81  hier=None,
82  model_ps=None,
83  target_fn='',
84  cutoff_dist_model_model=0.0,
85  cutoff_dist_model_data=0.0,
86  overlap_threshold=0.0,
87  target_mass_scale=1.0,
88  target_radii_scale=1.0,
89  model_radii_scale=1.0,
90  slope=0.0,
91  pointwise_restraint=False,
92  mm_container=None,
93  backbone_slope=False,
94  use_log_score=False,
95  orig_map_fn=None,
96  orig_map_res=None,
97  orig_map_apix=None,
98  orig_map_origin=None,
99  label=""):
100 
101  # some parameters
102  self.label = label
103  sigmaissampled = False
104  sigmamaxtrans = 0.3
105  sigmamin = 1.0
106  sigmamax = 100.0
107  sigmainit = 2.0
108  self.weight=1
109 
110  # setup target GMM
111  self.model_ps=[]
112  if hier is not None:
113  self.model_ps+=IMP.atom.get_leaves(hier)
114  if model_ps is not None:
115  self.model_ps+=model_ps
116  if len(self.model_ps)==0:
117  raise Exception("GaussianEM: must provide hier or model_ps")
118  self.m = self.model_ps[0].get_model()
119  print('will scale target mass by', target_mass_scale)
120  target_ps=[]
121  if target_fn != '':
122  IMP.isd_emxl.gmm_tools.decorate_gmm_from_text(
123  target_fn,
124  target_ps,
125  self.m,
126  radius_scale=target_radii_scale,
127  mass_scale=target_mass_scale)
128  else:
129  print('Gaussian EM restraint: must provide target density file')
130  return
131 
132  self.get_cc=False
133  if orig_map_fn is not None:
134  self.dmap=IMP.em.read_map(orig_map_fn,IMP.em.MRCReaderWriter())
135  self.dmap.update_voxel_size(orig_map_apix)
136  dh = self.dmap.get_header()
137  dh.set_resolution(orig_map_res)
138  self.fr = IMP.em.FitRestraint(self.model_ps,self.dmap,[0.,0.],
139  IMP.atom.Mass.get_mass_key(),1.0,
140  False)
141  #self.dmap.set_origin(-100/1.06,-100/1.06,-100/1.06)
142  if orig_map_origin is not None:
143  self.dmap.set_origin(-orig_map_origin[0]/orig_map_apix,
144  -orig_map_origin[1]/orig_map_apix,
145  -orig_map_origin[2]/orig_map_apix)
146 
147  frscore = self.fr.unprotected_evaluate(None)
148  print('init CC eval!',1.0-frscore)
149  self.get_cc=True
150 
151  # setup model GMM
152  if model_radii_scale != 1.0:
153  for p in self.model_ps:
154  rmax = sqrt(max(IMP.core.Gaussian(p).get_variances())) * \
155  model_radii_scale
156  if not IMP.core.XYZR.get_is_setup(p):
157  IMP.core.XYZR.setup_particle(p, rmax)
158  else:
159  IMP.core.XYZR(p).set_radius(rmax)
160 
161 
162  # sigma particle
163  sigmaglobal = tools.SetupNuisance(self.m, sigmainit,
164  sigmamin, sigmamax,
165  sigmaissampled).get_particle()
166 
167  # create restraint
168  print('target num particles', len(target_ps), \
169  'total weight', sum([IMP.atom.Mass(p).get_mass()
170  for p in target_ps]))
171  print('model num particles', len(self.model_ps), \
172  'total weight', sum([IMP.atom.Mass(p).get_mass()
173  for p in self.model_ps]))
174 
175  if not pointwise_restraint:
176  self.gaussianEM_restraint = \
177  IMP.isd.GaussianEMRestraint(self.m,
178  IMP.get_indexes(self.model_ps),
179  IMP.get_indexes(target_ps),
180  sigmaglobal.get_particle().get_index(),
181  cutoff_dist_model_model,
182  cutoff_dist_model_data,
183  slope,
184  backbone_slope)
185  else:
186  print('USING POINTWISE RESTRAINT')
187  print('mm_container',mm_container)
188  self.gaussianEM_restraint = \
189  IMP.isd_emxl.PointwiseGaussianEMRestraint(self.m,
190  IMP.get_indexes(self.model_ps),
191  IMP.get_indexes(target_ps),
192  sigmaglobal.get_particle().get_index(),
193  cutoff_dist_model_model,
194  cutoff_dist_model_data,
195  slope,
196  use_log_score,
197  mm_container)
198 
199  print('done EM setup')
200  self.rs = IMP.RestraintSet(self.m, 'GaussianEMRestraint')
201  self.rs.add_restraint(self.gaussianEM_restraint)
202 
203  def set_weight(self,weight):
204  self.weight = weight
205  self.rs.set_weight(weight)
206 
207  def set_label(self, label):
208  self.label = label
209 
210  def add_to_model(self):
211  self.m.add_restraint(self.rs)
212 
213  def get_restraint_set(self):
214  return self.rs
215 
216  def get_output(self):
217  self.m.update()
218  output = {}
219  score = self.weight * self.rs.unprotected_evaluate(None)
220  output["_TotalScore"] = str(score)
221  output["GaussianEMRestraint_" +
222  self.label] = str(score)
223  if self.get_cc:
224  frscore = self.fr.unprotected_evaluate(None)
225  output["CrossCorrelation"] = str(1.0-frscore)
226  #dd = self.fr.get_model_dens_map()
227  #IMP.em.write_map(dd,'test_map.mrc')
228  return output
229 
230  def evaluate(self):
231  return self.weight * self.rs.unprotected_evaluate(None)
232 
233 class GaussianPointRestraint(object):
234  def __init__(self,
235  hier,
236  target_map_fn,
237  map_res,
238  lmda_init,
239  sigma_init=None,
240  voxel_size=None,
241  origin=None,
242  cutoff_dist=0.0,
243  target_mass_scale=1.0,
244  slope=0.0,
245  mass_fraction=1.0,
246  backbone_slope=False,
247  label=""):
248 
249  # some parameters
250  self.label = label
251  self.weight = 1
252  self.cutoff_dist = cutoff_dist
253  sigma_min = 0.01
254  sigma_max = 10.
255  lmda_min = 0.01
256  lmda_max = 100
257  if lmda_init<lmda_min:
258  lmda_init=lmda_min
259 
260  # setup target GMM
261  self.model_ps=IMP.atom.get_leaves(hier)
262  self.m = self.model_ps[0].get_model()
263  IMP.atom.add_radii(hier)
264 
265  if sigma_init is None:
266  radius = IMP.core.XYZR(self.model_ps[0]).get_radius()
267  sigma_init = (0.42466090014400953*map_res)**2 + float(radius**2)/2
268 
269  for p in self.model_ps:
270  trans=IMP.algebra.Transformation3D(IMP.algebra.get_identity_rotation_3d(),
271  IMP.core.XYZ(p).get_coordinates())
272  shape = IMP.algebra.Gaussian3D(IMP.algebra.ReferenceFrame3D(trans),
273  [sigma_init]*3)
274  IMP.core.Gaussian.setup_particle(p,shape)
275 
276 
277  self.dmap=IMP.em.read_map(target_map_fn,IMP.em.MRCReaderWriter())
278  if voxel_size is not None:
279  self.dmap.update_voxel_size(voxel_size)
280  dh = self.dmap.get_header()
281  dh.set_resolution(map_res)
282  if origin is not None:
283  self.dmap.set_origin(-origin[0]/voxel_size,
284  -origin[1]/voxel_size,
285  -origin[2]/voxel_size)
286 
287  self.fr = IMP.em.FitRestraint(self.model_ps,self.dmap,[0.,0.],
288  IMP.atom.Mass.get_mass_key(),1.0,
289  False)
290  frscore = self.fr.unprotected_evaluate(None)
291  print('init CC eval!',1.0-frscore)
292  self.get_cc=True
293 
294 
295  self.sigma = tools.SetupNuisance(self.m, sigma_init,
296  sigma_min, sigma_max,
297  isoptimized=True).get_particle()
298  self.lmda = tools.SetupNuisance(self.m, lmda_init,
299  lmda_min,lmda_max,
300  isoptimized=True).get_particle()
301  self.gaussian_restraint = IMP.isd_emxl.GaussianPointRestraint(self.m,
302  IMP.get_indexes(self.model_ps),
303  self.sigma.get_particle_index(),
304  self.lmda.get_particle_index(),
305  self.dmap,
306  0,0,False,1.0)
307 
308  self.rs = IMP.RestraintSet(self.m, 'GaussianPointRestraint')
309  self.rs.add_restraint(self.gaussian_restraint)
310  self.rs.add_restraint(IMP.isd.JeffreysRestraint(self.m, self.sigma))
311  self.rs.add_restraint(IMP.isd.JeffreysRestraint(self.m, self.lmda))
312  print('done EM setup')
313 
314  def set_weight(self,weight):
315  self.weight = weight
316  self.rs.set_weight(weight)
317 
318  def set_label(self, label):
319  self.label = label
320 
321  def add_to_model(self):
322  self.m.add_restraint(self.rs)
323 
324  def get_restraint_set(self):
325  return self.rs
326 
327  def get_output(self):
328  self.m.update()
329  output = {}
330  score = self.weight * self.rs.unprotected_evaluate(None)
331  output["_TotalScore"] = str(score)
332  output["GaussianPointRestraint_" +
333  self.label] = str(score)
334  if self.get_cc:
335  frscore = self.fr.unprotected_evaluate(None)
336  output["CrossCorrelation"] = str(1.0-frscore)
337  output["GaussianPointRestraint_lambda"]=self.lmda.get_scale()
338  output["GaussianPointRestraint_sigma"]=self.sigma.get_scale()
339  #dd = self.fr.get_model_dens_map()
340  #IMP.em.write_map(dd,'test_map.mrc')
341  return output
342 
343  def evaluate(self):
344  return self.weight * self.rs.unprotected_evaluate(None)
345 
346  def get_mc_sample_objects(self,max_step):
347  """ HACK! Make a SampleObjects class that can be used with PMI::samplers"""
348  ps=[[self.sigma,self.lmda],max_step]
349  return [sampling_tools.SampleObjects('Nuisances',ps)]
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::Mass::get_mass_key
static FloatKey get_mass_key()
IMP::domino::get_index
Ints get_index(const kernel::ParticlesTemp &particles, const Subset &subset, const Subsets &excluded)
IMP::atom::get_mass
double get_mass(ResidueType c)
Get the mass from the residue type.
IMP::atom::add_radii
void add_radii(Hierarchy d, const ForceFieldParameters *ffp=get_all_atom_CHARMM_parameters(), FloatKey radius_key=FloatKey("radius"))
IMP::pmi.sampling_tools
Useful tools for setting up sampling.
Definition: sampling_tools.py:1
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::algebra::ReferenceFrame3D
A reference frame in 3D.
Definition: ReferenceFrame3D.h:20
IMP::core::XYZR::get_is_setup
static bool get_is_setup(const IMP::kernel::ParticleAdaptor &p)
Definition: XYZR.h:47
IMP::algebra::Gaussian3D
A Gaussian distribution in 3D.
Definition: Gaussian3D.h:24
IMP::core::XYZR::setup_particle
static XYZR setup_particle(kernel::Model *m, ParticleIndex pi)
Definition: XYZR.h:48
IMP::isd::GaussianEMRestraint
Creates a restraint between two Gaussian Mixture Models, "model" and "density".
Definition: GaussianEMRestraint.h:77
IMP::em
Basic utilities for handling cryo-electron microscopy 3D density maps.
IMP::core::XYZ
A decorator for a particle with x,y,z coordinates.
Definition: XYZ.h:30
IMP::em::read_map
DensityMap * read_map(std::string filename)
Read a density map from a file and return it.
IMP::em::FitRestraint
Calculate score based on fit to EM map.
Definition: FitRestraint.h:32
IMP::algebra::get_identity_rotation_3d
Rotation3D get_identity_rotation_3d()
Return a rotation that does not do anything.
Definition: Rotation3D.h:236
IMP::atom
Functionality for loading, creating, manipulating and scoring atomic structures.
IMP::atom::get_leaves
Hierarchies get_leaves(const Selection &h)
IMP::atom::Selection
Select hierarchy particles identified by the biological name.
Definition: Selection.h:62
IMP::core::Gaussian::setup_particle
static Gaussian setup_particle(kernel::Model *m, ParticleIndex pi)
Definition: Gaussian.h:48
IMP::core::XYZR
A decorator for a particle with x,y,z coordinates and a radius.
Definition: XYZR.h:27