Merge pull request deepmodeling#421 from amcadmus/api

deepmd namespace for the C++ interface
iProzd · Mar 16, 2021 · bd04db7 · bd04db7
2 parents fe1b76c + 48962fb
commit bd04db7
Show file tree

Hide file tree

Showing 22 changed files with 348 additions and 254 deletions.
diff --git a/source/api_cc/include/AtomMap.h b/source/api_cc/include/AtomMap.h
@@ -4,6 +4,7 @@
 
 // using namespace std;
 
+namespace deepmd{
 template <typename VALUETYPE>
 class AtomMap 
 {
@@ -25,3 +26,4 @@ class AtomMap
   std::vector<int> fwd_idx_map;
   std::vector<int> atype;
 };
+}
diff --git a/source/api_cc/include/DataModifier.h b/source/api_cc/include/DataModifier.h
@@ -2,6 +2,7 @@
 
 #include "DeepPot.h"
 
+namespace deepmd{
 class DipoleChargeModifier
 {
 public:
@@ -28,10 +29,10 @@ class DipoleChargeModifier
   int numb_types () const {assert(inited); return ntypes;};
   std::vector<int> sel_types () const {assert(inited); return sel_type;};
 private:
-  Session* session;
+  tensorflow::Session* session;
   std::string name_scope, name_prefix;
   int num_intra_nthreads, num_inter_nthreads;
-  GraphDef graph_def;
+  tensorflow::GraphDef graph_def;
   bool inited;
   VALUETYPE rcut;
   VALUETYPE cell_size;
@@ -42,9 +43,10 @@ class DipoleChargeModifier
   template<class VT> void get_vector(std::vector<VT> & vec, const std::string & name) const;
   void run_model (std::vector<VALUETYPE> &		dforce,
 		  std::vector<VALUETYPE> &		dvirial,
-		  Session *			session,
-		  const std::vector<std::pair<std::string, Tensor>> & input_tensors,
+		  tensorflow::Session *			session,
+		  const std::vector<std::pair<std::string, tensorflow::Tensor>> & input_tensors,
 		  const AtomMap<VALUETYPE> &	atommap,
 		  const int			nghost);
 };
+}
 
diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h
@@ -2,8 +2,8 @@
 
 #include "common.h"
 #include "neighbor_list.h"
-typedef double compute_t;
 
+namespace deepmd{
 class DeepPot 
 {
 public:
@@ -61,9 +61,9 @@ class DeepPot
   int dim_aparam () const {assert(inited); return daparam;};
   void get_type_map (std::string & type_map);
 private:
-  Session* session;
+  tensorflow::Session* session;
   int num_intra_nthreads, num_inter_nthreads;
-  GraphDef graph_def;
+  tensorflow::GraphDef graph_def;
   bool inited;
   template<class VT> VT get_scalar(const std::string & name) const;
   // VALUETYPE get_rcut () const;
@@ -92,7 +92,6 @@ class DeepPot
   // copy neighbor list info from host
   bool init_nbor;
   std::vector<int> sec_a;
-  compute_t *array_double;
   NeighborListData nlist_data;
   InputNlist nlist;
   AtomMap<VALUETYPE> atommap;
@@ -156,9 +155,9 @@ class DeepPotModelDevi
 		      const VALUETYPE eps);
 private:
   unsigned numb_models;
-  std::vector<Session*> sessions;
+  std::vector<tensorflow::Session*> sessions;
   int num_intra_nthreads, num_inter_nthreads;
-  std::vector<GraphDef> graph_defs;
+  std::vector<tensorflow::GraphDef> graph_defs;
   bool inited;
   template<class VT> VT get_scalar(const std::string name) const;
   // VALUETYPE get_rcut () const;
@@ -176,15 +175,15 @@ class DeepPotModelDevi
 
   // copy neighbor list info from host
   bool init_nbor;
-  compute_t *array_double;
   std::vector<std::vector<int> > sec;
-  AtomMap<VALUETYPE> atommap;
+  deepmd::AtomMap<VALUETYPE> atommap;
   NeighborListData nlist_data;
   InputNlist nlist;
 
   // function used for nborlist copy
   std::vector<std::vector<int> > get_sel() const;
-  void cum_sum(const std::vector<std::vector<int32> > n_sel);
+  void cum_sum(const std::vector<std::vector<tensorflow::int32> > n_sel);
 };
+}
 
 
diff --git a/source/api_cc/include/DeepTensor.h b/source/api_cc/include/DeepTensor.h
@@ -1,7 +1,9 @@
 #pragma once
 
-#include "DeepPot.h"
+#include "common.h"
+#include "neighbor_list.h"
 
