Refactor box bonding

5x speedup of `find_bonds_for_molecules_using_boxes`

propka/bonds.py

@@ -81,10 +81,6 @@ class BondMaker:
         self.num_pi_elec_conj_bonds_ligands = {'N.am': 1, 'N.pl3': 1}
         self.backbone_atoms = list(self.intra_residue_backbone_bonds.keys())
         self.terminal_oxygen_names = ['OXT', 'O\'\'']
-        self.boxes = {}
-        self.num_box_x = None
-        self.num_box_y = None
-        self.num_box_z = None
 
     def find_bonds_for_protein(self, protein):
         """Bonds proteins based on the way atoms normally bond.
@@ -286,14 +282,29 @@ class BondMaker:
         no_atoms = len(atoms)
         for i in range(no_atoms):
             for j in range(i+1, no_atoms):
-                if atoms[i] in atoms[j].bonded_atoms:
+                self._find_bonds_for_atoms(atoms[i], atoms[j])
-                    continue
+
-                if self.check_distance(atoms[i], atoms[j]):
+    def find_bonds_for_atoms_disjoint(self, atoms1, atoms2):
-                    self.make_bond(atoms[i], atoms[j])
+        """Finds all bonds between two disjoint sets of atoms.
-                    # di-sulphide bonds
+
-                    if atoms[i].element == 'S' and atoms[j].element == 'S':
+        Args:
-                        atoms[i].cysteine_bridge = True
+            atoms1:  list of atoms
-                        atoms[j].cysteine_bridge = True
+            atoms2:  list of atoms
+        """
+        for atom1 in atoms1:
+            for atom2 in atoms2:
+                self._find_bonds_for_atoms(atom1, atom2)
+
+    def _find_bonds_for_atoms(self, atom1, atom2):
+        assert atom1 is not atom2
+        if atom1 in atom2.bonded_atoms:
+            return
+        if self.check_distance(atom1, atom2):
+            self.make_bond(atom1, atom2)
+            # di-sulphide bonds
+            if atom1.element == 'S' and atom2.element == 'S':
+                atom1.cysteine_bridge = True
+                atom2.cysteine_bridge = True
 
     def check_distance(self, atom1, atom2):
         """Check distance between two atoms
@@ -344,75 +355,38 @@ class BondMaker:
         Args:
             atoms:  list of atoms for finding bonds
         """
-        box_size = BOX_SIZE
+        box_size = max(BOX_SIZE, self.max_sq_distance**0.5 + 0.01)
-        xmin = POS_MAX
+        boxes = {}
-        xmax = -1.0 * POS_MAX
+
-        ymin = POS_MAX
-        ymax = -1.0 * POS_MAX
-        zmin = POS_MAX
-        zmax = -1.0 * POS_MAX
         for atom in atoms:
-            if atom.x > xmax:
+            x = math.floor(atom.x / box_size)
-                xmax = atom.x
+            y = math.floor(atom.y / box_size)
-            if atom.y > ymax:
+            z = math.floor(atom.z / box_size)
-                ymax = atom.y
+            boxes.setdefault((x, y, z), []).append(atom)
-            if atom.z > zmax:
+
-                zmax = atom.z
+        for (x, y, z), value in boxes.items():
-            if atom.x < xmin:
-                xmin = atom.x
-            if atom.y < ymin:
-                ymin = atom.y
-            if atom.z < zmin:
-                zmin = atom.z
-        xlen = xmax - xmin
-        ylen = ymax - ymin
-        zlen = zmax - zmin
-        self.num_box_x = max(1, int(math.ceil(xlen/box_size)))
-        self.num_box_y = max(1, int(math.ceil(ylen/box_size)))
-        self.num_box_z = max(1, int(math.ceil(zlen/box_size)))
-        self.boxes = {}
-        for x in range(self.num_box_x):
-            for y in range(self.num_box_y):
-                for z in range(self.num_box_z):
-                    self.boxes[(x, y, z)] = []
-        for atom in atoms:
-            x = math.floor((atom.x - xmin)/box_size)
-            y = math.floor((atom.y - ymin)/box_size)
-            z = math.floor((atom.z - zmin)/box_size)
-            self.put_atom_in_box(x, y, z, atom)
-        for value in self.boxes.values():
             self.find_bonds_for_atoms(value)
 
-    def put_atom_in_box(self, x, y, z, atom):
+            for (dx, dy, dz) in [
-        """Put an atom in a box.
+                (-1, -1, -1),
-
+                (-1, -1, 0),
-        Args:
+                (-1, -1, 1),
-            x:  box x-coordinates
+                (-1, 0, -1),
-            y:  box y-coordinates
+                (-1, 0, 0),
-            z:  box z-coordinates
+                (-1, 0, 1),
-            atom:  the atom to place in a box
+                (-1, 1, -1),
-        """
+                (-1, 1, 0),
-        for (dx, dy, dz) in [
+                (-1, 1, 1),
-            (-1, -1, -1),
+                (0, -1, -1),
-            (-1, -1, 0),
+                (0, -1, 0),
-            (-1, -1, 1),
+                (0, -1, 1),
-            (-1, 0, -1),
+                (0, 0, -1),
-            (-1, 0, 0),
+            ]:
-            (-1, 0, 1),
+                try:
-            (-1, 1, -1),
+                    value2 = boxes[x + dx, y + dy, z + dz]
-            (-1, 1, 0),
+                except KeyError:
-            (-1, 1, 1),
+                    continue
-            (0, -1, -1),
+                self.find_bonds_for_atoms_disjoint(value, value2)
-            (0, -1, 0),
-            (0, -1, 1),
-            (0, 0, -1),
-            (0, 0, 0),
-        ]:
-            key = (x + dx, y + dy, z + dz)
-            try:
-                self.boxes[key].append(atom)
-            except KeyError:
-                pass
 
     @staticmethod
     def has_bond(atom1, atom2):