Create ctypes only neighbor list API variant

This moves the lammps.get_neighlist() method to lammps.numpy.get_neighlist().
lammps.get_neighlist() now returns a NeighList object, while the NumPy variants
returns a NumPyNeighList object. The main difference between the two is that while
the ctypes variant returns neighlist elements as

atom idx (int), numneighs (int), neighbors (POINTER(c_int))

the NumPy variant returns

atom idx (int), neighbors (numpy.array)

doc/src/Python_module.rst

@@ -152,3 +152,7 @@ Classes representing internal objects
 .. autoclass:: lammps.NeighList
    :members:
    :no-undoc-members:
+
+.. autoclass:: lammps.NumPyNeighList
+   :members:
+   :no-undoc-members:

examples/python/in.fix_python_invoke_neighlist

@@ -32,8 +32,9 @@ def post_force_callback(lmp, v):
     t = L.extract_global("ntimestep", 0)
     print(pid_prefix, "### POST_FORCE ###", t)
 
-    #mylist = L.get_neighlist(0)
+    #mylist = L.numpy.get_neighlist(0)
-    mylist = L.find_pair_neighlist("lj/cut", request=0)
+    idx = L.find_pair_neighlist("lj/cut", request=0)
+    mylist = L.numpy.get_neighlist(idx)
     print(pid_prefix, mylist)
     nlocal = L.extract_global("nlocal")
     nghost = L.extract_global("nghost")
@@ -43,8 +44,8 @@ def post_force_callback(lmp, v):
     v = L.numpy.extract_atom("v", nelem=nlocal+nghost, dim=3)
     f = L.numpy.extract_atom("f", nelem=nlocal+nghost, dim=3)
 
-    for iatom, numneigh, neighs in mylist:
+    for iatom, neighs in mylist:
-      print(pid_prefix, "- {}".format(iatom), x[iatom], v[iatom], f[iatom], " : ",  numneigh, "Neighbors")
+      print(pid_prefix, "- {}".format(iatom), x[iatom], v[iatom], f[iatom], " : ",  len(neighs), "Neighbors")
       for jatom in neighs:
         if jatom < nlocal:
             print(pid_prefix, "    *  ", jatom, x[jatom], v[jatom], f[jatom])

python/lammps.py

@@ -81,7 +81,12 @@ class MPIAbortException(Exception):
 class NeighList:
     """This is a wrapper class that exposes the contents of a neighbor list.
 
-    It can be used like a regular Python list.
+    It can be used like a regular Python list. Each element is a tuple of:
+
+    * the atom local index
+    * its number of neighbors
+    * and a pointer to an c_int array containing local atom indices of its
+      neighbors
 
     Internally it uses the lower-level LAMMPS C-library interface.
 
@@ -109,8 +114,8 @@ class NeighList:
 
     def get(self, element):
         """
-        :return: tuple with atom local index, number of neighbors and array of neighbor local atom indices
+        :return: tuple with atom local index, numpy array of neighbor local atom indices
-        :rtype:  (int, int, numpy.array)
+        :rtype:  (int, int, ctypes.POINTER(c_int))
         """
         iatom, numneigh, neighbors = self.lmp.get_neighlist_element_neighbors(self.idx, element)
         return iatom, numneigh, neighbors
@@ -129,6 +134,35 @@ class NeighList:
         for ii in range(inum):
             yield self.get(ii)
 
+# -------------------------------------------------------------------------
+
+class NumPyNeighList(NeighList):
+    """This is a wrapper class that exposes the contents of a neighbor list.
+
+    It can be used like a regular Python list. Each element is a tuple of:
+
+    * the atom local index
+    * a NumPy array containing the local atom indices of its neighbors
+
+    Internally it uses the lower-level LAMMPS C-library interface.
+
+    :param lmp: reference to instance of :py:class:`lammps`
+    :type  lmp: lammps
+    :param idx: neighbor list index
+    :type  idx: int
+    """
+    def __init__(self, lmp, idx):
+      super(NumPyNeighList, self).__init__(lmp, idx)
+
+    def get(self, element):
+        """
+        :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
+
+
 # -------------------------------------------------------------------------
 # -------------------------------------------------------------------------
 
