From 5bf8e1bc5b249a14ed5dd0a4cbb294fb2b3a3c1d Mon Sep 17 00:00:00 2001
From: Axel Kohlmeyer <akohlmey@gmail.com>
Date: Thu, 19 Apr 2018 18:37:15 -0400
Subject: [PATCH] add ability to retrieve the number of local rows for computes
 returning local data

this is done supporting the combination of type == 2 with style == 0,
i.e. a local scalar, which is not available in C++ (but there one can
access the compute style data member directly. for the python interface,
the pointer is automatically dereferenced and returned as a c_int.
---
 python/lammps.py | 11 ++++++++---
 src/library.cpp  | 10 +++++++++-
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/python/lammps.py b/python/lammps.py
index f5f252c45c..417427eb4b 100644
--- a/python/lammps.py
+++ b/python/lammps.py
@@ -384,9 +384,14 @@ class lammps(object):
       ptr = self.lib.lammps_extract_compute(self.lmp,id,style,type)
       return ptr
     if type == 2:
-      self.lib.lammps_extract_compute.restype = POINTER(POINTER(c_double))
-      ptr = self.lib.lammps_extract_compute(self.lmp,id,style,type)
-      return ptr
+      if style == 0:
+        self.lib.lammps_extract_compute.restype = POINTER(c_int)
+        ptr = self.lib.lammps_extract_compute(self.lmp,id,style,type)
+        return ptr[0]
+      else:
+        self.lib.lammps_extract_compute.restype = POINTER(POINTER(c_double))
+        ptr = self.lib.lammps_extract_compute(self.lmp,id,style,type)
+        return ptr
     return None
 
   # extract fix info
diff --git a/src/library.cpp b/src/library.cpp
index 07c4cae9c1..b7a0c0b5be 100644
--- a/src/library.cpp
+++ b/src/library.cpp
@@ -480,10 +480,13 @@ void *lammps_extract_atom(void *ptr, char *name)
      compute's internal data structure for the entity
      caller should cast it to (double *) for a scalar or vector
      caller should cast it to (double **) for an array
-   for per-atom or local data, returns a pointer to the
+   for per-atom or local vector/array data, returns a pointer to the
      compute's internal data structure for the entity
      caller should cast it to (double *) for a vector
      caller should cast it to (double **) for an array
+   for local data, accessing scalar data for the compute (type = 0),
+   returns a pointer that should be cast to (int *) which points to
+   an int with the number of local rows, i.e. the length of the local array.
    returns a void pointer to the compute's internal data structure
      for the entity which the caller can cast to the proper data type
    returns a NULL if id is not recognized or style/type not supported
@@ -541,6 +544,11 @@ void *lammps_extract_compute(void *ptr, char *id, int style, int type)
 
     if (style == 2) {
       if (!compute->local_flag) return NULL;
+      if (type == 0) {
+        if (compute->invoked_local != lmp->update->ntimestep)
+          compute->compute_local();
+        return (void *) &compute->size_local_rows;
+      }
       if (type == 1) {
         if (compute->invoked_local != lmp->update->ntimestep)
           compute->compute_local();