Add support for atom element

examples/simple-docking.py

@@ -96,9 +96,9 @@
 # Finally we can add coordinates for the deposited models. Typically these
 # will be stored in our own software's data structures somewhere (for this
 # example in simple lists 'atoms' and 'spheres'):
-atoms = [('A', 1, 'CA', 1., 2., 3.),
-         ('A', 2, 'CA', 4., 5., 6.),
-         ('A', 3, 'CA', 7., 8., 9.)]
+atoms = [('A', 1, 'C', 'CA', 1., 2., 3.),
+         ('A', 2, 'C', 'CA', 4., 5., 6.),
+         ('A', 3, 'C', 'CA', 7., 8., 9.)]
 spheres = [('B', 1, 2, 1., 2., 3., 1.2),
            ('B', 3, 4, 4., 5., 6., 1.2),
            ('B', 5, 6, 7., 8., 9., 1.2)]
@@ -114,9 +114,10 @@ class MyModel(ihm.model.Model):
     asym_unit_map = {'A': asymA, 'B': asymB}
 
     def get_atoms(self):
-        for asym, seq_id, atom_id, x, y, z in atoms:
+        for asym, seq_id, type_symbol, atom_id, x, y, z in atoms:
             yield ihm.model.Atom(asym_unit=self.asym_unit_map[asym],
-                                 seq_id=seq_id, atom_id=atom_id, x=x, y=y, z=z)
+                                 type_symbol=type_symbol, seq_id=seq_id,
+                                 atom_id=atom_id, x=x, y=y, z=z)
     def get_spheres(self):
         for asym, seq_id_start, seq_id_end, x, y, z, radius in spheres:
             yield ihm.model.Sphere(asym_unit=self.asym_unit_map[asym],

ihm/dumper.py

@@ -545,6 +545,7 @@ def dump_coords(self, system, writer):
                     l.write(starting_model_id=model._id,
                             group_PDB='HETATM' if atom.het else 'ATOM',
                             id=natom+1,
+                            type_symbol=atom.type_symbol,
                             atom_id=atom.atom_id,
                             comp_id=_amino_acids[oneletter],
                             asym_id=atom.asym_unit._id,
@@ -682,7 +683,8 @@ def dump_model_list(self, system, writer):
     def dump_atoms(self, system, writer):
         ordinal = 1
         with writer.loop("_atom_site",
-                         ["group_PDB", "id", "label_atom_id", "label_comp_id",
+                         ["group_PDB", "id", "type_symbol",
+                          "label_atom_id", "label_comp_id",
                           "label_seq_id",
                           "label_asym_id", "Cartn_x",
                           "Cartn_y", "Cartn_z", "label_entity_id",
@@ -691,6 +693,7 @@ def dump_atoms(self, system, writer):
                 for atom in model.get_atoms():
                     oneletter = atom.asym_unit.entity.sequence[atom.seq_id-1]
                     l.write(id=ordinal,
+                            type_symbol=atom.type_symbol,
                             group_PDB='HETATM' if atom.het else 'ATOM',
                             label_atom_id=atom.atom_id,
                             label_comp_id=_amino_acids[oneletter],

ihm/model.py

@@ -33,6 +33,7 @@ class Atom(object):
        :type asym_unit: :class:`ihm.AsymUnit`
        :param int seq_id: The residue index represented by this atom
        :param str atom_id: The name of the atom in the residue
+       :param str type_symbol: Element name
        :param float x: x coordinate of the atom
        :param float y: y coordinate of the atom
        :param float z: z coordinate of the atom
@@ -41,13 +42,14 @@ class Atom(object):
     """
 
     # Reduce memory usage
-    __slots__ = ['asym_unit', 'seq_id', 'atom_id', 'x', 'y', 'z', 'het',
-                 'biso']
+    __slots__ = ['asym_unit', 'seq_id', 'atom_id', 'type_symbol',
+                 'x', 'y', 'z', 'het', 'biso']
 
-    def __init__(self, asym_unit, seq_id, atom_id, x, y, z, het=False,
-                 biso=None):
+    def __init__(self, asym_unit, seq_id, atom_id, type_symbol, x, y, z,
+                 het=False, biso=None):
         self.asym_unit = asym_unit
         self.seq_id, self.atom_id = seq_id, atom_id
+        self.type_symbol = type_symbol
         self.x, self.y, self.z = x, y, z
         self.het, self.biso = het, biso
 

test/test_dumper.py

@@ -694,7 +694,8 @@ class TestStartingModel(ihm.startmodel.StartingModel):
             def get_atoms(self):
                 asym = self.asym_unit
                 return [ihm.model.Atom(asym_unit=asym, seq_id=1, atom_id='CA',
-                                       x=-8.0, y=-5.0, z=91.0, biso=42.)]
+                                       type_symbol='C', x=-8.0, y=-5.0, z=91.0,
+                                       biso=42.)]
             def get_seq_dif(self):
                 return [ihm.startmodel.MSESeqDif(db_seq_id=5, seq_id=7)]
 
@@ -777,7 +778,7 @@ def get_seq_dif(self):
 _ihm_starting_model_coord.Cartn_z
 _ihm_starting_model_coord.B_iso_or_equiv
 _ihm_starting_model_coord.ordinal_id
-1 ATOM 1 . CA ALA 42 99 1 -8.000 -5.000 91.000 42.000 1
+1 ATOM 1 C CA ALA 42 99 1 -8.000 -5.000 91.000 42.000 1
 #
 #
 loop_
@@ -1012,11 +1013,13 @@ def test_model_dumper_atoms(self):
         """Test ModelDumper with atoms"""
         system, model, asym = self._make_test_model()
         model._atoms = [ihm.model.Atom(asym_unit=asym, seq_id=1, atom_id='C',
-                                       x=1.0, y=2.0, z=3.0),
+                                       type_symbol='C', x=1.0, y=2.0, z=3.0),
                         ihm.model.Atom(asym_unit=asym, seq_id=1, atom_id='CA',
-                                       x=10.0, y=20.0, z=30.0, het=True),
+                                       type_symbol='C', x=10.0, y=20.0, z=30.0,
+                                       het=True),
                         ihm.model.Atom(asym_unit=asym, seq_id=2, atom_id='N',
-                                       x=4.0, y=5.0, z=6.0, biso=42.0)]
+                                       type_symbol='N', x=4.0, y=5.0, z=6.0,
+                                       biso=42.0)]
 
         dumper = ihm.dumper._ModelDumper()
         dumper.finalize(system) # assign model/group IDs
@@ -1038,6 +1041,7 @@ def test_model_dumper_atoms(self):
 loop_
 _atom_site.group_PDB
 _atom_site.id
+_atom_site.type_symbol
 _atom_site.label_atom_id
 _atom_site.label_comp_id
 _atom_site.label_seq_id
@@ -1048,9 +1052,9 @@ def test_model_dumper_atoms(self):
 _atom_site.label_entity_id
 _atom_site.B_iso_or_equiv
 _atom_site.model_id
-ATOM 1 C ALA 1 X 1.000 2.000 3.000 9 . 1
-HETATM 2 CA ALA 1 X 10.000 20.000 30.000 9 . 1
-ATOM 3 N CYS 2 X 4.000 5.000 6.000 9 42.000 1
+ATOM 1 C C ALA 1 X 1.000 2.000 3.000 9 . 1
+HETATM 2 C CA ALA 1 X 10.000 20.000 30.000 9 . 1
+ATOM 3 N N CYS 2 X 4.000 5.000 6.000 9 42.000 1
 #
 """)
 

test/test_examples.py

@@ -25,7 +25,7 @@ def test_simple_docking_example(self):
             # Make sure that a complete output file was produced
             with open(os.path.join(tmpdir, 'output.cif')) as fh:
                 contents = fh.readlines()
-            self.assertEqual(len(contents), 230)
+            self.assertEqual(len(contents), 231)
 
     def test_locations_example(self):
         """Test locations example"""

test/test_model.py

@@ -18,8 +18,8 @@ def test_sphere(self):
 
     def test_atom(self):
         """Test Atom class"""
-        s = ihm.model.Atom(asym_unit='foo', seq_id=1, atom_id='N', x=1.0,
-                             y=2.0, z=3.0)
+        s = ihm.model.Atom(asym_unit='foo', seq_id=1, atom_id='N',
+                           type_symbol='N', x=1.0, y=2.0, z=3.0)
         self.assertEqual(s.asym_unit, 'foo')
         self.assertEqual(s.seq_id, 1)