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

   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 "string.h"

#include "compute_attribute_atom.h"

#include "atom.h"

#include "domain.h"

#include "memory.h"

#include "error.h"

using namespace LAMMPS_NS;

enum{X,Y,Z,XU,YU,ZU,VX,VY,VZ,FX,FY,FZ,XYZ,V,F};

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

ComputeAttributeAtom::ComputeAttributeAtom(LAMMPS *lmp, int narg, char **arg) :

  Compute(lmp, narg, arg)

{

  if (narg != 4) error->all("Illegal compute ke/atom command");

  peratom_flag = 1;

  size_peratom = 0;

  allocate = 1;

  if (strcmp(arg[3],"x") == 0) which = X;

  else if (strcmp(arg[3],"y") == 0) which = Y;

  else if (strcmp(arg[3],"z") == 0) which = Z;

  else if (strcmp(arg[3],"xu") == 0) which = XU;

  else if (strcmp(arg[3],"yu") == 0) which = YU;

  else if (strcmp(arg[3],"zu") == 0) which = ZU;

  else if (strcmp(arg[3],"vx") == 0) which = VX;

  else if (strcmp(arg[3],"vy") == 0) which = VY;

  else if (strcmp(arg[3],"vz") == 0) which = VZ;

  else if (strcmp(arg[3],"fx") == 0) which = FX;

  else if (strcmp(arg[3],"fy") == 0) which = FY;

  else if (strcmp(arg[3],"fz") == 0) which = FZ;

  else if (strcmp(arg[3],"xyz") == 0) {

    which = XYZ;

    size_peratom = 3;

    allocate = 0;

  } else if (strcmp(arg[3],"v") == 0) {

    which = V;

    size_peratom = 3;

    allocate = 0;

  } else if (strcmp(arg[3],"f") == 0) {

    which = F;

    size_peratom = 3;

    allocate = 0;

  } else error->all("Illegal compute attribute/atom command");

  nmax = 0;

  s_attribute = NULL;

  v_attribute = NULL;

}

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

ComputeAttributeAtom::~ComputeAttributeAtom()

{

  if (allocate) {

    memory->sfree(s_attribute);

    memory->destroy_2d_double_array(v_attribute);

  }

}

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

void ComputeAttributeAtom::compute_peratom()

{

  // grow attribute array if necessary

  if (allocate && atom->nlocal > nmax) {

    if (size_peratom == 0) {

      memory->sfree(s_attribute);

      nmax = atom->nmax;

      s_attribute = (double *) 

	memory->smalloc(nmax*sizeof(double),

			"compute/attribute/atom:s_attribute");

    } else {

      memory->destroy_2d_double_array(v_attribute);

      nmax = atom->nmax;

      v_attribute =  

	memory->create_2d_double_array(nmax,size_peratom,

				       "compute/attribute/atom:v_attribute");

    }

  }

  // fill attribute vector with appropriate atom value

  // or simply set pointer to exisitng atom vector

  double xprd = domain->xprd;

  double yprd = domain->yprd;

  double zprd = domain->zprd;

  double **x = atom->x;

  double **v = atom->v;

  double **f = atom->f;

  int *mask = atom->mask;

  int *image = atom->image;

  int nlocal = atom->nlocal;

  if (which == X) {

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

      if (mask[i] & groupbit) s_attribute[i] = x[i][0];

      else s_attribute[i] = 0.0;

  } else if (which == Y) {

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

      if (mask[i] & groupbit) s_attribute[i] = x[i][1];

      else s_attribute[i] = 0.0;

  } else if (which == Z) {

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

      if (mask[i] & groupbit) s_attribute[i] = x[i][2];

      else s_attribute[i] = 0.0;

  } else if (which == XU) {

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

      if (mask[i] & groupbit)

	s_attribute[i] = x[i][0] + ((image[i] & 1023) - 512) * xprd;

      else s_attribute[i] = 0.0;

  } else if (which == YU) {

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

      if (mask[i] & groupbit)

	s_attribute[i] = x[i][1] + ((image[i] >> 10 & 1023) - 512) * yprd;

      else s_attribute[i] = 0.0;

  } else if (which == ZU) {

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

      if (mask[i] & groupbit)

	s_attribute[i] = x[i][2] + ((image[i] >> 20) - 512) * zprd;

      else s_attribute[i] = 0.0;

  } else if (which == VX) {

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

      if (mask[i] & groupbit) s_attribute[i] = v[i][0];

      else s_attribute[i] = 0.0;

  } else if (which == VY) {

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

      if (mask[i] & groupbit) s_attribute[i] = v[i][1];

      else s_attribute[i] = 0.0;

  } else if (which == VZ) {

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

      if (mask[i] & groupbit) s_attribute[i] = v[i][2];

      else s_attribute[i] = 0.0;

  } else if (which == FX) {

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

      if (mask[i] & groupbit) s_attribute[i] = f[i][0];

      else s_attribute[i] = 0.0;

  } else if (which == FY) {

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

      if (mask[i] & groupbit) s_attribute[i] = f[i][1];

      else s_attribute[i] = 0.0;

  } else if (which == FZ) {

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

      if (mask[i] & groupbit) s_attribute[i] = f[i][2];

      else s_attribute[i] = 0.0;

  } else if (which == XYZ) v_attribute = x;

  else if (which == V) v_attribute = v;

  else if (which == F) v_attribute = f;

  // set appropriate compute ptr to local array

  if (size_peratom == 0) scalar_atom = s_attribute;

  else vector_atom = v_attribute;

}

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

   memory usage of local atom-based array

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

