add find method to neighbor list wrapper classes

python/lammps/data.py

@@ -52,6 +52,8 @@ class NeighList:
 
     def get(self, element):
         """
+        Access a specific neighbor list entry. "element" must be a number from 0 to the size-1 of the list
+
         :return: tuple with atom local index, number of neighbors and ctypes pointer to neighbor's local atom indices
         :rtype:  (int, int, ctypes.POINTER(c_int))
         """
@@ -71,3 +73,20 @@ class NeighList:
 
         for ii in range(inum):
             yield self.get(ii)
+
+    def find(self, iatom):
+        """
+        Find the neighbor list for a specific (local) atom iatom.
+        If there is no list for iatom, (-1, None) is returned.
+
+        :return: tuple with number of neighbors and ctypes pointer to neighbor's local atom indices
+        :rtype:  (int, ctypes.POINTER(c_int))
+        """
+
+        inum = self.size
+        for ii in range(inum):
+            idx, numneigh, neighbors = self.get(ii)
+            if idx == iatom:
+                return numneigh, neighbors
+
+        return -1, None

python/lammps/numpy_wrapper.py

@@ -331,8 +331,25 @@ class NumPyNeighList(NeighList):
 
     def get(self, element):
         """
+        Access a specific neighbor list entry. "element" must be a number from 0 to the size-1 of the list
+
         :return: tuple with atom local index, numpy array of neighbor local atom indices
         :rtype:  (int, numpy.array)
         """
         iatom, neighbors = self.lmp.numpy.get_neighlist_element_neighbors(self.idx, element)
         return iatom, neighbors
+
+    def find(self, iatom):
+        """
+        Find the neighbor list for a specific (local) atom iatom.
+        If there is no list for iatom, None is returned.
+
+        :return: numpy array of neighbor local atom indices
+        :rtype:  numpy.array or None
+        """
+        inum = self.size
+        for ii in range(inum):
+          idx, neighbors = self.get(ii)
+          if idx == iatom:
+            return neighbors
+        return None

unittest/python/python-commands.py

@@ -122,9 +122,9 @@ create_atoms 1 single &
         run 0 post no""")
 
         idx = self.lmp.find_pair_neighlist("lj/cut")
-        self.assertEqual(0, 0)
+        self.assertNotEqual(idx, -1)
         self.assertEqual(self.lmp.find_pair_neighlist("morse"), -1)
-        nlist = self.lmp.get_neighlist(0)
+        nlist = self.lmp.get_neighlist(idx)
         self.assertEqual(len(nlist), 2)
         atom_i, numneigh_i, neighbors_i = nlist[0]
         atom_j, numneigh_j, _ = nlist[1]
@@ -167,6 +167,15 @@ create_atoms 1 single &
             self.assertEqual(idx,i)
             self.assertEqual(num,nlocal-1-i)
 
+        # look up neighbor list by atom index
+        num, neighs = nlist.find(2)
+        self.assertEqual(num,4)
+        self.assertIsNotNone(neighs,None)
+        # this one will fail
+        num, neighs = nlist.find(10)
+        self.assertEqual(num,-1)
+        self.assertIsNone(neighs,None)
+
     @unittest.skipIf(not has_manybody,"Full neighbor list test for manybody potential")
     def testNeighborListFull(self):
         self.lmp.commands_string("""