tx258/lammps-gran-kokkos
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first try at implementing lammps_extract_atom_size()

Browse Source
This commit is contained in:
Axel Kohlmeyer
2024-08-30 22:50:42 -04:00
parent e921af8efa
commit 9d9e591b54
13 changed files with 411 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
python/lammps/numpy_wrapper.py
View File

@ -54,7 +54,8 @@ class numpy_wrapper:
# -------------------------------------------------------------------------
def extract_atom(self, name, dtype=LAMMPS_AUTODETECT, nelem=LAMMPS_AUTODETECT, dim=LAMMPS_AUTODETECT):
def extract_atom(self, name, dtype=LAMMPS_AUTODETECT, nelem=LAMMPS_AUTODETECT,
dim=LAMMPS_AUTODETECT):
"""Retrieve per-atom properties from LAMMPS as NumPy arrays
This is a wrapper around the :py:meth:`lammps.extract_atom()` method.
@ -63,16 +64,16 @@ class numpy_wrapper:
.. note::
The returned arrays of per-atom data are by default dimensioned
for the range [0:nlocal] since that data is *always* valid. The
underlying storage for the data, however, is typically allocated
for the range of [0:nmax]. Whether there is valid data in the range
[nlocal:nlocal+nghost] depends on whether the property of interest
is also updated for ghost atoms. This is not often the case. In
some cases, it depends on a LAMMPS setting, see for example
:doc:`comm_modify vel yes <comm_modify>`. By using the optional
*nelem* parameter the size of the returned NumPy can be overridden.
There is no check whether the number of elements chosen is valid.
The returned vectors or arrays of per-atom data are dimensioned
according to the return value of :py:meth:`lammps.extract_atom_size()`.
Except for the "mass" property, the underlying storage will always be
dimensioned for the range [0:nmax]. The actual usable data may be
only in the range [0:nlocal] or [0:nlocal][0:dim]. Whether there is
valid data in the range [nlocal:nlocal+nghost] or [nlocal:local+nghost][0:dim]
depends on whether the property of interest is also updated for ghost atoms.
Also the value of *dim* depends on the value of *name*. By using the optional
*nelem* and *dim* parameters the dimensions of the returned NumPy array can
be overridden. There is no check whether the number of elements chosen is valid.
:param name: name of the property
:type name: string
@ -89,21 +90,10 @@ class numpy_wrapper:
dtype = self.lmp.extract_atom_datatype(name)
if nelem == LAMMPS_AUTODETECT:
if name == "mass":
nelem = self.lmp.extract_global("ntypes") + 1
else:
nelem = self.lmp.extract_global("nlocal")
nelem = self.lmp.extract_atom_size(name, LMP_SIZE_ROWS)
if dim == LAMMPS_AUTODETECT:
if dtype in (LAMMPS_INT_2D, LAMMPS_DOUBLE_2D, LAMMPS_INT64_2D):
# TODO add other fields
if name in ("x", "v", "f", "x0","omega", "angmom", "torque", "csforce", "vforce", "vest"):
dim = 3
elif name == "smd_data_9":
dim = 9
elif name == "smd_stress":
dim = 6
else:
dim = 2
dim = self.lmp.extract_atom_size(name, LMP_SIZE_COLS)
else:
dim = 1
@ -119,37 +109,6 @@ class numpy_wrapper:
# -------------------------------------------------------------------------
def extract_atom_iarray(self, name, nelem, dim=1):
warnings.warn("deprecated, use extract_atom instead", DeprecationWarning)
if name in ['id', 'molecule']:
c_int_type = self.lmp.c_tagint
elif name in ['image']:
c_int_type = self.lmp.c_imageint
else:
c_int_type = c_int
if dim == 1:
raw_ptr = self.lmp.extract_atom(name, LAMMPS_INT)
else:
raw_ptr = self.lmp.extract_atom(name, LAMMPS_INT_2D)
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)
if dim == 1:
raw_ptr = self.lmp.extract_atom(name, LAMMPS_DOUBLE)
else:
raw_ptr = self.lmp.extract_atom(name, LAMMPS_DOUBLE_2D)
return self.darray(raw_ptr, nelem, dim)
# -------------------------------------------------------------------------
def extract_compute(self, cid, cstyle, ctype):
"""Retrieve data from a LAMMPS compute