int ComputeAttributeAtom::memory_usage()

{

  int bytes = 0;

  if (allocate) {

    if (size_peratom == 0) bytes = nmax * sizeof(double);

    else bytes = size_peratom * nmax * sizeof(double);

  }

  return bytes;

}

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

   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 COMPUTE_ATTRIBUTE_ATOM_H

#define COMPUTE_ATTRIBUTE_ATOM_H

#include "compute.h"

namespace LAMMPS_NS {

class ComputeAttributeAtom : public Compute {

 public:

  ComputeAttributeAtom(class LAMMPS *, int, char **);

  ~ComputeAttributeAtom();

  void init() {}

  void compute_peratom();

  int memory_usage();

 private:

  int which,allocate,nmax;

  double *s_attribute,**v_attribute;

};

}

#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 "string.h"

#include "compute_sum_atom.h"

#include "atom.h"

#include "modify.h"

#include "force.h"

#include "comm.h"

#include "memory.h"

#include "error.h"

using namespace LAMMPS_NS;

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

ComputeSumAtom::ComputeSumAtom(LAMMPS *lmp, int narg, char **arg) :

  Compute(lmp, narg, arg)

{

  if (narg < 5) error->all("Illegal compute sum/atom command");

  // store pre-compute IDs

  npre = narg - 3;

  id_pre = new char*[npre];

  for (int i = 0; i < npre; i++) {

    int iarg = i + 3;

    int n = strlen(arg[iarg]) + 1;

    id_pre[i] = new char[n];

    strcpy(id_pre[i],arg[iarg]);

  }

  compute = new Compute*[npre];

  // check consistency of set of pre-computes for scalar & vector output

  peratom_flag = 1;

  int icompute = modify->find_compute(id_pre[0]);

  size_peratom = modify->compute[icompute]->size_peratom;

  for (int i = 1; i < npre; i++) {

    icompute = modify->find_compute(id_pre[i]);

    if (icompute < 0)

      error->all("Could not find compute sum/atom pre-compute ID");

    if (modify->compute[icompute]->peratom_flag == 0)

      error->all("Compute sum/atom compute does not compute vector per atom");

    if (modify->compute[icompute]->size_peratom != size_peratom)

      error->all("Inconsistent sizes of compute sum/atom compute vectors");

  }

  nmax = 0;

  s_value = NULL;

  v_value = NULL;

}

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

ComputeSumAtom::~ComputeSumAtom()

{

  delete [] compute;

  memory->sfree(s_value);

  memory->destroy_2d_double_array(v_value);

}

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

void ComputeSumAtom::init()

{

  // set ptrs to Computes used as pre-computes by this compute

  for (int i = 0; i < npre; i++) {

    int icompute = modify->find_compute(id_pre[i]);

    if (icompute < 0)

      error->all("Could not find compute sum/atom pre-compute ID");

    compute[i] = modify->compute[icompute];

  }

}

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

