support mixed numbers of atoms with mixed_type

deepmd/descriptor/se_atten.py

@@ -414,6 +414,8 @@ def build(
         tf.summary.histogram("nlist", self.nlist)
 
         self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt])
+        # prevent lookup error; the actual atype already used for nlist
+        atype = tf.clip_by_value(atype, 0, self.ntypes - 1)
         self.atype_nloc = tf.reshape(
             tf.slice(atype, [0, 0], [-1, natoms[0]]), [-1]
         )  ## lammps will have error without this

deepmd/fit/ener.py

@@ -559,6 +559,10 @@ def build(
                 aparam = tf.reshape(aparam, [-1, self.numb_aparam * natoms[0]])
 
         atype_nall = tf.reshape(atype, [-1, natoms[1]])
+        atype_filter = tf.cast(atype_nall >= 0, GLOBAL_TF_FLOAT_PRECISION)
+        # prevent embedding_lookup error,
+        # but the filter will be applied anyway
+        atype_nall = tf.clip_by_value(atype_nall, 0, self.ntypes - 1)
         self.atype_nloc = tf.reshape(
             tf.slice(atype_nall, [0, 0], [-1, natoms[0]]), [-1]
         )  ## lammps will make error
@@ -647,12 +651,13 @@ def build(
                 final_layer -= zero_layer
             outs = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[0]])
         # add bias
-        self.atom_ener_before = outs
+        self.atom_ener_before = outs * atype_filter
         self.add_type = tf.reshape(
             tf.nn.embedding_lookup(self.t_bias_atom_e, self.atype_nloc),
             [tf.shape(inputs)[0], natoms[0]],
         )
         outs = outs + self.add_type
+        outs *= atype_filter
         self.atom_ener_after = outs
 
         if self.tot_ener_zero:

source/lib/include/neighbor_list.h

@@ -142,6 +142,15 @@ int build_nlist_gpu(InputNlist& nlist,
                     const int& mem_size,
                     const float& rcut);
 
+/**
+ * @brief Filter the fake atom type.
+ * @details If >=0, set to 0; if <0, set to -1.
+ * @param ftype_out The output filtered atom type.
+ * @param ftype_in The input atom type.
+ * @param nloc The number of atoms.
+ */
+void filter_ftype_gpu_cuda(int* ftype_out, const int* ftype_in, const int nloc);
+
 void use_nei_info_gpu(int* nlist,
                       int* ntype,
                       bool* nmask,
@@ -175,6 +184,14 @@ int build_nlist_gpu_rocm(InputNlist& nlist,
                          const int& nall,
                          const int& mem_size,
                          const float& rcut);
+/**
+ * @brief Filter the fake atom type.
+ * @details If >=0, set to 0; if <0, set to -1.
+ * @param ftype_out The output filtered atom type.
+ * @param ftype_in The input atom type.
+ * @param nloc The number of atoms.
+ */
+void filter_ftype_gpu_rocm(int* ftype_out, const int* ftype_in, const int nloc);
 
 void use_nei_info_gpu_rocm(int* nlist,
                            int* ntype,

source/lib/src/cuda/neighbor_list.cu

@@ -276,4 +276,23 @@ template int build_nlist_gpu<double>(InputNlist &nlist,
                                      const int &nall,
                                      const int &mem_size,
                                      const float &rcut);
+
+__global__ void map_filter_ftype(int *ftype_out,
+                                 const int *ftype_in,
+                                 const int nloc) {
+  int ii = blockIdx.x * blockDim.x + threadIdx.x;
+  if (ii < nloc) {
+    ftype_out[ii] = ftype_in[ii] >= 0 ? 0 : -1;
+  }
+}
+
+void filter_ftype_gpu_cuda(int *ftype_out,
+                           const int *ftype_in,
+                           const int nloc) {
+  int nblock = (nloc + TPB - 1) / TPB;
+  map_filter_ftype<<<nblock, TPB>>>(ftype_out, ftype_in, nloc);
+  DPErrcheck(cudaGetLastError());
+  DPErrcheck(cudaDeviceSynchronize());
+}
+
 }  // namespace deepmd

source/lib/src/cuda/prod_env_mat.cu

@@ -129,6 +129,7 @@ __global__ void format_nlist_fill_a(uint_64* key,
   uint_64* key_in = key + idx * MAX_NBOR_SIZE;
   FPTYPE diff[3];
   const int& j_idx = nei_idx[idy];
+  if (type[j_idx] < 0) return;
   for (int dd = 0; dd < 3; dd++) {
     diff[dd] = coord[j_idx * 3 + dd] - coord[idx * 3 + dd];
   }
@@ -357,6 +358,7 @@ __global__ void compute_env_mat_a(FPTYPE* em,
   // <<<nloc, TPB>>>
   const int_64 bid = blockIdx.x;
   const unsigned int tid = threadIdx.x;
+  if (type[bid] < 0) return;
   if (tid >= nnei) {
     return;
   }

source/lib/src/fmt_nlist.cc

@@ -114,6 +114,7 @@ int format_nlist_i_cpu(std::vector<int> &fmt_nei_idx_a,
     // rcut is float in this function, so float rr is enough
     float diff[3];
     const int &j_idx = nei_idx[kk];
+    if (type[j_idx] < 0) continue;
     for (int dd = 0; dd < 3; ++dd) {
       diff[dd] = (float)posi[j_idx * 3 + dd] - (float)posi[i_idx * 3 + dd];
     }

source/lib/src/prod_env_mat.cc

@@ -82,12 +82,20 @@ void deepmd::prod_env_mat_a_cpu(FPTYPE *em,
     assert(fmt_nlist_a.size() == nnei);
     // record outputs
     for (int jj = 0; jj < nem; ++jj) {
-      em[ii * nem + jj] =
-          (d_em_a[jj] - avg[type[ii] * nem + jj]) / std[type[ii] * nem + jj];
+      if (type[ii] >= 0) {
+        em[ii * nem + jj] =
+            (d_em_a[jj] - avg[type[ii] * nem + jj]) / std[type[ii] * nem + jj];
+      } else {
+        em[ii * nem + jj] = 0;
+      }
     }
     for (int jj = 0; jj < nem * 3; ++jj) {
-      em_deriv[ii * nem * 3 + jj] =
-          d_em_a_deriv[jj] / std[type[ii] * nem + jj / 3];
+      if (type[ii] >= 0) {
+        em_deriv[ii * nem * 3 + jj] =
+            d_em_a_deriv[jj] / std[type[ii] * nem + jj / 3];
+      } else {
+        em_deriv[ii * nem * 3 + jj] = 0;
+      }
     }
     for (int jj = 0; jj < nnei * 3; ++jj) {
       rij[ii * nnei * 3 + jj] = d_rij_a[jj];

source/lib/src/rocm/neighbor_list.hip.cu

@@ -276,4 +276,21 @@ template int build_nlist_gpu_rocm<double>(InputNlist &nlist,
                                           const int &nall,
                                           const int &mem_size,
                                           const float &rcut);
+__global__ void map_filter_ftype(int *ftype_out,
+                                 const int *ftype_in,
+                                 const int nloc) {
+  int ii = blockIdx.x * blockDim.x + threadIdx.x;
+  if (ii < nloc) {
+    ftype_out[ii] = ftype_in[ii] >= 0 ? 0 : -1;
+  }
+}
+
+void filter_ftype_gpu_rocm(int *ftype_out,
+                           const int *ftype_in,
+                           const int nloc) {
+  int nblock = (nloc + TPB - 1) / TPB;
+  map_filter_ftype<<<nblock, TPB>>>(ftype_out, ftype_in, nloc);
+  DPErrcheck(hipGetLastError());
+  DPErrcheck(hipDeviceSynchronize());
+}
 }  // namespace deepmd

source/lib/src/rocm/prod_env_mat.hip.cu

@@ -126,6 +126,7 @@ __global__ void format_nlist_fill_a(uint_64* key,
   uint_64* key_in = key + idx * MAX_NBOR_SIZE;
   FPTYPE diff[3];
   const int& j_idx = nei_idx[idy];
+  if (type[j_idx] < 0) return;
   for (int dd = 0; dd < 3; dd++) {
     diff[dd] = coord[j_idx * 3 + dd] - coord[idx * 3 + dd];
   }
@@ -364,6 +365,7 @@ __global__ void compute_env_mat_a(FPTYPE* em,
   // <<<nloc, TPB>>>
   const int_64 bid = blockIdx.x;
   const unsigned int tid = threadIdx.x;
+  if (type[bid] < 0) return;
   if (tid >= nnei) {
     return;
   }

source/op/prod_env_mat_multi_device.cc

@@ -1176,7 +1176,7 @@ class ProdEnvMatAMixOp : public OpKernel {
         fake_type_shape.AddDim(nall);
         OP_REQUIRES_OK(context, context->allocate_temp(
                                     DT_INT32, fake_type_shape, &fake_type));
-        deepmd::memset_device_memory(fake_type.flat<int>().data(), 0, nall);
+        deepmd::filter_ftype_gpu_cuda(fake_type.flat<int>().data(), type, nall);
         const int* f_type = fake_type.flat<int>().data();
         // prepare coord and nlist
         _prepare_coord_nlist_gpu<FPTYPE>(
@@ -1227,7 +1227,7 @@ class ProdEnvMatAMixOp : public OpKernel {
         fake_type_shape.AddDim(nall);
         OP_REQUIRES_OK(context, context->allocate_temp(
                                     DT_INT32, fake_type_shape, &fake_type));
-        deepmd::memset_device_memory(fake_type.flat<int>().data(), 0, nall);
+        deepmd::filter_ftype_gpu_rocm(fake_type.flat<int>().data(), type, nall);
         const int* f_type = fake_type.flat<int>().data();
         // prepare coord and nlist
         _prepare_coord_nlist_gpu_rocm<FPTYPE>(
@@ -1272,6 +1272,11 @@ class ProdEnvMatAMixOp : public OpKernel {
         std::vector<int> type_cpy;
         int frame_nall = nall;
         std::vector<int> fake_type(nall, 0);
+        for (int ii = 0; ii < nall; ii++) {
+          if (type[ii] < 0) {
+            fake_type[ii] = -1;
+          }
+        }
         const int* f_type = &fake_type[0];
         // prepare coord and nlist
         _prepare_coord_nlist_cpu<FPTYPE>(