ComputeSnap is working with, matching results FitSNAP3 A matrix

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

/* IDEAS

-Need to define a local array for snad on local and ghost atoms, in addition

a local vector of (1+6)*size_array_cols scalars.  

-Reverse communicate local array

-MPI Allreduce on vector 

-Copy vector and array into output array

-size_array_cols = ncoeff

-size_array_rows = total number of atoms

-Boom!

-DONE: Need to define a local peratom array for snad and snad on local and ghost atoms

-DONE: Reverse communicate local peratom array

-DONE: Copy peratom array into output array

-DONE: size_array_cols = nperdim (ncoeff [+quadratic])

-DONE: size_array_rows = 1 + total number of atoms + 6

-DONE: size_peratom = (3+6)*nperdim*ntypes

INCOMPLETE: Mappy from local to global

INCOMPLETE: modify->find_compute() 

INCOMPLETE: eliminate local peratom array for viral, replace with fdotr 

 */

#include "compute_snap.h"

#include <cstring>

using namespace LAMMPS_NS;

enum{SCALAR,VECTOR,ARRAY};

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

  Compute(lmp, narg, arg), cutsq(NULL), list(NULL), snap(NULL),

  radelem(NULL), wjelem(NULL)

  radelem(NULL), wjelem(NULL), snap_peratom(NULL)

{

  array_flag = 1;

  extarray = 0;

  double rfac0, rmin0;

  int twojmax, switchflag, bzeroflag;

  radelem = NULL;

  if (quadraticflag) nperdim += (ncoeff*(ncoeff+1))/2;

  yoffset = nperdim;

  zoffset = 2*nperdim;

  size_array_rows = total number of atoms

  size_array_cols = 3*nperdim*atom->ntypes;

  comm_reverse = size_peratom_cols;

  virialoffset = 3*nperdim;

  natoms = atom->natoms;

  size_array_rows = 1+3*natoms+6;

  size_array_cols = nperdim*atom->ntypes+1; // extra col for reference potential

  ndims_peratom = 9;

  size_peratom = ndims_peratom*nperdim*atom->ntypes; // local atom force and virial data

  comm_reverse = size_peratom;

  nmax = 0;

  snap = NULL;

}

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

ComputeSnap::~ComputeSnap()

{

  memory->destroy(snap);

  memory->destroy(snap_peratom);

  memory->destroy(radelem);

  memory->destroy(wjelem);

  memory->destroy(cutsq);

  if (count > 1 && comm->me == 0)

    error->warning(FLERR,"More than one compute snap");

  snaptr->init();

  // allocate memory for global array

  //  printf("allocate memory for global array rows = %d cols = %d\n",

  //         size_array_rows,size_array_cols);

  memory->create(snap,size_array_rows,size_array_cols,

                 "snap:snap");

  array = snap;

  // INCOMPLETE: modify->find_compute() 

  // was called 223960 times by snappy Ta example

  // that is over 600 times per config?

  // how is this possible???

  // find compute for reference energy

  char *id_pe = (char *) "thermo_pe";

  int ipe = modify->find_compute(id_pe);

  c_pe = modify->compute[ipe];

  // add compute for reference virial tensor

  char *id_virial = (char *) "snap_press";

  int ivirial = modify->find_compute(id_virial);

  if (ivirial == -1)

    error->all(FLERR,"compute snap requires that compute snap_press exists!");

  c_virial = modify->compute[ivirial];

}

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

void ComputeSnap::init_list(int /*id*/, NeighList *ptr)

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

void ComputeSnap::compute()

void ComputeSnap::compute_array()

{

  int ntotal = atom->nlocal + atom->nghost;

  invoked_peratom = update->ntimestep;

  invoked_array = update->ntimestep;

  // grow snap array if necessary

  //  printf("Invoking compute snap on timestep %d\n",invoked_array);

  // grow snap_peratom array if necessary

  if (atom->nmax > nmax) {

    memory->destroy(snap);

    memory->destroy(snap_peratom);

    nmax = atom->nmax;

    memory->create(snap,size_array_rows,size_array_cols,

                   "snap:snap");

    array = snap;

    memory->create(snap_peratom,nmax,size_peratom,

                   "snap:snap_peratom");

  }

  // clear local array

  // clear global array

  // only need to zero out first row

  for (int icoeff = 0; icoeff < size_array_cols; icoeff++)

    snap[0][icoeff] = 0.0;

  // clear peratom array

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

    for (int icoeff = 0; icoeff < size_peratom_cols; icoeff++) {

      snap[i][icoeff] = 0.0;

    for (int icoeff = 0; icoeff < size_peratom; icoeff++) {

      snap_peratom[i][icoeff] = 0.0;

    }

  // invoke full neighbor list (will copy or build if necessary)

      const double radi = radelem[itype];

      const int* const jlist = firstneigh[i];

      const int jnum = numneigh[i];

      // const int typeoffset = threencoeff*(atom->type[i]-1);

      // const int quadraticoffset = threencoeff*atom->ntypes +

      //   threencoeffq*(atom->type[i]-1);

      const int typeoffset = 3*nperdim*(atom->type[i]-1);

      const int typeoffset = ndims_peratom*nperdim*(atom->type[i]-1);

      // insure rij, inside, and typej  are of size jnum

      snaptr->compute_ui(ninside);

      snaptr->compute_zi();

      if (quadraticflag) {

      //      if (quadraticflag) {

      snaptr->compute_bi();

      }

      //      }

      for (int jj = 0; jj < ninside; jj++) {

        const int j = snaptr->inside[jj];

        // Accumulate -dBi/dRi, -dBi/dRj

        double *snapi = snap[i]+typeoffset;

        double *snapj = snap[j]+typeoffset;

        double *snadi = snap_peratom[i]+typeoffset;

        double *snadj = snap_peratom[j]+typeoffset;

        double *snavi = snadi+virialoffset;

        double *snavj = snadj+virialoffset;

        for (int icoeff = 0; icoeff < ncoeff; icoeff++) {

          snapi[icoeff] += snaptr->dblist[icoeff][0];

          snapi[icoeff+yoffset] += snaptr->dblist[icoeff][1];

          snapi[icoeff+zoffset] += snaptr->dblist[icoeff][2];

          snapj[icoeff] -= snaptr->dblist[icoeff][0];

          snapj[icoeff+yoffset] -= snaptr->dblist[icoeff][1];

          snapj[icoeff+zoffset] -= snaptr->dblist[icoeff][2];

          snadi[icoeff] += snaptr->dblist[icoeff][0];

          snadi[icoeff+yoffset] += snaptr->dblist[icoeff][1];

          snadi[icoeff+zoffset] += snaptr->dblist[icoeff][2];

          snadj[icoeff] -= snaptr->dblist[icoeff][0];

          snadj[icoeff+yoffset] -= snaptr->dblist[icoeff][1];

          snadj[icoeff+zoffset] -= snaptr->dblist[icoeff][2];

          snavi[icoeff]           += snaptr->dblist[icoeff][0]*xtmp;

          snavi[icoeff+nperdim]   += snaptr->dblist[icoeff][1]*ytmp;

          snavi[icoeff+2*nperdim] += snaptr->dblist[icoeff][2]*ztmp;

          snavi[icoeff+3*nperdim] += snaptr->dblist[icoeff][1]*ztmp;

          snavi[icoeff+4*nperdim] += snaptr->dblist[icoeff][0]*ztmp;

          snavi[icoeff+5*nperdim] += snaptr->dblist[icoeff][0]*ytmp;

          snavj[icoeff]           -= snaptr->dblist[icoeff][0]*x[j][0];

          snavj[icoeff+nperdim]   -= snaptr->dblist[icoeff][1]*x[j][1];

          snavj[icoeff+2*nperdim] -= snaptr->dblist[icoeff][2]*x[j][2];

          snavj[icoeff+3*nperdim] -= snaptr->dblist[icoeff][1]*x[j][2];

          snavj[icoeff+4*nperdim] -= snaptr->dblist[icoeff][0]*x[j][2];

          snavj[icoeff+5*nperdim] -= snaptr->dblist[icoeff][0]*x[j][1];

        }

        if (quadraticflag) {

          const int quadraticoffset = ncoeff;

          snapi += quadraticoffset;

          snapj += quadraticoffset;

          snadi += quadraticoffset;

          snadj += quadraticoffset;

          snavi += quadraticoffset;

          snavj += quadraticoffset;

          int ncount = 0;

          for (int icoeff = 0; icoeff < ncoeff; icoeff++) {

            double bi = snaptr->blist[icoeff];

            double dbytmp = bi*biy;

            double dbztmp = bi*biz;

            snapi[ncount] +=         dbxtmp;

            snapi[ncount+yoffset] += dbytmp;

            snapi[ncount+zoffset] += dbztmp;

            snapj[ncount] -=         dbxtmp;

            snapj[ncount+yoffset] -= dbytmp;

            snapj[ncount+zoffset] -= dbztmp;

            snadi[ncount] +=         dbxtmp;

            snadi[ncount+yoffset] += dbytmp;

            snadi[ncount+zoffset] += dbztmp;

            snadj[ncount] -=         dbxtmp;

            snadj[ncount+yoffset] -= dbytmp;

            snadj[ncount+zoffset] -= dbztmp;

            snavi[ncount] +=           dbxtmp*xtmp;

            snavi[ncount+nperdim] +=   dbytmp*ytmp;

            snavi[ncount+2*nperdim] += dbztmp*ztmp;

            snavi[ncount+3*nperdim] += dbytmp*ztmp;

            snavi[ncount+4*nperdim] += dbxtmp*ztmp;

            snavi[ncount+5*nperdim] += dbxtmp*ytmp;

            snavj[ncount] -=            dbxtmp*x[j][0];

            snavj[ncount+nperdim] -=    dbytmp*x[j][1];

            snavj[ncount+2*nperdim] -=  dbztmp*x[j][2];

            snavj[ncount+3*nperdim] -=  dbytmp*x[j][2];

            snavj[ncount+4*nperdim] -=  dbxtmp*x[j][2];

            snavj[ncount+5*nperdim] -=  dbxtmp*x[j][1];

            ncount++;

            // upper-triangular elements of quadratic matrix

              double dbztmp = bi*snaptr->dblist[jcoeff][2]

                + biz*snaptr->blist[jcoeff];

              snapi[ncount] +=         dbxtmp;

              snapi[ncount+yoffset] += dbytmp;

              snapi[ncount+zoffset] += dbztmp;

              snapj[ncount] -=         dbxtmp;

              snapj[ncount+yoffset] -= dbytmp;

              snapj[ncount+zoffset] -= dbztmp;

              snadi[ncount] +=         dbxtmp;

              snadi[ncount+yoffset] += dbytmp;

              snadi[ncount+zoffset] += dbztmp;

              snadj[ncount] -=         dbxtmp;

              snadj[ncount+yoffset] -= dbytmp;

              snadj[ncount+zoffset] -= dbztmp;

              snavi[ncount] +=           dbxtmp*xtmp;

              snavi[ncount+nperdim] +=   dbytmp*ytmp;

              snavi[ncount+2*nperdim] += dbztmp*ztmp;

              snavi[ncount+3*nperdim] += dbytmp*ztmp;

              snavi[ncount+4*nperdim] += dbxtmp*ztmp;

              snavi[ncount+5*nperdim] += dbxtmp*ytmp;

              snavj[ncount] -=           dbxtmp*x[j][0];

              snavj[ncount+nperdim] -=   dbytmp*x[j][1];

              snavj[ncount+2*nperdim] -= dbztmp*x[j][2];

              snavj[ncount+3*nperdim] -= dbytmp*x[j][2];

              snavj[ncount+4*nperdim] -= dbxtmp*x[j][2];

              snavj[ncount+5*nperdim] -= dbxtmp*x[j][1];

              ncount++;

            }

          }

        }

      }

      // Accumulate Bi

      // linear contributions

      for (int icoeff = 0; icoeff < ncoeff; icoeff++)

        snap[0][icoeff] += snaptr->blist[icoeff];

      // quadratic contributions

      if (quadraticflag) {

        for (int icoeff = 0; icoeff < ncoeff; icoeff++) {

          double bveci = snaptr->blist[icoeff];

          snap[0][icoeff] += 0.5*bveci*bveci;

          for (int jcoeff = icoeff+1; jcoeff < ncoeff; jcoeff++) {

            double bvecj = snaptr->blist[jcoeff];

            snap[0][icoeff] += bveci*bvecj;

          }

        }

      }

    }

  }

  // INCOMPLETE

  // can get rid of virial from snap_peratom by doing

  // equivalent of Pair::virial_fdotr_compute()

  // before reverse communicate of snap_peratom

  // communicate snap contributions between neighbor procs

  comm->reverse_comm_compute(this);

  // construct global array

  for (int itype = 0; itype < atom->ntypes; itype++) {

    const int typeoffset = 3*nperdim*itype;

    for (int icoeff = 0; icoeff < nperdim; icoeff++) {

      // assign force rows

      // INCOMPLETE ignore local-global mapping for now

      int irow = 1;

      for (int i = 0; i < atom->nlocal; i++) {

        double *snadi = snap_peratom[i]+typeoffset;

        snap[irow++][icoeff+typeoffset] = snadi[icoeff];

        snap[irow++][icoeff+typeoffset] = snadi[icoeff+yoffset];

        snap[irow++][icoeff+typeoffset] = snadi[icoeff+zoffset];

      }

      // assign virial row

      int irow0 = irow;

      snap[irow++][icoeff+typeoffset] = 0.0; 

      snap[irow++][icoeff+typeoffset] = 0.0; 

      snap[irow++][icoeff+typeoffset] = 0.0; 

      snap[irow++][icoeff+typeoffset] = 0.0; 

      snap[irow++][icoeff+typeoffset] = 0.0; 

      snap[irow++][icoeff+typeoffset] = 0.0;

      for (int i = 0; i < atom->nlocal; i++) {

        double *snavi = snap_peratom[i]+typeoffset+virialoffset;

        irow = irow0;

        snap[irow++][icoeff+typeoffset] += snavi[icoeff];

        snap[irow++][icoeff+typeoffset] += snavi[icoeff+1*nperdim];

        snap[irow++][icoeff+typeoffset] += snavi[icoeff+2*nperdim];

        snap[irow++][icoeff+typeoffset] += snavi[icoeff+3*nperdim];

        snap[irow++][icoeff+typeoffset] += snavi[icoeff+4*nperdim];

        snap[irow++][icoeff+typeoffset] += snavi[icoeff+5*nperdim];

      }

    }

  }

  // assign energy row

  int icol = size_array_cols-1;

  int irow = 0;

  double reference_energy = c_pe->compute_scalar();

  snap[irow++][icol] = reference_energy; 

  // assign force rows

  // INCOMPLETE ignore local-global mapping for now

  for (int i = 0; i < atom->nlocal; i++) {

    snap[irow++][icol] = atom->f[i][0];

    snap[irow++][icol] = atom->f[i][1];

    snap[irow++][icol] = atom->f[i][2];

  }

  // assign virial row

  // switch to Voigt notation

  c_virial->compute_vector();

  snap[irow++][icol] = c_virial->vector[0]; 

  snap[irow++][icol] = c_virial->vector[1]; 

  snap[irow++][icol] = c_virial->vector[2]; 

  snap[irow++][icol] = c_virial->vector[5]; 

  snap[irow++][icol] = c_virial->vector[4]; 

  snap[irow++][icol] = c_virial->vector[3]; 

}

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

  m = 0;

  last = first + n;

  for (i = first; i < last; i++)

    for (icoeff = 0; icoeff < size_peratom_cols; icoeff++)

      buf[m++] = snap[i][icoeff];

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

      buf[m++] = snap_peratom[i][icoeff];

  return m;

}

  m = 0;

  for (i = 0; i < n; i++) {

    j = list[i];

    for (icoeff = 0; icoeff < size_peratom_cols; icoeff++)

      snap[j][icoeff] += buf[m++];

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

      snap_peratom[j][icoeff] += buf[m++];

  }

}

double ComputeSnap::memory_usage()

{

  double bytes = nmax*size_peratom_cols * sizeof(double); // snap

  double bytes = size_array_rows*size_array_cols * 

    sizeof(double);                                     // snap

  bytes += nmax*size_peratom * sizeof(double);          // snap_peratom

  bytes += snaptr->memory_usage();                      // SNA object

  return bytes;

  ~ComputeSnap();

  void init();

  void init_list(int, class NeighList *);

  void compute();

  void compute_array();

  int pack_reverse_comm(int, int, double *);

  void unpack_reverse_comm(int, int *, double *);

  double memory_usage();

 private:

  int nmax;

  int natoms, nmax, size_peratom;

  int ncoeff, nperdim, yoffset, zoffset;

  int virialoffset, ndims_peratom;

  double **cutsq;

  class NeighList *list;

  double **snap;

  double **snap_peratom;

  double rcutfac;

  double *radelem;

  double *wjelem;

  class SNA* snaptr;

  double cutmax;

  int quadraticflag;

  Compute *c_pe;

  Compute *c_virial;

};

}