void ComputeSumAtom::compute_peratom()

{

  int i,j,m;

  // grow sum array if necessary

  if (atom->nlocal > nmax) {

    nmax = atom->nmax;

    if (size_peratom == 0) {

      memory->sfree(s_value);

      s_value = (double *) 

	memory->smalloc(nmax*sizeof(double),"compute/sum/atom:s_value");

      scalar_atom = s_value;

    } else {

      memory->destroy_2d_double_array(v_value);

      v_value = memory->create_2d_double_array(nmax,size_peratom,

					       "compute/sum/atom:v_value");

      vector_atom = v_value;

    }

  }

  // sum over pre-computes

  // pre-computes of the pre-computes are not invoked

  int *mask = atom->mask;

  int nlocal = atom->nlocal;

  if (size_peratom == 0) {

    double *scalar = compute[0]->scalar_atom;

    for (i = 0; i < nlocal; i++)

      if (mask[i] & groupbit) s_value[i] = scalar[i];

      else s_value[i] = 0.0;

    for (m = 1; m < npre; m++) {

      scalar = compute[m]->scalar_atom;

      for (i = 0; i < nlocal; i++)

	if (mask[i] & groupbit) s_value[i] += scalar[i];

    }

  } else {

    double **vector = compute[0]->vector_atom;

    for (i = 0; i < nlocal; i++)

      if (mask[i] & groupbit) 

	for (j = 0; j < size_peratom; j++)

	  v_value[i][j] = vector[i][j];

      else 

	for (j = 0; j < size_peratom; j++)

	  v_value[i][j] = 0.0;

    for (m = 1; m < npre; m++) {

      vector = compute[m]->vector_atom;

      for (j = 0; j < size_peratom; j++)

	if (mask[i] & groupbit) v_value[i][j] += vector[i][j];

    }

  }

}

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

   memory usage of local atom-based array

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

int ComputeSumAtom::memory_usage()

{

  int bytes = 0;

  if (size_peratom == 0) bytes = nmax * sizeof(double);

  else bytes = nmax*size_peratom * sizeof(double);

  return bytes;

}

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

   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 COMPUTE_SUM_ATOM_H

#define COMPUTE_SUM_ATOM_H

#include "compute.h"

namespace LAMMPS_NS {

class ComputeSumAtom : public Compute {

 public:

  ComputeSumAtom(class LAMMPS *, int, char **);

  ~ComputeSumAtom();

  void init();

  void compute_peratom();

  int memory_usage();

 private:

  int nmax;

  class Compute **compute;

  double *s_value,**v_value;

};

}

#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 "stdlib.h"

#include "string.h"

#include "fix_ave_atom.h"

#include "atom.h"

#include "update.h"

#include "modify.h"

#include "compute.h"

#include "memory.h"

#include "error.h"

using namespace LAMMPS_NS;

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

FixAveAtom::FixAveAtom(LAMMPS *lmp, int narg, char **arg) :

  Fix(lmp, narg, arg)

{

  if (narg != 7) error->all("Illegal fix ave/atom command");

  nevery = atoi(arg[3]);

  nrepeat = atoi(arg[4]);

  peratom_freq = atoi(arg[5]);

  int n = strlen(arg[6]) + 1;

  id_compute = new char[n];

  strcpy(id_compute,arg[6]);

  // setup and error check

  if (nevery <= 0) error->all("Illegal fix ave/atom command");

  if (peratom_freq < nevery || peratom_freq % nevery ||

      (nrepeat-1)*nevery >= peratom_freq)

    error->all("Illegal fix ave/atom command");

  int icompute = modify->find_compute(id_compute);

  if (icompute < 0) error->all("Compute ID for fix ave/atom does not exist");

  if (modify->compute[icompute]->peratom_flag == 0)

    error->all("Fix ave/atom compute does not calculate per-atom info");

  if (modify->compute[icompute]->pressflag) pressure_every = nevery;

  peratom_flag = 1;

  // setup list of computes to call, including pre-computes

  ncompute = 1 + modify->compute[icompute]->npre;

  compute = new Compute*[ncompute];

  // perform initial allocation of atom-based array

  // register with Atom class

  size_peratom = modify->compute[icompute]->size_peratom;

  scalar = NULL;

  vector = NULL;

  grow_arrays(atom->nmax);

  atom->add_callback(0);

  // zero the array since dump may access it on timestep 0

  int nlocal = atom->nlocal;

  if (size_peratom == 0)

    for (int i = 0; i < nlocal; i++) scalar[i] = 0.0;

  else

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

      for (int m = 0; m < size_peratom; m++)

	vector[i][m] = 0.0;

  // nvalid = next step on which end_of_step does something

  irepeat = 0;

  nvalid = (update->ntimestep/peratom_freq)*peratom_freq + peratom_freq;

  nvalid -= (nrepeat-1)*nevery;

  if (nvalid <= update->ntimestep)

    error->all("Fix ave/atom cannot be started on this timestep");

}

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

