Merge pull request #87 from peastman/molecules

Eliminated a potential stack overflow when identifying molecules

openmmapi/include/openmm/internal/ContextImpl.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2012 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2013 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -250,9 +250,13 @@ public:
     void integratorDeleted() {
         integratorIsDeleted = true;
     }
+    /**
+     * This is the routine that actually computes the list of molecules returned by getMolecules().  Normally
+     * you should never call it.  It is exposed here because the same logic is useful to other classes too.
+     */
+    static std::vector<std::vector<int> > findMolecules(int numParticles, std::vector<std::vector<int> >& particleBonds);
 private:
     friend class Context;
-    static void tagParticlesInMolecule(int particle, int molecule, std::vector<int>& particleMolecule, std::vector<std::vector<int> >& particleBonds);
     Context& owner;
     const System& system;
     Integrator& integrator;

openmmapi/src/ContextImpl.cpp

@@ -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) 2008-2012 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2013 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -288,26 +288,54 @@ const vector<vector<int> >& ContextImpl::getMolecules() const {
         particleBonds[bonds[i].second].push_back(bonds[i].first);
     }
 
-    // Now tag particles by which molecule they belong to.
+    // Now identify particles by which molecule they belong to.
+
+    molecules = findMolecules(numParticles, particleBonds);
+    return molecules;
+}
+
+vector<vector<int> > ContextImpl::findMolecules(int numParticles, vector<vector<int> >& particleBonds) {
+    // This is essentially a recursive algorithm, but it is reformulated as a loop to avoid
+    // stack overflows.  It selects a particle, marks it as a new molecule, then recursively
+    // marks every particle bonded to it as also being in that molecule.
     
     vector<int> particleMolecule(numParticles, -1);
     int numMolecules = 0;
     for (int i = 0; i < numParticles; i++)
-        if (particleMolecule[i] == -1)
-            tagParticlesInMolecule(i, numMolecules++, particleMolecule, particleBonds);
-    molecules.resize(numMolecules);
-    for (int i = 0; i < numParticles; i++)
-        molecules[particleMolecule[i]].push_back(i);
-    return molecules;
-}
+        if (particleMolecule[i] == -1) {
+            // Start a new molecule.
+            
+            vector<int> particleStack;
+            vector<int> neighborStack;
+            particleStack.push_back(i);
+            neighborStack.push_back(0);
+            int molecule = numMolecules++;
             
-void ContextImpl::tagParticlesInMolecule(int particle, int molecule, vector<int>& particleMolecule, vector<vector<int> >& particleBonds) {
-    // Recursively tag particles as belonging to a particular molecule.
+            // Recursively tag all the bonded particles.
             
+            while (particleStack.size() > 0) {
+                int particle = particleStack.back();
                 particleMolecule[particle] = molecule;
-    for (int i = 0; i < (int) particleBonds[particle].size(); i++)
-        if (particleMolecule[particleBonds[particle][i]] == -1)
-            tagParticlesInMolecule(particleBonds[particle][i], molecule, particleMolecule, particleBonds);
+                int& neighbor = neighborStack.back();
+                while (neighbor < particleBonds[particle].size() && particleMolecule[particleBonds[particle][neighbor]] != -1)
+                    neighbor++;
+                if (neighbor < particleBonds[particle].size()) {
+                    particleStack.push_back(particleBonds[particle][neighbor]);
+                    neighborStack.push_back(0);
+                }
+                else {
+                    particleStack.pop_back();
+                    neighborStack.pop_back();
+                }
+            }
+        }
+    
+    // Build the final output vector.
+    
+    vector<vector<int> > molecules(numMolecules);
+    for (int i = 0; i < numParticles; i++)
+        molecules[particleMolecule[i]].push_back(i);
+    return molecules;
 }
 
 static void writeString(ostream& stream, string str) {

platforms/cuda/include/CudaContext.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2009-2012 Stanford University and the Authors.      *
+ * Portions copyright (c) 2009-2013 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -489,7 +489,6 @@ private:
     struct MoleculeGroup;
     class VirtualSiteInfo;
     void findMoleculeGroups();
-    static void tagAtomsInMolecule(int atom, int molecule, std::vector<int>& atomMolecule, std::vector<std::vector<int> >& atomBonds);
     /**
      * Ensure that all molecules marked as "identical" really are identical.  This should be
      * called whenever force field parameters change.  If necessary, it will rebuild the list

platforms/cuda/src/CudaContext.cpp

@@ -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) 2009-2012 Stanford University and the Authors.      *
+ * Portions copyright (c) 2009-2013 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -42,6 +42,7 @@
 #include "openmm/System.h"
 #include "openmm/VirtualSite.h"
 #include "CudaExpressionUtilities.h"
+#include "openmm/internal/ContextImpl.h"
 #include <algorithm>
 #include <cstdlib>
 #include <fstream>
@@ -647,15 +648,6 @@ void CudaContext::clearAutoclearBuffers() {
     }
 }
 
-void CudaContext::tagAtomsInMolecule(int atom, int molecule, vector<int>& atomMolecule, vector<vector<int> >& atomBonds) {
-    // Recursively tag atoms as belonging to a particular molecule.
-
-    atomMolecule[atom] = molecule;
-    for (int i = 0; i < (int) atomBonds[atom].size(); i++)
-        if (atomMolecule[atomBonds[atom][i]] == -1)
-            tagAtomsInMolecule(atomBonds[atom][i], molecule, atomMolecule, atomBonds);
-}
-
 /**
  * This class ensures that atom reordering doesn't break virtual sites.
  */
@@ -750,16 +742,14 @@ void CudaContext::findMoleculeGroups() {
             }
         }
 
