From 3db52f1aa16a6f252a671440496a4e24ae6ed507 Mon Sep 17 00:00:00 2001
From: Peter Eastman <peastman@stanford.edu>
Date: Thu, 16 Feb 2023 10:55:15 -0800
Subject: [PATCH] Optimizations to _findNearestDistance() (#268)

* Optimizations to _findNearestDistance()

* Minor changes based on review comments
---
 pdbfixer/pdbfixer.py | 41 ++++++++++++++++++++++++++++-------------
 1 file changed, 28 insertions(+), 13 deletions(-)

diff --git a/pdbfixer/pdbfixer.py b/pdbfixer/pdbfixer.py
index 26be392..431b021 100644
--- a/pdbfixer/pdbfixer.py
+++ b/pdbfixer/pdbfixer.py
@@ -6,7 +6,7 @@
 Biological Structures at Stanford, funded under the NIH Roadmap for
 Medical Research, grant U54 GM072970. See https://simtk.org.
 
-Portions copyright (c) 2013-2022 Stanford University and the Authors.
+Portions copyright (c) 2013-2023 Stanford University and the Authors.
 Authors: Peter Eastman
 Contributors:
 
@@ -954,8 +954,18 @@ def addMissingAtoms(self, seed=None):
             mm.LocalEnergyMinimizer.minimize(context)
             state = context.getState(getPositions=True)
             if newTopology.getNumResidues() > 1:
-                nearest = self._findNearestDistance(context, newTopology, newAtoms)
-                if nearest < 0.13:
+                # When looking for pairs of atoms that are too close to each other, exclude pairs that
+                # are in the same residue or are directly bonded to each other.
+
+                exclusions = dict((atom, {a.index for a in atom.residue.atoms()}) for atom in newAtoms)
+                for a1, a2 in newTopology.bonds():
+                    if a1 in exclusions:
+                        exclusions[a1].add(a2.index)
+                    if a2 in exclusions:
+                        exclusions[a2].add(a1.index)
+                cutoff = 0.13
+                nearest = self._findNearestDistance(context, newAtoms, cutoff, exclusions)
+                if nearest < cutoff:
 
                     # Some atoms are very close together.  Run some dynamics while slowly increasing the strength of the
                     # repulsive interaction to try to improve the result.
@@ -963,11 +973,11 @@ def addMissingAtoms(self, seed=None):
                     for i in range(10):
                         context.setParameter('C', 0.15*(i+1))
                         integrator.step(200)
-                        d = self._findNearestDistance(context, newTopology, newAtoms)
+                        d = self._findNearestDistance(context, newAtoms, cutoff, exclusions)
                         if d > nearest:
                             nearest = d
                             state = context.getState(getPositions=True)
-                            if nearest >= 0.13:
+                            if nearest >= cutoff:
                                 break
                     context.setState(state)
                     context.setParameter('C', 1.0)
@@ -1195,18 +1205,23 @@ def _createForceField(self, newTopology, water):
                 forcefield._templateSignatures[signature] = [template]
         return forcefield
 
-    def _findNearestDistance(self, context, topology, newAtoms):
+    def _findNearestDistance(self, context, newAtoms, cutoff, exclusions):
         """Given a set of newly added atoms, find the closest distance between one of those atoms and another atom."""
 
         positions = context.getState(getPositions=True).getPositions(asNumpy=True).value_in_unit(unit.nanometer)
-        atomResidue = [atom.residue for atom in topology.atoms()]
-        nearest = sys.float_info.max
+        boxSize = np.max(positions, axis=0)-np.min(positions, axis=0)
+        boxVectors = [(boxSize[0], 0, 0), (0, boxSize[1], 0), (0, 0, boxSize[2])]
+        cells = app.modeller._CellList(positions, cutoff, boxVectors, False)
+        nearest_squared = sys.float_info.max
         for atom in newAtoms:
-            p = positions-positions[atom.index]
-            dist = math.sqrt(min(np.dot(p[i], p[i]) for i in range(len(atomResidue)) if atomResidue[i] != atom.residue))
-            if dist < nearest:
-                nearest = dist
-        return nearest
+            excluded = exclusions[atom]
+            for i in cells.neighbors(positions[atom.index]):
+                if i not in excluded:
+                    p = positions[atom.index]-positions[i]
+                    dist_squared = np.dot(p, p)
+                    if dist_squared < nearest_squared:
+                        nearest_squared = dist_squared
+        return np.sqrt(nearest_squared)
 
 
 def main():