From 34b34d8410f74da95db4e95a48183127bfaefef8 Mon Sep 17 00:00:00 2001
From: Axel Kohlmeyer <akohlmey@gmail.com>
Date: Wed, 3 Aug 2016 13:43:39 -0400
Subject: [PATCH] complete implementation for group based imbalance class

(cherry picked from commit 8ff0085cba09cf13799b9cc280446de50f963a6d)
---
 src/imbalance.h         |  2 +-
 src/imbalance_group.cpp | 39 ++++++++++++++++++++++++++++++++++++++-
 src/imbalance_group.h   | 14 ++++++++------
 3 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/src/imbalance.h b/src/imbalance.h
index 55aa9d6d07..ec3fce0e0b 100644
--- a/src/imbalance.h
+++ b/src/imbalance.h
@@ -31,7 +31,7 @@ class Imbalance {
  public:
   // parse options. return number of arguments consumed.
   virtual int options(LAMMPS *lmp, int narg, char **arg) = 0;
-  // compute per-atom imbalance and apply to weight array
+  // compute and apply weigh factors to local atom array
   virtual void compute(LAMMPS *lmp, double *weights) = 0;
 };
 
diff --git a/src/imbalance_group.cpp b/src/imbalance_group.cpp
index 97e0a50b6a..cc95c282a6 100644
--- a/src/imbalance_group.cpp
+++ b/src/imbalance_group.cpp
@@ -14,14 +14,51 @@
 
 #include "pointers.h"
 #include "imbalance_group.h"
+#include "atom.h"
+#include "error.h"
+#include "force.h"
+#include "group.h"
 
 using namespace LAMMPS_NS;
 
 int ImbalanceGroup::options(LAMMPS *lmp, int narg, char **arg)
 {
-  return 0;
+  Error *error = lmp->error;
+  Force *force = lmp->force;
+  Group *group = lmp->group;
+
+  if (narg < 3) error->all(FLERR,"Illegal balance weight command");
+
+  _num = force->inumeric(FLERR,arg[0]);
+  if (_num < 1) error->all(FLERR,"Illegal balance weight command");
+  if (2*_num+1 > narg) error->all(FLERR,"Illegal balance weight command");
+
+  _id = new int[_num];
+  _factor = new double[_num];
+  for (int i = 0; i < _num; ++i) {
+    _id[i] = group->find(arg[2*i+1]);
+    if (_id[i] < 0)
+      error->all(FLERR,"Unknown group in balance weight command");
+    _factor[i] = force->numeric(FLERR,arg[2*i+2]);
+  }
+  return _num;
 }
  
 void ImbalanceGroup::compute(LAMMPS *lmp, double *weight)
 {
+  const int * const mask = lmp->atom->mask;
+  const int * const bitmask = lmp->group->bitmask;
+  const int nlocal = lmp->atom->nlocal;
+
+  if (_num == 0) return;
+
+  for (int i = 0; i < nlocal; ++i) {
+    const int imask = mask[i];
+    double iweight = weight[i];
+    for (int j = 0; j < _num; ++j) {
+      if (imask & bitmask[_id[j]])
+        iweight *= _factor[j];
+    }
+    weight[i] = iweight;
+  }
 }
diff --git a/src/imbalance_group.h b/src/imbalance_group.h
index 620996d8f9..52869f1c61 100644
--- a/src/imbalance_group.h
+++ b/src/imbalance_group.h
@@ -19,14 +19,16 @@
 namespace LAMMPS_NS {
 
 class ImbalanceGroup : public Imbalance {
- public:
-  ImbalanceGroup() : Imbalance() {};
-  virtual ~ImbalanceGroup() {};
 
-  // disallow copy constructor and assignment operator
+ public:
+  ImbalanceGroup() : Imbalance(), _num(0), _id(0), _factor(0) {};
+  virtual ~ImbalanceGroup() { delete[] _id; delete[] _factor; };
+
+  // internal data members
  private:
-  ImbalanceGroup(const ImbalanceGroup &) {};
-  ImbalanceGroup &operator=(const ImbalanceGroup &) {return *this;};
+  int _num;                     // number of groups with weights
+  int *_id;                     // list numerical id's of groups
+  double *_factor;              // list if group weight factors
 
   // required member functions
  public: