fix(dprc): count local atoms in qm-only maps

source/lib/src/pairwise.cc

@@ -90,22 +90,24 @@ 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<int>(
+      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<size_t>(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) {
       if (fragments[0][jj] < nloc) {
         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) {

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<int> idxs = {0, -1, 0, -1, 0, 0};
+  std::vector<std::vector<int>> fragments;
+  deepmd::group_atoms_cpu(fragments, idxs);
+  ASSERT_EQ(fragments.size(), 1);
+
+  std::vector<int> 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);
+}

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."""