From 14321d1fa0590ce737ae50e8667eb1ab6670bd8d Mon Sep 17 00:00:00 2001
From: Axel Kohlmeyer <akohlmey@gmail.com>
Date: Fri, 3 Jul 2020 22:19:44 -0400
Subject: [PATCH] add option to compute chunk/atom to access the number of
 chunks as a global scalar

---
 doc/src/compute_chunk_atom.rst | 10 ++++++----
 src/compute_chunk_atom.cpp     | 16 ++++++++++++++++
 src/compute_chunk_atom.h       |  1 +
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/doc/src/compute_chunk_atom.rst b/doc/src/compute_chunk_atom.rst
index 2ea12fb90c..a3d915b8a4 100644
--- a/doc/src/compute_chunk_atom.rst
+++ b/doc/src/compute_chunk_atom.rst
@@ -622,14 +622,16 @@ cylinder, x for a y-axis cylinder, and x for a z-axis cylinder.
 
 **Output info:**
 
-This compute calculates a per-atom vector, which can be accessed by
-any command that uses per-atom values from a compute as input.  See
-the :doc:`Howto output <Howto_output>` doc page for an overview of
+This compute calculates a per-atom vector (the chunk ID), which can
+be accessed by any command that uses per-atom values from a compute
+as input.  It also calculates a global scalar (the number of chunks),
+which can be similarly accessed everywhere outside of a per-atom context.
+See the :doc:`Howto output <Howto_output>` doc page for an overview of
 LAMMPS output options.
 
 The per-atom vector values are unitless chunk IDs, ranging from 1 to
 *Nchunk* (inclusive) for atoms assigned to chunks, and 0 for atoms not
-belonging to a chunk.
+belonging to a chunk.  The scalar contains the value of *Nchunk*.
 
 Restrictions
 """"""""""""
diff --git a/src/compute_chunk_atom.cpp b/src/compute_chunk_atom.cpp
index f7e2db69b9..c8e4269b56 100644
--- a/src/compute_chunk_atom.cpp
+++ b/src/compute_chunk_atom.cpp
@@ -64,6 +64,8 @@ ComputeChunkAtom::ComputeChunkAtom(LAMMPS *lmp, int narg, char **arg) :
   if (narg < 4) error->all(FLERR,"Illegal compute chunk/atom command");
 
   peratom_flag = 1;
+  scalar_flag = 1;
+  extscalar = 0;
   size_peratom_cols = 0;
   create_attribute = 1;
 
@@ -639,6 +641,20 @@ void ComputeChunkAtom::compute_peratom()
   for (int i = 0; i < nlocal; i++) chunk[i] = ichunk[i];
 }
 
+
+/* ----------------------------------------------------------------------
+   to return the number of chunks, we first need to make certain
+   that compute_peratom() has been called.
+------------------------------------------------------------------------- */
+double ComputeChunkAtom::compute_scalar()
+{
+  if (invoked_peratom != update->ntimestep)
+    compute_peratom();
+  invoked_scalar = update->ntimestep;
+
+  return (scalar = nchunk);
+}
+
 /* ----------------------------------------------------------------------
    set lock, so that nchunk will not change from startstep to stopstep
    called by fix for duration of time it requires lock
diff --git a/src/compute_chunk_atom.h b/src/compute_chunk_atom.h
index 5f7e165f98..98a3ea9a2b 100644
--- a/src/compute_chunk_atom.h
+++ b/src/compute_chunk_atom.h
@@ -39,6 +39,7 @@ class ComputeChunkAtom : public Compute {
   void init();
   void setup();
   void compute_peratom();
+  double compute_scalar();
   void set_arrays(int);
   double memory_usage();