Update DeepPotPT.cc

Author: iProzd

source/api_cc/src/DeepPotPT.cc

@@ -174,63 +174,37 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
     nlist_data.padding();
     if (do_message_passing) {
       int nswap = lmp_list.nswap;
-
-      std::vector<int> sendnum_new(nswap, 0);
-      std::vector<int> sendlist_new;
-      sendlist_new.reserve(
-          std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap, 0));
+      std::vector<std::vector<int>> sendlist_new;
+      sendlist_new.resize(nswap);
+      // select real atoms in sendlist
       for (int s = 0; s < nswap; ++s) {
         int cnt = 0;
+        sendlist_new[s].reserve(lmp_list.sendnum[s]);
         for (int k = 0; k < lmp_list.sendnum[s]; ++k) {
           const int old_idx = lmp_list.sendlist[s][k];
           int mapped = (old_idx >= 0 && old_idx < (int)fwd_map.size())
                            ? fwd_map[old_idx]
                            : -1;
           if (mapped >= 0) {
-            sendlist_new.push_back(mapped);
+            sendlist_new[s].push_back(mapped);
             ++cnt;
           }
         }
-        sendnum_new[s] = cnt;
-      }
-
-      std::vector<int> recvnum_new(nswap, 0);
-      // need check
-      for (int s = 0; s < nswap; ++s) {
-        recvnum_new[s] = sendnum_new[s];
-      }
-
-      std::vector<int> firstrecv_new(nswap, 0);
-      int acc = 0;
-      for (int s = 0; s < nswap; ++s) {
-        firstrecv_new[s] = acc;
-        acc += recvnum_new[s];
+        std::memcpy(lmp_list.sendlist[s], sendlist_new[s].data(), cnt * sizeof(int));
+        lmp_list.sendnum[s] = cnt;
+        lmp_list.recvnum[s] = cnt;
       }
-
       torch::Tensor sendproc_tensor =
           torch::from_blob(lmp_list.sendproc, {nswap}, int32_option);
       torch::Tensor recvproc_tensor =
           torch::from_blob(lmp_list.recvproc, {nswap}, int32_option);
 
       torch::Tensor firstrecv_tensor =
-          torch::from_blob(firstrecv_new.data(), {nswap}, int32_option).clone();
+            torch::from_blob(lmp_list.firstrecv, {nswap}, int32_option);
       torch::Tensor recvnum_tensor =
-          torch::from_blob(recvnum_new.data(), {nswap}, int32_option).clone();
+            torch::from_blob(lmp_list.recvnum, {nswap}, int32_option);
       torch::Tensor sendnum_tensor =
-          torch::from_blob(sendnum_new.data(), {nswap}, int32_option).clone();
-
-      torch::Tensor sendlist_tensor =
-          torch::from_blob(sendlist_new.data(),
-                           {static_cast<long>(sendlist_new.size())},
-                           int32_option)
-              .clone();
-
-      //      torch::Tensor firstrecv_tensor =
-      //          torch::from_blob(lmp_list.firstrecv, {nswap}, int32_option);
-      //      torch::Tensor recvnum_tensor =
-      //          torch::from_blob(lmp_list.recvnum, {nswap}, int32_option);
-      //      torch::Tensor sendnum_tensor =
-      //          torch::from_blob(lmp_list.sendnum, {nswap}, int32_option);
+            torch::from_blob(lmp_list.sendnum, {nswap}, int32_option);
       torch::Tensor communicator_tensor;
       if (lmp_list.world == 0) {
         communicator_tensor = torch::empty({1}, torch::kInt64);
@@ -240,12 +214,13 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
       }
 
       torch::Tensor nswap_tensor = torch::tensor(nswap, int32_option);
-      //      int total_send =
-      //          std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap,
-      //          0);
-      //      torch::Tensor sendlist_tensor =
-      //          torch::from_blob(lmp_list.sendlist, {total_send},
-      //          int32_option);
+
+      int total_send =
+            std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap,
+            0);
+      torch::Tensor sendlist_tensor =
+            torch::from_blob(lmp_list.sendlist, {total_send},
+            int32_option);
       comm_dict.insert_or_assign("send_list", sendlist_tensor);
       comm_dict.insert_or_assign("send_proc", sendproc_tensor);
       comm_dict.insert_or_assign("recv_proc", recvproc_tensor);