Add coordinates for models

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 = []

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

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__':