From 4952fa9a5495fdf1662c467b6f8905bb98a1d7bf Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 04:21:33 +0800 Subject: [PATCH 1/2] fix(dprc): count local atoms in qm-only maps Count local QM atoms independently of the MM-fragment loop so single-fragment pairwise maps report consistent metadata and cannot overrun the forward map. Cover direct map construction and mixed-frame TensorFlow padding. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lib/src/pairwise.cc | 10 +++-- source/lib/tests/test_pairwise.cc | 24 ++++++++++++ source/tests/tf/test_pairwise_dprc.py | 54 +++++++++++++++++++++++++++ 3 files changed, 84 insertions(+), 4 deletions(-) diff --git a/source/lib/src/pairwise.cc b/source/lib/src/pairwise.cc index b4a68b00b7..aa85d124ab 100644 --- a/source/lib/src/pairwise.cc +++ b/source/lib/src/pairwise.cc @@ -90,12 +90,17 @@ void deepmd::dprc_pairwise_map_cpu( } int max_fragment_size = max_fragment_real_size + max_fragment_ghost_size; int map_size = nqm + max_fragment_real_size + max_fragment_ghost_size; + // Count local QM atoms independently of the MM-fragment loop. A frame may + // contain only its QM fragment plus ignored -1 placeholders, in which case + // there are no QMMM map rows to perform this count later. + int nqm_real = static_cast( + std::count_if(fragments[0].begin(), fragments[0].end(), + [nloc](int atom_index) { return atom_index < nloc; })); // (3, 4, 0, 1, 2, 10, 11), // (3, 4, 5, 6, 7, 10, -1), // (3, 4, 8, 9, -1, 10, -1) forward_qmmm_map.resize(static_cast(nfragments - 1) * map_size); std::fill(forward_qmmm_map.begin(), forward_qmmm_map.end(), -1); - int nqm_real = nloc; // init for nfragments = 1 for (int ii = 0; ii < nfragments - 1; ++ii) { // real for (int jj = 0, kk = 0; jj < nqm; ++jj) { @@ -103,9 +108,6 @@ void deepmd::dprc_pairwise_map_cpu( forward_qmmm_map[ii * map_size + kk] = fragments[0][jj]; kk++; } - if (jj == nqm - 1) { - nqm_real = kk; - } } for (int jj = 0, kk = 0; jj < fragments[ii + 1].size(); ++jj) { if (fragments[ii + 1][jj] < nloc) { diff --git a/source/lib/tests/test_pairwise.cc b/source/lib/tests/test_pairwise.cc index 5f5d39d4fa..2e399f6920 100644 --- a/source/lib/tests/test_pairwise.cc +++ b/source/lib/tests/test_pairwise.cc @@ -41,3 +41,27 @@ TEST(TestPairwiseMap, pairwise_map) { EXPECT_EQ(nall_qm, 3); EXPECT_EQ(nall_qmmm, 7); } + +TEST(TestPairwiseMap, single_fragment_with_placeholders) { + // Local atoms 1 and 3 are placeholders, while atoms 4 and 5 are QM ghosts. + std::vector idxs = {0, -1, 0, -1, 0, 0}; + std::vector> fragments; + deepmd::group_atoms_cpu(fragments, idxs); + ASSERT_EQ(fragments.size(), 1); + + std::vector forward_qm_map, backward_qm_map, forward_qmmm_map, + backward_qmmm_map; + int nloc_qm, nloc_qmmm, nall_qm, nall_qmmm; + deepmd::dprc_pairwise_map_cpu( + forward_qm_map, backward_qm_map, forward_qmmm_map, backward_qmmm_map, + nloc_qm, nloc_qmmm, nall_qm, nall_qmmm, fragments, 4, 6); + + ASSERT_THAT(forward_qm_map, testing::ElementsAre(0, 2, 4, 5)); + ASSERT_THAT(backward_qm_map, testing::ElementsAre(0, -1, 1, -1, 2, 3)); + EXPECT_TRUE(forward_qmmm_map.empty()); + EXPECT_TRUE(backward_qmmm_map.empty()); + EXPECT_EQ(nloc_qm, 2); + EXPECT_EQ(nloc_qmmm, 2); + EXPECT_EQ(nall_qm, 4); + EXPECT_EQ(nall_qmmm, 4); +} diff --git a/source/tests/tf/test_pairwise_dprc.py b/source/tests/tf/test_pairwise_dprc.py index 450134a36c..b6d27ed76b 100644 --- a/source/tests/tf/test_pairwise_dprc.py +++ b/source/tests/tf/test_pairwise_dprc.py @@ -100,6 +100,60 @@ def test_op_single_frame(self) -> None: np.testing.assert_array_equal(natoms_qmmm, np.array([5, 7, 5], dtype=int)) np.testing.assert_array_equal(qmmm_frame_idx, np.array([0, 0, 0], dtype=int)) + def test_op_single_fragment_with_placeholders(self) -> None: + """Count a QM-only frame and pad its maps against a larger frame.""" + idxs = np.array( + [ + [0, -1, 0, -1, -1, 0, -1], + [0, 0, 0, 1, 1, -1, 1], + ], + dtype=int, + ) + natoms = np.array([5, 7, 5], dtype=int) + with self.cached_session() as sess: + t_outputs = op_module.dprc_pairwise_idx( + tf.convert_to_tensor(idxs, dtype=tf.int32), + tf.convert_to_tensor(natoms, dtype=tf.int32), + ) + ( + forward_qm_map, + backward_qm_map, + forward_qmmm_map, + backward_qmmm_map, + natoms_qm, + natoms_qmmm, + qmmm_frame_idx, + ) = run_sess(sess, t_outputs) + + # Frame 0 has two local QM atoms and one ghost. The local section is + # padded before its ghost to match frame 1's three local QM atoms. + np.testing.assert_array_equal( + forward_qm_map, + np.array([[0, 2, -1, 5], [0, 1, 2, -1]], dtype=int), + ) + np.testing.assert_array_equal( + backward_qm_map, + np.array( + [ + [0, -1, 1, -1, -1, 3, -1], + [0, 1, 2, -1, -1, -1, -1], + ], + dtype=int, + ), + ) + # Only frame 1 contains an MM fragment, so it contributes the sole + # QMMM row and retains its original frame index. + np.testing.assert_array_equal( + forward_qmmm_map, np.array([[0, 1, 2, 3, 4, 6]], dtype=int) + ) + np.testing.assert_array_equal( + backward_qmmm_map, + np.array([[0, 1, 2, 3, 4, -1, 5]], dtype=int), + ) + np.testing.assert_array_equal(natoms_qm, np.array([3, 4, 3], dtype=int)) + np.testing.assert_array_equal(natoms_qmmm, np.array([5, 6, 5], dtype=int)) + np.testing.assert_array_equal(qmmm_frame_idx, np.array([1], dtype=int)) + class TestConvertForwardMapOP(tf.test.TestCase): """Test convert_forward_map OP.""" From 65f3bf1e4f037d5a6bbd043e23d8ebc74b55da66 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 20:23:16 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/lib/tests/test_pairwise.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/source/lib/tests/test_pairwise.cc b/source/lib/tests/test_pairwise.cc index 2e399f6920..9a88092ea8 100644 --- a/source/lib/tests/test_pairwise.cc +++ b/source/lib/tests/test_pairwise.cc @@ -52,9 +52,9 @@ TEST(TestPairwiseMap, single_fragment_with_placeholders) { std::vector forward_qm_map, backward_qm_map, forward_qmmm_map, backward_qmmm_map; int nloc_qm, nloc_qmmm, nall_qm, nall_qmmm; - deepmd::dprc_pairwise_map_cpu( - forward_qm_map, backward_qm_map, forward_qmmm_map, backward_qmmm_map, - nloc_qm, nloc_qmmm, nall_qm, nall_qmmm, fragments, 4, 6); + deepmd::dprc_pairwise_map_cpu(forward_qm_map, backward_qm_map, + forward_qmmm_map, backward_qmmm_map, nloc_qm, + nloc_qmmm, nall_qm, nall_qmmm, fragments, 4, 6); ASSERT_THAT(forward_qm_map, testing::ElementsAre(0, 2, 4, 5)); ASSERT_THAT(backward_qm_map, testing::ElementsAre(0, -1, 1, -1, 2, 3));