feat(pt): implement comprehensive neighbor list support in DeepTensorPT with proper inheritance

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>

Author: Copilot

source/api_cc/include/DeepTensorPT.h

@@ -246,6 +246,8 @@ class DeepTensorPT : public DeepTensorBase {
   int gpu_id;
   bool gpu_enabled;
   NeighborListData nlist_data;
+  // Neighbor list tensors for efficient computation
+  at::Tensor firstneigh_tensor;
 
   /**
    * @brief Translate PyTorch exceptions to the DeePMD-kit exception.

source/api_cc/src/DeepTensorPT.cc

@@ -13,6 +13,23 @@
 
 using namespace deepmd;
 
+torch::Tensor createNlistTensor(const std::vector<std::vector<int>>& data) {
+  size_t total_size = 0;
+  for (const auto& row : data) {
+    total_size += row.size();
+  }
+  std::vector<int> flat_data;
+  flat_data.reserve(total_size);
+  for (const auto& row : data) {
+    flat_data.insert(flat_data.end(), row.begin(), row.end());
+  }
+
+  torch::Tensor flat_tensor = torch::tensor(flat_data, torch::kInt32);
+  int nloc = data.size();
+  int nnei = nloc > 0 ? total_size / nloc : 0;
+  return flat_tensor.view({1, nloc, nnei});
+}
+
 void DeepTensorPT::translate_error(std::function<void()> f) {
   try {
     f();
@@ -434,13 +451,157 @@ void DeepTensorPT::compute_inner(std::vector<VALUETYPE>& global_tensor,
                                  const std::vector<int>& atype,
                                  const std::vector<VALUETYPE>& box,
                                  const int nghost,
-                                 const InputNlist& inlist,
+                                 const InputNlist& lmp_list,
                                  const bool request_deriv) {
-  // Implement neighbor list support following DeepPotPT pattern
-  // For now, use the simple compute_inner approach
-  // TODO: Add full neighbor list optimization for better performance
-  compute_inner(global_tensor, force, virial, atom_tensor, atom_virial, coord,
-                atype, box, request_deriv);
+  torch::Device device(torch::kCUDA, gpu_id);
+  if (!gpu_enabled) {
+    device = torch::Device(torch::kCPU);
+  }
+
+  int natoms = atype.size();
+  auto options = torch::TensorOptions().dtype(torch::kFloat64);
+  torch::ScalarType floatType = torch::kFloat64;
+  if (std::is_same<VALUETYPE, float>::value) {
+    options = torch::TensorOptions().dtype(torch::kFloat32);
+    floatType = torch::kFloat32;
+  }
+  auto int32_option =
+      torch::TensorOptions().device(torch::kCPU).dtype(torch::kInt32);
+  auto int_option =
+      torch::TensorOptions().device(torch::kCPU).dtype(torch::kInt64);
+
+  // Select real atoms following DeepPotPT pattern
+  std::vector<VALUETYPE> dcoord, aparam_;
+  std::vector<int> datype, fwd_map, bkw_map;
+  int nghost_real, nall_real, nloc_real;
+  int nall = natoms;
+  int nframes = 1;
+  std::vector<VALUETYPE> aparam;  // Empty for tensor models
+  select_real_atoms_coord(dcoord, datype, aparam_, nghost_real, fwd_map,
+                          bkw_map, nall_real, nloc_real, coord, atype, aparam,
+                          nghost, ntypes, nframes, 0, nall, false);
+  int nloc = nall_real - nghost_real;
+
+  std::vector<VALUETYPE> coord_wrapped = dcoord;
+  at::Tensor coord_wrapped_Tensor =
+      torch::from_blob(coord_wrapped.data(), {1, nall_real, 3}, options)
+          .to(device);
+  std::vector<std::int64_t> atype_64(datype.begin(), datype.end());
+  at::Tensor atype_Tensor =
+      torch::from_blob(atype_64.data(), {1, nall_real}, int_option).to(device);
+
+  // Process neighbor list following DeepPotPT pattern
+  nlist_data.copy_from_nlist(lmp_list, nall - nghost);
+  nlist_data.shuffle_exclude_empty(fwd_map);
+  nlist_data.padding();
+
+  at::Tensor firstneigh = createNlistTensor(nlist_data.jlist);
+  firstneigh_tensor = firstneigh.to(torch::kInt64).to(device);
+
+  // Prepare box tensor
+  std::vector<VALUETYPE> box_wrapped = box;
+  at::Tensor box_tensor =
+      torch::from_blob(box_wrapped.data(), {1, 9}, options).to(device);
+
+  // Create input vector for model
+  std::vector<torch::jit::IValue> inputs;
+  inputs.push_back(coord_wrapped_Tensor);
+  inputs.push_back(atype_Tensor);
+  inputs.push_back(firstneigh_tensor);
+  inputs.push_back(box_tensor);
+
+  bool do_atom_virial_tensor = request_deriv;
+  inputs.push_back(do_atom_virial_tensor);
+
+  // Forward pass through model
+  c10::Dict<c10::IValue, c10::IValue> outputs =
+      module.forward(inputs).toGenericDict();
+
+  // Process global tensor
+  if (outputs.contains("global_tensor") || outputs.contains("dipole") ||
+      outputs.contains("global_dipole")) {
+    c10::IValue tensor_out;
+    if (outputs.contains("global_tensor")) {
+      tensor_out = outputs.at("global_tensor");
+    } else if (outputs.contains("global_dipole")) {
+      tensor_out = outputs.at("global_dipole");
+    } else {
+      tensor_out = outputs.at("dipole");
+    }
+
+    torch::Tensor flat_tensor = tensor_out.toTensor().view({-1}).to(floatType);
+    torch::Tensor cpu_tensor = flat_tensor.to(torch::kCPU);
+    global_tensor.assign(cpu_tensor.data_ptr<VALUETYPE>(),
+                         cpu_tensor.data_ptr<VALUETYPE>() + cpu_tensor.numel());
+  }
+
+  // Process force if available
+  if (outputs.contains("force") || outputs.contains("extended_force")) {
+    c10::IValue force_out = outputs.contains("extended_force")
+                                ? outputs.at("extended_force")
+                                : outputs.at("force");
+    torch::Tensor flat_force = force_out.toTensor().view({-1}).to(floatType);
+    torch::Tensor cpu_force = flat_force.to(torch::kCPU);
+    std::vector<VALUETYPE> dforce;
+    dforce.assign(cpu_force.data_ptr<VALUETYPE>(),
+                  cpu_force.data_ptr<VALUETYPE>() + cpu_force.numel());
+
+    // Map back to original atom order using select_map
+    force.resize(static_cast<size_t>(nframes) * fwd_map.size() * odim * 3);
+    select_map<VALUETYPE>(force, dforce, bkw_map, odim * 3, nframes,
+                          fwd_map.size(), nall_real);
+  }
+
+  // Process virial if available
+  if (outputs.contains("virial")) {
+    c10::IValue virial_out = outputs.at("virial");
+    torch::Tensor flat_virial = virial_out.toTensor().view({-1}).to(floatType);
+    torch::Tensor cpu_virial = flat_virial.to(torch::kCPU);
+    virial.assign(cpu_virial.data_ptr<VALUETYPE>(),
+                  cpu_virial.data_ptr<VALUETYPE>() + cpu_virial.numel());
+  }
+
+  // Process atom tensor if available
+  if (outputs.contains("atom_tensor")) {
+    c10::IValue atom_tensor_out = outputs.at("atom_tensor");
+    torch::Tensor flat_atom_tensor =
+        atom_tensor_out.toTensor().view({-1}).to(floatType);
+    torch::Tensor cpu_atom_tensor = flat_atom_tensor.to(torch::kCPU);
+    std::vector<VALUETYPE> datom_tensor_tmp;
+    datom_tensor_tmp.assign(
+        cpu_atom_tensor.data_ptr<VALUETYPE>(),
+        cpu_atom_tensor.data_ptr<VALUETYPE>() + cpu_atom_tensor.numel());
+
+    // Map back to original atom order using select_map
+    atom_tensor.resize(static_cast<size_t>(nframes) * fwd_map.size() * odim);
+    select_map<VALUETYPE>(atom_tensor, datom_tensor_tmp, bkw_map, odim, nframes,
+                          fwd_map.size(), nall_real);
+  }
+
+  // Process atomic virial if requested and available
+  if (request_deriv && (outputs.contains("atom_virial") ||
+                        outputs.contains("extended_virial"))) {
+    c10::IValue atom_virial_out = outputs.contains("extended_virial")
+                                      ? outputs.at("extended_virial")
+                                      : outputs.at("atom_virial");
+    torch::Tensor flat_atom_virial =
+        atom_virial_out.toTensor().view({-1}).to(floatType);
+    torch::Tensor cpu_atom_virial = flat_atom_virial.to(torch::kCPU);
+    std::vector<VALUETYPE> datom_virial_tmp;
+    datom_virial_tmp.assign(
+        cpu_atom_virial.data_ptr<VALUETYPE>(),
+        cpu_atom_virial.data_ptr<VALUETYPE>() + cpu_atom_virial.numel());
+
+    // Map back to original atom order using select_map
+    atom_virial.resize(static_cast<size_t>(nframes) * fwd_map.size() * odim *
+                       9);
+    select_map<VALUETYPE>(atom_virial, datom_virial_tmp, bkw_map, odim * 9,
+                          nframes, fwd_map.size(), nall_real);
+  } else if (request_deriv) {
+    // Fill with zeros if atomic virial not available but requested
+    atom_virial.assign(static_cast<size_t>(natoms) * odim * 9,
+                       static_cast<VALUETYPE>(0.0));
+  }
 }
 
 // Public wrapper functions