Use cell list to speed up addHydrogens() (#4816)

wrappers/python/openmm/app/modeller.py

@@ -871,6 +871,9 @@ class Modeller(object):
         newResidueTemplates = {}
         newIndices = []
         acceptors = [atom for atom in self.topology.atoms() if atom.element in (elem.oxygen, elem.nitrogen)]
+        positions = self.positions.value_in_unit(nanometer)
+        acceptorPositions = [positions[a.index] for a in acceptors]
+        cells = _CellList(acceptorPositions, 0.35, None, False)
         for chain in self.topology.chains():
             newChain = newTopology.addChain(chain.id)
             for residue in chain.residues():
@@ -927,14 +930,21 @@ class Modeller(object):
 
                                 nd1IsBonded = False
                                 ne2IsBonded = False
-                                for acceptor in acceptors:
+                                for acceptorIndex in cells.neighbors(nd1Pos.value_in_unit(nanometer)):
+                                    acceptor = acceptors[acceptorIndex]
                                     if acceptor.residue != residue:
                                         acceptorPos = self.positions[acceptor.index]
                                         if isHbond(nd1Pos, hd1Pos, acceptorPos):
                                             nd1IsBonded = True
                                             break
-                                        if isHbond(ne2Pos, he2Pos, acceptorPos):
-                                            ne2IsBonded = True
+                                if not nd1IsBonded:
+                                    for acceptorIndex in cells.neighbors(ne2Pos.value_in_unit(nanometer)):
+                                        acceptor = acceptors[acceptorIndex]
+                                        if acceptor.residue != residue:
+                                            acceptorPos = self.positions[acceptor.index]
+                                            if isHbond(ne2Pos, he2Pos, acceptorPos):
+                                                nd1IsBonded = True
+                                                break
                                 if ne2IsBonded and not nd1IsBonded:
                                     variant = 'HIE'
                                 else:
@@ -1664,6 +1674,14 @@ class _CellList(object):
     """This class organizes atom positions into cells, so the neighbors of a point can be quickly retrieved"""
 
     def __init__(self, positions, maxCutoff, vectors, periodic):
+        if vectors is None:
+            if len(positions) == 0:
+                vectors = (Vec3(maxCutoff, 0, 0), Vec3(0, maxCutoff, 0), Vec3(0, 0, maxCutoff))
+            else:
+                minPos = [min((pos[i] for pos in positions)) for i in range(3)]
+                maxPos = [max((pos[i] for pos in positions)) for i in range(3)]
+                width = [max(maxPos[i]-minPos[i], maxCutoff) for i in range(3)]
+                vectors = (Vec3(width[0], 0, 0), Vec3(0, width[1], 0), Vec3(0, 0, width[2]))
         self.positions = deepcopy(positions)
         self.cells = {}
         self.numCells = tuple((max(1, int(floor(vectors[i][i]/maxCutoff))) for i in range(3)))