diff --git a/python/lammps/data.py b/python/lammps/data.py index 1613936ddf..892d628356 100644 --- a/python/lammps/data.py +++ b/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 diff --git a/python/lammps/numpy_wrapper.py b/python/lammps/numpy_wrapper.py index ce64d68c90..1ecaf3163b 100644 --- a/python/lammps/numpy_wrapper.py +++ b/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 diff --git a/unittest/python/python-commands.py b/unittest/python/python-commands.py index 1c388460d1..530cea8b13 100644 --- a/unittest/python/python-commands.py +++ b/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("""