enable relancing to work with distributed grids

src/AMOEBA/amoeba_convolution.cpp

@@ -75,7 +75,34 @@ AmoebaConvolution::AmoebaConvolution(LAMMPS *lmp, Pair *pair,
   if (which == POLAR_GRIDC || which == INDUCE_GRIDC) flag3d = 0;
 
   nfft_global = (bigint) nx * ny * nz;
+}
 
+/* ----------------------------------------------------------------------
+   free all memory
+------------------------------------------------------------------------- */
+
+AmoebaConvolution::~AmoebaConvolution()
+{
+  deallocate_grid();
+}
+
+/* ----------------------------------------------------------------------
+   subset of FFT grids assigned to each proc may have changed
+   called by load balancer when proc subdomains are adjusted
+------------------------------------------------------------------------- */
+
+void AmoebaConvolution::reset_grid()
+{
+  deallocate_grid();
+  allocate_grid();
+}
+
+/* ----------------------------------------------------------------------
+   allocate all local grid data structs: FFT, Grid3d, Remap
+------------------------------------------------------------------------- */
+
+void AmoebaConvolution::allocate_grid()
+{
   // global indices of grid range from 0 to N-1
   // nlo_in,nhi_in = lower/upper limits of the 3d sub-brick of
   //   global grid that I own without ghost cells
@@ -247,11 +274,9 @@ AmoebaConvolution::AmoebaConvolution(LAMMPS *lmp, Pair *pair,
   memory->create(remap_buf,nqty*nfft_owned,"amoeba:remap_buf");
 }
 
-/* ----------------------------------------------------------------------
+/* ---------------------------------------------------------------------- */
-   free all memory
-------------------------------------------------------------------------- */
 
-AmoebaConvolution::~AmoebaConvolution()
+void AmoebaConvolution::deallocate_grid()
 {
   memory->destroy3d_offset(grid_brick,nzlo_out,nylo_out,nxlo_out);
   memory->destroy4d_offset_last(cgrid_brick,nzlo_out,nylo_out,nxlo_out);

src/AMOEBA/amoeba_convolution.h

@@ -42,6 +42,7 @@ class AmoebaConvolution : protected Pointers {
 
   AmoebaConvolution(class LAMMPS *, class Pair *, int, int, int, int, int);
   ~AmoebaConvolution();
+  void reset_grid();
   void *zero();
   FFT_SCALAR *pre_convolution();
   void *post_convolution();
@@ -67,6 +68,8 @@ class AmoebaConvolution : protected Pointers {
   double *gc_buf1, *gc_buf2;    // buffers for GridComm
   double *remap_buf;            // buffer for Remap
 
+  void allocate_grid();
+  void deallocate_grid();
   void *zero_3d();
   void *zero_4d();
   FFT_SCALAR *pre_convolution_3d();

src/AMOEBA/pair_amoeba.cpp

@@ -1409,6 +1409,33 @@ void PairAmoeba::unpack_reverse_comm(int n, int *list, double *buf)
   }
 }
 
+/* ----------------------------------------------------------------------
+   subset of FFT grids assigned to each proc may have changed
+   notify each instance of AmoebaConvolution class
+   called by load balancer when proc subdomains are adjusted
+------------------------------------------------------------------------- */
+
+void PairAmoeba::reset_grid()
+{
+  if (use_ewald) {
+    m_kspace->reset_grid();
+    p_kspace->reset_grid();
+    pc_kspace->reset_grid();
+    i_kspace->reset_grid();
+    ic_kspace->reset_grid();
+  }
+  if (use_dewald) d_kspace->reset_grid();
+
+  // qfac is shared by induce and polar
+  // gridfft1 is copy of FFT grid used within polar
+
+  memory->destroy(qfac);
+  memory->destroy(gridfft1);
+  int nmine = p_kspace->nfft_owned;
+  memory->create(qfac,nmine,"ameoba/induce:qfac");
+  memory->create(gridfft1,2*nmine,"amoeba/polar:gridfft1");
+}
+
 /* ----------------------------------------------------------------------
    pack own values to buf to send to another proc
 ------------------------------------------------------------------------- */

src/AMOEBA/pair_amoeba.h

@@ -44,6 +44,8 @@ class PairAmoeba : public Pair {
   int pack_reverse_comm(int, int, double *) override;
   void unpack_reverse_comm(int, int *, double *) override;
 
+  void reset_grid() override;
+
   void pack_forward_grid(int, void *, int, int *) override;
   void unpack_forward_grid(int, void *, int, int *) override;
   void pack_reverse_grid(int, void *, int, int *) override;

src/ELECTRODE/pppm_electrode.cpp

@@ -378,7 +378,7 @@ void PPPMElectrode::setup()
    called by fix balance b/c it changed sizes of processor sub-domains
 ------------------------------------------------------------------------- */
 
-void PPPMElectrode::setup_grid()
+void PPPMElectrode::reset_grid()
 {
   // free all arrays previously allocated
 

src/ELECTRODE/pppm_electrode.h

@@ -35,7 +35,7 @@ class PPPMElectrode : public PPPM, public ElectrodeKSpace {
   ~PPPMElectrode() override;
   void init() override;
   void setup() override;
-  void setup_grid() override;
+  void reset_grid() override;
   void compute(int, int) override;
 
   void compute_vector(double *, int, int, bool) override;

src/EXTRA-FIX/fix_ttm_grid.cpp

@@ -388,6 +388,17 @@ void FixTTMGrid::write_electron_temperatures(const std::string &filename)
   if (comm->me == 0) fclose(FPout);
 }
 
+/* ----------------------------------------------------------------------
+   subset of grid assigned to each proc may have changed
+   called by load balancer when proc subdomains are adjusted
+   not supported for now, b/c requires T_electron to persist, i.e. a remap()
+------------------------------------------------------------------------- */
+
+void FixTTMGrid::reset_grid()
+{
+  error->all(FLERR,"Fix ttm/grid does not support load balancing (yet)");
+}
+
 /* ----------------------------------------------------------------------
    pack own values to buf to send to another proc
 ------------------------------------------------------------------------- */

src/EXTRA-FIX/fix_ttm_grid.h

@@ -35,6 +35,8 @@ class FixTTMGrid : public FixTTM {
 
   // grid communication
 
+  void reset_grid() override;
+
   void pack_forward_grid(int, void *, int, int *) override;
   void unpack_forward_grid(int, void *, int, int *) override;
   void pack_reverse_grid(int, void *, int, int *) override;

src/KOKKOS/pppm_kokkos.cpp

@@ -551,7 +551,7 @@ void PPPMKokkos<DeviceType>::operator()(TagPPPM_setup_triclinic2, const int &n)
 ------------------------------------------------------------------------- */
 
 template<class DeviceType>
-void PPPMKokkos<DeviceType>::setup_grid()
+void PPPMKokkos<DeviceType>::reset_grid()
 {
   // free all arrays previously allocated
 

src/KOKKOS/pppm_kokkos.h

@@ -124,7 +124,7 @@ class PPPMKokkos : public PPPM, public KokkosBaseFFT {
   ~PPPMKokkos() override;
   void init() override;
   void setup() override;
-  void setup_grid() override;
+  void reset_grid() override;
   void settings(int, char **) override;
   void compute(int, int) override;
   int timing_1d(int, double &) override;

src/KSPACE/fix_tune_kspace.cpp

@@ -278,9 +278,13 @@ void FixTuneKspace::update_kspace_style(const std::string &new_kspace_style,
   force->init();
 
   // set up grid
-  force->kspace->setup_grid();
 
-  // Re-init neighbor list. Probably only needed when redefining the pair style. Should happen after pair->init() to get pair style neighbor list request registered
+  force->kspace->reset_grid();
+
+  // re-init neighbor list
+  // probably only needed when redefining the pair style
+  // should happen after pair->init() to get pair style 
+  //   neighbor list request registered
 
   neighbor->init();
 

src/KSPACE/msm.cpp

@@ -1379,7 +1379,7 @@ void MSM::set_proc_grid(int n)
    called by fix balance b/c it changed sizes of processor sub-domains
 ------------------------------------------------------------------------- */
 
-void MSM::setup_grid()
+void MSM::reset_grid()
 {
   // free all arrays previously allocated
   // pre-compute volume-dependent coeffs

src/KSPACE/msm.h

@@ -100,7 +100,7 @@ class MSM : public KSpace {
   void set_grid_global();
   void set_proc_grid(int);
   void set_grid_local();
-  void setup_grid() override;
+  void reset_grid() override;
   double estimate_1d_error(double, double);
   double estimate_3d_error();
   double estimate_total_error();

src/KSPACE/pppm.cpp

@@ -549,7 +549,7 @@ void PPPM::setup_triclinic()
    called by fix balance b/c it changed sizes of processor sub-domains
 ------------------------------------------------------------------------- */
 
-void PPPM::setup_grid()
+void PPPM::reset_grid()
 {
   // free all arrays previously allocated
 

src/KSPACE/pppm.h

@@ -32,7 +32,7 @@ class PPPM : public KSpace {
   void settings(int, char **) override;
   void init() override;
   void setup() override;
-  void setup_grid() override;
+  void reset_grid() override;
   void compute(int, int) override;
   int timing_1d(int, double &) override;
   int timing_3d(int, double &) override;

src/KSPACE/pppm_dipole.cpp

@@ -356,7 +356,7 @@ void PPPMDipole::setup()
    called by fix balance b/c it changed sizes of processor sub-domains
 ------------------------------------------------------------------------- */
 
-void PPPMDipole::setup_grid()
+void PPPMDipole::reset_grid()
 {
   // free all arrays previously allocated
 

src/KSPACE/pppm_dipole.h

@@ -30,7 +30,7 @@ class PPPMDipole : public PPPM {
   ~PPPMDipole() override;
   void init() override;
   void setup() override;
-  void setup_grid() override;
+  void reset_grid() override;
   void compute(int, int) override;
   int timing_1d(int, double &) override;
   int timing_3d(int, double &) override;

src/KSPACE/pppm_disp.cpp

@@ -794,7 +794,7 @@ void PPPMDisp::setup()
    called by fix balance b/c it changed sizes of processor sub-domains
 ------------------------------------------------------------------------- */
 
-void PPPMDisp::setup_grid()
+void PPPMDisp::reset_grid()
 {
   // free all arrays previously allocated
 

src/KSPACE/pppm_disp.h

@@ -34,7 +34,7 @@ class PPPMDisp : public KSpace {
   ~PPPMDisp() override;
   void init() override;
   void setup() override;
-  void setup_grid() override;
+  void reset_grid() override;
   void settings(int, char **) override;
   void compute(int, int) override;
   int timing_1d(int, double &) override;

src/MISC/fix_ipi.cpp

@@ -374,9 +374,9 @@ void FixIPI::initial_integrate(int /*vflag*/)
     //   kspace->setup() is in some cases not enough since, e.g., g_ewald needs
     //   to be reestimated due to changes in box dimensions.
     force->init();
-    // setup_grid() is necessary for pppm since init() is not calling
+    // reset_grid() is necessary for pppm since init() is not calling
-    //   setup() nor setup_grid() upon calling init().
+    //   setup() nor reset_grid() upon calling init().
-    if (force->kspace->pppmflag) force->kspace->setup_grid();
+    if (force->kspace->pppmflag) force->kspace->reset_grid();
     // other kspace styles might need too another setup()?
   } else if (!reset_flag && kspace_flag) {
     // original version

src/balance.cpp

@@ -27,6 +27,7 @@
 #include "comm.h"
 #include "domain.h"
 #include "fix_store_peratom.h"
+#include "force.h"
 #include "imbalance.h"
 #include "imbalance_group.h"
 #include "imbalance_neigh.h"
@@ -36,6 +37,7 @@
 #include "irregular.h"
 #include "memory.h"
 #include "modify.h"
+#include "pair.h"
 #include "rcb.h"
 #include "error.h"
 
@@ -366,6 +368,13 @@ void Balance::command(int narg, char **arg)
 
   if (outflag) dumpout(update->ntimestep);
 
+  // notify all classes that store distributed grids
+  // so they can adjust to new proc sub-domains
+  // no need to invoke kspace->reset_grid() b/c it does this in its init()
+
+  modify->reset_grid();
+  if (force->pair) force->pair->reset_grid();
+
   // check if any particles were lost
 
   bigint natoms;

src/compute.h

@@ -130,6 +130,8 @@ class Compute : protected Pointers {
   virtual int pack_reverse_comm(int, int, double *) { return 0; }
   virtual void unpack_reverse_comm(int, int *, double *) {}
 
+  virtual void reset_grid(){};
+
   virtual int get_grid_by_name(char *, int &) { return -1; };
   virtual void *get_grid_by_index(int) { return nullptr; };
   virtual int get_griddata_by_name(int, char *, int &) { return -1; };

src/compute_property_grid.cpp

@@ -41,14 +41,14 @@ ComputePropertyGrid::ComputePropertyGrid(LAMMPS *lmp, int narg, char **arg) :
 
   dimension = domain->dimension;
 
-  nx = utils::inumeric(FLERR,arg[3],false,lmp);
+  nxgrid = utils::inumeric(FLERR,arg[3],false,lmp);
-  ny = utils::inumeric(FLERR,arg[4],false,lmp);
+  nygrid = utils::inumeric(FLERR,arg[4],false,lmp);
-  nz = utils::inumeric(FLERR,arg[5],false,lmp);
+  nzgrid = utils::inumeric(FLERR,arg[5],false,lmp);
 
-  if (dimension == 2 && nz != 1) 
+  if (dimension == 2 && nzgrid != 1) 
     error->all(FLERR,"Compute property/grid for 2d requires nz = 1");
 
-  if (nx <= 0 || ny <= 0 || nz <= 0) 
+  if (nxgrid <= 0 || nygrid <= 0 || nzgrid <= 0) 
     error->all(FLERR, "Illegal compute property/grid command");
  
   nvalues = narg - 6;
@@ -108,36 +108,9 @@ ComputePropertyGrid::ComputePropertyGrid(LAMMPS *lmp, int narg, char **arg) :
     } else error->all(FLERR, "Illegal compute property/grid command");
   }
 
-  // instantiate the Grid class and allocate per-grid memory
+  // initial setup of distributed grid
 
-  if (dimension == 2) {
+  allocate_grid();
-    grid2d = new Grid2d(lmp, world, nx, ny, 0, 0.0, 0.0,
-                        nxlo_in, nxhi_in, nylo_in, nyhi_in,
-                        nxlo_out, nxhi_out, nylo_out, nyhi_out);
-    if (nvalues == 1)
-      memory->create2d_offset(vec2d, nylo_out, nyhi_out, nxlo_out, nxhi_out, 
-                              "property/grid:vec2d");
-    else
-      memory->create3d_offset_last(array2d, nylo_out, nyhi_out, nxlo_out,
-                                   nxhi_out, nvalues, "property/grid:array2d");
-    ngridout = (nxhi_out - nxlo_out + 1) * (nyhi_out - nylo_out + 1);
-    
-  } else {
-    grid3d = new Grid3d(lmp, world, nx, ny, nz, 0, 0.0, 0.0,
-                        nxlo_in, nxhi_in, nylo_in, nyhi_in, nzlo_in, nzhi_in, 
-                        nxlo_out, nxhi_out, nylo_out, nyhi_out, 
-                        nzlo_out, nzhi_out);
-    if (nvalues == 1)
-      memory->create3d_offset(vec3d, nzlo_out, nzhi_out, nylo_out, 
-                              nyhi_out, nxlo_out,
-                              nxhi_out, "property/grid:vec3d");
-    else
-      memory->create4d_offset_last(array3d, nzlo_out, nzhi_out, nylo_out, 
-                                   nyhi_out, nxlo_out,
-                                   nxhi_out, nvalues, "property/grid:array3d");
-    ngridout = (nxhi_out - nxlo_out + 1) * (nyhi_out - nylo_out + 1) * 
-      (nzhi_out - nzlo_out + 1);
-  }
 }
 
 /* ---------------------------------------------------------------------- */
@@ -146,12 +119,7 @@ ComputePropertyGrid::~ComputePropertyGrid()
 {
   delete[] pack_choice;
 
-  delete grid2d;
+  deallocate_grid();
-  delete grid3d;
-  memory->destroy2d_offset(vec2d,nylo_out,nxlo_out);
-  memory->destroy2d_offset(array2d,nylo_out,nxlo_out);
-  memory->destroy3d_offset(vec3d,nzlo_out,nylo_out,nxlo_out);
-  memory->destroy4d_offset_last(array3d,nzlo_out,nylo_out,nxlo_out);
 }
 
 /* ---------------------------------------------------------------------- */
@@ -160,19 +128,6 @@ void ComputePropertyGrid::compute_pergrid()
 {
   invoked_pergrid = update->ntimestep;
 
-  // set current size for portion of grid on each proc
-  // may change between compute invocations due to load balancing
-  
-  if (dimension == 2)
-    grid2d->get_bounds(nxlo_in,nxhi_in,nylo_in,nyhi_in);
-  else
-    grid3d->get_bounds(nxlo_in,nxhi_in,nylo_in,nyhi_in,nzlo_in,nzhi_in);
-
-  // reallocate data vector or array if changed
-  // NOTE: still need to implement
-
-
-
   // fill data vector or array with values for my grid pts
 
   if (nvalues == 1) {
@@ -182,6 +137,18 @@ void ComputePropertyGrid::compute_pergrid()
   }
 }
 
+/* ----------------------------------------------------------------------
+   subset of grid assigned to each proc may have changed
+   deallocate and reallocate Grid class and local data structs
+   called by load balancer when proc subdomains are adjusted
+---------------------------------------------------------------------- */
+
+void ComputePropertyGrid::reset_grid()
+{
+  deallocate_grid();
+  allocate_grid();
+}
+
 /* ----------------------------------------------------------------------
    return index of grid associated with name
    this class can store M named grids, indexed 0 to M-1
@@ -257,6 +224,54 @@ void *ComputePropertyGrid::get_griddata_by_index(int index)
   return nullptr;
 }
 
+/* ----------------------------------------------------------------------
+   instantiate the Grid class and allocate local per-grid memory
+---------------------------------------------------------------------- */
+
+void ComputePropertyGrid::allocate_grid()
+{
+  if (dimension == 2) {
+    grid2d = new Grid2d(lmp, world, nxgrid, nygrid, 0, 0.0, 0.0,
+                        nxlo_in, nxhi_in, nylo_in, nyhi_in,
+                        nxlo_out, nxhi_out, nylo_out, nyhi_out);
+    if (nvalues == 1)
+      memory->create2d_offset(vec2d, nylo_out, nyhi_out, nxlo_out, nxhi_out, 
+                              "property/grid:vec2d");
+    else
+      memory->create3d_offset_last(array2d, nylo_out, nyhi_out, nxlo_out,
+                                   nxhi_out, nvalues, "property/grid:array2d");
+    ngridout = (nxhi_out - nxlo_out + 1) * (nyhi_out - nylo_out + 1);
+    
+  } else {
+    grid3d = new Grid3d(lmp, world, nxgrid, nygrid, nzgrid, 0, 0.0, 0.0,
+                        nxlo_in, nxhi_in, nylo_in, nyhi_in, nzlo_in, nzhi_in, 
+                        nxlo_out, nxhi_out, nylo_out, nyhi_out, 
+                        nzlo_out, nzhi_out);
+    if (nvalues == 1)
+      memory->create3d_offset(vec3d, nzlo_out, nzhi_out, nylo_out, 
+                              nyhi_out, nxlo_out,
+                              nxhi_out, "property/grid:vec3d");
+    else
+      memory->create4d_offset_last(array3d, nzlo_out, nzhi_out, nylo_out, 
+                                   nyhi_out, nxlo_out,
+                                   nxhi_out, nvalues, "property/grid:array3d");
+    ngridout = (nxhi_out - nxlo_out + 1) * (nyhi_out - nylo_out + 1) * 
+      (nzhi_out - nzlo_out + 1);
+  }
+}
+
+/* ---------------------------------------------------------------------- */
+
+void ComputePropertyGrid::deallocate_grid()
+{
+  delete grid2d;
+  delete grid3d;
+  memory->destroy2d_offset(vec2d,nylo_out,nxlo_out);
+  memory->destroy2d_offset(array2d,nylo_out,nxlo_out);
+  memory->destroy3d_offset(vec3d,nzlo_out,nylo_out,nxlo_out);
+  memory->destroy4d_offset_last(array3d,nzlo_out,nylo_out,nxlo_out);
+}
+
 /* ----------------------------------------------------------------------
    memory usage of grid data
 ------------------------------------------------------------------------- */
@@ -277,23 +292,23 @@ void ComputePropertyGrid::pack_id(int n)
     if (nvalues == 0) {
       for (int iy = nylo_in; iy <= nyhi_in; iy++)
         for (int ix = nxlo_in; ix <= nxhi_in; ix++)
-          vec2d[iy][ix] = iy*nx + ix + 1;
+          vec2d[iy][ix] = iy*nxgrid + ix + 1;
     } else {
       for (int iy = nylo_in; iy <= nyhi_in; iy++)
         for (int ix = nxlo_in; ix <= nxhi_in; ix++)
-          array2d[iy][ix][n] = iy*nx + ix + 1;
+          array2d[iy][ix][n] = iy*nxgrid + ix + 1;
     }
   } else if (dimension == 3) {
     if (nvalues == 0) {
       for (int iz = nzlo_in; iz <= nzhi_in; iz++)
         for (int iy = nylo_in; iy <= nyhi_in; iy++)
           for (int ix = nxlo_in; ix <= nxhi_in; ix++)
-            vec3d[iz][iy][ix] = iz*ny*nx + iy*nx + ix + 1;
+            vec3d[iz][iy][ix] = iz*nygrid*nxgrid + iy*nxgrid + ix + 1;
     } else {
       for (int iz = nzlo_in; iz <= nzhi_in; iz++)
         for (int iy = nylo_in; iy <= nyhi_in; iy++)
           for (int ix = nxlo_in; ix <= nxhi_in; ix++)
-            array3d[iz][iy][ix][n] = iz*ny*nx + iy*nx + ix + 1;
+            array3d[iz][iy][ix][n] = iz*nygrid*nxgrid + iy*nxgrid + ix + 1;
     }
   }
 }
@@ -463,7 +478,7 @@ void ComputePropertyGrid::pack_z(int n)
 
 void ComputePropertyGrid::pack_xs(int n)
 {
-  double dx = 1.0/nx;
+  double dx = 1.0/nxgrid;
 
   if (dimension == 2) {
     if (nvalues == 0) {
@@ -494,7 +509,7 @@ void ComputePropertyGrid::pack_xs(int n)
 
 void ComputePropertyGrid::pack_ys(int n)
 {
-  double dy = 1.0/ny;
+  double dy = 1.0/nygrid;
 
   if (dimension == 2) {
     if (nvalues == 0) {
@@ -525,7 +540,7 @@ void ComputePropertyGrid::pack_ys(int n)
 
 void ComputePropertyGrid::pack_zs(int n)
 {
-  double dz = 1.0/nz;
+  double dz = 1.0/nzgrid;
 
   if (nvalues == 0) {
     for (int iz = nzlo_in; iz <= nzhi_in; iz++)
@@ -630,7 +645,7 @@ void ComputePropertyGrid::pack_zc(int n)
 
 void ComputePropertyGrid::pack_xsc(int n)
 {
-  double dx = 1.0/nx;
+  double dx = 1.0/nxgrid;
 
   if (dimension == 2) {
     if (nvalues == 0) {
@@ -661,7 +676,7 @@ void ComputePropertyGrid::pack_xsc(int n)
 
 void ComputePropertyGrid::pack_ysc(int n)
 {
-  double dy = 1.0/ny;
+  double dy = 1.0/nygrid;
 
   if (dimension == 2) {
     if (nvalues == 0) {
@@ -692,7 +707,7 @@ void ComputePropertyGrid::pack_ysc(int n)
 
 void ComputePropertyGrid::pack_zsc(int n)
 {
-  double dz = 1.0/nz;
+  double dz = 1.0/nzgrid;
 
   if (nvalues == 0) {
     for (int iz = nzlo_in; iz <= nzhi_in; iz++)

src/compute_property_grid.h

@@ -31,15 +31,17 @@ class ComputePropertyGrid : public Compute {
   void init() override {}
   void compute_pergrid() override;
 
-  int get_grid_by_name(char *, int &);
+  void reset_grid() override;
-  void *get_grid_by_index(int);
+
-  int get_griddata_by_name(int, char *, int &);
+  int get_grid_by_name(char *, int &) override;
-  void *get_griddata_by_index(int);
+  void *get_grid_by_index(int) override;
+  int get_griddata_by_name(int, char *, int &) override;
+  void *get_griddata_by_index(int) override;
 
   double memory_usage() override;
 
  private:
-  int nx,ny,nz;
+  int nxgrid,nygrid,nzgrid;
   int nvalues;
   int dimension;
 
@@ -53,6 +55,11 @@ class ComputePropertyGrid : public Compute {
   double **vec2d,***vec3d;
   double ***array2d,****array3d;
 
+  // local methods
+
+  void allocate_grid();
+  void deallocate_grid();
+
   typedef void (ComputePropertyGrid::*FnPtrPack)(int);
   FnPtrPack *pack_choice;    // ptrs to pack functions
 

src/dump_custom.cpp

@@ -1149,6 +1149,7 @@ int DumpCustom::count()
 void DumpCustom::pack(tagint *ids)
 {
   for (int n = 0; n < size_one; n++) (this->*pack_choice[n])(n);
+
   if (ids) {
     tagint *tag = atom->tag;
     for (int i = 0; i < nchoose; i++)

src/dump_grid.cpp

@@ -43,7 +43,6 @@ enum{COMPUTE,FIX};
 DumpGrid::DumpGrid(LAMMPS *lmp, int narg, char **arg) :
   Dump(lmp, narg, arg), idregion(nullptr), earg(nullptr), vtype(nullptr), 
   vformat(nullptr), columns(nullptr), columns_default(nullptr), 
-  choose(nullptr), dchoose(nullptr), clist(nullptr),
   field2index(nullptr), field2grid(nullptr), field2data(nullptr),
   argindex(nullptr), id_compute(nullptr), compute(nullptr),
   id_fix(nullptr), fix(nullptr), pack_choice(nullptr)
@@ -91,7 +90,8 @@ DumpGrid::DumpGrid(LAMMPS *lmp, int narg, char **arg) :
 
   if (ioptional < nfield &&
       strcmp(style,"image") != 0 && strcmp(style,"movie") != 0)
-    error->all(FLERR,"Invalid attribute {} in dump {} command",earg[ioptional],style);
+    error->all(FLERR,"Invalid attribute {} in dump {} command",
+               earg[ioptional],style);
 
   // noptional = # of optional args
   // reset nfield to subtract off optional args
@@ -272,7 +272,7 @@ void DumpGrid::init_style()
   Grid2d *grid2d;
   Grid3d *grid3d;
   
-  int nx,ny,nz,nxtmp,nytmp,nztmp;
+  int nxtmp,nytmp,nztmp;
   
   for (int i = 0; i < nfield; i++) {
      if (dimension == 2) {
@@ -283,10 +283,10 @@ void DumpGrid::init_style()
         ifix = fix[field2index[i]];
         grid2d = (Grid2d *) ifix->get_grid_by_index(field2grid[i]);
       }
-      if (i == 0) grid2d->get_size(nx,ny);
+      if (i == 0) grid2d->get_size(nxgrid,nygrid);
       else {
         grid2d->get_size(nxtmp,nytmp);
-        if (nxtmp != nx || nytmp != ny) 
+        if (nxtmp != nxgrid || nytmp != nygrid) 
           error->all(FLERR,"Dump grid field grid sizes do not match");
       }
 
@@ -298,10 +298,10 @@ void DumpGrid::init_style()
         ifix = fix[field2index[i]];
         grid3d = (Grid3d *) ifix->get_grid_by_index(field2grid[i]);
       }
-      if (i == 0) grid3d->get_size(nx,ny,nz);
+      if (i == 0) grid3d->get_size(nxgrid,nygrid,nzgrid);
       else {
         grid3d->get_size(nxtmp,nytmp,nztmp);
-        if (nxtmp != nx || nytmp != ny || nztmp != nz)
+        if (nxtmp != nxgrid || nytmp != nygrid || nztmp != nzgrid)
           error->all(FLERR,"Dump grid field grid sizes do not match");
       }
     }
@@ -505,23 +505,6 @@ int DumpGrid::count()
 {
   int i;
 
-  // grow choose arrays if needed
-  // NOTE: needs to change
-
-  /*
-  const int nlocal = atom->nlocal;
-  if (atom->nmax > maxlocal) {
-    maxlocal = atom->nmax;
-
-    memory->destroy(choose);
-    memory->destroy(dchoose);
-    memory->destroy(clist);
-    memory->create(choose,maxlocal,"dump:choose");
-    memory->create(dchoose,maxlocal,"dump:dchoose");
-    memory->create(clist,maxlocal,"dump:clist");
-  }
-  */
-
   // set current size for portion of grid on each proc
   // may change between dump snapshots due to load balancing
 
@@ -575,38 +558,6 @@ int DumpGrid::count()
     ngrid = (nxhi_in-nxlo_in+1) * (nyhi_in-nylo_in+1) * (nzhi_in-nzlo_in+1);
 
   return ngrid;
-
-  // choose all local grid pts for output
-  // NOTE: this needs to change
-
-  //for (i = 0; i < nlocal; i++) choose[i] = 1;
-
-  // un-choose if not in region
-  // NOTE: this needs to change
-
-  /*
-  if (idregion) {
-    auto region = domain->get_region_by_id(idregion);
-    region->prematch();
-    double **x = atom->x;
-    for (i = 0; i < nlocal; i++)
-      if (choose[i] && region->match(x[i][0],x[i][1],x[i][2]) == 0)
-        choose[i] = 0;
-  }
-  */
-
-  // compress choose flags into clist
-  // nchoose = # of selected atoms
-  // clist[i] = local index of each selected atom
-  // NOTE: this neds to change
-
-  /*
-  nchoose = 0;
-  for (i = 0; i < nlocal; i++)
-    if (choose[i]) clist[nchoose++] = i;
-
-  return nchoose;
-  */
 }
 
 /* ---------------------------------------------------------------------- */
@@ -615,14 +566,21 @@ void DumpGrid::pack(tagint *ids)
 {
   for (int n = 0; n < size_one; n++) (this->*pack_choice[n])(n);
 
-  // NOTE: this needs to be grid IDs ?
+  // ids = list of my grid IDs
-  /*
+
   if (ids) {
-    tagint *tag = atom->tag;
+    int m = 0;
-    for (int i = 0; i < nchoose; i++)
+    if (dimension == 2) {
-      ids[i] = tag[clist[i]];
+        for (int iy = nylo_in; iy <= nyhi_in; iy++)
+          for (int ix = nxlo_in; ix <= nxhi_in; ix++)
+            ids[m++] = iy*nxgrid + ix + 1;
+    } else if (dimension == 3) {
+        for (int iz = nzlo_in; iz <= nzhi_in; iz++)
+          for (int iy = nylo_in; iy <= nyhi_in; iy++)
+            for (int ix = nxlo_in; ix <= nxhi_in; ix++)
+              ids[m++] = iz*nygrid*nxgrid + iy*nxgrid + ix + 1;
+    }
   }
-  */
 }
 
 /* ----------------------------------------------------------------------
@@ -963,10 +921,6 @@ int DumpGrid::modify_param(int narg, char **arg)
 double DumpGrid::memory_usage()
 {
   double bytes = Dump::memory_usage();
-  //NOTE: restore if use choose
-  //bytes += memory->usage(choose,maxlocal);
-  //bytes += memory->usage(dchoose,maxlocal);
-  //bytes += memory->usage(clist,maxlocal);
   return bytes;
 }
 

src/dump_grid.h

@@ -49,11 +49,7 @@ class DumpGrid : public Dump {
 
   int dimension;
 
-  int nchoose;        // # of selected atoms
+  int nxgrid,nygrid,nzgrid;   // global grid size
-  int maxlocal;       // size of atom selection and variable arrays
-  int *choose;        // local indices of selected atoms
-  double *dchoose;    // value for each atom to threshold against
-  int *clist;         // compressed list of indices of selected atoms
 
   int nfield;                 // # of keywords listed by user
   int ioptional;              // index of start of optional args

src/fix.h

@@ -211,6 +211,8 @@ class Fix : protected Pointers {
   virtual int pack_reverse_comm(int, int, double *) { return 0; }
   virtual void unpack_reverse_comm(int, int *, double *) {}
 
+  virtual void reset_grid(){};
+
   virtual void pack_forward_grid(int, void *, int, int *){};
   virtual void unpack_forward_grid(int, void *, int, int *){};
   virtual void pack_reverse_grid(int, void *, int, int *){};

src/fix_ave_grid.cpp

@@ -691,6 +691,17 @@ void FixAveGrid::end_of_step()
   }
 }
 
+/* ----------------------------------------------------------------------
+   subset of grid assigned to each proc may have changed
+   called by load balancer when proc subdomains are adjusted
+   not supported for now, b/c requires per-grid values to persist, i.e. a remap()
+------------------------------------------------------------------------- */
+
+void FixAveGrid::reset_grid()
+{
+  error->all(FLERR,"Fix ave/grid does not support load balancing (yet)");
+}
+
 /* ----------------------------------------------------------------------
    sum per-atom contributions to owned+ghost grid cells
    sets one of vec2d,array2d,vec3d,array3d

src/fix_ave_grid.h

@@ -32,6 +32,9 @@ class FixAveGrid : public Fix {
   void init() override;
   void setup(int) override;
   void end_of_step() override;
+
+  void reset_grid() override;
+
   int get_grid_by_name(char *, int &) override;
   void *get_grid_by_index(int) override;
   int get_griddata_by_name(int, char *, int &) override;

src/fix_balance.cpp

@@ -25,6 +25,7 @@
 #include "kspace.h"
 #include "modify.h"
 #include "neighbor.h"
+#include "pair.h"
 #include "rcb.h"
 #include "update.h"
 
@@ -155,9 +156,6 @@ void FixBalance::post_constructor()
 
 void FixBalance::init()
 {
-  if (force->kspace) kspace_flag = 1;
-  else kspace_flag = 0;
-
   balance->init_imbalance(1);
 }
 
@@ -278,11 +276,13 @@ void FixBalance::rebalance()
   }
 
   // reset proc sub-domains
-  // check and warn if any proc's subbox is smaller than neigh skin
-  //   since may lead to lost atoms in comm->exchange()
 
   if (domain->triclinic) domain->set_lamda_box();
   domain->set_local_box();
+
+  // check and warn if any proc's subbox is smaller than neigh skin
+  //   since may lead to lost atoms in comm->exchange()
+
   domain->subbox_too_small_check(neighbor->skin);
 
   // output of new decomposition
@@ -303,9 +303,12 @@ void FixBalance::rebalance()
   else if (irregular->migrate_check()) irregular->migrate_atoms();
   if (domain->triclinic) domain->lamda2x(atom->nlocal);
 
-  // invoke KSpace setup_grid() to adjust to new proc sub-domains
+  // notify all classes that store distributed grids
+  // so they can adjust to new proc sub-domains
 
-  if (kspace_flag) force->kspace->setup_grid();
+  modify->reset_grid();
+  if (force->pair) force->pair->reset_grid();
+  if (force->kspace) force->kspace->reset_grid();
 
   // pending triggers pre_neighbor() to compute final imbalance factor
   // can only be done after atoms migrate in comm->exchange()

src/fix_balance.h

@@ -43,16 +43,15 @@ class FixBalance : public Fix {
   int nevery, lbstyle, nitermax;
   double thresh, stopthresh;
   char bstr[4];
-  int wtflag;    // 1 for weighted balancing
+  int wtflag;               // 1 for weighted balancing
 
   double imbnow;            // current imbalance factor
   double imbprev;           // imbalance factor before last rebalancing
   double imbfinal;          // imbalance factor after last rebalancing
   double maxloadperproc;    // max load on any processor
   int itercount;            // iteration count of last call to Balance
-  int kspace_flag;          // 1 if KSpace solver defined
   int pending;
-  bigint lastbalance;    // last timestep balancing was attempted
+  bigint lastbalance;       // last timestep balancing was attempted
 
   class Balance *balance;
   class Irregular *irregular;

src/kspace.h

@@ -124,7 +124,7 @@ class KSpace : protected Pointers {
   virtual void settings(int, char **){};
   virtual void init() = 0;
   virtual void setup() = 0;
-  virtual void setup_grid(){};
+  virtual void reset_grid(){};
   virtual void compute(int, int) = 0;
   virtual void compute_group_group(int, int, int){};
 

src/modify.cpp

@@ -791,6 +791,19 @@ int Modify::min_reset_ref()
   return itmpall;
 }
 
+/* ----------------------------------------------------------------------
+   reset grids for any Fix or Compute that uses distributed grids
+   called by load balancer when proc sub-domains change
+------------------------------------------------------------------------- */
+
+void Modify::reset_grid()
+{
+  for (int i = 0; i < nfix; i++)
+    if (fix[i]->pergrid_flag) fix[i]->reset_grid();
+  for (int i = 0; i < ncompute; i++)
+    if (compute[i]->pergrid_flag) compute[i]->reset_grid();
+}
+
 /* ----------------------------------------------------------------------
    add a new fix or replace one with same ID
 ------------------------------------------------------------------------- */

src/modify.h

@@ -101,6 +101,8 @@ class Modify : protected Pointers {
   virtual int min_dof();
   virtual int min_reset_ref();
 
+  void reset_grid();
+
   Fix *add_fix(int, char **, int trysuffix = 1);
   Fix *add_fix(const std::string &, int trysuffix = 1);
   Fix *replace_fix(const char *, int, char **, int trysuffix = 1);

src/pair.h

@@ -203,6 +203,8 @@ class Pair : protected Pointers {
   virtual int pack_reverse_comm(int, int, double *) { return 0; }
   virtual void unpack_reverse_comm(int, int *, double *) {}
 
+  virtual void reset_grid() {}
+
   virtual void pack_forward_grid(int, void *, int, int *) {}
   virtual void unpack_forward_grid(int, void *, int, int *) {}
   virtual void pack_reverse_grid(int, void *, int, int *) {}