Add ANISOU saving in .pdb files

atomium/files/pdb2pdbdict.py

@@ -74,7 +74,7 @@ def atom_to_atom_dict(atom):
      "residue_name": residue_name, "full_id": id_,
      "chain_id": chain_id, "residue_id": residue_id, "insert_code": insert_code,
      "x": atom.x, "y": atom.y, "z": atom.z,
-     "occupancy": 1.0,
+     "occupancy": 1.0, "anisotropy": atom.anisotropy,
      "element": atom.element, "charge": atom.charge,
      "temp_factor": atom.bfactor if atom.bfactor else None,
     }

atomium/files/pdbdict2pdbstring.py

@@ -147,9 +147,13 @@ def pack_model(lines, model_dict, multi=0):
         for residue in chain["residues"]:
             for atom in residue["atoms"]:
                 lines.append(atom_dict_to_atom_line(atom, hetero=False))
+                if atom["anisotropy"] != [0] * 6:
+                    lines.append(atom_dict_to_anisou_line(atom))
     for molecule in model_dict["molecules"]:
         for atom in molecule["atoms"]:
             lines.append(atom_dict_to_atom_line(atom, hetero=True))
+            if atom["anisotropy"] != [0] * 6:
+                lines.append(atom_dict_to_anisou_line(atom))
     if multi > 0:
         lines.append("ENDMDL".ljust(80))
 
@@ -188,6 +192,37 @@ def atom_dict_to_atom_line(d, hetero=False):
     return line
 
 
+def atom_dict_to_anisou_line(d):
+    """Converts an atom ``dict`` to an ANISOU record.
+
+    :param dict d: The atom dictionary to pack."""
+
+    line = "{:6}{:5} {:4}{:1}{:3} {:1}{:4}{:1} "
+    line += "{:>7}{:>7}{:>7}{:>7}{:>7}{:>7}      {:>2}{:2}"
+    atom_name = d["atom_name"] if d["atom_name"] else ""
+    atom_name = " " + atom_name if len(atom_name) < 4 else atom_name
+    anisotropy = [round(x * 10000 )for x in d["anisotropy"]]
+    line = line.format(
+     "ANISOU",
+     d["atom_id"],
+     atom_name,
+     d["alt_loc"] if d["alt_loc"] else "",
+     d["residue_name"] if d["residue_name"] else "",
+     d["chain_id"],
+     d["residue_id"] if d["residue_id"] else "",
+     d["insert_code"],
+     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 "",
+     d["element"] if d["element"] else "",
+     str(d["charge"])[::-1] if d["charge"] else "",
+    )
+    return line
+
+
 def pack_connections(lines, pdb_dict):
     """Adds CONECT records to a list of lines.
 

tests/integration/files/1lol_output.pdb

@@ -10,6 +10,7 @@ REMARK   2 RESOLUTION.    1.90 ANGSTROMS.
 REMARK   3
 REMARK   3   R VALUE            (WORKING SET) : 0.193
 ATOM      1  N   VAL A  11       3.696  33.898  63.219  1.00  21.5           N
+ANISOU    1  N   VAL A  11     2406   1892   1614    198    519   -328       N
 ATOM      2  CA  VAL A  11       3.198  33.218  61.983  1.00 19.76           C
 ATOM      3  C   VAL A  11       3.914  31.863  61.818  1.00 19.29           C
 ATOM      4  O   VAL A  11       5.132  31.792  61.932  1.00 19.78           O

tests/integration/files/chaina_output.pdb

@@ -1,4 +1,5 @@
 ATOM      1  N   VAL A  11       3.696  33.898  63.219  1.00  21.5           N  
+ANISOU    1  N   VAL A  11     2406   1892   1614    198    519   -328       N
 ATOM      2  CA  VAL A  11       3.198  33.218  61.983  1.00 19.76           C  
 ATOM      3  C   VAL A  11       3.914  31.863  61.818  1.00 19.29           C  
 ATOM      4  O   VAL A  11       5.132  31.792  61.932  1.00 19.78           O  

tests/unit/files_tests/test_pdb_dict_to_pdb_string.py

@@ -226,48 +226,57 @@ def test_can_pack_structure_multiple_models(self, mock_con, mock_model):
 class ModelPackingTests(TestCase):
 
     def setUp(self):
+        self.atoms =  [{"id": c + i + n, "anisotropy": [0] * 6}
+         for c in "rm" for i in "1234" for n in "12"]
+        self.atoms[1]["anisotropy"][0] = 1
+        self.atoms[-2]["anisotropy"][0] = 1
+        self.lines = []
         self.model_dict = {
          "chains": [
-          {"residues": [{"atoms": ["r11", "r12"]}, {"atoms": ["r21", "r22"]}]},
-          {"residues": [{"atoms": ["r31", "r32"]}, {"atoms": ["r41", "r42"]}]},
+          {"residues": [{"atoms": self.atoms[:2]}, {"atoms": self.atoms[2:4]}]},
+          {"residues": [{"atoms": self.atoms[4:6]}, {"atoms": self.atoms[6:8]}]},
          ],
          "molecules": [
-          {"atoms": ["m11", "m12"]}, {"atoms": ["m21", "m22"]},
-          {"atoms": ["m31", "m32"]}, {"atoms": ["m41", "m42"]}
+          {"atoms": self.atoms[8:10]}, {"atoms": self.atoms[10:12]},
+          {"atoms": self.atoms[12:14]}, {"atoms": self.atoms[14:16]}
          ]
         }
-        self.lines = []
+
 
 
     @patch("atomium.files.pdbdict2pdbstring.atom_dict_to_atom_line")
-    def test_can_pack_sole_model(self, mock_line):
+    @patch("atomium.files.pdbdict2pdbstring.atom_dict_to_anisou_line")
+    def test_can_pack_sole_model(self, mock_an, mock_line):
         mock_line.side_effect = ["a" + str(i) for i in range(16)]
+        mock_an.return_value = "AN"
         pack_model(self.lines, self.model_dict, multi=0)
-        for char in ["r", "m"]:
-            for num1 in ["1", "2", "3", "4"]:
-                for num2 in ["1", "2"]:
-                    mock_line.assert_any_call(
-                     char + num1 + num2, hetero=char == "m"
-                    )
+        for atom in self.atoms:
+            mock_line.assert_any_call(
+             atom, hetero=atom["id"].startswith("m")
+            )
+        mock_an.assert_any_call(self.atoms[1])
+        mock_an.assert_any_call(self.atoms[-2])
         self.assertEqual(self.lines, [
-         "a0", "a1", "a2", "a3", "a4", "a5", "a6", "a7",
-         "a8", "a9", "a10", "a11", "a12", "a13", "a14", "a15"
+         "a0", "a1", "AN", "a2", "a3", "a4", "a5", "a6", "a7",
+         "a8", "a9", "a10", "a11", "a12", "a13", "a14", "AN", "a15"
         ])
 
 
     @patch("atomium.files.pdbdict2pdbstring.atom_dict_to_atom_line")
-    def test_can_pack_model_in_series(self, mock_line):
+    @patch("atomium.files.pdbdict2pdbstring.atom_dict_to_anisou_line")
+    def test_can_pack_model_in_series(self, mock_an, mock_line):
         mock_line.side_effect = ["a" + str(i) for i in range(16)]
+        mock_an.return_value = "AN"
         pack_model(self.lines, self.model_dict, multi=5)
-        for char in ["r", "m"]:
-            for num1 in ["1", "2", "3", "4"]:
-                for num2 in ["1", "2"]:
-                    mock_line.assert_any_call(
-                     char + num1 + num2, hetero=char == "m"
-                    )
+        for atom in self.atoms:
+            mock_line.assert_any_call(
+             atom, hetero=atom["id"].startswith("m")
+            )
+        mock_an.assert_any_call(self.atoms[1])
+        mock_an.assert_any_call(self.atoms[-2])
         self.assertEqual(self.lines, [
-         "MODEL        5".ljust(80), "a0", "a1", "a2", "a3", "a4", "a5", "a6", "a7",
-         "a8", "a9", "a10", "a11", "a12", "a13", "a14", "a15", "ENDMDL".ljust(80)
+         "MODEL        5".ljust(80), "a0", "a1", "AN", "a2", "a3", "a4", "a5", "a6", "a7",
+         "a8", "a9", "a10", "a11", "a12", "a13", "a14", "AN", "a15", "ENDMDL".ljust(80)
         ])
 
 
@@ -350,6 +359,55 @@ def test_can_convert_heteroatom_dict_to_line(self):
 
 
 
+class AtomDictToAnisouLineTests(TestCase):
+
+    def setUp(self):
+        self.atom_dict = {
+         "atom_id": 107, "atom_name": "N", "alt_loc": "A",
+         "residue_name": "GLY", "anisotropy": [0.34, -0.3456, 0.098, 0, -0.1231343, 0.9],
+         "chain_id": "B", "residue_id": 13, "insert_code": "C",
+         "x": 12.681, "y": 7.302, "z": -25.21,
+         "occupancy": 0.5, "temp_factor": 15.5,
+         "element": "N", "charge": -2,
+        }
+
+
+    def test_can_convert_empty_atom_dict_to_line(self):
+        for key in self.atom_dict:
+            self.atom_dict[key] = None
+        self.atom_dict["atom_id"] = 0
+        self.atom_dict["chain_id"], self.atom_dict["insert_code"] = "", ""
+        self.atom_dict["occupancy"], self.atom_dict["charge"] = 1, 0
+        self.atom_dict["anisotropy"] = [0, 0, 0, 0, 0, 0]
+        line = atom_dict_to_anisou_line(self.atom_dict)
+        self.assertEqual(line[:6], "ANISOU")
+        self.assertEqual(line[6:11], "    0")
+        self.assertEqual(line[11:].strip(), "")
+
+
+    def test_can_convert_full_atom_dict_to_line(self):
+        line = atom_dict_to_anisou_line(self.atom_dict)
+        self.assertEqual(line[:6], "ANISOU")
+        self.assertEqual(line[6:11], "  107")
+        self.assertEqual(line[11], " ")
+        self.assertEqual(line[12:16], " N  ")
+        self.assertEqual(line[16], "A")
+        self.assertEqual(line[17:20], "GLY")
+        self.assertEqual(line[20], " ")
+        self.assertEqual(line[21], "B")
+        self.assertEqual(line[22:26], "  13")
+        self.assertEqual(line[26], "C")
+        self.assertEqual(line[28:35], "   3400")
+        self.assertEqual(line[35:42], "  -3456")
+        self.assertEqual(line[42:49], "    980")
+        self.assertEqual(line[49:56], "       ")
+        self.assertEqual(line[56:63], "  -1231")
+        self.assertEqual(line[63:70], "   9000")
+        self.assertEqual(line[76:78], " N")
+        self.assertEqual(line[78:], "2-")
+
+
+
 class ConnectionsPackingTests(TestCase):
 
     def test_can_pack_connections(self):

tests/unit/files_tests/test_pdb_to_pdb_dict.py

@@ -115,6 +115,7 @@ def setUp(self):
         self.atom.element = "N"
         self.atom.charge = -2
         self.atom.bfactor = 12.5
+        self.atom.anisotropy = [1, 2, 0, 3, 3, 4]
         self.residue = Mock()
         self.residue.name = "GLY"
         self.residue.id = "A13B"
@@ -131,7 +132,7 @@ def test_can_convert_full_atom_to_dict(self):
          "residue_name": "GLY",
          "chain_id": "A", "residue_id": 13, "insert_code": "B", "full_id": "A13B",
          "x": 12.681, "y": 37.302, "z": -25.211,
-         "occupancy": 1.0, "temp_factor": 12.5,
+         "occupancy": 1.0, "temp_factor": 12.5, "anisotropy": [1, 2, 0, 3, 3, 4],
          "element": "N", "charge": -2,
         })
 
@@ -149,7 +150,7 @@ def test_can_convert_full_heteroatom_to_dict(self):
          "residue_name": "SUC",
          "chain_id": "A", "residue_id": 200, "insert_code": "", "full_id": "A200",
          "x": 12.681, "y": 37.302, "z": -25.211,
-         "occupancy": 1.0, "temp_factor": 12.5,
+         "occupancy": 1.0, "temp_factor": 12.5, "anisotropy": [1, 2, 0, 3, 3, 4],
          "element": "N", "charge": -2,
         })