+namespace deepmd{
 class DeepTensor
 {
 public:
@@ -29,10 +31,10 @@ class DeepTensor
   int output_dim () const {assert(inited); return odim;};
   const std::vector<int> & sel_types () const {assert(inited); return sel_type;};
 private:
-  Session* session;
+  tensorflow::Session* session;
   std::string name_scope;
   int num_intra_nthreads, num_inter_nthreads;
-  GraphDef graph_def;
+  tensorflow::GraphDef graph_def;
   bool inited;
   VALUETYPE rcut;
   VALUETYPE cell_size;
@@ -44,8 +46,8 @@ class DeepTensor
   template<class VT> VT get_scalar(const std::string & name) const;
   template<class VT> void get_vector (std::vector<VT> & vec, const std::string & name) const;
   void run_model (std::vector<VALUETYPE> &		d_tensor_,
-		  Session *				session, 
-		  const std::vector<std::pair<std::string, Tensor>> & input_tensors,
+		  tensorflow::Session *			session, 
+		  const std::vector<std::pair<std::string, tensorflow::Tensor>> & input_tensors,
 		  const AtomMap<VALUETYPE> &		atommap, 
 		  const int				nghost = 0);
   void compute_inner (std::vector<VALUETYPE> &		value,
@@ -59,4 +61,5 @@ class DeepTensor
 		      const int				nghost,
 		      const InputNlist&			inlist);
 };
+}
 
diff --git a/source/api_cc/include/common.h b/source/api_cc/include/common.h
@@ -13,7 +13,9 @@
 #include <tensorflow/core/graph/default_device.h>
 #include <tensorflow/core/graph/graph_def_builder.h>
 
-using namespace tensorflow;
+
+namespace deepmd{
+
 #if TF_MAJOR_VERSION >= 2 && TF_MINOR_VERSION >= 2
 typedef tensorflow::tstring STRINGTYPE;
 #else
@@ -37,7 +39,7 @@ struct NeighborListData
 public:
   void copy_from_nlist(const InputNlist & inlist);
   void shuffle(const std::vector<int> & fwd_map);
-  void shuffle(const AtomMap<VALUETYPE> & map);
+  void shuffle(const deepmd::AtomMap<VALUETYPE> & map);
   void shuffle_exclude_empty(const std::vector<int> & fwd_map);
   void make_inlist(InputNlist & inlist);
 };
@@ -76,136 +78,52 @@ get_env_nthreads(int & num_intra_nthreads,
 		 int & num_inter_nthreads);
 
 void
-checkStatus(const tensorflow::Status& status);
+check_status(
+    const tensorflow::Status& status);
 
-std::string name_prefix(const std::string & name_scope);
+std::string 
+name_prefix(
+    const std::string & name_scope);
 
 template<typename VT>
 VT
-session_get_scalar(Session* session, const std::string name, const std::string scope = "");
+session_get_scalar(
+    tensorflow::Session* session, 
+    const std::string name, 
+    const std::string scope = "");
 
 template<typename VT>
 void
-session_get_vector(std::vector<VT> & o_vec, Session* session, const std::string name_, const std::string scope = "");
+session_get_vector(
+    std::vector<VT> & o_vec, 
+    tensorflow::Session* session, 
+    const std::string name_, 
+    const std::string scope = "");
 
 int
