Add .pdb saving

atomium/mmtf.py

@@ -498,4 +498,4 @@ def add_het_to_groups(het, group_type_list, group_id_list, group_list, ins_list)
         group_type_list.append(len(group_list) - 1)
     id_, insert = split_residue_id(atoms[0])
     group_id_list.append(id_)
-    ins_list.append(insert)
+    ins_list.append(insert if insert != "?" else "")

atomium/pdb.py

@@ -3,7 +3,10 @@
 from datetime import datetime
 import re
 from itertools import groupby, chain
+import valerius
+from math import ceil
 from .data import CODES
+from .structures import Residue, Ligand
 
 def pdb_string_to_pdb_dict(filestring):
     """Takes a .pdb filestring and turns into a ``dict`` which represents its
@@ -464,3 +467,90 @@ def merge_lines(lines, start, join=" "):
 
     string = join.join([line[start:].strip() for line in lines])
     return string
+
+
+def structure_to_pdb_string(structure):
+    """Converts a :py:class:`.AtomStructure` to a .pdb filestring.
+
+    :param AtomStructure structure: the structure to convert.
+    :rtype: ``str``"""
+
+    lines = []
+    pack_sequences(structure, lines)
+    atoms = sorted(structure.atoms(), key=lambda a: a.id)
+    for i, atom in enumerate(atoms):
+        atom_to_atom_line(atom, lines)
+        if isinstance(atom.structure, Residue) and (
+         atom is atoms[-1] or atoms[i + 1].chain is not atom.chain or
+          isinstance(atoms[i + 1].structure, Ligand)):
+            lines.append("TER")
+    return "\n".join(lines)
+
+
+def pack_sequences(structure, lines):
+    try:
+        for chain in sorted(structure.chains(), key=lambda c: c.id):
+            residues = valerius.from_string(chain.sequence).codes
+            length = len(residues)
+            line_count = ceil(length / 13)
+            for line_num in range(line_count):
+                lines += ["SEQRES {:>3} {} {:>4}  {}".format(
+                 line_num + 1, chain.id, length,
+                 " ".join(residues[line_num * 13: (line_num + 1) * 13])
+                )]
+    except AttributeError: pass
+
+
+def atom_to_atom_line(a, lines):
+    line = "{:6}{:5} {:4} {:3} {:1}{:4}{:1}   "
+    line += "{:>8}{:>8}{:>8}  1.00{:6}          {:>2}{:2}"
+    id_, residue_name, chain_id, residue_id, insert_code = "", "", "", "", ""
+    if a.structure:
+        id_, residue_name = a.structure.id, a.structure._name
+        chain_id = a.chain.id if a.chain is not None else ""
+        residue_id = int("".join([c for c in id_ if c.isdigit()]))
+        insert_code = id_[-1] if id_ and id_[-1].isalpha() else ""
+    atom_name = a._name or ""
+    atom_name = " " + atom_name if len(atom_name) < 4 else atom_name
+    occupancy = "  1.00"
+    line = line.format(
+     "HETATM" if isinstance(a.structure, Ligand) else "ATOM",
+     a.id, atom_name, residue_name, chain_id, residue_id, insert_code,
+     "{:.3f}".format(a.x) if a.x is not None else "",
+     "{:.3f}".format(a.y) if a.y is not None else "",
+     "{:.3f}".format(a.z) if a.z is not None else "",
+     a.bvalue if a.bvalue is not None else "", a.element or "",
+     str(int(a.charge))[::-1] if a.charge else "",
+    )
+    lines.append(line)
+    if a.anisotropy != [0, 0, 0, 0, 0, 0]:
+        lines.append(atom_to_anisou_line(a, atom_name,
+         residue_name, chain_id, residue_id, insert_code))
+
+
+def atom_to_anisou_line(a, name, res_name, chain_id, res_id, insert):
+    """Converts an :py:class:`.Atom` to an ANISOU record.
+
+    :param Atom a: The Atom to pack.
+    :param str name: The atom name to use.
+    :param str res_name: The residue name to use.
+    :param str chain_id: The chain ID to use.
+    :param str res_id: The residue ID to use.
+    :param str insert: The residue insert code to use.
+    :rtype: ``str``"""
+
+    line = "ANISOU{:5} {:4} {:3} {:1}{:4}{:1} "
+    line += "{:>7}{:>7}{:>7}{:>7}{:>7}{:>7}      {:>2}{:2}"
+    anisotropy = [round(x * 10000 )for x in a.anisotropy]
+    line = line.format(
+     a.id, name, res_name, chain_id, res_id, insert,
+     anisotropy[0] if anisotropy[0] is not 0 else "",
+     anisotropy[1] if anisotropy[1] is not 0 else "",
+     anisotropy[2] if anisotropy[2] is not 0 else "",
+     anisotropy[3] if anisotropy[3] is not 0 else "",
+     anisotropy[4] if anisotropy[4] is not 0 else "",
+     anisotropy[5] if anisotropy[5] is not 0 else "",
+     a.element if a.element else "",
+     str(int(a.charge))[::-1] if a.charge else "",
+    )
+    return line

atomium/structures.py

@@ -306,6 +306,11 @@ def save(self, path):
         elif ext == "mmtf":
             from .mmtf import structure_to_mmtf_string
             string = structure_to_mmtf_string(self)
+        elif ext == "pdb":
+            from .pdb import structure_to_pdb_string
+            string = structure_to_pdb_string(self)
+        else:
+            raise ValueError("Unsupported file extension: " + ext)
         save(string, path)
 
 

tests/integration/test_file_saving.py

@@ -20,6 +20,17 @@ def check_file_saving(self, filename):
         f.model.save("tests/integration/files/saved_" + filename)
         f2 = atomium.open("tests/integration/files/saved_" + filename)
         self.assertTrue(f.model.equivalent_to(f2.model))
+        self.assertEqual(len(f.model.chains()), len(f2.model.chains()))
+        for chain1, chain2 in zip(sorted(f.model.chains(), key=lambda c: c.id),
+         sorted(f2.model.chains(), key=lambda c: c.id)):
+            self.assertEqual(chain1.sequence, chain2.sequence)
+            self.assertEqual(chain1.id, chain2.id)
+            self.assertTrue(chain1.equivalent_to(chain2))
+        for lig1, lig2 in zip(sorted(f.model.ligands(), key=lambda c: c.id),
+         sorted(f2.model.ligands(), key=lambda c: c.id)):
+            self.assertEqual(lig1.name, lig2.name)
+            self.assertEqual(lig1.id, lig2.id)
+            self.assertTrue(lig1.equivalent_to(lig2))
 
 
 
@@ -49,6 +60,13 @@ def test_can_save_4y60(self):
         self.check_file_saving("4y60.cif")
 
 
+    def test_chain(self):
+        f = atomium.open("tests/integration/files/1lol.cif")
+        f.model.chain("A").save("tests/integration/files/chaina.cif")
+        chain = atomium.open("tests/integration/files/chaina.cif").model
+        self.assertTrue(f.model.chain("A").equivalent_to(chain))
+
+
 
 class MmtfFileSavingTests(SavingTest):
 
@@ -74,3 +92,44 @@ def test_can_save_5xme(self):
 
     def test_can_save_4y60(self):
         self.check_file_saving("4y60.mmtf")
+
+
+    def test_chain(self):
+        f = atomium.open("tests/integration/files/1lol.mmtf")
+        f.model.chain("A").save("tests/integration/files/chaina.mmtf")
+        chain = atomium.open("tests/integration/files/chaina.mmtf").model
+        self.assertTrue(f.model.chain("A").equivalent_to(chain))
+
+
+
+class PdbFileSavingTests(SavingTest):
+
+    def test_can_save_1lol(self):
+        self.check_file_saving("1lol.pdb")
+
+
+    def test_can_save_1cbn(self):
+        self.check_file_saving("1cbn.pdb")
+
+
+    def test_can_save_1m4x(self):
+        self.check_file_saving("1m4x.pdb")
+
+
+    def test_can_save_1xda(self):
+        self.check_file_saving("1xda.pdb")
+
+
+    def test_can_save_5xme(self):
+        self.check_file_saving("5xme.pdb")
+
+
+    def test_can_save_4y60(self):
+        self.check_file_saving("4y60.pdb")
+
+
+    def test_chain(self):
+        f = atomium.open("tests/integration/files/1lol.pdb")
+        f.model.chain("A").save("tests/integration/files/chaina.pdb")
+        chain = atomium.open("tests/integration/files/chaina.pdb").model
+        self.assertTrue(f.model.chain("A").equivalent_to(chain))

tests/unit/test_mmtf.py

@@ -602,7 +602,7 @@ def test_can_add_existing_het_to_group(self, mock_split):
          "elementList": ["B", "A"], "formalChargeList": [-2, -1]
         }], ["B"]
         het = Mock(_name="HET")
-        mock_split.return_value = "AB"
+        mock_split.return_value = "A?"
         het.atoms.return_value = [Mock(id=2, _name="AA", element="A", charge=-1), Mock(id=1, _name="BB", element="B", charge=-2)]
         add_het_to_groups(het, a, b, c, d)
         self.assertEqual(a, [0, 0])
@@ -611,4 +611,4 @@ def test_can_add_existing_het_to_group(self, mock_split):
          "groupName": "HET", "atomNameList": ["BB", "AA"],
          "elementList": ["B", "A"], "formalChargeList": [-2, -1]
         }])
-        self.assertEqual(d, ["B", "B"])
+        self.assertEqual(d, ["B", ""])

tests/unit/test_pdb.py

@@ -721,3 +721,133 @@ def test_can_vary_join(self):
          merge_lines(self.lines, 8, join="."),
          "89.ij.89"
         )
+
+
+
+class StructureToPdbStringTests(TestCase):
+
+    def setUp(self):
+        self.patch1 = patch("atomium.pdb.pack_sequences")
+        self.patch2 = patch("atomium.pdb.atom_to_atom_line")
+        self.mock_pk = self.patch1.start()
+        self.mock_at = self.patch2.start()
+        self.mock_at.side_effect = lambda a, l: l.append(str(a.id))
+
+
+    def tearDown(self):
+        self.patch1.stop()
+        self.patch2.stop()
+
+
+    def test_can_convert_structure_with_one_chain_to_lines(self):
+        structure = Mock()
+        structure.atoms.return_value = [
+         Mock(id=1, structure=Mock(Residue), chain=1), Mock(id=2, structure=Mock(Residue), chain=1)
+        ]
+        s = structure_to_pdb_string(structure)
+        self.assertEqual(self.mock_pk.call_args_list[0][0][0], structure)
+        self.assertEqual(s, "1\n2\nTER")
+
+
+
+    def test_can_convert_structure_with_two_chains_to_lines(self):
+        structure = Mock()
+        structure.atoms.return_value = [
+         Mock(id=1, structure=Mock(Residue), chain=1), Mock(id=2, structure=Mock(Residue), chain=1),
+         Mock(id=3, structure=Mock(Residue), chain=2), Mock(id=4, structure=Mock(Residue), chain=2)
+        ]
+        s = structure_to_pdb_string(structure)
+        self.assertEqual(self.mock_pk.call_args_list[0][0][0], structure)
+        self.assertEqual(s, "1\n2\nTER\n3\n4\nTER")
+
+
+    def test_can_convert_structure_with_one_chain_and_ligands_to_lines(self):
+        structure = Mock()
+        structure.atoms.return_value = [
+         Mock(id=1, structure=Mock(Residue), chain=1), Mock(id=2, structure=Mock(Residue), chain=1),
+         Mock(id=3, structure=Mock(Ligand), chain=1), Mock(id=4, structure=Mock(Ligand), chain=1)
+        ]
+        s = structure_to_pdb_string(structure)
+        self.assertEqual(self.mock_pk.call_args_list[0][0][0], structure)
+        self.assertEqual(s, "1\n2\nTER\n3\n4")
+
+
+    def test_can_convert_structure_with_two_chains_and_ligands_to_lines(self):
+        structure = Mock()
+        structure.atoms.return_value = [
+         Mock(id=1, structure=Mock(Residue), chain=1), Mock(id=2, structure=Mock(Residue), chain=1),
+         Mock(id=3, structure=Mock(Residue), chain=2), Mock(id=4, structure=Mock(Residue), chain=2),
+         Mock(id=5, structure=Mock(Ligand), chain=1), Mock(id=6, structure=Mock(Ligand), chain=1)
+        ]
+        s = structure_to_pdb_string(structure)
+        self.assertEqual(self.mock_pk.call_args_list[0][0][0], structure)
+        self.assertEqual(s, "1\n2\nTER\n3\n4\nTER\n5\n6")
+
+
+
+class SequencePackingTests(TestCase):
+
+    def test_can_pack_sequence(self):
+        structure = Mock()
+        structure.chains.return_value = [Mock(sequence="ABC" * 10, id=2), Mock(sequence="GA", id=1)]
+        lines = []
+        pack_sequences(structure, lines)
+        self.assertEqual(lines, [
+         "SEQRES   1 1    2  DG DA",
+         "SEQRES   1 2   30  ALA XXX CYS ALA XXX CYS ALA XXX CYS ALA XXX CYS ALA",
+         "SEQRES   2 2   30  XXX CYS ALA XXX CYS ALA XXX CYS ALA XXX CYS ALA XXX",
+         "SEQRES   3 2   30  CYS ALA XXX CYS"
+        ])
+
+
+    def test_can_handle_no_chains(self):
+        structure = Mock()
+        structure.chains.side_effect = AttributeError
+        lines = []
+        pack_sequences(structure, lines)
+        self.assertEqual(lines, [])
+
+
+
+class AtomToAtomLineTests(TestCase):
+
+    @patch("atomium.pdb.atom_to_anisou_line")
+    def test_can_convert_atom_to_line_residue(self, mock_an):
+        atom = Mock(structure=Mock(Residue, id="A100B", _name="RES"),
+         chain=Mock(id="A"), _name="CD", id=100, x=1.2, y=2.7, z=-9, bvalue=12.1,
+          charge=-2, element="C")
+        lines = []
+        atom_to_atom_line(atom, lines)
+        self.assertEqual(lines[0], "ATOM    100  CD  RES A 100B      1.200   2.700  -9.000  1.00  12.1           C2-")
+        self.assertEqual(lines[1], mock_an.return_value)
+        mock_an.assert_called_with(atom, " CD", "RES", "A", 100, "B")
+
+
+    @patch("atomium.pdb.atom_to_anisou_line")
+    def test_can_convert_atom_to_line_ligand(self, mock_an):
+        atom = Mock(structure=Mock(Ligand, id="A100B", _name="RES"),
+         chain=Mock(id="A"), _name="CD", id=100, x=1.2, y=2.7, z=-9, bvalue=12.1,
+          charge=-2, element="C", anisotropy=[0, 0, 0, 0, 0, 0])
+        lines = []
+        atom_to_atom_line(atom, lines)
+        self.assertEqual(lines[0], "HETATM  100  CD  RES A 100B      1.200   2.700  -9.000  1.00  12.1           C2-")
+        self.assertFalse(mock_an.called)
+
+
+    @patch("atomium.pdb.atom_to_anisou_line")
+    def test_can_convert_atom_to_line_bare(self, mock_an):
+        atom = Mock(_name="CD", id=100, x=1.2, y=2.7, z=-9, bvalue=12.1,
+          charge=-2, element="C", anisotropy=[0, 0, 0, 0, 0, 0], structure=None)
+        lines = []
+        atom_to_atom_line(atom, lines)
+        self.assertEqual(lines[0], "ATOM    100  CD                  1.200   2.700  -9.000  1.00  12.1           C2-")
+        self.assertFalse(mock_an.called)
+
+
+
+class AtomToAnisouLinetests(TestCase):
+
+    def test_can_convert_atom_to_anisou_line(self):
+        atom = Mock(id=4, anisotropy=[0.12, 0.24, 0.36, 0.48, 0.512, 0.1], charge=-1, element="P")
+        line = atom_to_anisou_line(atom, "A", "RES", "C", "101", "D")
+        self.assertEqual(line, "ANISOU    4 A    RES C101 D    1200   2400   3600   4800   5120   1000       P1-")