/* -*- c++ -*- ----------------------------------------------------------

   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator

   http://lammps.sandia.gov, Sandia National Laboratories

   Steve Plimpton, sjplimp@sandia.gov

   Copyright (2003) Sandia Corporation.  Under the terms of Contract

   DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains

   certain rights in this software.  This software is distributed under

   the GNU General Public License.

   See the README file in the top-level LAMMPS directory.

------------------------------------------------------------------------- */

#include "imbalance.h"

using namespace LAMMPS_NS;

/* ---------------------------------------------------------------------- */

Imbalance::Imbalance(LAMMPS *lmp) : Pointers(lmp) {}

/* -*- c++ -*- ----------------------------------------------------------

   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator

   http://lammps.sandia.gov, Sandia National Laboratories

   Steve Plimpton, sjplimp@sandia.gov

   Copyright (2003) Sandia Corporation.  Under the terms of Contract

   DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains

   certain rights in this software.  This software is distributed under

   the GNU General Public License.

   See the README file in the top-level LAMMPS directory.

------------------------------------------------------------------------- */

#ifndef LMP_IMBALANCE_H

#define LMP_IMBALANCE_H

#include <stdio.h>

#include "pointers.h"

namespace LAMMPS_NS {

class Imbalance : protected Pointers {

 public:

  Imbalance(class LAMMPS *);

  virtual ~Imbalance() {};

  // parse options. return number of arguments consumed (required)

  virtual int options(int narg, char **arg) = 0;

  // reinitialize internal data (needed for fix balance) (optional)

  virtual void init() {};

  // compute and apply weight factors to local atom array (required)

  virtual void compute(double *weights) = 0;

  // print information about the state of this imbalance compute (required)

  virtual void info(FILE *fp) = 0;

  // disallow default and copy constructor, assignment operator

  // private:

  //Imbalance() {};

  //Imbalance(const Imbalance &) {};

  //Imbalance &operator=(const Imbalance &) {return *this;};

};

}

#endif

/* ----------------------------------------------------------------------

   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator

   http://lammps.sandia.gov, Sandia National Laboratories

   Steve Plimpton, sjplimp@sandia.gov

   Copyright (2003) Sandia Corporation.  Under the terms of Contract

   DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains

   certain rights in this software.  This software is distributed under

   the GNU General Public License.

   See the README file in the top-level LAMMPS directory.

------------------------------------------------------------------------- */

#include "imbalance_group.h"

#include "atom.h"

#include "force.h"

#include "group.h"

#include "error.h"

using namespace LAMMPS_NS;

/* -------------------------------------------------------------------- */

ImbalanceGroup::ImbalanceGroup(LAMMPS *lmp) : Imbalance(lmp), id(0), factor(0)

{}

/* -------------------------------------------------------------------- */

ImbalanceGroup::~ImbalanceGroup()

{

  delete [] id;

  delete [] factor;

}

/* -------------------------------------------------------------------- */

int ImbalanceGroup::options(int narg, char **arg)

{

  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 2*num+1;

}

/* -------------------------------------------------------------------- */

void ImbalanceGroup::compute(double *weight)

{

  const int * const mask = atom->mask;

  const int * const bitmask = group->bitmask;

  const int nlocal = 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;

  }

}

/* -------------------------------------------------------------------- */

void ImbalanceGroup::info(FILE *fp)

{

  if (num > 0) {

    const char * const * const names = group->names;

    fprintf(fp,"  group weights:");

    for (int i = 0; i < num; ++i)

      fprintf(fp," %s=%g",names[id[i]],factor[i]);

    fputs("\n",fp);

  }

}

/* -*- c++ -*- ----------------------------------------------------------

   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator

   http://lammps.sandia.gov, Sandia National Laboratories

   Steve Plimpton, sjplimp@sandia.gov

   Copyright (2003) Sandia Corporation.  Under the terms of Contract

   DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains

   certain rights in this software.  This software is distributed under

   the GNU General Public License.

   See the README file in the top-level LAMMPS directory.

------------------------------------------------------------------------- */

