modified versions of creating atoms on subset of lattice, ditto for set type/fraction

src/random_mars.cpp

@@ -21,6 +21,8 @@
 
 using namespace LAMMPS_NS;
 
+enum{ADD,SUBTRACT};
+
 /* ---------------------------------------------------------------------- */
 
 RanMars::RanMars(LAMMPS *lmp, int seed) : Pointers(lmp),
@@ -149,7 +151,7 @@ double RanMars::besselexp(double theta, double alpha, double cp)
 {
   double first,v1,v2;
 
-  if (theta < 0.0 || alpha < 0.0 || alpha > 1.0)
+  if (theta < 0.0 || alpha < 0.0 || alpha < 1.0)
     error->all(FLERR,"Invalid Bessel exponential distribution parameters");
 
   v1 = uniform();
@@ -166,3 +168,131 @@ double RanMars::besselexp(double theta, double alpha, double cp)
 
   return first;
 }
+
+/* ----------------------------------------------------------------------
+   select random subset of size Ntarget out of Ntotal items
+   Ntotal = sum of Nmine across all procs
+   mark,next = vectors of length Nmine
+   return mark = 0 for unselected item, 1 for selected item
+   next = work vector used to store linked lists for 2 active sets of items
+   IMPORTANT: this method must be called simultaneously by all procs
+------------------------------------------------------------------------- */
+
+void RanMars::select_subset(bigint ntarget, int nmine, int *mark, int *next)
+{
+  int mode,index,oldindex,newvalue,nflip,which,niter;
+  int active[2],first[2],last[2];
+  int newactive[2],newfirst[2],newlast[2];
+  bigint nmark,nactive,nactiveall,nflipall,bnflip;
+  bigint activeall[2],bsum[3],bsumall[3];
+  double thresh;
+
+  active[0] = nmine;
+  active[1] = 0;
+  first[0] = 0;
+  first[1] = -1;
+  last[0] = nmine-1;
+  last[1] = -1;
+
+  bigint bnmine = nmine;
+  bigint bnall;
+  MPI_Allreduce(&bnmine,&bnall,1,MPI_LMP_BIGINT,MPI_SUM,world);
+  activeall[0] = bnall;
+
+  for (int i = 0; i < nmine; i++) mark[i] = 0;
+  for (int i = 0; i < nmine; i++) next[i] = i+1;
+  next[nmine-1] = -1;
+
+  nmark = 0;
+  niter = 0;
+
+  while (nmark != ntarget) {
+
+    // choose to ADD or SUBTRACT from current nmark
+    // thresh = desired flips / size of active set
+    // nactive = size of current active set, only for debug output below
+
+    if (ntarget-nmark > 0) {
+      mode = ADD;
+      // nactive = active[mode];
+      thresh = 1.0 * (ntarget-nmark) / activeall[mode];
+    } else {
+      mode = SUBTRACT;
+      // nactive = active[mode];
+      thresh = 1.0 * (nmark-ntarget) / activeall[mode];
+    }
+
+    // bound thresh for RNG accuracy
+  
+    thresh = MAX(thresh,0.01);
+    thresh = MIN(thresh,0.99);
+
+    // new empty active sets for next iteration
+  
+    newactive[0] = newactive[1] = 0;
+    newfirst[0] = newfirst[1] = -1;
+    newlast[0] = newlast[1] = -1;
+
+    // index = first value in ADD or SUBTRACT set
+  
+    if (mode == ADD) newvalue = 1;
+    else if (mode == SUBTRACT) newvalue = 0;
+    index = first[mode];
+
+    // flip marks from 0 -> 1 (ADD) or 1 -> 0 (SUBTRACT)
+    // loop over active set via next vector = linked list
+    // flip each value based on RN < thresh
+  
+    nflip = 0;
+    while (index >= 0) {
+      if (uniform() < thresh) {
+        mark[index] = newvalue;
+        nflip++;
+      }
+      oldindex = index;
+      index = next[index];
+
+      // oldindex can now be appended to a new active set
+      // which = which of two new active sets to append to
+    
+      which = mark[oldindex];
+      newactive[which]++;
+      if (newfirst[which] < 0) newfirst[which] = oldindex;
+      if (newlast[which] >= 0) next[newlast[which]] = oldindex;
+      newlast[which] = oldindex;
+      next[oldindex] = -1;
+
+      // set active sets for next iteration to the new ones
+      // next vector is already updated
+  
+      active[0] = newactive[0];
+      active[1] = newactive[1];  
+      first[0] = newfirst[0];
+      first[1] = newfirst[1];  
+      last[0] = newlast[0];
+      last[1] = newlast[1];  
+    }
+
+    // update nmark and activeall
+
+    bsum[0] = nflip;
+    bsum[1] = active[0];
+    bsum[2] = active[1];
+    bsum[3] = nactive;
+    MPI_Allreduce(&bsum,&bsumall,4,MPI_LMP_BIGINT,MPI_SUM,world);
+    nflipall = bsumall[0];
+    activeall[0] = bsumall[1];
+    activeall[1] = bsumall[2];
+    nactiveall = bsumall[3];
+
+    if (mode == ADD) nmark += nflipall;
+    else if (mode == SUBTRACT) nmark -= nflipall;
+
+    niter++;
+    
+    // DEBUG output of stats
+
+    //if (comm->me == 0) printf("%d %ld %ld %g %ld\n",
+    //                          niter,nmark,nactiveall,thresh,nflipall);
+  }
+}