@@ -1537,11 +1571,15 @@ class lammps(object):
     self.callback[fix_name] = { 'function': cFunc, 'caller': caller }
     self.lib.lammps_set_fix_external_callback(self.lmp, fix_name.encode(), cFunc, cCaller)
 
+
   # -------------------------------------------------------------------------
 
   def get_neighlist(self, idx):
     """Returns an instance of :class:`NeighList` which wraps access to the neighbor list with the given index
 
+    See :py:meth:`lammps.numpy.get_neighlist() <lammps.numpy_wrapper.get_neighlist()>` if you want to use
+    NumPy arrays instead of ``c_int`` pointers.
+
     :param idx: index of neighbor list
     :type  idx: int
     :return: an instance of :class:`NeighList` wrapping access to neighbor list data
@@ -1549,7 +1587,37 @@ class lammps(object):
     """
     if idx < 0:
         return None
-    return NeighList(self, idx)
+    return NeighList(self.lmp, idx)
+
+  # -------------------------------------------------------------------------
+
+  def get_neighlist_size(self, idx):
+    """Return the number of elements in neighbor list with the given index
+
+    :param idx: neighbor list index
+    :type  idx: int
+    :return: number of elements in neighbor list with index idx
+    :rtype:  int
+     """
+    return self.lib.lammps_neighlist_num_elements(self.lmp, idx)
+
+  # -------------------------------------------------------------------------
+
+  def get_neighlist_element_neighbors(self, idx, element):
+    """Return data of neighbor list entry
+
+    :param element: neighbor list index
+    :type  element: int
+    :param element: neighbor list element index
+    :type  element: int
+    :return: tuple with atom local index, number of neighbors and array of neighbor local atom indices
+    :rtype:  (int, int, POINTER(c_int))
+    """
+    c_iatom = c_int()
+    c_numneigh = c_int()
+    c_neighbors = POINTER(c_int)()
+    self.lib.lammps_neighlist_element_neighbors(self.lmp, idx, element, byref(c_iatom), byref(c_numneigh), byref(c_neighbors))
+    return c_iatom.value, c_numneigh.value, c_neighbors
 
   # -------------------------------------------------------------------------
 
@@ -1579,7 +1647,7 @@ class lammps(object):
     style = style.encode()
     exact = int(exact)
     idx = self.lib.lammps_find_pair_neighlist(self.lmp, style, exact, nsub, request)
-    return self.get_neighlist(idx)
+    return idx
 
   # -------------------------------------------------------------------------
 
@@ -1595,7 +1663,7 @@ class lammps(object):
      """
     fixid = fixid.encode()
     idx = self.lib.lammps_find_fix_neighlist(self.lmp, fixid, request)
-    return self.get_neighlist(idx)
+    return idx
 
   # -------------------------------------------------------------------------
 
@@ -1611,38 +1679,7 @@ class lammps(object):
      """
     computeid = computeid.encode()
     idx = self.lib.lammps_find_compute_neighlist(self.lmp, computeid, request)
-    return self.get_neighlist(idx)
+    return idx
-
-  # -------------------------------------------------------------------------
-
-  def get_neighlist_size(self, idx):
-    """Return the number of elements in neighbor list with the given index
-
-    :param idx: neighbor list index
-    :type  idx: int
-    :return: number of elements in neighbor list with index idx
-    :rtype:  int
-     """
-    return self.lib.lammps_neighlist_num_elements(self.lmp, idx)
-
-  # -------------------------------------------------------------------------
-
-  def get_neighlist_element_neighbors(self, idx, element):
-    """Return data of neighbor list entry
-
-    :param element: neighbor list index
-    :type  element: int
-    :param element: neighbor list element index
-    :type  element: int
-    :return: tuple with atom local index, number of neighbors and array of neighbor local atom indices
-    :rtype:  (int, int, numpy.array)
-    """
-    c_iatom = c_int()
-    c_numneigh = c_int()
-    c_neighbors = POINTER(c_int)()
-    self.lib.lammps_neighlist_element_neighbors(self.lmp, idx, element, byref(c_iatom), byref(c_numneigh), byref(c_neighbors))
-    neighbors = self.numpy.iarray(c_int, c_neighbors, c_numneigh.value, 1)
-    return c_iatom.value, c_numneigh.value, neighbors
 
 # -------------------------------------------------------------------------
 
@@ -1664,6 +1701,8 @@ class numpy_wrapper:
   def __init__(self, lmp):
     self.lmp = lmp
 
