From e160ca82d411846b41cd000e9d7766fb69950866 Mon Sep 17 00:00:00 2001 From: Ben Webb Date: Fri, 8 Sep 2023 17:49:30 -0700 Subject: [PATCH] Add coordinates for models --- modules/mmcif/pyext/src/data.py | 39 +++++++++++++++++++++++++++++- modules/mmcif/pyext/src/util.py | 2 ++ modules/mmcif/test/test_convert.py | 38 +++++++++++++++++++++++------ 3 files changed, 71 insertions(+), 8 deletions(-) diff --git a/modules/mmcif/pyext/src/data.py b/modules/mmcif/pyext/src/data.py index 691acaf165..28cf3222cd 100644 --- a/modules/mmcif/pyext/src/data.py +++ b/modules/mmcif/pyext/src/data.py @@ -597,7 +597,8 @@ def _add_protocol(self, prot): # with existing protocols. This should still be performant as # we generally don't have more than one or two protocols. def step_equal(x, y): - return type(x) == type(y) and x.__dict__ == y.__dict__ + return (type(x) == type(y) # noqa: E721 + and x.__dict__ == y.__dict__) def analysis_equal(x, y): return (len(x.steps) == len(y.steps) @@ -641,6 +642,12 @@ def _add_hierarchy(self, h, top_h, modeled_assembly, all_software): class _CoordinateHandler(object): def __init__(self): self._representation = ihm.representation.Representation() + # IHM atoms/spheres corresponding to IMP beads/residues/atoms + # We build them up front (rather than on the fly) as the original + # IMP objects may have been destroyed or changed (e.g. if we read + # multiple frames from an RMF file) by the time we write the mmCIF. + self._atoms = [] + self._spheres = [] def add_chain(self, c, asym): ps = self._get_structure_particles(c) @@ -650,10 +657,40 @@ def add_chain(self, c, asym): seg = segfactory.add(p, None) if seg: self._representation.append(seg) + self._add_atom_or_sphere(p, asym) last = segfactory.get_last() if last: self._representation.append(last) + def _add_atom_or_sphere(self, p, asym): + if isinstance(p, IMP.atom.Atom): + residue = IMP.atom.get_residue(p) + xyz = IMP.core.XYZ(p).get_coordinates() + element = p.get_element() + element = IMP.atom.get_element_table().get_name(element) + atom_name = p.get_atom_type().get_string() + het = atom_name.startswith('HET:') + if het: + atom_name = atom_name[4:] + self._atoms.append(ihm.model.Atom( + asym_unit=asym, seq_id=residue.get_index(), + atom_id=atom_name, type_symbol=element, + x=xyz[0], y=xyz[1], z=xyz[2], het=het, + biso=p.get_temperature_factor(), + occupancy=p.get_occupancy())) + else: + if isinstance(p, IMP.atom.Fragment): + resinds = p.get_residue_indexes() + sbegin = resinds[0] + send = resinds[-1] + else: # residue + sbegin = send = p.get_index() + xyzr = IMP.core.XYZR(p) + xyz = xyzr.get_coordinates() + self._spheres.append(ihm.model.Sphere( + asym_unit=asym, seq_id_range=(sbegin, send), + x=xyz[0], y=xyz[1], z=xyz[2], radius=xyzr.get_radius())) + def _get_structure_particles(self, h): """Return particles sorted by residue index""" ps = [] diff --git a/modules/mmcif/pyext/src/util.py b/modules/mmcif/pyext/src/util.py index 8bf4df4483..eee2c46ae4 100644 --- a/modules/mmcif/pyext/src/util.py +++ b/modules/mmcif/pyext/src/util.py @@ -467,6 +467,8 @@ def _add_hierarchy(self, h, top_h, state, name, ensemble): self._external_files.add_hierarchy(h, top_h) model = ihm.model.Model(assembly=assembly, protocol=protocol, representation=representation) + model._atoms = ch._atoms + model._spheres = ch._spheres ensemble.model_group.append(model) ensemble.num_models += 1 return ensemble diff --git a/modules/mmcif/test/test_convert.py b/modules/mmcif/test/test_convert.py index 046ed3fb25..fb36c7e0cb 100644 --- a/modules/mmcif/test/test_convert.py +++ b/modules/mmcif/test/test_convert.py @@ -346,24 +346,48 @@ def test_different_assembly(self): self.assertEqual(len(c.system.orphan_assemblies[0]), 3) self.assertEqual(len(c.system.orphan_assemblies[1]), 2) - def test_representation(self): - """Test representation""" + def test_model_creation(self): + """Test creation of ihm Model objects""" m = IMP.Model() top = IMP.atom.Hierarchy.setup_particle(IMP.Particle(m)) self.add_chains(m, top) c = IMP.mmcif.Convert() chain0 = top.get_child(0).get_child(0) - chain0.add_child(IMP.atom.Residue.setup_particle(IMP.Particle(m), - IMP.atom.ALA, 1)) residue = IMP.atom.Residue.setup_particle(IMP.Particle(m), - IMP.atom.ALA, 2) + IMP.atom.ALA, 1) + IMP.core.XYZR.setup_particle( + residue, IMP.algebra.Sphere3D(IMP.algebra.Vector3D(1, 2, 3), 4)) + chain0.add_child(residue) + residue = IMP.atom.Residue.setup_particle(IMP.Particle(m), + IMP.atom.HIS, 2) atom = IMP.atom.Atom.setup_particle(IMP.Particle(m), IMP.atom.AT_CA) + IMP.core.XYZR.setup_particle( + atom, IMP.algebra.Sphere3D(IMP.algebra.Vector3D(5, 6, 7), 8)) residue.add_child(atom) chain0.add_child(residue) - chain0.add_child(IMP.atom.Fragment.setup_particle(IMP.Particle(m), - [3, 4])) + frag = IMP.atom.Fragment.setup_particle(IMP.Particle(m), [3, 4]) + chain0.add_child(frag) + IMP.core.XYZR.setup_particle( + frag, IMP.algebra.Sphere3D(IMP.algebra.Vector3D(9, 10, 11), 12)) c.add_model([top], []) + model, = c.system.state_groups[0][0][0] + # Representation should contain residue, atom, fragment + r1, r2, r3 = model.representation + self.assertIsInstance(r1, ihm.representation.ResidueSegment) + self.assertIsInstance(r2, ihm.representation.AtomicSegment) + self.assertIsInstance(r3, ihm.representation.FeatureSegment) + # Coordinates should contain one atom, two spheres + a1, = model._atoms + self.assertEqual(a1.seq_id, 2) + self.assertEqual(a1.atom_id, 'CA') + self.assertEqual(a1.type_symbol, 'C') + self.assertFalse(a1.het) + s1, s2 = model._spheres + self.assertEqual(s1.seq_id_range, (1, 1)) + self.assertAlmostEqual(s1.radius, 4.0, delta=1e-2) + self.assertEqual(s2.seq_id_range, (3, 4)) + self.assertAlmostEqual(s2.radius, 12.0, delta=1e-2) if __name__ == '__main__':