#ifndef LMP_IMBALANCE_GROUP_H

#define LMP_IMBALANCE_GROUP_H

#include "imbalance.h"

namespace LAMMPS_NS {

class ImbalanceGroup : public Imbalance {

 public:

  ImbalanceGroup(class LAMMPS *);

  virtual ~ImbalanceGroup();

  // parse options, return number of arguments consumed

  virtual int options(int, char **);

  // compute and apply weight factors to local atom array

  virtual void compute(double *);

  // print information about the state of this imbalance compute

  virtual void info(FILE *);

 private:

  int num;                     // number of groups with weights

  int *id;                     // numerical IDs of groups

  double *factor;              // group weight factors

};

}

#endif

/* ----------------------------------------------------------------------

   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator

   http://lammps.sandia.gov, Sandia National Laboratories

   Steve Plimpton, sjplimp@sandia.gov

   Copyright (2003) Sandia Corporation.  Under the terms of Contract

   DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains

   certain rights in this software.  This software is distributed under

   the GNU General Public License.

   See the README file in the top-level LAMMPS directory.

------------------------------------------------------------------------- */

#include <mpi.h>

#include "imbalance_neigh.h"

#include "atom.h"

#include "comm.h"

#include "force.h"

#include "neighbor.h"

#include "neigh_request.h"

#include "neigh_list.h"

#include "error.h"

using namespace LAMMPS_NS;

/* -------------------------------------------------------------------- */

ImbalanceNeigh::ImbalanceNeigh(LAMMPS *lmp) : Imbalance(lmp)

{

  did_warn = 0;

}

/* -------------------------------------------------------------------- */

int ImbalanceNeigh::options(int narg, char **arg)

{

  if (narg < 1) error->all(FLERR,"Illegal balance weight command");

  factor = force->numeric(FLERR,arg[0]);

  if (factor < 0.0) error->all(FLERR,"Illegal balance weight command");

  return 1;

}

/* -------------------------------------------------------------------- */

void ImbalanceNeigh::compute(double *weight)

{

  int req;

  if (factor == 0.0) return;

  // find suitable neighbor list

  // we can only make use of certain (conventional) neighbor lists

  for (req = 0; req < neighbor->old_nrequest; ++req) {

    if ((neighbor->old_requests[req]->half ||

         neighbor->old_requests[req]->gran ||

         neighbor->old_requests[req]->respaouter ||

         neighbor->old_requests[req]->half_from_full) &&

        neighbor->old_requests[req]->skip == 0 &&

        neighbor->lists[req] && neighbor->lists[req]->numneigh) break;

  }

  if (req >= neighbor->old_nrequest || neighbor->ago < 0) {

    if (comm->me == 0 && !did_warn)

      error->warning(FLERR,"No suitable neighbor list found. "

                     "Neighbor weighted balancing skipped");

    did_warn = 1;

    return;

  }

  NeighList *list = neighbor->lists[req];

  bigint neighsum = 0;

  const int inum = list->inum;

  const int * const ilist = list->ilist;

  const int * const numneigh = list->numneigh;

  // first pass: get local number of neighbors

  for (int i = 0; i < inum; ++i) neighsum += numneigh[ilist[i]];

  double allatoms = static_cast <double>(atom->natoms);

  if (allatoms == 0.0) allatoms = 1.0;

  double allavg;

  double myavg = static_cast<double>(neighsum)/allatoms;

  MPI_Allreduce(&myavg,&allavg,1,MPI_DOUBLE,MPI_SUM,world);

  // second pass: compute and apply weights

  double scale = 1.0/allavg;

  for (int ii = 0; ii < inum; ++ii) {

    const int i = ilist[ii];

    weight[i] *= (1.0-factor) + factor*scale*numneigh[i];

  }

}

/* -------------------------------------------------------------------- */

void ImbalanceNeigh::info(FILE *fp)

{

  fprintf(fp,"  neigh weight factor: %g\n",factor);

}