-        // Now tag atoms by which molecule they belong to.
+        // Now identify atoms by which molecule they belong to.
 
-        vector<int> atomMolecule(numAtoms, -1);
-        int numMolecules = 0;
-        for (int i = 0; i < numAtoms; i++)
-            if (atomMolecule[i] == -1)
-                tagAtomsInMolecule(i, numMolecules++, atomMolecule, atomBonds);
-        vector<vector<int> > atomIndices(numMolecules);
-        for (int i = 0; i < numAtoms; i++)
-            atomIndices[atomMolecule[i]].push_back(i);
+        vector<vector<int> > atomIndices = ContextImpl::findMolecules(numAtoms, atomBonds);
+        int numMolecules = atomIndices.size();
+        vector<int> atomMolecule(numAtoms);
+        for (int i = 0; i < (int) atomIndices.size(); i++)
+            for (int j = 0; j < (int) atomIndices[i].size(); j++)
+                atomMolecule[atomIndices[i][j]] = i;
 
         // Construct a description of each molecule.
 

platforms/opencl/include/OpenCLContext.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2009-2011 Stanford University and the Authors.      *
+ * Portions copyright (c) 2009-2013 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -589,7 +589,6 @@ private:
     struct MoleculeGroup;
     class VirtualSiteInfo;
     void findMoleculeGroups();
-    static void tagAtomsInMolecule(int atom, int molecule, std::vector<int>& atomMolecule, std::vector<std::vector<int> >& atomBonds);
     /**
      * Ensure that all molecules marked as "identical" really are identical.  This should be
      * called whenever force field parameters change.  If necessary, it will rebuild the list

platforms/opencl/src/OpenCLContext.cpp

@@ -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) 2009-2012 Stanford University and the Authors.      *
+ * Portions copyright (c) 2009-2013 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -39,6 +39,7 @@
 #include "openmm/Platform.h"
 #include "openmm/System.h"
 #include "openmm/VirtualSite.h"
+#include "openmm/internal/ContextImpl.h"
 #include <algorithm>
 #include <fstream>
 #include <iostream>
@@ -618,15 +619,6 @@ void OpenCLContext::reduceBuffer(OpenCLArray& array, int numBuffers) {
     executeKernel(reduceReal4Kernel, bufferSize, 128);
 }
 
-void OpenCLContext::tagAtomsInMolecule(int atom, int molecule, vector<int>& atomMolecule, vector<vector<int> >& atomBonds) {
-    // Recursively tag atoms as belonging to a particular molecule.
-
-    atomMolecule[atom] = molecule;
-    for (int i = 0; i < (int) atomBonds[atom].size(); i++)
-        if (atomMolecule[atomBonds[atom][i]] == -1)
-            tagAtomsInMolecule(atomBonds[atom][i], molecule, atomMolecule, atomBonds);
-}
-
 /**
  * This class ensures that atom reordering doesn't break virtual sites.
  */
@@ -722,16 +714,14 @@ void OpenCLContext::findMoleculeGroups() {
             }
         }
 