-session_input_tensors (std::vector<std::pair<std::string, Tensor>> & input_tensors,
+session_input_tensors (std::vector<std::pair<std::string, tensorflow::Tensor>> & input_tensors,
 		       const std::vector<VALUETYPE> &	dcoord_,
 		       const int &			ntypes,
 		       const std::vector<int> &		datype_,
 		       const std::vector<VALUETYPE> &	dbox, 
 		       const VALUETYPE &		cell_size,
 		       const std::vector<VALUETYPE> &	fparam_,
 		       const std::vector<VALUETYPE> &	aparam_,
-		       const AtomMap<VALUETYPE>&	atommap,
+		       const deepmd::AtomMap<VALUETYPE>&atommap,
 		       const std::string		scope = "");
 
 int
-session_input_tensors (std::vector<std::pair<std::string, Tensor>> & input_tensors,
+session_input_tensors (std::vector<std::pair<std::string, tensorflow::Tensor>> & input_tensors,
 		       const std::vector<VALUETYPE> &	dcoord_,
 		       const int &			ntypes,
 		       const std::vector<int> &		datype_,
 		       const std::vector<VALUETYPE> &	dbox,		    
 		       InputNlist &		dlist, 
 		       const std::vector<VALUETYPE> &	fparam_,
 		       const std::vector<VALUETYPE> &	aparam_,
-		       const AtomMap<VALUETYPE>&	atommap,
+		       const deepmd::AtomMap<VALUETYPE>&atommap,
 		       const int			nghost,
 		       const int			ago,
 		       const std::string		scope = "");
-
-// int
-// session_input_tensors (std::vector<std::pair<std::string, Tensor>> & input_tensors,
-// 		       const std::vector<VALUETYPE> &	dcoord_,
-// 		       const int &			ntypes,
-// 		       const std::vector<int> &		datype_,
-// 		       const std::vector<VALUETYPE> &	dbox,		    
-// 		       InputNlist &		dlist, 
-// 		       const std::vector<VALUETYPE> &	fparam_,
-// 		       const std::vector<VALUETYPE> &	aparam_,
-// 		       const AtomMap<VALUETYPE>&	atommap,
-// 		       const int			nghost,
-// 		       const std::string			scope = "");
-
-// int 
-// session_input_tensors (std::vector<std::pair<std::string, Tensor>>& input_tensors,
-// 		       const std::vector<VALUETYPE>          & dcoord_,
-// 		       const int                        & ntypes,
-// 		       const std::vector<int>                & atype_,
-// 		       const std::vector<VALUETYPE>          & dbox,
-// 		       const int                        * ilist, 
-// 		       const int                        * jrange,
-// 		       const int                        * jlist,
-// 		       const std::vector<VALUETYPE>		& fparam_,
-// 		       const std::vector<VALUETYPE>	        & aparam_,
-// 		       const AtomMap<VALUETYPE>      & atommap,
-// 		       const int			& nghost);
-
-
-template<typename VT>
-VT
-session_get_scalar(Session* session, const std::string name_, const std::string scope) 
-{
-  std::string name = name_;
-  if (scope != "") {
-    name = scope + "/" + name;
-  }
-  std::vector<Tensor> output_tensors;
-  checkStatus (session->Run(std::vector<std::pair<std::string, Tensor>> ({}), 
-			    {name.c_str()}, 
-			    {}, 
-			    &output_tensors));
-  Tensor output_rc = output_tensors[0];
-  auto orc = output_rc.flat <VT> ();
-  return orc(0);
-}
-
-template<typename VT>
-void
-session_get_vector(std::vector<VT> & o_vec, Session* session, const std::string name_, const std::string scope) 
-{
-  std::string name = name_;
-  if (scope != "") {
-    name = scope + "/" + name;
-  }
-  std::vector<Tensor> output_tensors;
-  checkStatus (session->Run(std::vector<std::pair<std::string, Tensor>> ({}), 
-			    {name.c_str()}, 
-			    {}, 
-			    &output_tensors));
-  Tensor output_rc = output_tensors[0];
-  assert(1 == output_rc.shape().dims());
-  int dof = output_rc.shape().dim_size(0);
-  o_vec.resize(dof);
-  auto orc = output_rc.flat <VT> ();
-  for (int ii = 0; ii < dof; ++ii){
-    o_vec[ii] = orc(ii);
-  }  
-}
-
-
-template<typename VT>
-void 
-select_map(std::vector<VT> & out,
-	   const std::vector<VT > & in,
-	   const std::vector<int > & idx_map, 
-	   const int & stride)
-{
-#ifdef DEBUG
-  assert(in.size() / stride * stride == in.size()), "in size should be multiples of stride"
-#endif
-  for (int ii = 0; ii < in.size() / stride; ++ii){
-#ifdef DEBUG
-    assert(ii < idx_map.size()), "idx goes over the idx map size";
-    assert(idx_map[ii] < out.size()), "mappped idx goes over the out size";
-#endif
-    if (idx_map[ii] >= 0) {
-      int to_ii = idx_map[ii];
-      for (int dd = 0; dd < stride; ++dd){
-	out[to_ii * stride + dd] = in[ii * stride + dd];
-      }
-    }
-  }
 }
 
-
diff --git a/source/api_cc/src/AtomMap.cc b/source/api_cc/src/AtomMap.cc
@@ -3,6 +3,8 @@
 #include <algorithm>
 #include <cassert>
 
+using namespace deepmd;
+
 template <typename VALUETYPE>
 AtomMap<VALUETYPE>::
 AtomMap() {}

diff --git a/source/api_cc/src/DataModifier.cc b/source/api_cc/src/DataModifier.cc
@@ -1,5 +1,8 @@
 #include "DataModifier.h"
 
+using namespace deepmd;
+using namespace tensorflow;
+
 DipoleChargeModifier::
 DipoleChargeModifier()
     : inited (false)
@@ -30,9 +33,9 @@ init (const std::string & model,
   get_env_nthreads(num_intra_nthreads, num_inter_nthreads);
   options.config.set_inter_op_parallelism_threads(num_inter_nthreads);
   options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
-  checkStatus(NewSession(options, &session));
-  checkStatus(ReadBinaryProto(Env::Default(), model, &graph_def));
-  checkStatus(session->Create(graph_def));  
+  deepmd::check_status(NewSession(options, &session));
+  deepmd::check_status(ReadBinaryProto(Env::Default(), model, &graph_def));
+  deepmd::check_status(session->Create(graph_def));  
   // int nnodes = graph_def.node_size();
   // for (int ii = 0; ii < nnodes; ++ii){
   //   cout << ii << " \t " << graph_def.node(ii).name() << endl;
@@ -80,7 +83,7 @@ run_model (std::vector<VALUETYPE> &		dforce,
   }
 
   std::vector<Tensor> output_tensors;
-  checkStatus (session->Run(input_tensors, 
+  deepmd::check_status (session->Run(input_tensors, 
 			    {"o_dm_force", "o_dm_virial", "o_dm_av"},
 			    {}, 
 			    &output_tensors));