ICODE-ADC/lammps
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Create ctypes only neighbor list API variant

Browse Source 
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)
This commit is contained in:
Richard Berger
2020-10-05 16:21:37 -04:00
parent 68147306e7
commit 02b10380bc
3 changed files with 137 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
doc/src/Python_module.rst
View File

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

9
examples/python/in.fix_python_invoke_neighlist
View File

@ -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.find_pair_neighlist("lj/cut", request=0)
#mylist = L.numpy.get_neighlist(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:
print(pid_prefix, "- {}".format(iatom), x[iatom], v[iatom], f[iatom], " : ", numneigh, "Neighbors")
for iatom, neighs in mylist:
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])

166
python/lammps.py
View File

@ -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
:rtype: (int, int, numpy.array)
:return: tuple with atom local index, numpy array of neighbor local atom indices
: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)
# -------------------------------------------------------------------------
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
return idx
# -------------------------------------------------------------------------
@ -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: