add lammps compute style for atomic deep tensor

doc/getting-started.md

@@ -417,9 +417,11 @@ and then run the program:
 Note that the model for MD simulations is required to be compatible with the DeePMD-kit package. See [Model compatibility](troubleshooting/model-compatability.md) for details. 
 
 ### Run MD with LAMMPS
-Include deepmd in the pair_style
 
-#### Syntax
+#### pair_style `deepmd`
+
+The DeePMD-kit package provides the pair_style `deepmd`
+
 ```
 pair_style deepmd models ... keyword value ...
 ```
@@ -439,14 +441,14 @@ pair_style deepmd models ... keyword value ...
         level = The level parameter for computing the relative model deviation
 </pre>
 
-#### Examples
+##### Examples
 ```
 pair_style deepmd graph.pb
 pair_style deepmd graph.pb fparam 1.2
 pair_style deepmd graph_0.pb graph_1.pb graph_2.pb out_file md.out out_freq 10 atomic relative 1.0
 ```
 
-#### Description
+##### Description
 Evaluate the interaction of the system by using [Deep Potential][DP] or [Deep Potential Smooth Edition][DP-SE]. It is noticed that deep potential is not a "pairwise" interaction, but a multi-body interaction. 
 
 This pair style takes the deep potential defined in a model file that usually has the .pb extension. The model can be trained and frozen by package [DeePMD-kit](https://github.com/deepmodeling/deepmd-kit).
@@ -461,11 +463,35 @@ Ef_i = -------------
 ```
 where `Df_i` is the absolute model deviation of the force on atom `i`, `|f_i|` is the norm of the the force and `level` is provided as the parameter of the keyword `relative`.
 
-
-#### Restrictions
+##### Restrictions
 - The `deepmd` pair style is provided in the USER-DEEPMD package, which is compiled from the DeePMD-kit, visit the [DeePMD-kit website](https://github.com/deepmodeling/deepmd-kit) for more information.
 
 
+#### Compute tensorial prperties
+
+The DeePMD-kit package provide the compute `deeptensor/atom` for computing atomic tensorial properties. 
+
+```
+compute ID group-ID deeptensor/atom model_file
+```
+- ID: user-assigned name of the computation
+- group-ID: ID of the group of atoms to compute
+- deeptensor/atom: the style of this compute
+- model_file: the name of the binary model file.
+
+##### Examples
+```
+compute         dipole all deeptensor/atom dipole.pb
+```
+The result of the compute can be dump to trajctory file by 
+```
+dump            1 all custom 100 water.dump id type c_dipole[1] c_dipole[2] c_dipole[3] 
+```
+
+##### Restrictions
+- The `deeptensor/atom` compute is provided in the USER-DEEPMD package, which is compiled from the DeePMD-kit, visit the [DeePMD-kit website](https://github.com/deepmodeling/deepmd-kit) for more information.
+
+
 #### Long-range interaction
 The reciprocal space part of the long-range interaction can be calculated by LAMMPS command `kspace_style`. To use it with DeePMD-kit, one writes 
 ```bash

source/lmp/compute_deeptensor_atom.cpp

@@ -0,0 +1,172 @@
+#include "compute_deeptensor_atom.h"
+#include <cstring>
+#include <algorithm>
+#include "atom.h"
+#include "update.h"
+#include "neighbor.h"
+#include "neigh_list.h"
+#include "neigh_request.h"
+#include "comm.h"
+#include "force.h"
+#include "pair.h"
+#include "fix.h"
+#include "memory.h"
+#include "error.h"
+
+#include "domain.h"
+#include "update.h"
+#include "modify.h"
+#include "fix.h"
+
+using namespace LAMMPS_NS;
+
+#ifdef HIGH_PREC
+#define VALUETYPE double
+#else
+#define VALUETYPE float
+#endif
+
+/* ---------------------------------------------------------------------- */
+
+ComputeDeeptensorAtom::ComputeDeeptensorAtom(LAMMPS *lmp, int narg, char **arg) :
+  Compute(lmp, narg, arg),
+  dp(lmp),
+  tensor(nullptr)
+{
+  if (narg < 4) error->all(FLERR,"Illegal compute deeptensor/atom command");
+
+  // parse args
+  std::string model_file = std::string(arg[3]);
+
+  // initialize deeptensor
+  int gpu_rank = dp.get_node_rank();
+  std::string model_file_content = dp.get_file_content(model_file);
+  dt.init(model_file, gpu_rank);
+  sel_types = dt.sel_types();
+  std::sort(sel_types.begin(), sel_types.end());
+
+  peratom_flag = 1;
+  size_peratom_cols = dt.output_dim();
+  pressatomflag = 0;
+  timeflag = 1;
+
+  nmax = 0;
+}
+
+/* ---------------------------------------------------------------------- */
+
+ComputeDeeptensorAtom::~ComputeDeeptensorAtom()
+{
+  memory->destroy(tensor);
+}
+
+/* ---------------------------------------------------------------------- */
+
+void ComputeDeeptensorAtom::init()
+{
+  // need an occasional full neighbor list
+
+  int irequest = neighbor->request(this,instance_me);
+  neighbor->requests[irequest]->half = 0;
+  neighbor->requests[irequest]->pair = 0;
+  neighbor->requests[irequest]->compute = 1;
+  // neighbor->requests[irequest]->full = 1;
+  neighbor->requests[irequest]->occasional = 1;
+}
+
+void ComputeDeeptensorAtom::init_list(int /*id*/, NeighList *ptr)
+{
+  list = ptr;
+}
+
+/* ---------------------------------------------------------------------- */
+
+void ComputeDeeptensorAtom::compute_peratom()
+{
+  invoked_peratom = update->ntimestep;
+
+  // grow local tensor array if necessary
+  // needs to be atom->nmax in length
+  if (atom->nmax > nmax) {
+    memory->destroy(tensor);
+    nmax = atom->nmax;
+    memory->create(tensor, nmax, size_peratom_cols, "deeptensor/atom:tensor");
+    array_atom = tensor;
+  }
+
+  double **x = atom->x;
+  double **f = atom->f;
+  int *type = atom->type;
+  int *mask = atom->mask;
+  int nlocal = atom->nlocal;
+  int nghost = atom->nghost;
+  int nall = nlocal + nghost;
+  int newton_pair = force->newton_pair;
+
+  std::vector<VALUETYPE > dcoord (nall * 3, 0.);
+  std::vector<VALUETYPE > dbox (9, 0) ;
+  std::vector<int > dtype (nall);
+  // get type
+  for (int ii = 0; ii < nall; ++ii){
+    dtype[ii] = type[ii] - 1;
+  }
+  // get box
+  dbox[0] = domain->h[0];	// xx
+  dbox[4] = domain->h[1];	// yy
+  dbox[8] = domain->h[2];	// zz
+  dbox[7] = domain->h[3];	// zy
+  dbox[6] = domain->h[4];	// zx
+  dbox[3] = domain->h[5];	// yx
+  // get coord
+  for (int ii = 0; ii < nall; ++ii){
+    for (int dd = 0; dd < 3; ++dd){
+      dcoord[ii*3+dd] = x[ii][dd] - domain->boxlo[dd];
+    }
+  }
+
+  // invoke full neighbor list (will copy or build if necessary)
+  neighbor->build_one(list);
+  deepmd::InputNlist lmp_list (list->inum, list->ilist, list->numneigh, list->firstneigh);
+
+  // declare outputs
+  std::vector<VALUETYPE > gtensor, force, virial, atensor, avirial;
+
+  // compute tensors
+  dt.compute (gtensor, force, virial, atensor, avirial,
+	      dcoord, dtype, dbox, nghost, lmp_list);
+
+  // store the result in tensor
+  int iter_tensor = 0;
+  for(int ii = 0; ii < nlocal; ++ii){
+    std::vector<int>::iterator _it =
+	std::find(sel_types.begin(), sel_types.end(), dtype[ii]);
+    bool selected = (_it != sel_types.end());    
+    bool ingroup = (mask[ii] & groupbit);    
+    // record when selected and in group
+    if (selected && ingroup){
+      for(int jj = 0; jj < size_peratom_cols; ++jj){
+	tensor[ii][jj] = atensor[iter_tensor+jj];
+      }
+    }
+    // if not selected or not in group set to 0.
+    else{
+      for(int jj = 0; jj < size_peratom_cols; ++jj){
+	tensor[ii][jj] = 0.0;
+      }
+    }    
+    if (selected) {
+      iter_tensor += size_peratom_cols;
+    }
+  }
+}
+
+
+/* ----------------------------------------------------------------------
+   memory usage of local atom-based array
+------------------------------------------------------------------------- */
+
+double ComputeDeeptensorAtom::memory_usage()
+{
+  double bytes = nmax*size_peratom_cols * sizeof(double);
+  return bytes;
+}

source/lmp/compute_deeptensor_atom.h

@@ -0,0 +1,38 @@
+#ifdef COMPUTE_CLASS
+
+ComputeStyle(deeptensor/atom,ComputeDeeptensorAtom)
+
+#else
+
+#ifndef LMP_COMPUTE_DEEPTENSOR_ATOM_H
+#define LMP_COMPUTE_DEEPTENSOR_ATOM_H
+
+#include "compute.h"
+#include "pair_deepmd.h"
+#include "deepmd/DeepTensor.h"
+
+namespace LAMMPS_NS {
+
+class ComputeDeeptensorAtom : public Compute {
+ public:
+  ComputeDeeptensorAtom(class LAMMPS *, int, char **);
+  ~ComputeDeeptensorAtom();
+  void init();
+  void compute_peratom();
+  double memory_usage();
+  void init_list(int, class NeighList *);
+
+ private:
+  int nmax;
+  double **tensor;
+  PairDeepMD dp;
+  class NeighList *list;
+  deepmd::DeepTensor dt;
+  std::vector<int > sel_types;
+};
+
+}
+
+#endif
+#endif
+