+  # -------------------------------------------------------------------------
+
   def _ctype_to_numpy_int(self, ctype_int):
     import numpy as np
     if ctype_int == c_int32:
@@ -1672,6 +1711,8 @@ class numpy_wrapper:
       return np.int64
     return np.intc
 
+  # -------------------------------------------------------------------------
+
   def extract_atom(self, name, dtype=LAMMPS_AUTODETECT, nelem=LAMMPS_AUTODETECT, dim=LAMMPS_AUTODETECT):
     """Retrieve per-atom properties from LAMMPS as NumPy arrays
 
@@ -1727,6 +1768,8 @@ class numpy_wrapper:
       return self.iarray(c_int64, raw_ptr, nelem, dim)
     return raw_ptr
 
+  # -------------------------------------------------------------------------
+
   def extract_atom_iarray(self, name, nelem, dim=1):
     warnings.warn("deprecated, use extract_atom instead", DeprecationWarning)
 
@@ -1744,6 +1787,8 @@ class numpy_wrapper:
 
     return self.iarray(c_int_type, raw_ptr, nelem, dim)
 
+  # -------------------------------------------------------------------------
+
   def extract_atom_darray(self, name, nelem, dim=1):
     warnings.warn("deprecated, use extract_atom instead", DeprecationWarning)
 
@@ -1754,6 +1799,8 @@ class numpy_wrapper:
 
     return self.darray(raw_ptr, nelem, dim)
 
+  # -------------------------------------------------------------------------
+
   def extract_compute(self, cid, style, type):
     """Retrieve data from a LAMMPS compute
 
@@ -1791,6 +1838,8 @@ class numpy_wrapper:
         return self.darray(value, nlocal, ncols)
     return value
 
+  # -------------------------------------------------------------------------
+
   def extract_fix(self, fid, style, type, nrow=0, ncol=0):
     """Retrieve data from a LAMMPS fix
 
@@ -1831,6 +1880,8 @@ class numpy_wrapper:
         return self.darray(value, nrows, ncols)
     return value
 
+  # -------------------------------------------------------------------------
+
   def extract_variable(self, name, group=None, vartype=LMP_VAR_EQUAL):
     """ Evaluate a LAMMPS variable and return its data
 
@@ -1854,6 +1905,43 @@ class numpy_wrapper:
       return np.ctypeslib.as_array(value)
     return value
 
+  # -------------------------------------------------------------------------
+
+  def get_neighlist(self, idx):
+    """Returns an instance of :class:`NumPyNeighList` which wraps access to the neighbor list with the given index
+
+    :param idx: index of neighbor list
+    :type  idx: int
+    :return: an instance of :class:`NumPyNeighList` wrapping access to neighbor list data
+    :rtype:  NumPyNeighList
+    """
+    if idx < 0:
+        return None
+    return NumPyNeighList(self.lmp, idx)
+
+  # -------------------------------------------------------------------------
+
+  def get_neighlist_element_neighbors(self, idx, element):
+    """Return data of neighbor list entry
+
+    This function is a wrapper around the function
+    :py:meth:`lammps.get_neighlist_element_neighbors() <lammps.lammps.get_neighlist_element_neighbors()>`
+    method. It behaves the same as the original method, but returns a NumPy array containing the neighbors
+    instead of a ``ctypes`` pointer.
+
+    :param element: neighbor list index
+    :type  element: int
+    :param element: neighbor list element index
+    :type  element: int
+    :return: tuple with atom local index and numpy array of neighbor local atom indices
+    :rtype:  (int, numpy.array)
+    """
+    iatom, numneigh, c_neighbors = self.lmp.get_neighlist_element_neighbors(idx, element)
+    neighbors = self.iarray(c_int, c_neighbors, numneigh, 1)
+    return iatom, neighbors
+
+  # -------------------------------------------------------------------------
+
   def iarray(self, c_int_type, raw_ptr, nelem, dim=1):
     import numpy as np
     np_int_type = self._ctype_to_numpy_int(c_int_type)
@@ -1867,6 +1955,8 @@ class numpy_wrapper:
     a.shape = (nelem, dim)
     return a
 
+  # -------------------------------------------------------------------------
+
   def darray(self, raw_ptr, nelem, dim=1):
     import numpy as np
     if dim == 1: