Refactorize the implementation of the neighbor list building.

.github/workflows/test_python.yml

@@ -5,7 +5,7 @@ name: Test Python
 jobs:
   testpython:
     name: Test Python
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-18.04
     strategy:  
       matrix:
         include:

source/api_cc/src/NNPInter.cc

@@ -190,7 +190,10 @@ void
 NNPInter::
 init (const std::string & model, const int & gpu_rank, const std::string & file_content)
 {
-  assert (!inited);
+  if (inited){
+    std::cerr << "WARNING: deepmd-kit should not be initialized twice, do nothing at the second call of initializer" << std::endl;
+    return ;
+  }
   SessionOptions options;
   options.config.set_inter_op_parallelism_threads(num_inter_nthreads);
   options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
@@ -500,7 +503,10 @@ void
 NNPInterModelDevi::
 init (const std::vector<std::string> & models, const int & gpu_rank, const std::vector<std::string> & file_contents)
 {
-  assert (!inited);
+  if (inited){
+    std::cerr << "WARNING: deepmd-kit should not be initialized twice, do nothing at the second call of initializer" << std::endl;
+    return ;
+  }
   numb_models = models.size();
   sessions.resize(numb_models);
   graph_defs.resize(numb_models);

source/lib/include/coord.h

@@ -0,0 +1,35 @@
+#pragma once
+
+#include "region.h"
+
+// normalize coords
+template <typename FPTYPE>
+void
+normalize_coord_cpu(
+    FPTYPE * coord,
+    const int natom,
+    const Region<FPTYPE> & region);
+
+// copy coordinates
+// outputs:
+//	out_c, out_t, mapping, nall
+// inputs:
+//	in_c, in_t, nloc, mem_nall, rc, region
+//	mem_nall is the size of allocated memory for out_c, out_t, mapping
+// returns
+//	0: succssful
+//	1: the memory is not large enough to hold all copied coords and types.
+//	   i.e. nall > mem_nall
+template <typename FPTYPE>
+int
+copy_coord_cpu(
+    FPTYPE * out_c,
+    int * out_t,
+    int * mapping,
+    int * nall,
+    const FPTYPE * in_c,
+    const int * in_t,
+    const int & nloc,
+    const int & mem_nall,
+    const float & rcut,
+    const Region<FPTYPE> & region);

source/lib/include/env_mat.h

@@ -24,7 +24,6 @@ void env_mat_a_cpu (
     std::vector<FPTYPE > &	        descrpt_a_deriv,
     std::vector<FPTYPE > &	        rij_a,
     const std::vector<FPTYPE > &	posi,
-    const int &				ntypes,
     const std::vector<int > &		type,
     const int &				i_idx,
     const std::vector<int > &		fmt_nlist,
@@ -53,7 +52,6 @@ void env_mat_r_cpu (
     std::vector<FPTYPE > &	        descrpt_a_deriv,
     std::vector<FPTYPE > &	        rij_a,
     const std::vector<FPTYPE > &	posi,
-    const int &				ntypes,
     const std::vector<int > &		type,
     const int &				i_idx,
     const std::vector<int > &		fmt_nlist_a,

source/lib/include/fmt_nlist.h

@@ -1,11 +1,23 @@
 #pragma once
 
 #include <vector>
+#include "neighbor_list.h"
 #include "SimulationRegion.h"
 
+template <typename FPTYPE>
+void format_nlist_cpu(
+    int * nlist,
+    const InputNlist & in_nlist,
+    const FPTYPE * coord, 
+    const int * type, 
+    const int nloc, 
+    const int nall, 
+    const float rcut, 
+    const std::vector<int> sec);
+
 // return:	-1	OK
 //		> 0	the type of unsuccessful neighbor list
-int format_nlist_fill_a (
+int format_nlist_i_fill_a (
     std::vector<int > &			fmt_nei_idx_a,
     std::vector<int > &			fmt_nei_idx_r,
     const std::vector<double > &	posi,
@@ -22,10 +34,9 @@ int format_nlist_fill_a (
 
 
 template<typename FPTYPE> 
-int format_nlist_cpu (
+int format_nlist_i_cpu (
     std::vector<int > &			fmt_nei_idx_a,
     const std::vector<FPTYPE > &	posi,
-    const int &				ntypes,
     const std::vector<int > &		type,
     const int &				i_idx,
     const std::vector<int > &		nei_idx_a, 

source/lib/include/neighbor_list.h

@@ -3,10 +3,59 @@
 #include <algorithm>
 #include <iterator>
 #include <cassert>
+#include <vector>
 
+#include "region.h"
 #include "utilities.h"
 #include "SimulationRegion.h"
 
+// format of the input neighbor list
+struct InputNlist
+{
+  int inum;
+  int * ilist;
+  int * numneigh;
+  int ** firstneigh;
+  InputNlist (
+      int inum_, 
+      int * ilist_,
+      int * numneigh_, 
+      int ** firstneigh_
+      ) 
+      : inum(inum_), ilist(ilist_), numneigh(numneigh_), firstneigh(firstneigh_)
+      {}
+};
+
+void convert_nlist(
+    InputNlist & to_nlist,
+    std::vector<std::vector<int> > & from_nlist
+    );
+
+
+// build neighbor list.
+// outputs
+
+//	nlist, max_list_size
+//	max_list_size is the maximal size of jlist.
+// inputs
+//	c_cpy, nloc, nall, mem_size, rcut, region
+//	mem_size is the size of allocated memory for jlist.
+// returns
+//	0: succssful
+//	1: the memory is not large enough to hold all neighbors.
+//	   i.e. max_list_size > mem_nall
+template <typename FPTYPE>
+int
+build_nlist_cpu(
+    InputNlist & nlist,
+    int * max_list_size,
+    const FPTYPE * c_cpy,
+    const int & nloc, 
+    const int & nall, 
+    const int & mem_size,
+    const float & rcut);
+
+
 // build nlist by an extended grid
 void
 build_nlist (std::vector<std::vector<int > > &	nlist0,

source/lib/include/prod_env_mat.h

@@ -18,7 +18,6 @@ void prod_env_mat_a_cpu(
     const FPTYPE * std, 
     const int nloc, 
     const int nall, 
-    const int ntypes, 
     const float rcut, 
     const float rcut_smth, 
     const std::vector<int> sec);
@@ -39,7 +38,6 @@ void prod_env_mat_r_cpu(
     const FPTYPE * std, 
     const int nloc, 
     const int nall, 
-    const int ntypes, 
     const float rcut, 
     const float rcut_smth, 
     const std::vector<int> sec);

source/lib/include/region.h

@@ -0,0 +1,31 @@
+#pragma once
+
+template<typename FPTYPE>
+struct Region
+{
+  FPTYPE * boxt;
+  FPTYPE * rec_boxt;
+  Region();
+  ~Region();
+};
+
+template<typename FPTYPE>
+void
+init_region_cpu(
+    Region<FPTYPE> & region,
+    const FPTYPE * boxt);
+
+template<typename FPTYPE>
+void
+convert_to_inter_cpu(
+    FPTYPE * ri, 
+    const Region<FPTYPE> & region,
+    const FPTYPE * rp);
+
+template<typename FPTYPE>
+void
+convert_to_phys_cpu(
+    FPTYPE * rp, 
+    const Region<FPTYPE> & region,
+    const FPTYPE * ri);
+

source/lib/src/coord.cc

@@ -0,0 +1,110 @@
+#include "coord.h"
+#include "neighbor_list.h"
+#include "SimulationRegion.h"
+#include <vector>
+
+// normalize coords
+template <typename FPTYPE>
+void
+normalize_coord_cpu(
+    FPTYPE * coord,
+    const int natom,
+    const Region<FPTYPE> & region)
+{
+  for(int ii = 0; ii < natom; ++ii){
+    FPTYPE ri[3];
+    convert_to_inter_cpu(ri, region, coord+3*ii);
+    for(int dd = 0; dd < 3; ++dd){
+      while(ri[dd] >= 1.) ri[dd] -= 1.;
+      while(ri[dd] <  0.) ri[dd] += 1.;
+    }
+    convert_to_phys_cpu(coord+3*ii, region, ri);
+  }
+}
+
+
+template <typename FPTYPE>
+int
+copy_coord_cpu(
+    FPTYPE * out_c,
+    int * out_t,
+    int * mapping,
+    int * nall,
+    const FPTYPE * in_c,
+    const int * in_t,
+    const int & nloc,
+    const int & mem_nall,
+    const float & rcut,
+    const Region<FPTYPE> & region)
+{
+  std::vector<double> coord(nloc * 3);
+  std::vector<int> atype(nloc);
+  std::copy(in_c, in_c+nloc*3, coord.begin());
+  std::copy(in_t, in_t+nloc, atype.begin());
+  SimulationRegion<double> tmpr;
+  double tmp_boxt[9];
+  std::copy(region.boxt, region.boxt+9, tmp_boxt);
+  tmpr.reinitBox(tmp_boxt);
+
+  std::vector<double > out_coord;
+  std::vector<int> out_atype, out_mapping, ncell, ngcell;
+  copy_coord(out_coord, out_atype, out_mapping, ncell, ngcell, coord, atype, rcut, tmpr);
+
+  *nall = out_atype.size();
+  if(*nall > mem_nall){
+    // size of the output arrays is not large enough
+    return 1;
+  }
+  else{
+    std::copy(out_coord.begin(), out_coord.end(), out_c);
+    std::copy(out_atype.begin(), out_atype.end(), out_t);
+    std::copy(out_mapping.begin(), out_mapping.end(), mapping);
+  }
+  return 0;
+}
+
+
+template
+void
+normalize_coord_cpu<double>(
+    double * coord,
+    const int natom,
+    const Region<double> & region);
+
+template
+void
+normalize_coord_cpu<float>(
+    float * coord,
+    const int natom,
+    const Region<float> & region);
+
+template
+int
+copy_coord_cpu<double>(
+    double * out_c,
+    int * out_t,
+    int * mapping,
+    int * nall,
+    const double * in_c,
+    const int * in_t,
+    const int & nloc,
+    const int & mem_nall,
+    const float & rcut,
+    const Region<double> & region);
+
+template
+int
+copy_coord_cpu<float>(
+    float * out_c,
+    int * out_t,
+    int * mapping,
+    int * nall,
+    const float * in_c,
+    const int * in_t,
+    const int & nloc,
+    const int & mem_nall,
+    const float & rcut,
+    const Region<float> & region);
+
+
+

source/lib/src/env_mat.cc

@@ -97,7 +97,6 @@ void env_mat_a_cpu (
     std::vector<FPTYPE > &	        descrpt_a_deriv,
     std::vector<FPTYPE > &	        rij_a,
     const std::vector<FPTYPE > &	posi,
-    const int &				ntypes,
     const std::vector<int > &		type,
     const int &				i_idx,
     const std::vector<int > &		fmt_nlist_a,
@@ -244,7 +243,6 @@ void env_mat_r_cpu (
     std::vector<FPTYPE > &	        descrpt_a_deriv,
     std::vector<FPTYPE > &	        rij_a,
     const std::vector<FPTYPE > &	posi,
-    const int &				ntypes,
     const std::vector<int > &		type,
     const int &				i_idx,
     const std::vector<int > &		fmt_nlist,
@@ -305,7 +303,6 @@ void env_mat_a_cpu<double> (
     std::vector<double > &	        descrpt_a_deriv,
     std::vector<double > &	        rij_a,
     const std::vector<double > &	posi,
-    const int &				ntypes,
     const std::vector<int > &		type,
     const int &				i_idx,
     const std::vector<int > &		fmt_nlist,
@@ -320,7 +317,6 @@ void env_mat_a_cpu<float> (
     std::vector<float > &	        descrpt_a_deriv,
     std::vector<float > &	        rij_a,
     const std::vector<float > &		posi,
-    const int &				ntypes,
     const std::vector<int > &		type,
     const int &				i_idx,
     const std::vector<int > &		fmt_nlist,
@@ -335,7 +331,6 @@ void env_mat_r_cpu<double> (
     std::vector<double > &	        descrpt_r_deriv,
     std::vector<double > &	        rij_r,
     const std::vector<double > &	posi,
-    const int &				ntypes,
     const std::vector<int > &		type,
     const int &				i_idx,
     const std::vector<int > &		fmt_nlist,
@@ -350,7 +345,6 @@ void env_mat_r_cpu<float> (
     std::vector<float > &	        descrpt_r_deriv,
     std::vector<float > &	        rij_r,
     const std::vector<float > &		posi,
-    const int &				ntypes,
     const std::vector<int > &		type,
     const int &				i_idx,
     const std::vector<int > &		fmt_nlist,