support dimension of aparam to be nall

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/descriptor/se_a_mask.py

@@ -301,10 +301,12 @@ def build(
         dstd = self.dstd
 
         """
-        ``aparam'' shape is [nframes, natoms]
+        ``aparam'' shape is [nframes, nall]
         aparam[:, :] is the real/virtual sign for each atom.
         """
         aparam = input_dict["aparam"]
+        with tf.variable_scope("fitting_attr" + suffix, reuse=reuse):
+            t_aparam_nall = tf.constant(True, name="aparam_nall", dtype=tf.bool)
         self.mask = tf.cast(aparam, tf.int32)
         self.mask = tf.reshape(self.mask, [-1, natoms[1]])
 

deepmd/entrypoints/freeze.py

@@ -224,6 +224,7 @@ def _make_node_names(
             "spin_attr/ntypes_spin",
             "fitting_attr/dfparam",
             "fitting_attr/daparam",
+            "fitting_attr/aparam_nall",
         ]
     elif model_type == "dos":
         nodes += [

deepmd/model/pairwise_dprc.py

@@ -125,6 +125,7 @@ def build(
         with tf.variable_scope("fitting_attr" + suffix, reuse=reuse):
             t_dfparam = tf.constant(0, name="dfparam", dtype=tf.int32)
             t_daparam = tf.constant(1, name="daparam", dtype=tf.int32)
+            t_aparam_nall = tf.constant(True, name="aparam_nall", dtype=tf.bool)
         with tf.variable_scope("descrpt_attr" + suffix, reuse=reuse):
             t_ntypes = tf.constant(self.ntypes, name="ntypes", dtype=tf.int32)
             t_rcut = tf.constant(

source/api_c/include/c_api_internal.h

@@ -41,6 +41,7 @@ struct DP_DeepPot {
   std::string exception;
   int dfparam;
   int daparam;
+  bool aparam_nall;
 };
 
 struct DP_DeepPotModelDevi {
@@ -51,6 +52,7 @@ struct DP_DeepPotModelDevi {
   std::string exception;
   int dfparam;
   int daparam;
+  bool aparam_nall;
 };
 
 struct DP_DeepTensor {

source/api_c/include/deepmd.hpp

@@ -1051,7 +1051,9 @@ class DeepPot {
     }
 
     if (aparam.size() != daparam * nloc &&
-        aparam.size() != nframes * daparam * nloc) {
+        aparam.size() != nframes * daparam * nloc &&
+        aparam.size() != daparam * nall &&
+        aparam.size() != nframes * daparam * nall) {
       throw deepmd::hpp::deepmd_exception(
           "the dim of atom parameter provided is not consistent with what the "
           "model uses");
@@ -1460,7 +1462,9 @@ class DeepPotModelDevi {
     }
 
     if (aparam.size() != daparam * nloc &&
-        aparam.size() != nframes * daparam * nloc) {
+        aparam.size() != nframes * daparam * nloc &&
+        aparam.size() != daparam * nall &&
+        aparam.size() != nframes * daparam * nall) {
       throw deepmd::hpp::deepmd_exception(
           "the dim of atom parameter provided is not consistent with what the "
           "model uses");

source/api_c/src/c_api.cc

@@ -29,6 +29,7 @@ DP_DeepPot::DP_DeepPot() {}
 DP_DeepPot::DP_DeepPot(deepmd::DeepPot& dp) : dp(dp) {
   dfparam = dp.dim_fparam();
   daparam = dp.dim_aparam();
+  aparam_nall = dp.is_aparam_all();
 }
 
 DP_DeepPot* DP_NewDeepPot(const char* c_model) {
@@ -65,6 +66,7 @@ DP_DeepPotModelDevi::DP_DeepPotModelDevi(deepmd::DeepPotModelDevi& dp)
     : dp(dp) {
   dfparam = dp.dim_fparam();
   daparam = dp.dim_aparam();
+  aparam_nall = dp.is_aparam_all();
 }
 
 DP_DeepPotModelDevi* DP_NewDeepPotModelDevi(const char** c_models,
@@ -249,7 +251,10 @@ inline void DP_DeepPotComputeNList_variant(DP_DeepPot* dp,
   }
   std::vector<VALUETYPE> aparam_;
   if (aparam) {
-    aparam_.assign(aparam, aparam + nframes * (natoms - nghost) * dp->daparam);
+    aparam_.assign(aparam,
+                   aparam + nframes *
+                                (dp->aparam_all ? natoms : (natoms - nghost)) *
+                                dp->daparam);
   }
   std::vector<double> e;
   std::vector<VALUETYPE> f, v, ae, av;
@@ -433,7 +438,9 @@ void DP_DeepPotModelDeviComputeNList_variant(DP_DeepPotModelDevi* dp,
   }
   std::vector<VALUETYPE> aparam_;
   if (aparam) {
-    aparam_.assign(aparam, aparam + (natoms - nghost) * dp->daparam);
+    aparam_.assign(
+        aparam,
+        aparam + (dp->aparam_all ? natoms : (natoms - nghost)) * dp->daparam);
   }
   // different from DeepPot
   std::vector<double> e;

source/api_cc/include/DeepPot.h

@@ -291,6 +291,16 @@ class DeepPot {
    **/
   void get_type_map(std::string& type_map);
 
+  /**
+   * @brief Get whether the atom dimension of aparam is nall instead of fparam.
+   * @param[out] aparam_nall whether the atom dimension of aparam is nall
+   *instead of fparam.
+   **/
+  bool is_aparam_nall() const {
+    assert(inited);
+    return aparam_nall;
+  };
+
  private:
   tensorflow::Session* session;
   int num_intra_nthreads, num_inter_nthreads;
@@ -309,6 +319,7 @@ class DeepPot {
   int ntypes_spin;
   int dfparam;
   int daparam;
+  bool aparam_nall;
   /**
    * @brief Validate the size of frame and atomic parameters.
    * @param[in] nframes The number of frames.
@@ -572,6 +583,15 @@ class DeepPotModelDevi {
   void compute_relative_std_f(std::vector<VALUETYPE>& std,
                               const std::vector<VALUETYPE>& avg,
                               const VALUETYPE eps);
+  /**
+   * @brief Get whether the atom dimension of aparam is nall instead of fparam.
+   * @param[out] aparam_nall whether the atom dimension of aparam is nall
+   *instead of fparam.
+   **/
+  bool is_aparam_nall() const {
+    assert(inited);
+    return aparam_nall;
+  };
 
  private:
   unsigned numb_models;
@@ -592,6 +612,7 @@ class DeepPotModelDevi {
   int ntypes_spin;
   int dfparam;
   int daparam;
+  bool aparam_nall;
   template <typename VALUETYPE>
   void validate_fparam_aparam(const int& nloc,
                               const std::vector<VALUETYPE>& fparam,

source/api_cc/include/common.h

@@ -227,6 +227,8 @@ int session_get_dtype(tensorflow::Session* session,
  * @param[in] fparam_ Frame parameters.
  * @param[in] aparam_ Atom parameters.
  * @param[in] atommap Atom map.
+ * @param[in] aparam_nall Whether the atomic dimesion of atomic parameters is
+ * nall.
  * @param[in] scope The scope of the tensors.
  */
 template <typename MODELTYPE, typename VALUETYPE>
@@ -240,6 +242,7 @@ int session_input_tensors(
     const std::vector<VALUETYPE>& fparam_,
     const std::vector<VALUETYPE>& aparam_,
     const deepmd::AtomMap& atommap,
+    const bool aparam_nall = false,
     const std::string scope = "");
 
 /**
@@ -254,6 +257,8 @@ int session_input_tensors(
  * @param[in] atommap Atom map.
  * @param[in] nghost Number of ghost atoms.
  * @param[in] ago Update the internal neighbour list if ago is 0.
+ * @param[in] aparam_nall Whether the atomic dimesion of atomic parameters is
+ * nall.
  * @param[in] scope The scope of the tensors.
  */
 template <typename MODELTYPE, typename VALUETYPE>
@@ -269,6 +274,7 @@ int session_input_tensors(
     const deepmd::AtomMap& atommap,
     const int nghost,
     const int ago,
+    const bool aparam_nall = false,
     const std::string scope = "");
 
 /**
@@ -284,6 +290,8 @@ int session_input_tensors(
  * @param[in] atommap Atom map.
  * @param[in] nghost Number of ghost atoms.
  * @param[in] ago Update the internal neighbour list if ago is 0.
+ * @param[in] aparam_nall Whether the atomic dimesion of atomic parameters is
+ * nall.
  * @param[in] scope The scope of the tensors.
  */
 template <typename MODELTYPE, typename VALUETYPE>
@@ -298,6 +306,7 @@ int session_input_tensors_mixed_type(
     const std::vector<VALUETYPE>& fparam_,
     const std::vector<VALUETYPE>& aparam_,
     const deepmd::AtomMap& atommap,
+    const bool aparam_nall = false,
     const std::string scope = "");
 
 /**