From b951d825b80245669c7e087b1479984a4f88a200 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 01:21:44 +0800 Subject: [PATCH 1/4] fix(dpmodel): mask virtual EnvMat centers Use safe type indices and neutral normalization values before zeroing virtual-center descriptors. Add radial, angular, and strict Array API regression coverage. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/utils/env_mat.py | 25 ++++++++++++-- .../common/dpmodel/array_api/test_env_mat.py | 19 +++++++++++ source/tests/common/dpmodel/test_env_mat.py | 34 +++++++++++++++++++ 3 files changed, 76 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py index 0b0bd18c35..0ea1e81aa3 100644 --- a/deepmd/dpmodel/utils/env_mat.py +++ b/deepmd/dpmodel/utils/env_mat.py @@ -159,10 +159,31 @@ def call( em, diff, sw = self._call(nlist, coord_ext, radial_only) nf, nloc, nnei = nlist.shape atype = xp_take_first_n(atype_ext, 1, nloc) + center_is_real = atype >= 0 + # Virtual atoms use a negative type sentinel. Never pass that sentinel to + # ``take``: NumPy treats -1 as the final real type, while stricter array + # namespaces may reject it. Type zero is only a safe placeholder because + # the corresponding rows are masked out after normalization below. + safe_atype = xp.where(center_is_real, atype, xp.zeros_like(atype)) + center_mask = xp.reshape(center_is_real, (nf, nloc, 1, 1)) if davg is not None: - em -= xp.reshape(xp.take(davg, xp.reshape(atype, (-1,)), axis=0), em.shape) + center_avg = xp.reshape( + xp.take(davg, xp.reshape(safe_atype, (-1,)), axis=0), em.shape + ) + center_avg = xp.where(center_mask, center_avg, xp.zeros_like(center_avg)) + em -= center_avg if dstd is not None: - em /= xp.reshape(xp.take(dstd, xp.reshape(atype, (-1,)), axis=0), em.shape) + center_std = xp.reshape( + xp.take(dstd, xp.reshape(safe_atype, (-1,)), axis=0), em.shape + ) + # A neutral scale avoids hidden divide-by-zero/NaN values in the + # masked branch, which is important for differentiable backends. + center_std = xp.where(center_mask, center_std, xp.ones_like(center_std)) + em /= center_std + # The neighbor-list contract leaves every slot empty for a virtual center, + # but mask the final descriptor explicitly so malformed or externally built + # lists cannot leak a fake center into downstream descriptor computations. + em = xp.where(center_mask, em, xp.zeros_like(em)) return em, diff, sw def _call( diff --git a/source/tests/common/dpmodel/array_api/test_env_mat.py b/source/tests/common/dpmodel/array_api/test_env_mat.py index 607b0515bf..204cf710f7 100644 --- a/source/tests/common/dpmodel/array_api/test_env_mat.py +++ b/source/tests/common/dpmodel/array_api/test_env_mat.py @@ -4,6 +4,7 @@ import array_api_strict as xp from deepmd.dpmodel.utils.env_mat import ( + EnvMat, compute_smooth_weight, ) @@ -23,3 +24,21 @@ def test_compute_smooth_weight(self) -> None: self.assert_namespace_equal(w, d) self.assert_device_equal(w, d) self.assert_dtype_equal(w, d) + + def test_virtual_center_uses_safe_normalization_indices(self) -> None: + """Strict array indexing must never receive the negative type sentinel.""" + coord = xp.asarray([[0.0, 0.0, 0.0]], dtype=xp.float64) + atype = xp.asarray([[-1]], dtype=xp.int64) + nlist = xp.asarray([[[-1]]], dtype=xp.int64) + davg = xp.asarray([[[11.0, 13.0, 17.0, 19.0]]], dtype=xp.float64) + # A zero placeholder scale verifies that masking happens before division; + # otherwise the virtual row can create hidden NaN or infinity values. + dstd = xp.zeros_like(davg) + + env_mat, diff, switch = EnvMat(2.0, 0.5).call(coord, atype, nlist, davg, dstd) + + for output in (env_mat, diff, switch): + self.assertTrue(bool(xp.all(output == xp.zeros_like(output)))) + self.assert_namespace_equal(output, coord) + self.assert_device_equal(output, coord) + self.assert_dtype_equal(output, coord) diff --git a/source/tests/common/dpmodel/test_env_mat.py b/source/tests/common/dpmodel/test_env_mat.py index 39f5003d98..58b95cb797 100644 --- a/source/tests/common/dpmodel/test_env_mat.py +++ b/source/tests/common/dpmodel/test_env_mat.py @@ -12,6 +12,7 @@ ) from .case_single_frame_with_nlist import ( TestCaseSingleFrameWithNlist, + TestCaseSingleFrameWithNlistWithVirtual, ) @@ -38,3 +39,36 @@ def test_self_consistency( np.testing.assert_allclose(mm0, mm1) np.testing.assert_allclose(diff0, diff1) np.testing.assert_allclose(ww0, ww1) + + +class TestEnvMatWithVirtualCenter( + unittest.TestCase, TestCaseSingleFrameWithNlistWithVirtual +): + def setUp(self) -> None: + TestCaseSingleFrameWithNlistWithVirtual.setUp(self) + + def test_normalization_keeps_virtual_centers_zero(self) -> None: + """Virtual centers must not borrow normalization data from a real type.""" + nf, nloc, nnei = self.nlist.shape + virtual_center = self.atype_ext[:, :nloc] < 0 + + for radial_only, width in ((False, 4), (True, 1)): + with self.subTest(radial_only=radial_only): + # Nonzero values make accidental ``-1`` indexing observable: NumPy + # would otherwise select the final real-type row silently. + davg = np.arange(1, self.nt * nnei * width + 1, dtype=np.float64) + davg = davg.reshape(self.nt, nnei, width) + dstd = np.full_like(davg, 2.0) + + env_mat, diff, switch = EnvMat(self.rcut, self.rcut_smth).call( + self.coord_ext, + self.atype_ext, + self.nlist, + davg, + dstd, + radial_only=radial_only, + ) + + np.testing.assert_allclose(env_mat[virtual_center], 0.0) + np.testing.assert_allclose(diff[virtual_center], 0.0) + np.testing.assert_allclose(switch[virtual_center], 0.0) From 6d67a1257ff030529219ef5bd26bb4d8d1da1f62 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 23 Jul 2026 20:11:43 +0800 Subject: [PATCH 2/4] fix(dpmodel): mask all virtual EnvMat outputs Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/utils/env_mat.py | 6 ++-- .../common/dpmodel/array_api/test_env_mat.py | 29 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py index 0ea1e81aa3..ead976cd1c 100644 --- a/deepmd/dpmodel/utils/env_mat.py +++ b/deepmd/dpmodel/utils/env_mat.py @@ -181,9 +181,11 @@ def call( center_std = xp.where(center_mask, center_std, xp.ones_like(center_std)) em /= center_std # The neighbor-list contract leaves every slot empty for a virtual center, - # but mask the final descriptor explicitly so malformed or externally built - # lists cannot leak a fake center into downstream descriptor computations. + # but mask all center-indexed outputs explicitly so malformed or externally + # built lists cannot leak a fake center into downstream descriptor paths. em = xp.where(center_mask, em, xp.zeros_like(em)) + diff = xp.where(center_mask, diff, xp.zeros_like(diff)) + sw = xp.where(center_mask, sw, xp.zeros_like(sw)) return em, diff, sw def _call( diff --git a/source/tests/common/dpmodel/array_api/test_env_mat.py b/source/tests/common/dpmodel/array_api/test_env_mat.py index 204cf710f7..b4680c0489 100644 --- a/source/tests/common/dpmodel/array_api/test_env_mat.py +++ b/source/tests/common/dpmodel/array_api/test_env_mat.py @@ -42,3 +42,32 @@ def test_virtual_center_uses_safe_normalization_indices(self) -> None: self.assert_namespace_equal(output, coord) self.assert_device_equal(output, coord) self.assert_dtype_equal(output, coord) + + def test_mixed_centers_mask_all_virtual_outputs(self) -> None: + """The shared JAX/pt_expt EnvMat keeps real rows and masks bad virtual rows.""" + coord = xp.asarray([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0]], dtype=xp.float64) + atype = xp.asarray([[1, -1]], dtype=xp.int64) + # Deliberately violate the normal virtual-center contract to verify that + # em, diff, and switch all provide the documented defense in depth. + nlist = xp.asarray([[[1], [0]]], dtype=xp.int64) + davg = xp.asarray( + [[[11.0, 13.0, 17.0, 19.0]], [[0.0, 0.0, 0.0, 0.0]]], + dtype=xp.float64, + ) + dstd = xp.asarray( + [[[0.0, 0.0, 0.0, 0.0]], [[1.0, 1.0, 1.0, 1.0]]], + dtype=xp.float64, + ) + + env_mat, diff, switch = EnvMat(2.0, 0.5).call(coord, atype, nlist, davg, dstd) + + for output in (env_mat, diff, switch): + real_output = output[:, :1, ...] + virtual_output = output[:, 1:, ...] + self.assertTrue(bool(xp.any(real_output != xp.zeros_like(real_output)))) + self.assertTrue( + bool(xp.all(virtual_output == xp.zeros_like(virtual_output))) + ) + self.assert_namespace_equal(output, coord) + self.assert_device_equal(output, coord) + self.assert_dtype_equal(output, coord) From 27f2a5a0f3c9f16fcf3f3e83cfd6f202e5a83084 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Mon, 27 Jul 2026 18:17:08 +0800 Subject: [PATCH 3/4] fix(dpmodel): neutralize normalization for virtual EnvMat centers _make_env_mat already zeroes em, diff and sw wherever nlist < 0, so a virtual center -- whose neighbor row is empty by the neighbor-list contract -- only leaves this function nonzero because normalization shifts it: xp.take with the negative sentinel silently selects the last real type's davg/dstd. Gather with a safe index and neutralize the offset and scale instead. Masking em/diff/sw again afterwards makes the descriptor depend on atype_ext, which the compiled pt_expt DPA2 lower turns into wrong forces (TestCompiledVaryingNatoms, 100% of force elements). A scalar no-op multiply in the same place is fine, so it is the atype dependency, not the extra op. --- deepmd/dpmodel/utils/env_mat.py | 15 ++++++++------- .../common/dpmodel/array_api/test_env_mat.py | 19 +++++++++++-------- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py index f97c15eac8..b765a61fec 100644 --- a/deepmd/dpmodel/utils/env_mat.py +++ b/deepmd/dpmodel/utils/env_mat.py @@ -204,9 +204,16 @@ def call( # Virtual atoms use a negative type sentinel. Never pass that sentinel to # ``take``: NumPy treats -1 as the final real type, while stricter array # namespaces may reject it. Type zero is only a safe placeholder because - # the corresponding rows are masked out after normalization below. + # the gathered rows are neutralized below. safe_atype = xp.where(center_is_real, atype, xp.zeros_like(atype)) center_mask = xp.reshape(center_is_real, (nf, nloc, 1, 1)) + # ``_make_env_mat`` already zeroes em, diff and sw wherever ``nlist < 0``, + # so a virtual center -- whose neighbor row is empty by the neighbor-list + # contract -- leaves this function at zero as long as normalization does + # not shift it. Neutralizing the offset and the scale is therefore the + # whole fix; masking em/diff/sw again afterwards would make the + # descriptor depend on ``atype_ext``, which the compiled pt_expt DPA2 + # lower miscompiles into wrong forces. if davg is not None: center_avg = xp.reshape( xp.take(davg, xp.reshape(safe_atype, (-1,)), axis=0), em.shape @@ -221,12 +228,6 @@ def call( # masked branch, which is important for differentiable backends. center_std = xp.where(center_mask, center_std, xp.ones_like(center_std)) em /= center_std - # The neighbor-list contract leaves every slot empty for a virtual center, - # but mask all center-indexed outputs explicitly so malformed or externally - # built lists cannot leak a fake center into downstream descriptor paths. - em = xp.where(center_mask, em, xp.zeros_like(em)) - diff = xp.where(center_mask, diff, xp.zeros_like(diff)) - sw = xp.where(center_mask, sw, xp.zeros_like(sw)) return em, diff, sw def _call( diff --git a/source/tests/common/dpmodel/array_api/test_env_mat.py b/source/tests/common/dpmodel/array_api/test_env_mat.py index b4680c0489..5314f1a1cb 100644 --- a/source/tests/common/dpmodel/array_api/test_env_mat.py +++ b/source/tests/common/dpmodel/array_api/test_env_mat.py @@ -43,19 +43,22 @@ def test_virtual_center_uses_safe_normalization_indices(self) -> None: self.assert_device_equal(output, coord) self.assert_dtype_equal(output, coord) - def test_mixed_centers_mask_all_virtual_outputs(self) -> None: - """The shared JAX/pt_expt EnvMat keeps real rows and masks bad virtual rows.""" + def test_mixed_centers_keep_virtual_rows_at_zero(self) -> None: + """A real center is normalized; a virtual one is left untouched at zero.""" coord = xp.asarray([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0]], dtype=xp.float64) - atype = xp.asarray([[1, -1]], dtype=xp.int64) - # Deliberately violate the normal virtual-center contract to verify that - # em, diff, and switch all provide the documented defense in depth. - nlist = xp.asarray([[[1], [0]]], dtype=xp.int64) + atype = xp.asarray([[0, -1]], dtype=xp.int64) + # The virtual center's neighbor row is empty, per the neighbor-list + # contract, so _make_env_mat leaves its outputs at zero. Only the + # normalization below could shift them off zero. + nlist = xp.asarray([[[1], [-1]]], dtype=xp.int64) + # The last row is what an unguarded ``take`` selects for atype -1, so + # keep it nonzero: borrowing it would shift the virtual row off zero. davg = xp.asarray( - [[[11.0, 13.0, 17.0, 19.0]], [[0.0, 0.0, 0.0, 0.0]]], + [[[0.0, 0.0, 0.0, 0.0]], [[11.0, 13.0, 17.0, 19.0]]], dtype=xp.float64, ) dstd = xp.asarray( - [[[0.0, 0.0, 0.0, 0.0]], [[1.0, 1.0, 1.0, 1.0]]], + [[[1.0, 1.0, 1.0, 1.0]], [[2.0, 2.0, 2.0, 2.0]]], dtype=xp.float64, ) From 8fe00baf8d31d49bd50ea6c50f68c52a492151a3 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 2 Aug 2026 21:59:41 +0800 Subject: [PATCH 4/4] docs(dpmodel): state virtual-center nlist contract in EnvMat.call The zero-output guarantee for virtual centers rests on the neighbor row being empty. Document that the in-tree neighbor-list builder fills a virtual atom's row with -1 by construction. Coding-Agent: opencode opencode-Version: 1.18.9 Model: ustc/deepseek-v4-flash Reasoning-Effort: max --- deepmd/dpmodel/utils/env_mat.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py index b765a61fec..158745cbad 100644 --- a/deepmd/dpmodel/utils/env_mat.py +++ b/deepmd/dpmodel/utils/env_mat.py @@ -172,7 +172,11 @@ def call( Parameters ---------- nlist - The neighbor list. shape: nf x nloc x nnei + The neighbor list. shape: nf x nloc x nnei. Entries equal to ``-1`` + mark empty neighbor slots, and a virtual center (whose atom type is + negative) must have an entire neighbor row of ``-1``; the in-tree + builder (``deepmd.dpmodel.utils.nlist.build_neighbor_list``) fills + the full row of a virtual atom with ``-1`` by construction. coord_ext The extended coordinates of atoms. shape: nf x (nallx3) atype_ext