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/src/DeepTensorPT.cc

@@ -5,6 +5,7 @@
 #include <torch/csrc/jit/runtime/jit_exception.h>
 
 #include <cstdint>
+#include <numeric>  // for std::iota
 #include <sstream>
 
 #include "common.h"
@@ -206,239 +207,19 @@ void DeepTensorPT::compute_inner(std::vector<VALUETYPE>& global_tensor,
                                  const std::vector<int>& atype,
                                  const std::vector<VALUETYPE>& box,
                                  const bool 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);
-  if (std::is_same<VALUETYPE, float>::value) {
-    options = torch::TensorOptions().dtype(torch::kFloat32);
-  }
-  auto int_option =
-      torch::TensorOptions().device(torch::kCPU).dtype(torch::kInt64);
-
-  // Convert inputs to tensors
-  std::vector<VALUETYPE> coord_wrapped = coord;
-  at::Tensor coord_tensor =
-      torch::from_blob(coord_wrapped.data(), {1, natoms, 3}, options)
-          .to(device);
-
-  std::vector<std::int64_t> atype_64(atype.begin(), atype.end());
-  at::Tensor atype_tensor =
-      torch::from_blob(atype_64.data(), {1, natoms}, int_option).to(device);
-
-  std::vector<VALUETYPE> box_wrapped = box;
-  at::Tensor box_tensor =
-      torch::from_blob(box_wrapped.data(), {1, 9}, options).to(device);
-
-  // Create input vector
-  std::vector<torch::jit::IValue> inputs;
-  inputs.push_back(coord_tensor);
-  inputs.push_back(atype_tensor);
-  inputs.push_back(box_tensor);
-
-  // Forward pass through model
-  torch::jit::IValue result;
-  if (request_deriv) {
-    inputs.push_back(torch::tensor(true));  // do_atomic_virial
-    result = module.forward(inputs);
-  } else {
-    result = module.forward(inputs);
-  }
-
-  auto result_dict = result.toGenericDict();
-
-  // Extract results - try common key names
-  torch::Tensor global_tensor_tensor, atom_tensor_tensor;
-
-  // Try different possible keys for global tensor
-  if (result_dict.contains("global_tensor")) {
-    global_tensor_tensor = result_dict.at("global_tensor").toTensor().cpu();
-  } else if (result_dict.contains("tensor")) {
-    global_tensor_tensor = result_dict.at("tensor").toTensor().cpu();
-  } else if (result_dict.contains("global_dipole")) {
-    global_tensor_tensor = result_dict.at("global_dipole").toTensor().cpu();
-  } else if (result_dict.contains("dipole")) {
-    // For models that only output atomic tensor, sum to get global
-    auto dipole_tensor = result_dict.at("dipole").toTensor().cpu();
-    global_tensor_tensor =
-        torch::sum(dipole_tensor, 1, true);  // Sum over atoms, keep dims
-  } else {
-    throw deepmd::deepmd_exception(
-        "PyTorch tensor model output missing global tensor (expected "
-        "'global_tensor', 'tensor', 'global_dipole', or 'dipole' key)");
-  }
-
-  // Try different possible keys for atomic tensor
-  if (result_dict.contains("atomic_tensor")) {
-    atom_tensor_tensor = result_dict.at("atomic_tensor").toTensor().cpu();
-  } else if (result_dict.contains("atom_tensor")) {
-    atom_tensor_tensor = result_dict.at("atom_tensor").toTensor().cpu();
-  } else if (result_dict.contains("dipole")) {
-    atom_tensor_tensor = result_dict.at("dipole").toTensor().cpu();
-  } else {
-    throw deepmd::deepmd_exception(
-        "PyTorch tensor model output missing atomic tensor (expected "
-        "'atomic_tensor', 'atom_tensor', or 'dipole' key)");
-  }
-
-  // Determine task dimension if not already known
-  if (odim == -1) {
-    if (global_tensor_tensor.dim() >= 2) {
-      odim = global_tensor_tensor.size(-1);
-    } else if (atom_tensor_tensor.dim() >= 3) {
-      odim = atom_tensor_tensor.size(-1);
-    } else {
-      throw deepmd::deepmd_exception(
-          "Unable to determine task dimension from model output");
-    }
-  }
-
-  // Copy global tensor - convert to desired type
-  global_tensor.resize(odim);
-  torch::Tensor global_tensor_converted;
-  if (std::is_same<VALUETYPE, float>::value) {
-    global_tensor_converted = global_tensor_tensor.to(torch::kFloat32);
-    auto global_tensor_acc = global_tensor_converted.accessor<float, 2>();
-    for (int i = 0; i < odim; ++i) {
-      global_tensor[i] = global_tensor_acc[0][i];
-    }
-  } else {
-    global_tensor_converted = global_tensor_tensor.to(torch::kFloat64);
-    auto global_tensor_acc = global_tensor_converted.accessor<double, 2>();
-    for (int i = 0; i < odim; ++i) {
-      global_tensor[i] = global_tensor_acc[0][i];
-    }
-  }
-
-  // Copy atom tensor - convert to desired type
-  atom_tensor.resize(static_cast<size_t>(natoms) * static_cast<size_t>(odim));
-  torch::Tensor atom_tensor_converted;
-  if (std::is_same<VALUETYPE, float>::value) {
-    atom_tensor_converted = atom_tensor_tensor.to(torch::kFloat32);
-    auto atom_tensor_acc = atom_tensor_converted.accessor<float, 3>();
-    for (int i = 0; i < natoms; ++i) {
-      for (int j = 0; j < odim; ++j) {
-        atom_tensor[i * odim + j] = atom_tensor_acc[0][i][j];
-      }
-    }
-  } else {
-    atom_tensor_converted = atom_tensor_tensor.to(torch::kFloat64);
-    auto atom_tensor_acc = atom_tensor_converted.accessor<double, 3>();
-    for (int i = 0; i < natoms; ++i) {
-      for (int j = 0; j < odim; ++j) {
-        atom_tensor[i * odim + j] = atom_tensor_acc[0][i][j];
-      }
-    }
-  }
-
-  if (request_deriv) {
-    // Try to get derivative tensors with error handling
-    torch::Tensor force_tensor, virial_tensor, atom_virial_tensor;
-
-    if (result_dict.contains("force")) {
-      force_tensor = result_dict.at("force").toTensor().cpu();
-    } else {
-      throw deepmd::deepmd_exception(
-          "PyTorch tensor model output missing force tensor when derivatives "
-          "requested");
-    }
-
-    if (result_dict.contains("virial")) {
-      virial_tensor = result_dict.at("virial").toTensor().cpu();
-    } else {
-      throw deepmd::deepmd_exception(
-          "PyTorch tensor model output missing virial tensor when derivatives "
-          "requested");
-    }
-
-    if (result_dict.contains("atomic_virial")) {
-      atom_virial_tensor = result_dict.at("atomic_virial").toTensor().cpu();
-    } else if (result_dict.contains("atom_virial")) {
-      atom_virial_tensor = result_dict.at("atom_virial").toTensor().cpu();
-    } else {
-      // Fill with zeros when atomic virial is not available
-      // This may happen with some models that don't compute atomic virial
-      atom_virial_tensor =
-          torch::zeros({1, odim, natoms, 9}, virial_tensor.options());
-    }
-
-    // Copy force - convert to desired type
-    force.resize(static_cast<size_t>(natoms) * 3 * static_cast<size_t>(odim));
-    torch::Tensor force_converted;
-    if (std::is_same<VALUETYPE, float>::value) {
-      force_converted = force_tensor.to(torch::kFloat32);
-      auto force_acc = force_converted.accessor<float, 4>();
-      for (int d = 0; d < odim; ++d) {
-        for (int i = 0; i < natoms; ++i) {
-          for (int j = 0; j < 3; ++j) {
-            force[d * natoms * 3 + i * 3 + j] = force_acc[0][d][i][j];
-          }
-        }
-      }
-    } else {
-      force_converted = force_tensor.to(torch::kFloat64);
-      auto force_acc = force_converted.accessor<double, 4>();
-      for (int d = 0; d < odim; ++d) {
-        for (int i = 0; i < natoms; ++i) {
-          for (int j = 0; j < 3; ++j) {
-            force[d * natoms * 3 + i * 3 + j] = force_acc[0][d][i][j];
-          }
-        }
-      }
-    }
-
-    // Copy virial - convert to desired type
-    virial.resize(odim * 9);
-    torch::Tensor virial_converted;
-    if (std::is_same<VALUETYPE, float>::value) {
-      virial_converted = virial_tensor.to(torch::kFloat32);
-      auto virial_acc = virial_converted.accessor<float, 3>();
-      for (int d = 0; d < odim; ++d) {
-        for (int i = 0; i < 9; ++i) {
-          virial[d * 9 + i] = virial_acc[0][d][i];
-        }
-      }
-    } else {
-      virial_converted = virial_tensor.to(torch::kFloat64);
-      auto virial_acc = virial_converted.accessor<double, 3>();
-      for (int d = 0; d < odim; ++d) {
-        for (int i = 0; i < 9; ++i) {
-          virial[d * 9 + i] = virial_acc[0][d][i];
-        }
-      }
-    }
-
-    // Copy atom virial - convert to desired type
-    atom_virial.resize(static_cast<size_t>(natoms) * 9 *
-                       static_cast<size_t>(odim));
-    torch::Tensor atom_virial_converted;
-    if (std::is_same<VALUETYPE, float>::value) {
-      atom_virial_converted = atom_virial_tensor.to(torch::kFloat32);
-      auto atom_virial_acc = atom_virial_converted.accessor<float, 4>();
-      for (int d = 0; d < odim; ++d) {
-        for (int i = 0; i < natoms; ++i) {
-          for (int j = 0; j < 9; ++j) {
-            atom_virial[d * natoms * 9 + i * 9 + j] =
-                atom_virial_acc[0][d][i][j];
-          }
-        }
-      }
-    } else {
-      atom_virial_converted = atom_virial_tensor.to(torch::kFloat64);
-      auto atom_virial_acc = atom_virial_converted.accessor<double, 4>();
-      for (int d = 0; d < odim; ++d) {
-        for (int i = 0; i < natoms; ++i) {
-          for (int j = 0; j < 9; ++j) {
-            atom_virial[d * natoms * 9 + i * 9 + j] =
-                atom_virial_acc[0][d][i][j];
-          }
-        }
-      }
-    }
-  }
+  // This is the simpler version without neighbor list optimization
+  // Use a dummy neighbor list and call the full version
+  deepmd::InputNlist dummy_nlist;
+  // Initialize dummy neighbor list with empty data
+  dummy_nlist.inum = atype.size();
+  dummy_nlist.ilist.resize(dummy_nlist.inum);
+  std::iota(dummy_nlist.ilist.begin(), dummy_nlist.ilist.end(), 0);
+  dummy_nlist.numneigh.resize(dummy_nlist.inum, 0);
+  dummy_nlist.firstneigh.resize(dummy_nlist.inum);
+
+  // Call the neighbor list version with nghost=0 and empty neighbor list
+  compute_inner(global_tensor, force, virial, atom_tensor, atom_virial, coord,
+                atype, box, 0, dummy_nlist, request_deriv);
 }
 
 template <typename VALUETYPE>