FixAveAtom::~FixAveAtom()

{

  // unregister callback to this fix from Atom class

  atom->delete_callback(id,0);

  delete [] id_compute;

  delete [] compute;

  memory->sfree(scalar);

  memory->destroy_2d_double_array(vector);

}

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

int FixAveAtom::setmask()

{

  int mask = 0;

  mask |= END_OF_STEP;

  return mask;

}

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

void FixAveAtom::init()

{

  // set ptrs to one or more computes called each end-of-step

  int icompute = modify->find_compute(id_compute);

  if (icompute < 0) error->all("Compute ID for fix ave/atom does not exist");

  ncompute = 0;

  if (modify->compute[icompute]->npre)

    for (int i = 0; i < modify->compute[icompute]->npre; i++) {

      int ic = modify->find_compute(modify->compute[icompute]->id_pre[i]);

      if (ic < 0)

	error->all("Precompute ID for fix ave/atom does not exist");

      compute[ncompute++] = modify->compute[ic];

    }

  compute[ncompute++] = modify->compute[icompute];

}

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

void FixAveAtom::end_of_step()

{

  int i,m;

  // skip if not step which requires doing something

  if (update->ntimestep != nvalid) return;

  // zero if first step

  int nlocal = atom->nlocal;

  if (irepeat == 0) {

    if (size_peratom == 0)

      for (i = 0; i < nlocal; i++) scalar[i] = 0.0;

    else

      for (i = 0; i < nlocal; i++)

	for (m = 0; m < size_peratom; m++)

	  vector[i][m] = 0.0;

  }

  // accumulate results of compute to local copy

  for (i = 0; i < ncompute; i++) compute[i]->compute_peratom();

  int *mask = atom->mask;

  if (size_peratom == 0) {

    double *compute_scalar = compute[ncompute-1]->scalar_atom;

    for (i = 0; i < nlocal; i++)

      if (mask[i] & groupbit) scalar[i] += compute_scalar[i];

  } else {

    double **compute_vector = compute[ncompute-1]->vector_atom;

    for (i = 0; i < nlocal; i++)

      if (mask[i] & groupbit)

	for (m = 0; m < size_peratom; m++)

	  vector[i][m] += compute_vector[i][m];

  }

  irepeat++;

  nvalid += nevery;

  // divide by nrepeat if final step

  // reset irepeat and nvalid

  if (irepeat == nrepeat) {

    double repeat = nrepeat;

    if (size_peratom == 0)

      for (i = 0; i < nlocal; i++)

	scalar[i] /= repeat;

    else

      for (i = 0; i < nlocal; i++)

	for (m = 0; m < size_peratom; m++)

	  vector[i][m] /= repeat;

    irepeat = 0;

    nvalid = update->ntimestep+peratom_freq - (nrepeat-1)*nevery;

  }

}

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

   memory usage of local atom-based array

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

int FixAveAtom::memory_usage()

{

  int bytes;

  if (size_peratom == 0) bytes = atom->nmax * sizeof(double);

  else bytes = atom->nmax*size_peratom * sizeof(double);

  return bytes;

}

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

   allocate atom-based array

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

void FixAveAtom::grow_arrays(int nmax)

{

  if (size_peratom == 0) {

    scalar = (double *) memory->srealloc(scalar,nmax,"fix_ave/atom:scalar");