1 """Interface between XML Restraint and IMP Restraint."""
8 class _RestraintSets(object):
10 self._restraint_sets = dict()
11 self.set_by_restraint = dict()
12 self.set_by_name = dict()
13 def add(self, restraint_type, weight, name, restraint, model = None):
14 if restraint_type
is None:
17 restraint_type =
'%s|%s' % (restraint_type, name)
19 set_by_weight = self._restraint_sets[weight]
21 set_by_weight = self._restraint_sets[weight] = dict()
23 current_set = set_by_weight[restraint_type]
27 model.add_restraint(current_set)
28 current_set.add_restraint(restraint)
29 self.set_by_restraint[restraint] = current_set
30 self.set_by_name[name] = current_set
31 def set_model(self, model):
32 for weight, set_by_weight
in self._restraint_sets.iteritems():
33 for rest_set
in set_by_weight.itervalues():
34 model.add_restraint(rest_set)
35 rest_set.set_weight(float(weight))
36 def remove_restraint(self, restraint):
38 current_set = self.set_by_restraint[restraint]
41 current_set.remove_restraint(restraint)
42 del self.set_by_restraint[restraint]
43 def get_weight(self, restraint):
45 current_set = self.set_by_restraint[restraint]
48 return current_set.get_weight()
49 def remove_all(self, model):
50 for set_by_weight
in self._restraint_sets.itervalues():
51 for rest_set
in set_by_weight.itervalues():
52 model.remove_restraint(rest_set)
53 self._restraint_sets.clear()
59 self._children = list()
60 self._restraint_sets = _RestraintSets()
62 def add_to_representation(self, repr):
63 """ Place restraint into IMP Model."""
65 if repr._model
is None:
68 self._model = repr._model
69 for child
in self._children:
70 child.create_restraint(repr)
71 self.set_include(include=
True)
74 def add_to_model(self):
75 self._restraint_sets.remove_all(self._model)
76 for child
in self._children:
77 child.include_restraint(self._restraint_sets)
78 self._restraint_sets.set_model(self._model)
80 def set_include(self, name=None, include=False, weight=-1.0):
82 def set_include_rec(node):
83 for child
in node._children:
84 set_include_rec(child)
85 node.include_me = include
90 for child
in self._children:
91 set_include_rec(child)
93 r = self.get_restraint_by_name(name)
97 r.include_me = include
102 node.include_me =
False
103 for child
in node._children:
106 for child
in self._children:
108 self._restraint_sets.remove_all(self._model)
111 def show_all_restraints(self):
112 """Show restraint name, initial weight, and score for the current state of the model"""
115 if node.imp_restraint
is not None:
116 print "=== ", node.name,
" ==="
117 print "Restraint initial weight = ", node.root._restraint_sets.get_weight(node.imp_restraint)
118 node.imp_restraint.evaluate(
False)
120 for child
in node._children:
124 for child
in self._children:
129 def get_all_restraints_by_name(self, name):
130 """Assuming there are many restraint objects with the same name."""
133 if node.name == name:
135 for child
in node._children:
139 for child
in self._children:
143 def get_restraint_by_name(self, name):
144 """Assuming there is just one restraint object with the same name."""
147 if node.name == name:
149 for child
in node._children:
155 for child
in self._children:
161 def get_restraint_set_by_name(self, name):
162 """Get an IMP::RestraintSet by name specified in restraint XML file."""
163 return self._restraint_sets.set_by_name.get(name)
165 def get_rigid_body_by_id(self, id):
166 """Get rigid body by particle id."""
170 return node.rigid_body
171 for child
in node._children:
177 for child
in self._children:
184 def get_all_rigid_bodies(self):
185 """Get all rigid bodies."""
189 all_bodies.append(node.rigid_body)
190 for child
in node._children:
194 for child
in self._children:
201 def set_root_rec(node):
203 for child
in node._children:
206 for child
in self._children:
210 class _RestraintNode(object):
213 def __init__(self, attributes, restraint_type=None):
214 self.imp_restraint =
None
215 self.rigid_body =
None
217 name = attributes.get(
'name')
221 self.name =
'object_%d' % _RestraintNode.counter
222 _RestraintNode.counter += 1
223 self._children = list()
224 self.child_restraints = list()
225 self.restraint_type = restraint_type
226 self.include_me =
False
227 self.repr_particle =
None
229 def get_particle(self):
230 particle_list = list()
231 for child
in self._children:
232 if isinstance(child, _RestraintParticle):
233 particle_list.append(child)
236 def get_source(self):
238 for child
in self._children:
240 if isinstance(child, _RestraintSource):
241 source_list.append(child)
242 source_list[counter].author_list = child.get_author()
243 source_list[counter].journal = child.get_journal()
244 source_list[counter].title = child.get_title()
245 source_list[counter].year = child.get_year()
249 def append_restraint_with_weight(self, restraint, restraint_sets):
251 self.child_restraints.append(restraint)
252 if self.restraint_type:
253 restraint_sets.add(self.restraint_type, self.weight, self.name, restraint)
256 self.child_restraints = list()
257 for child
in self._children:
258 restraint = child.create_restraint(repr)
260 self.child_restraints.append(restraint)
262 def include_restraint(self, restraint_sets):
263 if self.restraint_type
and self.include_me:
264 for child
in self.child_restraints:
265 restraint_sets.add(self.restraint_type, self.weight, self.name,
267 for child
in self._children:
268 child.include_restraint(restraint_sets)
271 class _RestraintNodeWithWeight(_RestraintNode):
272 def __init__(self, attributes, type):
273 _RestraintNode.__init__(self, attributes, type)
274 self.weight = float(attributes.get(
'weight', 1))
277 class _RigidBodyRestraintNode(_RestraintNodeWithWeight):
278 def __init__(self, attributes):
279 _RestraintNodeWithWeight.__init__(self, attributes,
'RigidBody')
282 for child
in self._children:
283 restraint = child.create_rigid_body_restraint(repr)
286 class _ExcludedVolumeRestraintNode(_RestraintNodeWithWeight):
287 def __init__(self, attributes, restraint_type):
288 _RestraintNodeWithWeight.__init__(self, attributes, restraint_type)
291 self.child_restraints = list()
292 for child
in self._children:
293 restraint = child.create_excluded_volume_restraint(repr)
296 self.child_restraints.append(restraint)
299 class _ConnectivityRestraintNode(_RestraintNodeWithWeight):
300 def __init__(self, attributes, restraint_type):
301 _RestraintNodeWithWeight.__init__(self, attributes, restraint_type)
304 self.child_restraints = list()
305 for child
in self._children:
306 restraint = child.create_connectivity_restraint(repr)
309 self.child_restraints.append(restraint)
312 class _CrosslinkRestraintNode(_RestraintNodeWithWeight):
313 def __init__(self, attributes, restraint_type):
314 _RestraintNodeWithWeight.__init__(self, attributes, restraint_type)
317 self.child_restraints = list()
318 for child
in self._children:
319 restraint = child.create_crosslink_restraint(repr)
322 self.child_restraints.append(restraint)
325 class _DistanceRestraintNode(_RestraintNodeWithWeight):
326 def __init__(self, attributes, restraint_type):
327 _RestraintNodeWithWeight.__init__(self, attributes, restraint_type)
330 self.child_restraints = list()
331 for child
in self._children:
332 restraint = child.create_distance_restraint(repr)
335 self.child_restraints.append(restraint)
338 class _DiameterRestraintNode(_RestraintNodeWithWeight):
339 def __init__(self, attributes, restraint_type):
340 _RestraintNodeWithWeight.__init__(self, attributes, restraint_type)
343 self.child_restraints = list()
344 for child
in self._children:
345 restraint = child.create_diameter_restraint(repr)
348 self.child_restraints.append(restraint)
351 class _SAXSRestraintNode(_RestraintNodeWithWeight):
352 def __init__(self, attributes, restraint_type):
353 _RestraintNodeWithWeight.__init__(self, attributes, restraint_type)
356 self.child_restraints = list()
357 for child
in self._children:
358 restraint = child.create_saxs_restraint(repr)
361 self.child_restraints.append(restraint)
364 class _SANSRestraintNode(_RestraintNodeWithWeight):
365 def __init__(self, attributes, restraint_type):
366 _RestraintNodeWithWeight.__init__(self, attributes, restraint_type)
369 self.child_restraints = list()
371 class _EMRestraintNode(_RestraintNodeWithWeight):
372 def __init__(self, attributes, restraint_type):
373 _RestraintNodeWithWeight.__init__(self, attributes, restraint_type)
376 self.child_restraints = list()
377 for child
in self._children:
378 restraint = child.create_em_restraint(repr)
381 self.child_restraints.append(restraint)
384 class _RestraintY2H(_ConnectivityRestraintNode):
385 def __init__(self, attributes):
386 _ConnectivityRestraintNode.__init__(self, attributes,
"Y2H")
388 class _RestraintExcludedVolume(_ExcludedVolumeRestraintNode):
389 def __init__(self, attributes):
390 _ExcludedVolumeRestraintNode.__init__(self, attributes,
"ExcludedVolume")
392 class _RestraintPulldown(_ConnectivityRestraintNode):
393 def __init__(self, attributes):
394 _ConnectivityRestraintNode.__init__(self, attributes,
"Pulldown")
396 class _RestraintXrayStruc(_ConnectivityRestraintNode):
397 def __init__(self, attributes):
398 _ConnectivityRestraintNode.__init__(self, attributes,
"XrayStruc")
400 class _RestraintMSMS(_ConnectivityRestraintNode):
401 def __init__(self, attributes):
402 _ConnectivityRestraintNode.__init__(self, attributes,
"MSMS")
404 class _RestraintArray(_ConnectivityRestraintNode):
405 def __init__(self, attributes):
406 _ConnectivityRestraintNode.__init__(self, attributes,
"Array")
408 class _RestraintCopurification(_ConnectivityRestraintNode):
409 def __init__(self, attributes):
410 _ConnectivityRestraintNode.__init__(self, attributes,
"Copurification")
412 class _RestraintCrossLink(_CrosslinkRestraintNode):
413 def __init__(self, attributes):
414 _CrosslinkRestraintNode.__init__(self, attributes,
"CrossLink")
416 class _RestraintDistance(_DistanceRestraintNode):
417 def __init__(self, attributes):
418 _DistanceRestraintNode.__init__(self, attributes,
"Distance")
420 class _RestraintDiameter(_DiameterRestraintNode):
421 def __init__(self, attributes):
422 _DiameterRestraintNode.__init__(self, attributes,
"Diameter")
424 class _RestraintSAXS(_SAXSRestraintNode):
425 def __init__(self, attributes):
426 _SAXSRestraintNode.__init__(self, attributes,
"SAXS")
428 class _RestraintSANS(_SANSRestraintNode):
429 def __init__(self, attributes):
430 _SANSRestraintNode.__init__(self, attributes,
"SANS")
432 class _RestraintEM(_EMRestraintNode):
433 def __init__(self, attributes):
434 _EMRestraintNode.__init__(self, attributes,
"EM")
436 class _RestraintRestraint(_RestraintNode):
437 def __init__(self, attributes):
438 _RestraintNode.__init__(self, attributes)
439 self.imp_restraint =
None
440 self.std_dev = float(attributes.get(
'std_dev', 1))
441 self.distance = float(attributes.get(
'distance', 0))
442 self.max_diameter = float(attributes.get(
'max_diameter', -1))
443 self.profile_filename = attributes.get(
'profile_filename',
'')
444 self.density_filename = attributes.get(
'density_filename',
'')
445 self.resolution = float(attributes.get(
'resolution', -1))
446 self.spacing = float(attributes.get(
'spacing', -1))
447 self.xorigin = float(attributes.get(
'xorigin', 0.0))
448 self.yorigin = float(attributes.get(
'yorigin', 0.0))
449 self.zorigin = float(attributes.get(
'zorigin', 0.0))
450 self.linker_length = float(attributes.get(
'linker_length', -1))
452 def include_restraint(self, restraint_sets):
455 def get_weight(self):
456 return self.root._restraint_sets.get_weight(self.imp_restraint)
458 def create_rigid_body_restraint(self, repr):
459 _RestraintNode.create_restraint(self, repr)
460 self.rigid_bodies = list()
461 for child, child_r
in zip(self._children, self.child_restraints):
462 if child.repr_particle:
463 child.repr_particle.force_field = 0
464 mhs = IMP.atom.Hierarchies()
467 self.rigid_bodies.append(rb[0])
468 child.rigid_body = rb[0]
471 def create_force(self):
472 for child
in self._children:
473 repr = child.repr_particle
476 print 'repr.force_field=', repr.force_field
477 if repr
and repr.force_field == 1:
478 print 'Use the force.'
479 hierarchy = repr.model_decorator
483 print 'hierarchy present'
487 topology = ff.create_topology(hierarchy)
488 topology.apply_default_patches(ff)
490 bonds = topology.add_bonds(hierarchy, ff)
491 angles = ff.create_angles(bonds)
492 dihedrals = ff.create_dihedrals(bonds)
493 impropers = topology.add_impropers(hierarchy, ff)
507 ff.add_well_depths(hierarchy)
512 pair_filter.set_bonds(bonds)
513 pair_filter.set_angles(angles)
514 pair_filter.set_dihedrals(dihedrals)
515 nbl.add_pair_filter(pair_filter)
522 _RestraintNode.create_restraint(self, repr)
525 k = 1.0/(2.0*self.std_dev*self.std_dev)
528 mhs = IMP.atom.Hierarchies()
529 if not self.child_restraints:
530 raise Exception,
"Restraint %s is empty" % self.id
531 for child
in self.child_restraints:
533 first_particle = self.child_restraints[0].get_particle()
537 rbs_tmp.append(mh.get_particle())
538 rbs = IMP.core.RigidBodies(rbs_tmp)
542 self.sc = IMP.restrainer.create_simple_connectivity_on_molecules(mhs)
544 connectivity_restraint = self.sc.get_restraint()
545 self.imp_restraint = connectivity_restraint
546 return connectivity_restraint
549 _RestraintNode.create_restraint(self, repr)
551 mhs = IMP.atom.Hierarchies()
552 if not self.child_restraints:
553 raise Exception,
"Restraint %s is empty" % self.id
554 for child
in self.child_restraints:
556 first_particle = self.child_restraints[0].get_particle()
560 rbs_tmp.append(mh.get_particle())
561 rbs = IMP.core.RigidBodies(rbs_tmp)
562 self.sev = IMP.restrainer.create_simple_excluded_volume_on_rigid_bodies(rbs)
564 self.sev = IMP.restrainer.create_simple_excluded_volume_on_molecules(mhs)
565 ev_restraint = self.sev.get_restraint()
566 self.imp_restraint = ev_restraint
569 def create_saxs_restraint(self, repr):
570 _RestraintNode.create_restraint(self, repr)
572 if self.profile_filename:
574 particles = IMP.atom.Hierarchies()
575 for child
in self.child_restraints:
577 particles.append(atoms)
580 self.imp_restraint = self.saxs_restraint
581 return self.saxs_restraint
583 def create_em_restraint(self, repr):
584 _RestraintNode.create_restraint(self, repr)
586 if self.density_filename:
587 mhs = IMP.atom.Hierarchies()
589 self.spacing, self.resolution)
590 self.dmap_header = self.dmap.get_header_writable()
592 self.dmap.set_origin (self.xorigin, self.yorigin, self.zorigin)
593 self.dmap.update_voxel_size(self.spacing)
594 self.dmap_header.set_resolution (self.resolution)
596 for child
in self.child_restraints:
599 self.sef = IMP.restrainer.create_simple_em_fit(mhs, self.dmap)
600 em_restraint = self.sef.get_restraint()
601 self.imp_restraint = em_restraint
604 def create_crosslink_restraint(self, repr):
605 _RestraintNode.create_restraint(self, repr)
609 _RestraintNode.create_restraint(self, repr)
612 k = 1.0/(2.0*self.std_dev*self.std_dev)
617 self.child_restraints[1])
618 return self.imp_restraint
620 def create_diameter_restraint(self, repr):
621 _RestraintNode.create_restraint(self, repr)
624 for child
in self.child_restraints:
625 ps.append(child.get_particle())
626 self.diameter_restraint = IMP.restrainer.create_simple_diameter(
627 ps, self.max_diameter)
628 self.imp_restraint = self.diameter_restraint.get_restraint()
629 return self.imp_restraint
631 class _RestraintSource(_RestraintNode):
632 def __init__(self, attributes):
633 _RestraintNode.__init__(self, attributes)
634 self.pubmed_id = int(attributes.get(
'pubmed_id', -1))
636 self.__author_list = list()
637 self.__journal =
None
641 def get_author(self):
642 if not self.__author_list:
643 for child
in self._children:
644 if isinstance(child, _RestraintAuthor):
645 self.__author_list.append(child)
647 return self.__author_list
649 def get_journal(self):
650 if self.__journal
is None:
651 for child
in self._children:
652 if isinstance(child, _RestraintJournal):
653 self.__journal = child.text
655 return self.__journal
658 if self.__title
is None:
659 for child
in self._children:
660 if isinstance(child, _RestraintTitle):
661 self.__title = child.text
666 if self.__year
is None:
667 for child
in self._children:
668 if isinstance(child, _RestraintYear):
669 self.__year = child.text
673 class _RestraintAuthor(_RestraintNode):
674 def __init__(self, attributes):
675 _RestraintNode.__init__(self, attributes)
676 self.first_name = attributes.get(
'first_name',
'')
677 self.last_name = attributes.get(
'last_name',
'')
679 class _RestraintJournal(_RestraintNode):
680 def __init__(self, attributes, text):
681 _RestraintNode.__init__(self, attributes)
684 class _RestraintTitle(_RestraintNode):
685 def __init__(self, attributes, text):
686 _RestraintNode.__init__(self, attributes)
689 class _RestraintYear(_RestraintNode):
690 def __init__(self, attributes, text):
691 _RestraintNode.__init__(self, attributes)
694 class _RestraintParticle(_RestraintNode):
695 def __init__(self, attributes):
696 _RestraintNode.__init__(self, attributes)
697 self.id = attributes.get(
'id',
'')
698 self.residue = attributes.get(
'residue',
'')
699 self.atom = attributes.get(
'atom',
'')
702 self.repr_particle = self.find_representation(repr)
703 if not self.repr_particle:
704 raise Exception,
"Particle with id=%s not found" % (self.id)
705 decorator = self.repr_particle.model_decorator
714 def find_representation(self, repr):
715 return repr.find_by_id(self.id)