From ad47c1c2142b9d37ade040e20226bce430776767 Mon Sep 17 00:00:00 2001 From: OutisLi <137472077+OutisLi@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:13:21 +0800 Subject: [PATCH 1/2] fix(pt_expt): make graph trace samples representative --- deepmd/pt_expt/utils/serialization.py | 20 ++- .../test_neighbor_preparation_compat.py | 123 +++++++++++++++ .../infer/test_deep_eval_pt_checkpoint.py | 38 +++++ .../pt_expt/utils/test_graph_trace_samples.py | 140 ++++++++++++++++++ 4 files changed, 318 insertions(+), 3 deletions(-) create mode 100644 source/tests/common/dpmodel/test_neighbor_preparation_compat.py create mode 100644 source/tests/pt_expt/utils/test_graph_trace_samples.py diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 22f5a2a494..1bec3c9b31 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -501,7 +501,7 @@ def build_synthetic_graph_inputs( Single source of the trace-time graph inputs, shared by ``.pt2`` export (:func:`_trace_and_export`) and compiled training (:func:`deepmd.pt_expt.train.training._trace_and_compile_graph`). Builds a - small random system and runs the carry-all + small deterministic system and runs the carry-all :func:`~deepmd.dpmodel.utils.neighbor_graph.build_neighbor_graph` with a padded ``GraphLayout(edge_capacity=e_max)`` trace sample. Export uses the destination-major graph with CSR metadata; compiled training selects the @@ -519,8 +519,10 @@ def build_synthetic_graph_inputs( (before ``fparam``/``aparam``); there is no ``charge_spin`` slot at all (native spin rejects ``add_chg_spin_ebd`` at build). - The system (``rng(42)``, ``box = rcut*3``, centered coords, ``atype[:, i] = - i % ntypes``) is identical for both callers. + The system is identical for both callers and independent of dense neighbor + capacities. For at least four atoms it includes an isolated center and a + connected three-atom cluster. The highest atom type occurs in the cluster + even when the type map is larger than the sample. Parameters ---------- @@ -587,9 +589,21 @@ def build_synthetic_graph_inputs( box_np = (np.eye(3, dtype=np.float64) * box_size).reshape(1, 9) rng = np.random.default_rng(42) coord_np = rng.random((nframes, nloc, 3)) * box_size * 0.5 + box_size * 0.25 + if nloc >= 4: + # Exercise multi-edge aggregation and empty neighborhoods in the same + # trace. Every other atom lies in [.75*rcut, 2.25*rcut]^3, so the + # origin is isolated even under periodic minimum-image distances. + coord_np[:, [0, 2, 3]] = rcut * np.array( + [[1.5, 1.5, 1.5], [1.6, 1.5, 1.5], [1.5, 1.65, 1.5]] + ) + # Keep the isolated center inside n_local (the final atom is a + # synthetic ghost), including the minimal four-atom fixture. + coord_np[:, 1] = 0.0 atype_np = np.zeros((nframes, nloc), dtype=np.int64) for i in range(nloc): atype_np[:, i] = i % ntypes + # Include the upper type-index boundary without growing the node axis. + atype_np[:, 0] = ntypes - 1 coord_t = torch.tensor(coord_np, dtype=dtype, device=device) atype_t = torch.tensor(atype_np, dtype=torch.int64, device=device) diff --git a/source/tests/common/dpmodel/test_neighbor_preparation_compat.py b/source/tests/common/dpmodel/test_neighbor_preparation_compat.py new file mode 100644 index 0000000000..383fdb4e3c --- /dev/null +++ b/source/tests/common/dpmodel/test_neighbor_preparation_compat.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Keep graph orchestration separate from descriptor preparation semantics.""" + +import copy +import unittest +from unittest.mock import ( + call, + patch, +) + +from deepmd.dpmodel.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.dpmodel.model.base_model import ( + BaseModel, +) +from deepmd.utils.argcheck import ( + normalize, +) + + +def normalized_model(descriptor: dict) -> dict: + """Exercise the same strict schema normalization as training input.""" + return normalize( + { + "model": { + "type_map": ["H"], + "descriptor": descriptor, + "fitting_net": {}, + }, + "learning_rate": {"start_lr": 0.001, "decay_steps": 10}, + "loss": {}, + "training": { + "numb_steps": 1, + "training_data": {"systems": ["unused"]}, + }, + }, + check=True, + )["model"] + + +class TestNeighborPreparationCompatibility(unittest.TestCase): + """Mock data measurement only, retaining model and descriptor dispatch.""" + + @patch("deepmd.utils.update_sel.BaseUpdateSel.get_nbor_stat") + def test_omitted_dpa1_sel_keeps_auto_statistics(self, stats) -> None: + stats.return_value = (0.5, [25]) + model = normalized_model({"type": "dpa1", "rcut": 6.0}) + self.assertEqual(model["descriptor"]["sel"], "auto") + before = copy.deepcopy(model) + updated, minimum = BaseModel.update_sel(None, ["H"], model) + stats.assert_called_once_with(None, ["H"], 6.0, mixed_type=True) + self.assertEqual(updated["descriptor"]["sel"], 28) + self.assertEqual(minimum, 0.5) + self.assertEqual(model, before) + + @patch("deepmd.utils.update_sel.BaseUpdateSel.get_nbor_stat") + def test_mixed_hybrid_prepares_and_constructs(self, stats) -> None: + stats.side_effect = [(0.5, [4]), (0.4, [4])] + model = normalized_model( + { + "type": "hybrid", + "list": [ + { + "type": "dpa1", + "rcut": 6.0, + "rcut_smth": 0.5, + "sel": 8, + "neuron": [8, 8], + "axis_neuron": 4, + "attn": 8, + "attn_layer": 0, + "tebd_dim": 4, + }, + { + "type": "se_e2_a", + "rcut": 6.0, + "rcut_smth": 0.5, + "sel": [8], + "neuron": [8, 8], + "axis_neuron": 4, + }, + ], + } + ) + updated, minimum = BaseModel.update_sel(None, ["H"], model) + self.assertEqual( + stats.call_args_list, + [ + call(None, ["H"], 6.0, mixed_type=True), + call(None, ["H"], 6.0, mixed_type=False), + ], + ) + self.assertEqual(minimum, 0.4) + descriptor = BaseDescriptor(**updated["descriptor"], ntypes=1) + self.assertEqual(descriptor.get_sel(), [8]) + + @patch("deepmd.utils.update_sel.BaseUpdateSel.get_nbor_stat") + def test_dpa4c_only_measures_minimum_distance(self, stats) -> None: + stats.return_value = (0.5, [0]) + model = normalized_model({"type": "dpa4c", "rcut": 6.0}) + before = copy.deepcopy(model) + updated, minimum = BaseModel.update_sel(None, ["H"], model) + # This tiny cutoff deliberately does not discover neighbor capacity. + stats.assert_called_once_with(None, None, 1e-6, mixed_type=True) + self.assertEqual(minimum, 0.5) + self.assertEqual(updated, before) + self.assertEqual(model, before) + self.assertNotIn("sel", updated["descriptor"]) + + @patch("deepmd.utils.update_sel.BaseUpdateSel.get_nbor_stat") + def test_dpa4_preserves_explicit_and_auto_capacity(self, stats) -> None: + stats.return_value = (0.5, [25]) + for sel, expected in [(32, 32), ("auto", 28), ("auto:1.5", 40)]: + with self.subTest(sel=sel): + stats.reset_mock() + model = normalized_model({"type": "dpa4", "rcut": 6.0, "sel": sel}) + before = copy.deepcopy(model) + updated, minimum = BaseModel.update_sel(None, ["H"], model) + stats.assert_called_once_with(None, ["H"], 6.0, mixed_type=True) + self.assertEqual(updated["descriptor"]["sel"], expected) + self.assertEqual(minimum, 0.5) + self.assertEqual(model, before) diff --git a/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py b/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py index 00eab95d62..29e3cc9172 100644 --- a/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py +++ b/source/tests/pt_expt/infer/test_deep_eval_pt_checkpoint.py @@ -636,6 +636,44 @@ def test_eval_matches_public_forward(self) -> None: err_msg=name, ) + def test_dense_introspection_survives_graph_checkpoint(self) -> None: + """Graph evaluation must not discard capacities used by introspection.""" + coords = np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]]) + atom_types = np.array([0, 1], dtype=np.int32) + coord_t = torch.tensor(coords, dtype=torch.float64, device=DEVICE) + atype_t = torch.tensor(atom_types[None], dtype=torch.int64, device=DEVICE) + # Construct the two-atom dense list independently of the checkpoint + # loader: each center has one neighbor and the remaining slots pad. + nlist = torch.full( + (1, 2, sum(self.model.get_sel())), -1, dtype=torch.int64, device=DEVICE + ) + nlist[0, :, 0] = torch.tensor([1, 0], device=DEVICE) + mapping = torch.arange(2, device=DEVICE).reshape(1, 2) + atomic_model = self.model.get_dp_atomic_model() + with torch.no_grad(): + descriptor, rot_mat, g2, h2, _ = atomic_model.descriptor( + coord_t, atype_t, nlist, mapping=mapping + ) + fitting = atomic_model.fitting_net + fitting.set_return_middle_output(True) + try: + middle = fitting(descriptor, atype_t, gr=rot_mat, g2=g2, h2=h2)[ + "middle_output" + ] + finally: + fitting.set_return_middle_output(False) + + for layout, path in self.pt_paths.items(): + with self.subTest(layout=layout): + dp = DeepPot(path, auto_batch_size=False) + for actual, expected in ( + (dp.eval_descriptor(coords, None, atom_types), descriptor), + (dp.eval_fitting_last_layer(coords, None, atom_types), middle), + ): + np.testing.assert_allclose( + actual, expected.detach().cpu().numpy(), rtol=1e-10, atol=1e-10 + ) + def test_nlist_controls_fail_without_changing_graph_semantics(self) -> None: output_def = ModelOutputDef(self.model.atomic_output_def()) path = self.pt_paths["plain"] diff --git a/source/tests/pt_expt/utils/test_graph_trace_samples.py b/source/tests/pt_expt/utils/test_graph_trace_samples.py new file mode 100644 index 0000000000..db877a49ae --- /dev/null +++ b/source/tests/pt_expt/utils/test_graph_trace_samples.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Graph trace coverage must not depend on dense neighbor capacities.""" + +from unittest.mock import patch + +import pytest +import torch + +from deepmd.pt_expt.utils.serialization import ( + build_synthetic_graph_inputs, + count_synthetic_graph_edges, +) + + +class _GraphSampleModel: + def __init__(self, ntypes=118): + self.ntypes = ntypes + + def get_rcut(self): + return 4.0 + + def get_type_map(self): + return [f"T{i}" for i in range(self.ntypes)] + + def get_sel(self): + raise AssertionError("Graph samples must not query dense capacities") + + def get_dim_fparam(self): + return 0 + + def get_dim_aparam(self): + return 0 + + def get_dim_chg_spin(self): + return 0 + + +@pytest.mark.parametrize("nloc", [4, 7]) +@pytest.mark.parametrize("ntypes", [1, 118]) +@pytest.mark.parametrize("canonicalize", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_graph_sample_covers_sparse_topology_and_high_types( + nloc, ntypes, canonicalize, dtype +): + model = _GraphSampleModel(ntypes) + kwargs = {"nframes": 2, "nloc": nloc, "dtype": dtype, "device": torch.device("cpu")} + with patch( + "deepmd.pt_expt.utils.serialization.build_neighbor_list", + side_effect=AssertionError("Graph samples must not build dense lists"), + ): + real = count_synthetic_graph_edges(model, **kwargs) + sample = build_synthetic_graph_inputs( + model, real + 5, canonicalize=canonicalize, **kwargs + ) + repeated = build_synthetic_graph_inputs( + model, real + 5, canonicalize=canonicalize, **kwargs + ) + for actual, expected in zip(sample, repeated, strict=True): + if actual is not None: + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + atype, n_node, n_local, edge_index, edge_vec, edge_mask = sample[:6] + assert atype.shape == (2 * nloc,) + assert atype[0] == atype[nloc] == ntypes - 1 + assert (atype < ntypes - 1).any() == (ntypes > 1) + assert n_node.tolist() == [nloc, nloc] + assert n_local.tolist() == [nloc - 1, nloc - 1] + assert edge_mask.sum() == real + assert (~edge_mask).sum() >= 5 + real_edges = edge_index[:, edge_mask] + # An owned center in each frame has no incoming or outgoing real edge. + assert not torch.isin(real_edges, torch.tensor([1, nloc + 1])).any() + assert 1 < n_local.min() + # The highest-type node has two guaranteed neighbors and occurs at + # both ends of directed edges. + assert (real_edges[0] == 0).sum() >= 2 + assert (real_edges[1] == 0).sum() >= 2 + assert torch.isfinite(edge_vec).all() + assert all(value is not None for value in sample[6:10]) == canonicalize + + +def test_graph_export_high_type_sparse_runtime(): + from deepmd.pt_expt.model import ( + get_model, + ) + from deepmd.pt_expt.utils import ( + env, + ) + from deepmd.pt_expt.utils.serialization import ( + _trace_and_export, + ) + + config = { + "type_map": [f"T{i}" for i in range(118)], + "descriptor": { + "type": "se_atten", + "sel": 416, + "rcut": 4.0, + "rcut_smth": 0.5, + "neuron": [3, 6], + "axis_neuron": 2, + "attn_layer": 0, + "seed": 42, + }, + "fitting_net": {"neuron": [8, 8], "seed": 42}, + } + model = get_model(config).to(env.DEVICE).eval() + exported, metadata, _, _ = _trace_and_export( + {"model": model.serialize()}, model_json_override=None, lower_kind="graph" + ) + # Dense compatibility metadata remains available to existing C++ readers. + assert metadata["sel"] == model.get_sel() + assert metadata["nnei"] == sum(model.get_sel()) + loaded = exported.module() + for nframes, nloc in ((1, 1), (1, 4), (3, 8)): + inputs = list( + build_synthetic_graph_inputs( + model, + None, + nframes=nframes, + nloc=nloc, + dtype=torch.float64, + device=env.DEVICE, + ) + ) + inputs[2] = inputs[1].clone() # Runtime may own every atom. + actual = loaded(*inputs) + reference = model.forward_common_lower_graph(*inputs[:10]) + torch.testing.assert_close( + actual["energy"], reference["energy_redu"], rtol=1e-10, atol=1e-10 + ) + for key, ref_key in ( + ("force", "energy_derv_r"), + ("virial", "energy_derv_c_redu"), + ): + torch.testing.assert_close( + actual[key], + reference[ref_key].reshape(actual[key].shape), + rtol=1e-10, + atol=1e-10, + ) From 16409639b843f10491dcea9946e9c7a43bf4c009 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:16:57 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/pt_expt/utils/test_graph_trace_samples.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/source/tests/pt_expt/utils/test_graph_trace_samples.py b/source/tests/pt_expt/utils/test_graph_trace_samples.py index db877a49ae..cc9f7df9aa 100644 --- a/source/tests/pt_expt/utils/test_graph_trace_samples.py +++ b/source/tests/pt_expt/utils/test_graph_trace_samples.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Graph trace coverage must not depend on dense neighbor capacities.""" -from unittest.mock import patch +from unittest.mock import ( + patch, +) import pytest import torch