-        // Now tag atoms by which molecule they belong to.
+        // Now identify atoms by which molecule they belong to.
 
-        vector<int> atomMolecule(numAtoms, -1);
-        int numMolecules = 0;
-        for (int i = 0; i < numAtoms; i++)
-            if (atomMolecule[i] == -1)
-                tagAtomsInMolecule(i, numMolecules++, atomMolecule, atomBonds);
-        vector<vector<int> > atomIndices(numMolecules);
-        for (int i = 0; i < numAtoms; i++)
-            atomIndices[atomMolecule[i]].push_back(i);
+        vector<vector<int> > atomIndices = ContextImpl::findMolecules(numAtoms, atomBonds);
+        int numMolecules = atomIndices.size();
+        vector<int> atomMolecule(numAtoms);
+        for (int i = 0; i < (int) atomIndices.size(); i++)
+            for (int j = 0; j < (int) atomIndices[i].size(); j++)
+                atomMolecule[atomIndices[i][j]] = i;
 
         // Construct a description of each molecule.
 

tests/TestFindMolecules.cpp

+/* -------------------------------------------------------------------------- *
+ *                                   OpenMM                                   *
+ * -------------------------------------------------------------------------- *
+ * This is part of the OpenMM molecular simulation toolkit originating from   *
+ * Simbios, the NIH National Center for Physics-Based Simulation of           *
+ * Biological Structures at Stanford, funded under the NIH Roadmap for        *
+ * Medical Research, grant U54 GM072970. See https://simtk.org.               *
+ *                                                                            *
+ * Portions copyright (c) 2013 Stanford University and the Authors.           *
+ * Authors: Peter Eastman                                                     *
+ * Contributors:                                                              *
+ *                                                                            *
+ * Permission is hereby granted, free of charge, to any person obtaining a    *
+ * copy of this software and associated documentation files (the "Software"), *
+ * to deal in the Software without restriction, including without limitation  *
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,   *
+ * and/or sell copies of the Software, and to permit persons to whom the      *
+ * Software is furnished to do so, subject to the following conditions:       *
+ *                                                                            *
+ * The above copyright notice and this permission notice shall be included in *
+ * all copies or substantial portions of the Software.                        *
+ *                                                                            *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   *
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    *
+ * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,    *
+ * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR      *
+ * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE  *
+ * USE OR OTHER DEALINGS IN THE SOFTWARE.                                     *
+ * -------------------------------------------------------------------------- */
+
+#include "openmm/internal/AssertionUtilities.h"
+#include "openmm/Context.h"
+#include "openmm/internal/ContextImpl.h"
+#include "openmm/HarmonicBondForce.h"
+#include "openmm/Platform.h"
+#include "openmm/System.h"
+#include "openmm/VerletIntegrator.h"
+#include <iostream>
+#include <vector>
+
+using namespace OpenMM;
+using namespace std;
+
+void testFindMolecules() {
+    const int numMolecules = 5;
+    const int moleculeSize[] = {1, 10, 100, 1000, 10000};
+    vector<int> particleMolecule;
+    System system;
+    HarmonicBondForce* bonds = new HarmonicBondForce();
+    system.addForce(bonds);
+    for (int i = 0; i < numMolecules; i++)
+        for (int j = 0; j < moleculeSize[i]; j++) {
+            int index = system.addParticle(1.0);
+            particleMolecule.push_back(i);
+            if (j > 0)
+                bonds->addBond(index, index-1, 1.0, 1.0);
+        }
+    VerletIntegrator integrator(1.0);
+    Context context(system, integrator, Platform::getPlatformByName("Reference"));
+    ContextImpl* contextImpl = *reinterpret_cast<ContextImpl**>(&context);
+    const vector<vector<int> >& molecules = contextImpl->getMolecules();
+    ASSERT_EQUAL(numMolecules, molecules.size());
+    for (int i = 0; i < numMolecules; i++) {
+        ASSERT_EQUAL(moleculeSize[i], molecules[i].size());
+        for (int j = 0; j < moleculeSize[i]; j++)
+            ASSERT_EQUAL(particleMolecule[molecules[i][j]], i);
+    }
+}
+
+int main() {
+    try {
+        testFindMolecules();
+    }
+    catch(const exception& e) {
+        cout << "exception: " << e.what() << endl;
+        return 1;
+    }
+    cout << "Done" << endl;
+    return 0;
+}