fix(pt): remap spins after null atom filtering

source/api_cc/src/DeepSpinPT.cc

@@ -160,7 +160,17 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener,
   at::Tensor coord_wrapped_Tensor =
       torch::from_blob(coord_wrapped.data(), {1, nall_real, 3}, options)
           .to(device);
-  std::vector<VALUETYPE> spin_wrapped = spin;
+  // ``select_real_atoms_coord`` compacts coordinates and atom types in
+  // ``bkw_map`` order.  Spin is another per-atom side channel, so it must use
+  // the same mapping; taking a prefix of the original spin buffer is only
+  // correct when every filtered NULL atom happens to be at the end.
+  std::vector<VALUETYPE> spin_wrapped(static_cast<size_t>(nall_real) * 3);
+  for (int ii = 0; ii < nall_real; ++ii) {
+    for (int dd = 0; dd < 3; ++dd) {
+      spin_wrapped[static_cast<size_t>(ii) * 3 + dd] =
+          spin[static_cast<size_t>(bkw_map[ii]) * 3 + dd];
+    }
+  }
   at::Tensor spin_wrapped_Tensor =
       torch::from_blob(spin_wrapped.data(), {1, nall_real, 3}, options)
           .to(device);

source/api_cc/tests/test_deeppot_dpa_pt_spin.cc

@@ -314,6 +314,103 @@ TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_lmp_nlist) {
   }
 }
 
+TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_lmp_nlist_null_atom_in_middle) {
+  using VALUETYPE = TypeParam;
+  const std::vector<VALUETYPE>& coord = this->coord;
+  const std::vector<VALUETYPE>& spin = this->spin;
+  const std::vector<int>& atype = this->atype;
+  const std::vector<VALUETYPE>& box = this->box;
+  deepmd::DeepSpin& dp = this->dp;
+
+  // This fixture deliberately uses the TorchScript ``.pth`` artifact, which
+  // dispatches through DeepSpinPT rather than the separate PTExpt ``.pt2``
+  // implementation.
+  ASSERT_EQ(deepmd::get_backend(kModelPath), deepmd::DPBackend::PyTorch);
+
+  auto compute_with_nlist = [&dp, &box](
+                                double& energy, std::vector<VALUETYPE>& force,
+                                std::vector<VALUETYPE>& force_mag,
+                                std::vector<VALUETYPE>& virial,
+                                const std::vector<VALUETYPE>& input_coord,
+                                const std::vector<VALUETYPE>& input_spin,
+                                const std::vector<int>& input_atype,
+                                std::vector<std::vector<int> >& nlist_data) {
+    const int natoms = static_cast<int>(input_atype.size());
+    std::vector<int> ilist(natoms), numneigh(natoms);
+    std::vector<int*> firstneigh(natoms);
+    deepmd::InputNlist inlist(natoms, ilist.data(), numneigh.data(),
+                              firstneigh.data());
+    convert_nlist(inlist, nlist_data);
+    dp.compute(energy, force, force_mag, virial, input_coord, input_spin,
+               input_atype, box, 0, inlist, 0);
+  };
+
+  // First evaluate the compact six-atom representation.  The padded system
+  // below describes the same physical atoms and neighbor graph, so all real
+  // atom predictions must remain unchanged after the NULL atom is filtered.
+  std::vector<std::vector<int> > compact_nlist = {
+      {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5},
+      {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}};
+  double compact_energy;
+  std::vector<VALUETYPE> compact_force, compact_force_mag, compact_virial;
+  compute_with_nlist(compact_energy, compact_force, compact_force_mag,
+                     compact_virial, coord, spin, atype, compact_nlist);
+
+  constexpr int null_index = 3;
+  std::vector<VALUETYPE> padded_coord = coord;
+  padded_coord.insert(padded_coord.begin() + null_index * 3,
+                      {static_cast<VALUETYPE>(8.1), static_cast<VALUETYPE>(7.2),
+                       static_cast<VALUETYPE>(6.3)});
+  std::vector<VALUETYPE> padded_spin = spin;
+  padded_spin.insert(
+      padded_spin.begin() + null_index * 3,
+      {static_cast<VALUETYPE>(0.91), static_cast<VALUETYPE>(-0.73),
+       static_cast<VALUETYPE>(0.57)});
+  std::vector<int> padded_atype = atype;
+  padded_atype.insert(padded_atype.begin() + null_index, -1);
+
+  // The NULL row is empty and no real row references it.  After fwd_map
+  // compaction this becomes exactly ``compact_nlist``; the distinct NULL spin
+  // makes an incorrectly prefix-truncated spin tensor observably different.
+  std::vector<std::vector<int> > padded_nlist = {
+      {1, 2, 4, 5, 6}, {0, 2, 4, 5, 6}, {0, 1, 4, 5, 6}, {},
+      {0, 1, 2, 5, 6}, {0, 1, 2, 4, 6}, {0, 1, 2, 4, 5}};
+  double padded_energy;
+  std::vector<VALUETYPE> padded_force, padded_force_mag, padded_virial;
+  compute_with_nlist(padded_energy, padded_force, padded_force_mag,
+                     padded_virial, padded_coord, padded_spin, padded_atype,
+                     padded_nlist);
+
+  // These calls describe identical real atoms and neighbor graphs, so use a
+  // strict equivalence tolerance instead of this file's loose float reference
+  // tolerance.  The old prefix-based spin mapping differs by less than 0.1 and
+  // would otherwise let the float regression pass despite the wrong inputs.
+  constexpr double equivalence_tolerance =
+      std::is_same<VALUETYPE, double>::value ? 1e-10 : 1e-5;
+  EXPECT_NEAR(padded_energy, compact_energy, equivalence_tolerance);
+  ASSERT_EQ(padded_virial.size(), compact_virial.size());
+  for (size_t ii = 0; ii < compact_virial.size(); ++ii) {
+    EXPECT_NEAR(padded_virial[ii], compact_virial[ii], equivalence_tolerance);
+  }
+
+  ASSERT_EQ(padded_force.size(), padded_atype.size() * 3);
+  ASSERT_EQ(padded_force_mag.size(), padded_atype.size() * 3);
+  const std::vector<int> compact_to_padded = {0, 1, 2, 4, 5, 6};
+  for (size_t ii = 0; ii < compact_to_padded.size(); ++ii) {
+    const int padded_index = compact_to_padded[ii];
+    for (int dd = 0; dd < 3; ++dd) {
+      EXPECT_NEAR(padded_force[padded_index * 3 + dd],
+                  compact_force[ii * 3 + dd], equivalence_tolerance);
+      EXPECT_NEAR(padded_force_mag[padded_index * 3 + dd],
+                  compact_force_mag[ii * 3 + dd], equivalence_tolerance);
+    }
+  }
+  for (int dd = 0; dd < 3; ++dd) {
+    EXPECT_EQ(padded_force[null_index * 3 + dd], static_cast<VALUETYPE>(0));
+    EXPECT_EQ(padded_force_mag[null_index * 3 + dd], static_cast<VALUETYPE>(0));
+  }
+}
+
 TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_lmp_nlist_atomic) {
   using VALUETYPE = TypeParam;
   const std::vector<VALUETYPE>& coord = this->coord;