fix(cc): fix message passing when nloc is 0 (deepmodeling#4021)

Fix deepmodeling#3732. Removes the special codes for `nloc==0` which has been
supported by deepmodeling#4005.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
	- Improved robustness in tensor handling when inputs are empty.
- Enhanced test coverage for parameters in the `test_pair_deepmd_si` and
`test_pair_deepmd_mpi` functions, testing scenarios with and without
balance arguments.
- **Bug Fixes**
- Adjusted control flow in the `compute` method to simplify handling of
zero atoms or local atoms.


<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

source/api_cc/src/DeepPotPT.cc

@@ -38,7 +38,12 @@ torch::Tensor createNlistTensor(const std::vector<std::vector<int>>& data) {
     row_tensors.push_back(row_tensor);
   }
 
-  torch::Tensor tensor = torch::cat(row_tensors, 0).unsqueeze(0);
+  torch::Tensor tensor;
+  if (row_tensors.size() > 0) {
+    tensor = torch::cat(row_tensors, 0).unsqueeze(0);
+  } else {
+    tensor = torch::empty({1, 0, 0}, torch::kInt32);
+  }
   return tensor;
 }
 DeepPotPT::DeepPotPT() : inited(false) {}
@@ -152,25 +157,6 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
                           nghost, ntypes, 1, daparam, nall, aparam_nall);
   int nloc = nall_real - nghost_real;
   int nframes = 1;
-  // TODO: dpa2 model may need a fake communication op to deal with nloc == 0.
-  // this should be fixed after wrapping comm op as a pure c++ implementation.
-  if (nloc == 0) {
-    // no backward map needed
-    ener.resize(nframes);
-    // dforce of size nall * 3
-    force.resize(static_cast<size_t>(nframes) * fwd_map.size() * 3);
-    fill(force.begin(), force.end(), (VALUETYPE)0.0);
-    // dvirial of size 9
-    virial.resize(static_cast<size_t>(nframes) * 9);
-    fill(virial.begin(), virial.end(), (VALUETYPE)0.0);
-    // datom_energy_ of size nall
-    atom_energy.resize(static_cast<size_t>(nframes) * fwd_map.size());
-    fill(atom_energy.begin(), atom_energy.end(), (VALUETYPE)0.0);
-    // datom_virial_ of size nall * 9
-    atom_virial.resize(static_cast<size_t>(nframes) * fwd_map.size() * 9);
-    fill(atom_virial.begin(), atom_virial.end(), (VALUETYPE)0.0);
-    return;
-  }
   std::vector<VALUETYPE> coord_wrapped = dcoord;
   at::Tensor coord_wrapped_Tensor =
       torch::from_blob(coord_wrapped.data(), {1, nall_real, 3}, options)
@@ -345,23 +331,6 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
   }
   auto int_options = torch::TensorOptions().dtype(torch::kInt64);
   int nframes = 1;
-  if (natoms == 0) {
-    // no backward map needed
-    ener.resize(nframes);
-    // dforce of size nall * 3
-    force.resize(static_cast<size_t>(nframes) * natoms * 3);
-    fill(force.begin(), force.end(), (VALUETYPE)0.0);
-    // dvirial of size 9
-    virial.resize(static_cast<size_t>(nframes) * 9);
-    fill(virial.begin(), virial.end(), (VALUETYPE)0.0);
-    // datom_energy_ of size nall
-    atom_energy.resize(static_cast<size_t>(nframes) * natoms);
-    fill(atom_energy.begin(), atom_energy.end(), (VALUETYPE)0.0);
-    // datom_virial_ of size nall * 9
-    atom_virial.resize(static_cast<size_t>(nframes) * natoms * 9);
-    fill(atom_virial.begin(), atom_virial.end(), (VALUETYPE)0.0);
-    return;
-  }
   std::vector<torch::jit::IValue> inputs;
   at::Tensor coord_wrapped_Tensor =
       torch::from_blob(coord_wrapped.data(), {1, natoms, 3}, options)

source/lmp/tests/test_lammps_dpa_pt.py

@@ -680,7 +680,7 @@ def test_pair_deepmd_si(lammps_si):
 )
 @pytest.mark.parametrize(
     ("balance_args",),
-    [(["--balance"],)],
+    [(["--balance"],), ([],)],
 )
 def test_pair_deepmd_mpi(balance_args: list):
     with tempfile.NamedTemporaryFile() as f:

source/lmp/tests/test_lammps_dpa_sel_pt.py

@@ -683,7 +683,7 @@ def test_pair_deepmd_si(lammps_si):
 )
 @pytest.mark.parametrize(
     ("balance_args",),
-    [(["--balance"],)],
+    [(["--balance"],), ([],)],
 )
 def test_pair_deepmd_mpi(balance_args: list):
     with tempfile.NamedTemporaryFile() as f: