From 088e103735059511ce650d7ad719c6368d374f74 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Wed, 15 Jul 2026 02:52:06 +0800 Subject: [PATCH] fix(pt): remap spins after null atom filtering Compact spin vectors with the same backward map used for coordinates and atom types in DeepSpinPT neighbor-list inference. Add strict compact-versus-padded regression coverage for NULL atoms in the middle of both float and double inputs. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/api_cc/src/DeepSpinPT.cc | 12 ++- .../api_cc/tests/test_deeppot_dpa_pt_spin.cc | 97 +++++++++++++++++++ 2 files changed, 108 insertions(+), 1 deletion(-) diff --git a/source/api_cc/src/DeepSpinPT.cc b/source/api_cc/src/DeepSpinPT.cc index 3114746caa..7537ca2588 100644 --- a/source/api_cc/src/DeepSpinPT.cc +++ b/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 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 spin_wrapped(static_cast(nall_real) * 3); + for (int ii = 0; ii < nall_real; ++ii) { + for (int dd = 0; dd < 3; ++dd) { + spin_wrapped[static_cast(ii) * 3 + dd] = + spin[static_cast(bkw_map[ii]) * 3 + dd]; + } + } at::Tensor spin_wrapped_Tensor = torch::from_blob(spin_wrapped.data(), {1, nall_real, 3}, options) .to(device); diff --git a/source/api_cc/tests/test_deeppot_dpa_pt_spin.cc b/source/api_cc/tests/test_deeppot_dpa_pt_spin.cc index 4432ddd02b..a81bd6ac45 100644 --- a/source/api_cc/tests/test_deeppot_dpa_pt_spin.cc +++ b/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& coord = this->coord; + const std::vector& spin = this->spin; + const std::vector& atype = this->atype; + const std::vector& 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& force, + std::vector& force_mag, + std::vector& virial, + const std::vector& input_coord, + const std::vector& input_spin, + const std::vector& input_atype, + std::vector >& nlist_data) { + const int natoms = static_cast(input_atype.size()); + std::vector ilist(natoms), numneigh(natoms); + std::vector 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 > 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 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 padded_coord = coord; + padded_coord.insert(padded_coord.begin() + null_index * 3, + {static_cast(8.1), static_cast(7.2), + static_cast(6.3)}); + std::vector padded_spin = spin; + padded_spin.insert( + padded_spin.begin() + null_index * 3, + {static_cast(0.91), static_cast(-0.73), + static_cast(0.57)}); + std::vector 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 > 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 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::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 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(0)); + EXPECT_EQ(padded_force_mag[null_index * 3 + dd], static_cast(0)); + } +} + TYPED_TEST(TestInferDeepSpinDpaPtNopbc, cpu_lmp_nlist_atomic) { using VALUETYPE = TypeParam; const std::vector& coord = this->coord;