fix(pt_expt): make graph trace samples representative

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)

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)

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

source/tests/pt_expt/utils/test_graph_trace_samples.py

@@ -0,0 +1,142 @@
+# 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,
+            )