From 0b2c374036756e54ad71941db5afbb6fe8e473cf Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Sun, 21 Jun 2026 09:38:12 +0000
Subject: [PATCH 1/8] fix: cover mixed types in stat sampling

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
---
 deepmd/pt/utils/stat.py                | 118 ++++++++++++++++++++-----
 deepmd/utils/model_stat.py             |  81 +++++++++++++++++
 source/tests/common/test_model_stat.py |  98 ++++++++++++++++++++
 source/tests/pt/test_observed_type.py  |  67 ++++++++++++++
 4 files changed, 344 insertions(+), 20 deletions(-)
 create mode 100644 source/tests/common/test_model_stat.py

diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py
index f8c3685b78..d586d4e679 100644
--- a/deepmd/pt/utils/stat.py
+++ b/deepmd/pt/utils/stat.py
@@ -75,38 +75,116 @@ def make_stat_input(
                 except StopIteration:
                     iterator = iter(dataloaders[i])
                     stat_data = next(iterator)
-                if (
-                    "find_fparam" in stat_data
-                    and "fparam" in stat_data
-                    and stat_data["find_fparam"] == 0.0
-                ):
-                    # for model using default fparam
-                    stat_data.pop("fparam")
-                    stat_data.pop("find_fparam")
-                for dd in stat_data:
-                    if stat_data[dd] is None:
-                        sys_stat[dd] = None
-                    elif isinstance(stat_data[dd], torch.Tensor):
-                        if dd not in sys_stat:
-                            sys_stat[dd] = []
-                        sys_stat[dd].append(stat_data[dd])
-                    elif isinstance(stat_data[dd], np.float32):
-                        sys_stat[dd] = stat_data[dd]
-                    else:
-                        pass
+                _append_stat_data(sys_stat, stat_data)
+            _append_missing_type_frames(sys_stat, datasets[i])
 
         for key in sys_stat:
             if isinstance(sys_stat[key], np.float32):
                 pass
             elif sys_stat[key] is None or sys_stat[key][0] is None:
                 sys_stat[key] = None
-            elif isinstance(stat_data[dd], torch.Tensor):
+            elif isinstance(sys_stat[key][0], torch.Tensor):
                 sys_stat[key] = torch.cat(sys_stat[key], dim=0)
         dict_to_device(sys_stat)
         lst.append(sys_stat)
     return lst
 
 
+def _append_stat_data(sys_stat: dict[str, Any], stat_data: dict[str, Any]) -> None:
+    """Append one statistics batch to the per-system accumulator."""
+    if (
+        "find_fparam" in stat_data
+        and "fparam" in stat_data
+        and stat_data["find_fparam"] == 0.0
+    ):
+        # for model using default fparam
+        stat_data.pop("fparam")
+        stat_data.pop("find_fparam")
+    for dd, value in stat_data.items():
+        if value is None:
+            sys_stat[dd] = None
+        elif isinstance(value, torch.Tensor):
+            if dd not in sys_stat:
+                sys_stat[dd] = []
+            sys_stat[dd].append(value)
+        elif isinstance(value, np.float32):
+            sys_stat[dd] = value
+
+
+def _append_missing_type_frames(sys_stat: dict[str, Any], dataset: Any) -> None:
+    """Add representative mixed-type frames for atom types missed by sampling.
+
+    Global output statistics solve a per-type linear regression from the sampled
+    frame compositions.  In mixed-type datasets a random small sample can miss a
+    type that exists elsewhere in the dataset, making that type's bias
+    unconstrained.  We therefore append a minimal set of real frames, one by one,
+    until every type present in the full system is also present in the statistics
+    sample.  Non-mixed systems have a fixed composition, so the initially sampled
+    frames already cover the system-level types.
+    """
+    if "real_natoms_vec" not in sys_stat or sys_stat["real_natoms_vec"] is None:
+        return
+    if not hasattr(dataset, "data_system"):
+        return
+    sampled_natoms_vec = sys_stat["real_natoms_vec"]
+    if len(sampled_natoms_vec) == 0:
+        return
+    sampled_counts = torch.cat(sampled_natoms_vec, dim=0)[:, 2:].sum(dim=0)
+    dataset_counts, first_frame_for_type = _mixed_type_coverage(dataset)
+    if dataset_counts is None or first_frame_for_type is None:
+        return
+
+    missing_types = np.flatnonzero((dataset_counts > 0) & (sampled_counts.numpy() == 0))
+    if len(missing_types) == 0:
+        return
+
+    # Import lazily to keep this utility independent from dataloader import time.
+    from deepmd.pt.utils.dataloader import (
+        collate_batch,
+    )
+
+    used_frames: set[int] = set()
+    while len(missing_types) > 0:
+        frame_idx = first_frame_for_type[int(missing_types[0])]
+        if frame_idx < 0 or frame_idx in used_frames:
+            break
+        used_frames.add(frame_idx)
+        # Reuse the dataset and collate path so that augmented frames have
+        # exactly the same tensor layout as ordinary DataLoader batches.
+        extra_batch = collate_batch([dataset[frame_idx]])
+        _append_stat_data(sys_stat, extra_batch)
+        sampled_counts += extra_batch["real_natoms_vec"][:, 2:].sum(dim=0)
+        missing_types = np.flatnonzero(
+            (dataset_counts > 0) & (sampled_counts.numpy() == 0)
+        )
+
+
+def _mixed_type_coverage(dataset: Any) -> tuple[np.ndarray | None, np.ndarray | None]:
+    """Return full-dataset type counts and one representative frame per type."""
+    data_system = dataset.data_system
+    if not getattr(data_system, "mixed_type", False):
+        return None, None
+    ntypes = data_system.get_ntypes()
+    counts = np.zeros(ntypes, dtype=np.int64)
+    first_frame_for_type = np.full(ntypes, -1, dtype=np.int64)
+    frame_offset = 0
+    for set_dir, frame_end in zip(
+        data_system.dirs, data_system.prefix_sum, strict=True
+    ):
+        type_path = set_dir / "real_atom_types.npy"
+        real_type = type_path.load_numpy()
+        if getattr(data_system, "enforce_type_map", False):
+            real_type = data_system.type_idx_map[real_type].astype(np.int32)
+        real_type = real_type.reshape(frame_end - frame_offset, data_system.natoms)
+        for type_i in range(ntypes):
+            frame_hits = np.flatnonzero((real_type == type_i).any(axis=1))
+            counts[type_i] += int((real_type == type_i).sum())
+            if first_frame_for_type[type_i] < 0 and len(frame_hits) > 0:
+                first_frame_for_type[type_i] = frame_offset + int(frame_hits[0])
+        frame_offset = frame_end
+    return counts, first_frame_for_type
+
+
 def _restore_from_file(
     stat_file_path: DPPath,
     keys: list[str] = ["energy"],
diff --git a/deepmd/utils/model_stat.py b/deepmd/utils/model_stat.py
index 33ebbcae57..0464758024 100644
--- a/deepmd/utils/model_stat.py
+++ b/deepmd/utils/model_stat.py
@@ -63,6 +63,7 @@ def collect_batches(
                 if dd == "natoms_vec":
                     stat_data[dd] = stat_data[dd].astype(np.int32)
                 sys_stat[dd].append(stat_data[dd])
+        _append_missing_type_frames(data, ii, sys_stat)
         for dd in sys_stat:
             if merge_sys:
                 for bb in sys_stat[dd]:
@@ -72,6 +73,86 @@ def collect_batches(
     return all_stat
 
 
+def _append_missing_type_frames(
+    data: Any, sys_idx: int, sys_stat: dict[str, list[Any]]
+) -> None:
+    """Append representative mixed-type frames for types missed by sampling.
+
+    Energy/output bias statistics regress one bias per atom type from the sampled
+    frame compositions.  Mixed-type systems can contain types that do not appear
+    in the small random statistics sample.  When that happens, append the first
+    frame containing each missing type so the regression is constrained for every
+    type that exists in the underlying system.  Standard (non-mixed) systems have
+    fixed composition and do not need augmentation.
+    """
+    if "real_natoms_vec" not in sys_stat or not hasattr(data, "data_systems"):
+        return
+    if getattr(data, "mixed_systems", False):
+        # In mixed-system batching sys_idx is intentionally ignored by get_batch;
+        # keep the historical sampling behaviour rather than guessing ownership.
+        return
+    data_system = data.data_systems[sys_idx]
+    dataset_counts, first_frame_for_type = _mixed_type_coverage(data_system)
+    if dataset_counts is None or first_frame_for_type is None:
+        return
+    sampled_counts = np.concatenate(sys_stat["real_natoms_vec"], axis=0)[:, 2:].sum(
+        axis=0
+    )
+    missing_types = np.flatnonzero((dataset_counts > 0) & (sampled_counts == 0))
+    if len(missing_types) == 0:
+        return
+
+    used_frames: set[int] = set()
+    while len(missing_types) > 0:
+        frame_idx = first_frame_for_type[int(missing_types[0])]
+        if frame_idx < 0 or frame_idx in used_frames:
+            break
+        used_frames.add(frame_idx)
+        extra_batch = data_system.get_single_frame(frame_idx, num_worker=1)
+        extra_batch["natoms_vec"] = data.natoms_vec[sys_idx].astype(np.int32)
+        extra_batch["default_mesh"] = data.default_mesh[sys_idx]
+        for key, value in extra_batch.items():
+            if key == "natoms_vec":
+                value = value.astype(np.int32)
+            if (
+                key not in {"natoms_vec", "default_mesh"}
+                and isinstance(value, np.ndarray)
+                and value.ndim >= 1
+            ):
+                value = value.reshape((1, *value.shape))
+            sys_stat[key].append(value)
+        sampled_counts += (
+            extra_batch["real_natoms_vec"].reshape(1, -1)[:, 2:].sum(axis=0)
+        )
+        missing_types = np.flatnonzero((dataset_counts > 0) & (sampled_counts == 0))
+
+
+def _mixed_type_coverage(
+    data_system: Any,
+) -> tuple[np.ndarray | None, np.ndarray | None]:
+    """Return full mixed-type counts and a representative frame per type."""
+    if not getattr(data_system, "mixed_type", False):
+        return None, None
+    ntypes = data_system.get_ntypes()
+    counts = np.zeros(ntypes, dtype=np.int64)
+    first_frame_for_type = np.full(ntypes, -1, dtype=np.int64)
+    frame_offset = 0
+    for set_dir, frame_end in zip(
+        data_system.dirs, data_system.prefix_sum, strict=True
+    ):
+        real_type = (set_dir / "real_atom_types.npy").load_numpy()
+        if getattr(data_system, "enforce_type_map", False):
+            real_type = data_system.type_idx_map[real_type].astype(np.int32)
+        real_type = real_type.reshape(frame_end - frame_offset, data_system.natoms)
+        for type_i in range(ntypes):
+            frame_hits = np.flatnonzero((real_type == type_i).any(axis=1))
+            counts[type_i] += int((real_type == type_i).sum())
+            if first_frame_for_type[type_i] < 0 and len(frame_hits) > 0:
+                first_frame_for_type[type_i] = frame_offset + int(frame_hits[0])
+        frame_offset = frame_end
+    return counts, first_frame_for_type
+
+
 def make_stat_input(
     data: Any,
     nbatches: int,
diff --git a/source/tests/common/test_model_stat.py b/source/tests/common/test_model_stat.py
new file mode 100644
index 0000000000..b108684c3b
--- /dev/null
+++ b/source/tests/common/test_model_stat.py
@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Tests for backend-agnostic statistics sampling helpers."""
+
+import unittest
+
+import numpy as np
+
+from deepmd.utils.model_stat import (
+    make_stat_input,
+)
+
+
+class _FakeTypePath:
+    def __init__(self, real_types: np.ndarray) -> None:
+        self.real_types = real_types
+
+    def load_numpy(self) -> np.ndarray:
+        return self.real_types
+
+
+class _FakeSetDir:
+    def __init__(self, real_types: np.ndarray) -> None:
+        self.real_types = real_types
+
+    def __truediv__(self, name: str) -> _FakeTypePath:
+        assert name == "real_atom_types.npy"
+        return _FakeTypePath(self.real_types)
+
+
+class _FakeMixedDataSystem:
+    mixed_type = True
+    enforce_type_map = False
+    natoms = 2
+    dirs: list[_FakeSetDir]
+    prefix_sum: list[int]
+
+    def __init__(self) -> None:
+        self.dirs = [_FakeSetDir(np.array([[0, -1], [1, -1]], dtype=np.int32))]
+        self.prefix_sum = [2]
+
+    def get_ntypes(self) -> int:
+        return 2
+
+    def get_single_frame(self, index: int, num_worker: int = 1) -> dict:
+        assert index == 1
+        return {
+            "coord": np.zeros((6,), dtype=np.float32),
+            "type": np.array([1, -1], dtype=np.int32),
+            "atype": np.array([1, -1], dtype=np.int32),
+            "box": np.eye(3, dtype=np.float32).reshape(-1),
+            "real_natoms_vec": np.array([2, 2, 0, 1], dtype=np.int32),
+            "find_energy": np.float32(1.0),
+            "energy": np.array([1.0], dtype=np.float64),
+        }
+
+
+class _FakeMixedData:
+    mixed_systems = False
+    natoms_vec: list[np.ndarray]
+    default_mesh: list[np.ndarray]
+
+    def __init__(self) -> None:
+        self.data_systems = [_FakeMixedDataSystem()]
+        self.natoms_vec = [np.array([2, 2, 1, 0], dtype=np.int32)]
+        self.default_mesh = [np.array([], dtype=np.int32)]
+
+    def get_nsystems(self) -> int:
+        return 1
+
+    def get_batch(self, sys_idx: int | None = None) -> dict:
+        assert sys_idx == 0
+        return {
+            "coord": np.zeros((1, 6), dtype=np.float32),
+            "type": np.array([[0, -1]], dtype=np.int32),
+            "atype": np.array([[0, -1]], dtype=np.int32),
+            "box": np.eye(3, dtype=np.float32).reshape(1, 9),
+            "real_natoms_vec": np.array([[2, 2, 1, 0]], dtype=np.int32),
+            "natoms_vec": np.array([2, 2, 1, 0], dtype=np.int32),
+            "default_mesh": np.array([], dtype=np.int32),
+            "find_energy": np.float32(1.0),
+            "energy": np.array([[0.0]], dtype=np.float64),
+        }
+
+
+class TestModelStatSamplingCoverage(unittest.TestCase):
+    """Mixed-type make_stat_input should cover types beyond initial batches."""
+
+    def test_make_stat_input_appends_missing_mixed_type_frame(self) -> None:
+        sampled = make_stat_input(_FakeMixedData(), nbatches=1)
+
+        self.assertEqual(len(sampled), 1)
+        counts = sampled[0]["real_natoms_vec"][:, 2:].sum(axis=0)
+        self.assertTrue(np.all(counts > 0))
+        self.assertEqual(sampled[0]["energy"].shape[0], 2)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/pt/test_observed_type.py b/source/tests/pt/test_observed_type.py
index c834ec567e..25cf58b07a 100644
--- a/source/tests/pt/test_observed_type.py
+++ b/source/tests/pt/test_observed_type.py
@@ -18,6 +18,7 @@
 import torch
 
 from deepmd.pt.utils.stat import (
+    _append_missing_type_frames,
     _restore_observed_type_from_file,
     _save_observed_type_to_file,
     collect_observed_types,
@@ -72,6 +73,72 @@ def test_out_of_range_index_ignored(self) -> None:
         self.assertEqual(result, ["O"])
 
 
+class _FakeTypePath:
+    def __init__(self, real_types: np.ndarray) -> None:
+        self.real_types = real_types
+
+    def load_numpy(self) -> np.ndarray:
+        return self.real_types
+
+
+class _FakeSetDir:
+    def __init__(self, real_types: np.ndarray) -> None:
+        self.real_types = real_types
+
+    def __truediv__(self, name: str) -> _FakeTypePath:
+        assert name == "real_atom_types.npy"
+        return _FakeTypePath(self.real_types)
+
+
+class _FakeMixedDataSystem:
+    mixed_type = True
+    enforce_type_map = False
+    natoms = 2
+    dirs: list[_FakeSetDir]
+    prefix_sum: list[int]
+
+    def __init__(self) -> None:
+        self.dirs = [_FakeSetDir(np.array([[0, -1], [1, -1]], dtype=np.int32))]
+        self.prefix_sum = [2]
+
+    def get_ntypes(self) -> int:
+        return 2
+
+
+class _FakeMixedDataset:
+    data_system: _FakeMixedDataSystem
+
+    def __init__(self) -> None:
+        self.data_system = _FakeMixedDataSystem()
+
+    def __getitem__(self, index: int) -> dict:
+        assert index == 1
+        return {
+            "coord": np.zeros((6,), dtype=np.float32),
+            "atype": np.array([1, -1], dtype=np.int32),
+            "box": np.eye(3, dtype=np.float32).reshape(-1),
+            "real_natoms_vec": np.array([2, 2, 0, 1], dtype=np.int32),
+        }
+
+
+class TestStatSamplingCoverage(unittest.TestCase):
+    """Mixed-type statistics samples should cover all dataset types."""
+
+    def test_append_missing_mixed_type_frame(self) -> None:
+        sys_stat = {
+            "coord": [torch.zeros((1, 6), dtype=torch.float32)],
+            "atype": [torch.tensor([[0, -1]], dtype=torch.int32)],
+            "box": [torch.eye(3, dtype=torch.float32).reshape(1, 9)],
+            "real_natoms_vec": [torch.tensor([[2, 2, 1, 0]], dtype=torch.int32)],
+        }
+
+        _append_missing_type_frames(sys_stat, _FakeMixedDataset())
+
+        self.assertEqual(len(sys_stat["real_natoms_vec"]), 2)
+        sampled_counts = torch.cat(sys_stat["real_natoms_vec"], dim=0)[:, 2:].sum(dim=0)
+        self.assertTrue(torch.all(sampled_counts > 0))
+
+
 class TestObservedTypeStatFile(unittest.TestCase):
     """Test stat file save/load round-trip for observed_type."""
 

From a739210e6159ec34b6c9e385c925a9df12d7fdb7 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Date: Fri, 26 Jun 2026 01:18:43 +0800
Subject: [PATCH 2/8] fix(pt): keep stat sampling coverage on cpu

---
 deepmd/pt/utils/stat.py               | 12 +++++++++---
 deepmd/utils/model_stat.py            | 11 +++++++----
 source/tests/pt/test_observed_type.py | 27 +++++++++++++++------------
 3 files changed, 31 insertions(+), 19 deletions(-)

diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py
index d586d4e679..97726b3891 100644
--- a/deepmd/pt/utils/stat.py
+++ b/deepmd/pt/utils/stat.py
@@ -145,13 +145,19 @@ def _append_missing_type_frames(sys_stat: dict[str, Any], dataset: Any) -> None:
 
     used_frames: set[int] = set()
     while len(missing_types) > 0:
-        frame_idx = first_frame_for_type[int(missing_types[0])]
-        if frame_idx < 0 or frame_idx in used_frames:
+        frame_idx: int | None = None
+        for type_i in missing_types:
+            candidate = int(first_frame_for_type[int(type_i)])
+            if candidate >= 0 and candidate not in used_frames:
+                frame_idx = candidate
+                break
+        if frame_idx is None:
             break
         used_frames.add(frame_idx)
         # Reuse the dataset and collate path so that augmented frames have
         # exactly the same tensor layout as ordinary DataLoader batches.
-        extra_batch = collate_batch([dataset[frame_idx]])
+        with torch.device("cpu"):
+            extra_batch = collate_batch([dataset[frame_idx]])
         _append_stat_data(sys_stat, extra_batch)
         sampled_counts += extra_batch["real_natoms_vec"][:, 2:].sum(dim=0)
         missing_types = np.flatnonzero(
diff --git a/deepmd/utils/model_stat.py b/deepmd/utils/model_stat.py
index 0464758024..e57a0ef3c5 100644
--- a/deepmd/utils/model_stat.py
+++ b/deepmd/utils/model_stat.py
@@ -104,16 +104,19 @@ def _append_missing_type_frames(
 
     used_frames: set[int] = set()
     while len(missing_types) > 0:
-        frame_idx = first_frame_for_type[int(missing_types[0])]
-        if frame_idx < 0 or frame_idx in used_frames:
+        frame_idx: int | None = None
+        for type_i in missing_types:
+            candidate = int(first_frame_for_type[int(type_i)])
+            if candidate >= 0 and candidate not in used_frames:
+                frame_idx = candidate
+                break
+        if frame_idx is None:
             break
         used_frames.add(frame_idx)
         extra_batch = data_system.get_single_frame(frame_idx, num_worker=1)
         extra_batch["natoms_vec"] = data.natoms_vec[sys_idx].astype(np.int32)
         extra_batch["default_mesh"] = data.default_mesh[sys_idx]
         for key, value in extra_batch.items():
-            if key == "natoms_vec":
-                value = value.astype(np.int32)
             if (
                 key not in {"natoms_vec", "default_mesh"}
                 and isinstance(value, np.ndarray)
diff --git a/source/tests/pt/test_observed_type.py b/source/tests/pt/test_observed_type.py
index 25cf58b07a..d53c49193d 100644
--- a/source/tests/pt/test_observed_type.py
+++ b/source/tests/pt/test_observed_type.py
@@ -125,18 +125,21 @@ class TestStatSamplingCoverage(unittest.TestCase):
     """Mixed-type statistics samples should cover all dataset types."""
 
     def test_append_missing_mixed_type_frame(self) -> None:
-        sys_stat = {
-            "coord": [torch.zeros((1, 6), dtype=torch.float32)],
-            "atype": [torch.tensor([[0, -1]], dtype=torch.int32)],
-            "box": [torch.eye(3, dtype=torch.float32).reshape(1, 9)],
-            "real_natoms_vec": [torch.tensor([[2, 2, 1, 0]], dtype=torch.int32)],
-        }
-
-        _append_missing_type_frames(sys_stat, _FakeMixedDataset())
-
-        self.assertEqual(len(sys_stat["real_natoms_vec"]), 2)
-        sampled_counts = torch.cat(sys_stat["real_natoms_vec"], dim=0)[:, 2:].sum(dim=0)
-        self.assertTrue(torch.all(sampled_counts > 0))
+        with torch.device("cpu"):
+            sys_stat = {
+                "coord": [torch.zeros((1, 6), dtype=torch.float32)],
+                "atype": [torch.tensor([[0, -1]], dtype=torch.int32)],
+                "box": [torch.eye(3, dtype=torch.float32).reshape(1, 9)],
+                "real_natoms_vec": [torch.tensor([[2, 2, 1, 0]], dtype=torch.int32)],
+            }
+
+            _append_missing_type_frames(sys_stat, _FakeMixedDataset())
+
+            self.assertEqual(len(sys_stat["real_natoms_vec"]), 2)
+            sampled_counts = torch.cat(sys_stat["real_natoms_vec"], dim=0)[:, 2:].sum(
+                dim=0
+            )
+            self.assertTrue(torch.all(sampled_counts > 0))
 
 
 class TestObservedTypeStatFile(unittest.TestCase):

From 3db3757bc07db49b9d4e50b1f523e45641fbb81c Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Date: Fri, 26 Jun 2026 01:27:18 +0800
Subject: [PATCH 3/8] docs(test): document mixed-type stat fixtures

---
 source/tests/common/test_model_stat.py | 15 +++++++++++++++
 source/tests/pt/test_observed_type.py  | 13 +++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/source/tests/common/test_model_stat.py b/source/tests/common/test_model_stat.py
index b108684c3b..512bd01fd3 100644
--- a/source/tests/common/test_model_stat.py
+++ b/source/tests/common/test_model_stat.py
@@ -11,23 +11,31 @@
 
 
 class _FakeTypePath:
+    """Fake path object that returns in-memory atom types."""
+
     def __init__(self, real_types: np.ndarray) -> None:
         self.real_types = real_types
 
     def load_numpy(self) -> np.ndarray:
+        """Return the stored atom-type array."""
         return self.real_types
 
 
 class _FakeSetDir:
+    """Fake set directory exposing ``real_atom_types.npy``."""
+
     def __init__(self, real_types: np.ndarray) -> None:
         self.real_types = real_types
 
     def __truediv__(self, name: str) -> _FakeTypePath:
+        """Return a fake path for the atom-type file."""
         assert name == "real_atom_types.npy"
         return _FakeTypePath(self.real_types)
 
 
 class _FakeMixedDataSystem:
+    """Minimal mixed-type data system for stat sampling tests."""
+
     mixed_type = True
     enforce_type_map = False
     natoms = 2
@@ -39,9 +47,11 @@ def __init__(self) -> None:
         self.prefix_sum = [2]
 
     def get_ntypes(self) -> int:
+        """Return the number of real atom types."""
         return 2
 
     def get_single_frame(self, index: int, num_worker: int = 1) -> dict:
+        """Return the representative frame containing the missing type."""
         assert index == 1
         return {
             "coord": np.zeros((6,), dtype=np.float32),
@@ -55,6 +65,8 @@ def get_single_frame(self, index: int, num_worker: int = 1) -> dict:
 
 
 class _FakeMixedData:
+    """Minimal multi-system data wrapper for ``make_stat_input``."""
+
     mixed_systems = False
     natoms_vec: list[np.ndarray]
     default_mesh: list[np.ndarray]
@@ -65,9 +77,11 @@ def __init__(self) -> None:
         self.default_mesh = [np.array([], dtype=np.int32)]
 
     def get_nsystems(self) -> int:
+        """Return the number of systems."""
         return 1
 
     def get_batch(self, sys_idx: int | None = None) -> dict:
+        """Return the initially sampled batch that misses one type."""
         assert sys_idx == 0
         return {
             "coord": np.zeros((1, 6), dtype=np.float32),
@@ -86,6 +100,7 @@ class TestModelStatSamplingCoverage(unittest.TestCase):
     """Mixed-type make_stat_input should cover types beyond initial batches."""
 
     def test_make_stat_input_appends_missing_mixed_type_frame(self) -> None:
+        """Append a representative frame when the first batch misses a type."""
         sampled = make_stat_input(_FakeMixedData(), nbatches=1)
 
         self.assertEqual(len(sampled), 1)
diff --git a/source/tests/pt/test_observed_type.py b/source/tests/pt/test_observed_type.py
index d53c49193d..f88aaae223 100644
--- a/source/tests/pt/test_observed_type.py
+++ b/source/tests/pt/test_observed_type.py
@@ -74,23 +74,31 @@ def test_out_of_range_index_ignored(self) -> None:
 
 
 class _FakeTypePath:
+    """Fake path object that returns in-memory atom types."""
+
     def __init__(self, real_types: np.ndarray) -> None:
         self.real_types = real_types
 
     def load_numpy(self) -> np.ndarray:
+        """Return the stored atom-type array."""
         return self.real_types
 
 
 class _FakeSetDir:
+    """Fake set directory exposing ``real_atom_types.npy``."""
+
     def __init__(self, real_types: np.ndarray) -> None:
         self.real_types = real_types
 
     def __truediv__(self, name: str) -> _FakeTypePath:
+        """Return a fake path for the atom-type file."""
         assert name == "real_atom_types.npy"
         return _FakeTypePath(self.real_types)
 
 
 class _FakeMixedDataSystem:
+    """Minimal mixed-type data system for stat sampling tests."""
+
     mixed_type = True
     enforce_type_map = False
     natoms = 2
@@ -102,16 +110,20 @@ def __init__(self) -> None:
         self.prefix_sum = [2]
 
     def get_ntypes(self) -> int:
+        """Return the number of real atom types."""
         return 2
 
 
 class _FakeMixedDataset:
+    """Minimal PyTorch dataset wrapper for mixed-type stat sampling."""
+
     data_system: _FakeMixedDataSystem
 
     def __init__(self) -> None:
         self.data_system = _FakeMixedDataSystem()
 
     def __getitem__(self, index: int) -> dict:
+        """Return the representative frame containing the missing type."""
         assert index == 1
         return {
             "coord": np.zeros((6,), dtype=np.float32),
@@ -125,6 +137,7 @@ class TestStatSamplingCoverage(unittest.TestCase):
     """Mixed-type statistics samples should cover all dataset types."""
 
     def test_append_missing_mixed_type_frame(self) -> None:
+        """Append a representative frame under an explicit CPU device."""
         with torch.device("cpu"):
             sys_stat = {
                 "coord": [torch.zeros((1, 6), dtype=torch.float32)],

From deafe7a5bbb888c50e1c7b87e53f1bf4577e70d3 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz-bot@users.noreply.github.com>
Date: Tue, 11 Aug 2026 03:58:24 +0800
Subject: [PATCH 4/8] fix(data): preserve mixed statistics coverage

Keep representative frames aligned with sampled schemas, cover rare types in mixed-system batches, and honor minimum pair-distance filtering.

Coding-Agent: Codex
Codex-Version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning-Effort: xhigh
---
 deepmd/pt/utils/stat.py                |  63 +++++++----
 deepmd/utils/model_stat.py             | 146 ++++++++++++++++++++-----
 source/tests/common/test_model_stat.py | 103 ++++++++++++++++-
 source/tests/pt/test_observed_type.py  |  58 ++++++++--
 4 files changed, 313 insertions(+), 57 deletions(-)

diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py
index dfd7f744ee..58ae05452c 100644
--- a/deepmd/pt/utils/stat.py
+++ b/deepmd/pt/utils/stat.py
@@ -176,7 +176,11 @@ def make_stat_input(
                     stat_data = select_batch_frames(stat_data, frame_mask)
                 accepted_batches += 1
                 _append_stat_data(sys_stat, stat_data)
-            _append_missing_type_frames(sys_stat, datasets[system_index])
+            _append_missing_type_frames(
+                sys_stat,
+                datasets[system_index],
+                min_pair_dist=min_pair_dist,
+            )
 
         if not sys_stat:
             if min_pair_dist > 0.0:
@@ -226,7 +230,9 @@ def _append_stat_data(sys_stat: dict[str, Any], stat_data: dict[str, Any]) -> No
             sys_stat[dd] = value
 
 
-def _append_missing_type_frames(sys_stat: dict[str, Any], dataset: Any) -> None:
+def _append_missing_type_frames(
+    sys_stat: dict[str, Any], dataset: Any, min_pair_dist: float = 0.0
+) -> None:
     """Add representative mixed-type frames for atom types missed by sampling.
 
     Global output statistics solve a per-type linear regression from the sampled
@@ -245,8 +251,8 @@ def _append_missing_type_frames(sys_stat: dict[str, Any], dataset: Any) -> None:
     if len(sampled_natoms_vec) == 0:
         return
     sampled_counts = torch.cat(sampled_natoms_vec, dim=0)[:, 2:].sum(dim=0)
-    dataset_counts, first_frame_for_type = _mixed_type_coverage(dataset)
-    if dataset_counts is None or first_frame_for_type is None:
+    dataset_counts, candidate_frames = _mixed_type_coverage(dataset)
+    if dataset_counts is None or candidate_frames is None:
         return
 
     missing_types = np.flatnonzero((dataset_counts > 0) & (sampled_counts.numpy() == 0))
@@ -260,19 +266,33 @@ def _append_missing_type_frames(sys_stat: dict[str, Any], dataset: Any) -> None:
 
     used_frames: set[int] = set()
     while len(missing_types) > 0:
-        frame_idx: int | None = None
+        extra_batch: dict[str, Any] | None = None
         for type_i in missing_types:
-            candidate = int(first_frame_for_type[int(type_i)])
-            if candidate >= 0 and candidate not in used_frames:
-                frame_idx = candidate
+            for frame_idx in candidate_frames[int(type_i)]:
+                if frame_idx in used_frames:
+                    continue
+                used_frames.add(frame_idx)
+                # Reuse the dataset and collate path so augmented frames have
+                # exactly the same tensor layout as DataLoader batches.
+                with torch.device("cpu"):
+                    candidate_batch = collate_batch([dataset[frame_idx]])
+                frame_mask = min_pair_dist_frame_mask(candidate_batch, min_pair_dist)
+                if frame_mask is not None and not torch.any(frame_mask):
+                    continue
+                if frame_mask is not None:
+                    candidate_batch = select_batch_frames(candidate_batch, frame_mask)
+                extra_batch = candidate_batch
+                break
+            if extra_batch is not None:
                 break
-        if frame_idx is None:
+        if extra_batch is None:
+            log.warning(
+                "No frame containing atom types %s satisfies "
+                "min_pair_dist=%s; statistics will not cover these types.",
+                missing_types.tolist(),
+                min_pair_dist,
+            )
             break
-        used_frames.add(frame_idx)
-        # Reuse the dataset and collate path so that augmented frames have
-        # exactly the same tensor layout as ordinary DataLoader batches.
-        with torch.device("cpu"):
-            extra_batch = collate_batch([dataset[frame_idx]])
         _append_stat_data(sys_stat, extra_batch)
         sampled_counts += extra_batch["real_natoms_vec"][:, 2:].sum(dim=0)
         missing_types = np.flatnonzero(
@@ -280,14 +300,16 @@ def _append_missing_type_frames(sys_stat: dict[str, Any], dataset: Any) -> None:
         )
 
 
-def _mixed_type_coverage(dataset: Any) -> tuple[np.ndarray | None, np.ndarray | None]:
-    """Return full-dataset type counts and one representative frame per type."""
+def _mixed_type_coverage(
+    dataset: Any,
+) -> tuple[np.ndarray | None, list[list[int]] | None]:
+    """Return full-dataset type counts and candidate frames for every type."""
     data_system = dataset.data_system
     if not getattr(data_system, "mixed_type", False):
         return None, None
     ntypes = data_system.get_ntypes()
     counts = np.zeros(ntypes, dtype=np.int64)
-    first_frame_for_type = np.full(ntypes, -1, dtype=np.int64)
+    candidate_frames: list[list[int]] = [[] for _ in range(ntypes)]
     frame_offset = 0
     for set_dir, frame_end in zip(
         data_system.dirs, data_system.prefix_sum, strict=True
@@ -300,10 +322,11 @@ def _mixed_type_coverage(dataset: Any) -> tuple[np.ndarray | None, np.ndarray |
         for type_i in range(ntypes):
             frame_hits = np.flatnonzero((real_type == type_i).any(axis=1))
             counts[type_i] += int((real_type == type_i).sum())
-            if first_frame_for_type[type_i] < 0 and len(frame_hits) > 0:
-                first_frame_for_type[type_i] = frame_offset + int(frame_hits[0])
+            candidate_frames[type_i].extend(
+                frame_offset + int(frame_idx) for frame_idx in frame_hits
+            )
         frame_offset = frame_end
-    return counts, first_frame_for_type
+    return counts, candidate_frames
 
 
 def _restore_from_file(
diff --git a/deepmd/utils/model_stat.py b/deepmd/utils/model_stat.py
index 01bb6563f1..cdb52cbff0 100644
--- a/deepmd/utils/model_stat.py
+++ b/deepmd/utils/model_stat.py
@@ -115,13 +115,19 @@ def _append_missing_type_frames(
     """
     if "real_natoms_vec" not in sys_stat or not hasattr(data, "data_systems"):
         return
-    if getattr(data, "mixed_systems", False):
-        # In mixed-system batching sys_idx is intentionally ignored by get_batch;
-        # keep the historical sampling behaviour rather than guessing ownership.
-        return
-    data_system = data.data_systems[sys_idx]
-    dataset_counts, first_frame_for_type = _mixed_type_coverage(data_system)
-    if dataset_counts is None or first_frame_for_type is None:
+    mixed_systems = getattr(data, "mixed_systems", False)
+    if mixed_systems:
+        dataset_counts, candidate_frames = _mixed_system_coverage(data)
+    else:
+        data_system = data.data_systems[sys_idx]
+        dataset_counts, first_frame_for_type = _mixed_type_coverage(data_system)
+        if dataset_counts is None or first_frame_for_type is None:
+            return
+        candidate_frames = [
+            [] if frame_idx < 0 else [(sys_idx, int(frame_idx))]
+            for frame_idx in first_frame_for_type
+        ]
+    if dataset_counts is None or candidate_frames is None:
         return
     sampled_counts = np.concatenate(sys_stat["real_natoms_vec"], axis=0)[:, 2:].sum(
         axis=0
@@ -130,34 +136,118 @@ def _append_missing_type_frames(
     if len(missing_types) == 0:
         return
 
-    used_frames: set[int] = set()
+    used_frames: set[tuple[int, int]] = set()
     while len(missing_types) > 0:
-        frame_idx: int | None = None
+        appended = False
         for type_i in missing_types:
-            candidate = int(first_frame_for_type[int(type_i)])
-            if candidate >= 0 and candidate not in used_frames:
-                frame_idx = candidate
+            for system_idx, frame_idx in candidate_frames[int(type_i)]:
+                candidate = (system_idx, frame_idx)
+                if candidate in used_frames:
+                    continue
+                used_frames.add(candidate)
+                extra_batch = _get_representative_batch(
+                    data,
+                    system_idx,
+                    frame_idx,
+                    merge_mixed=mixed_systems,
+                )
+                if not _append_compatible_batch(sys_stat, extra_batch):
+                    continue
+                sampled_counts += (
+                    extra_batch["real_natoms_vec"].reshape(1, -1)[:, 2:].sum(axis=0)
+                )
+                appended = True
+                break
+            if appended:
                 break
-        if frame_idx is None:
+        if not appended:
+            if mixed_systems:
+                log.warning(
+                    "Cannot add shape-compatible representative frames for atom "
+                    "types %s in mixed-system statistics.",
+                    missing_types.tolist(),
+                )
             break
-        used_frames.add(frame_idx)
-        extra_batch = data_system.get_single_frame(frame_idx, num_worker=1)
-        extra_batch["natoms_vec"] = data.natoms_vec[sys_idx].astype(np.int32)
-        extra_batch["default_mesh"] = data.default_mesh[sys_idx]
-        for key, value in extra_batch.items():
-            if (
-                key not in {"natoms_vec", "default_mesh"}
-                and isinstance(value, np.ndarray)
-                and value.ndim >= 1
-            ):
-                value = value.reshape((1, *value.shape))
-            sys_stat[key].append(value)
-        sampled_counts += (
-            extra_batch["real_natoms_vec"].reshape(1, -1)[:, 2:].sum(axis=0)
-        )
         missing_types = np.flatnonzero((dataset_counts > 0) & (sampled_counts == 0))
 
 
+def _get_representative_batch(
+    data: Any,
+    system_idx: int,
+    frame_idx: int,
+    *,
+    merge_mixed: bool,
+) -> dict[str, Any]:
+    """Load one frame and give it the same layout as an ordinary batch."""
+    data_system = data.data_systems[system_idx]
+    extra_batch = data_system.get_single_frame(frame_idx, num_worker=1)
+    extra_batch["natoms_vec"] = data.natoms_vec[system_idx].astype(np.int32)
+    extra_batch["default_mesh"] = data.default_mesh[system_idx]
+    for key, value in list(extra_batch.items()):
+        if (
+            key not in {"natoms_vec", "default_mesh"}
+            and isinstance(value, np.ndarray)
+            and value.ndim >= 1
+        ):
+            extra_batch[key] = value.reshape((1, *value.shape))
+    if not merge_mixed:
+        return extra_batch
+
+    # Reuse the production mixed-batch merger so atomic labels and input arrays
+    # follow exactly the same padding rules as randomly sampled mixed batches.
+    merged_batch = data._merge_batch_data([extra_batch])
+    if "real_natoms_vec" in extra_batch:
+        # _merge_batch_data derives this field from fixed system composition;
+        # preserve the actual per-frame composition for mixed-type systems.
+        merged_batch["real_natoms_vec"] = extra_batch["real_natoms_vec"]
+    return merged_batch
+
+
+def _append_compatible_batch(
+    sys_stat: dict[str, list[Any]], extra_batch: dict[str, Any]
+) -> bool:
+    """Append a batch without changing the sampled schema or array widths."""
+    selected: dict[str, Any] = {}
+    for key, value in extra_batch.items():
+        if key not in sys_stat:
+            continue
+        reference = sys_stat[key][0]
+        if isinstance(reference, np.ndarray) and isinstance(value, np.ndarray):
+            reference_shape = (
+                reference.shape[1:] if reference.ndim >= 2 else reference.shape
+            )
+            value_shape = value.shape[1:] if value.ndim >= 2 else value.shape
+            if reference_shape != value_shape:
+                return False
+        selected[key] = value
+    if "real_natoms_vec" not in selected:
+        return False
+    for key, value in selected.items():
+        sys_stat[key].append(value)
+    return True
+
+
+def _mixed_system_coverage(
+    data: Any,
+) -> tuple[np.ndarray | None, list[list[tuple[int, int]]] | None]:
+    """Aggregate type counts and representative frames across mixed systems."""
+    if not hasattr(data, "data_systems") or len(data.data_systems) == 0:
+        return None, None
+    ntypes = int(data.get_ntypes())
+    counts = np.zeros(ntypes, dtype=np.int64)
+    candidate_frames: list[list[tuple[int, int]]] = [[] for _ in range(ntypes)]
+    for system_idx, data_system in enumerate(data.data_systems):
+        system_counts, first_frame_for_type = _mixed_type_coverage(data_system)
+        if system_counts is None or first_frame_for_type is None:
+            system_counts = np.asarray(data.natoms_vec[system_idx][2:], dtype=np.int64)
+            first_frame_for_type = np.where(system_counts > 0, 0, -1)
+        counts[: len(system_counts)] += system_counts
+        for type_i, frame_idx in enumerate(first_frame_for_type):
+            if frame_idx >= 0:
+                candidate_frames[type_i].append((system_idx, int(frame_idx)))
+    return counts, candidate_frames
+
+
 def _mixed_type_coverage(
     data_system: Any,
 ) -> tuple[np.ndarray | None, np.ndarray | None]:
diff --git a/source/tests/common/test_model_stat.py b/source/tests/common/test_model_stat.py
index 512bd01fd3..80d3a41d2d 100644
--- a/source/tests/common/test_model_stat.py
+++ b/source/tests/common/test_model_stat.py
@@ -61,6 +61,7 @@ def get_single_frame(self, index: int, num_worker: int = 1) -> dict:
             "real_natoms_vec": np.array([2, 2, 0, 1], dtype=np.int32),
             "find_energy": np.float32(1.0),
             "energy": np.array([1.0], dtype=np.float64),
+            "fid": index,
         }
 
 
@@ -86,7 +87,6 @@ def get_batch(self, sys_idx: int | None = None) -> dict:
         return {
             "coord": np.zeros((1, 6), dtype=np.float32),
             "type": np.array([[0, -1]], dtype=np.int32),
-            "atype": np.array([[0, -1]], dtype=np.int32),
             "box": np.eye(3, dtype=np.float32).reshape(1, 9),
             "real_natoms_vec": np.array([[2, 2, 1, 0]], dtype=np.int32),
             "natoms_vec": np.array([2, 2, 1, 0], dtype=np.int32),
@@ -96,17 +96,116 @@ def get_batch(self, sys_idx: int | None = None) -> dict:
         }
 
 
+class _FakeFixedDataSystem:
+    """One-frame fixed-composition system used by mixed-batch tests."""
+
+    mixed_type = False
+
+    def __init__(self, atom_type: int) -> None:
+        self.atom_type = atom_type
+
+    def get_single_frame(self, index: int, num_worker: int = 1) -> dict:
+        """Return the only frame in this fixed-composition system."""
+        assert index == 0
+        return {
+            "coord": np.full((6,), self.atom_type, dtype=np.float32),
+            "type": np.full((2,), self.atom_type, dtype=np.int32),
+            "atype": np.full((2,), self.atom_type, dtype=np.int32),
+            "box": np.eye(3, dtype=np.float32).reshape(-1),
+            "find_energy": np.float32(1.0),
+            "energy": np.array([float(self.atom_type)], dtype=np.float64),
+            "fid": index,
+        }
+
+
+class _FakeMixedSystemsData:
+    """Mixed-batch wrapper whose random sample always misses the rare system."""
+
+    mixed_systems = True
+
+    def __init__(self) -> None:
+        self.data_systems = [_FakeFixedDataSystem(0), _FakeFixedDataSystem(1)]
+        self.natoms_vec = [
+            np.array([2, 2, 2, 0], dtype=np.int32),
+            np.array([2, 2, 0, 2], dtype=np.int32),
+        ]
+        self.default_mesh = [
+            np.array([], dtype=np.int32),
+            np.array([], dtype=np.int32),
+        ]
+
+    def get_nsystems(self) -> int:
+        """Return the number of underlying systems."""
+        return 2
+
+    def get_ntypes(self) -> int:
+        """Return the global number of atom types."""
+        return 2
+
+    def get_batch(self, sys_idx: int | None = None) -> dict:
+        """Mimic a low-probability rare system by always sampling type zero."""
+        raw_batch = self.data_systems[0].get_single_frame(0)
+        raw_batch = {
+            key: value.reshape((1, *value.shape))
+            if isinstance(value, np.ndarray) and value.ndim >= 1
+            else value
+            for key, value in raw_batch.items()
+            if key not in {"atype", "fid"}
+        }
+        raw_batch["natoms_vec"] = self.natoms_vec[0]
+        raw_batch["default_mesh"] = self.default_mesh[0]
+        return self._merge_batch_data([raw_batch])
+
+    def _merge_batch_data(self, batch_data: list[dict]) -> dict:
+        """Merge same-sized frames using the production mixed-batch schema."""
+        return {
+            "natoms_vec": np.array([2, 2, 2, 0], dtype=np.int32),
+            "real_natoms_vec": np.vstack([batch["natoms_vec"] for batch in batch_data]),
+            "type": np.concatenate([batch["type"] for batch in batch_data]),
+            "default_mesh": np.array([], dtype=np.int32),
+            "coord": np.concatenate([batch["coord"] for batch in batch_data]),
+            "box": np.concatenate([batch["box"] for batch in batch_data]),
+            "find_energy": batch_data[0]["find_energy"],
+            "energy": np.concatenate([batch["energy"] for batch in batch_data]),
+        }
+
+
 class TestModelStatSamplingCoverage(unittest.TestCase):
     """Mixed-type make_stat_input should cover types beyond initial batches."""
 
     def test_make_stat_input_appends_missing_mixed_type_frame(self) -> None:
         """Append a representative frame when the first batch misses a type."""
-        sampled = make_stat_input(_FakeMixedData(), nbatches=1)
+        data = _FakeMixedData()
+        ordinary_batch = data.get_batch(sys_idx=0)
+        self.assertIn("type", ordinary_batch)
+        self.assertNotIn("atype", ordinary_batch)
+        self.assertNotIn("fid", ordinary_batch)
+
+        raw_frame = data.data_systems[0].get_single_frame(1)
+        self.assertIn("type", raw_frame)
+        self.assertIn("atype", raw_frame)
+        self.assertIn("fid", raw_frame)
+
+        sampled = make_stat_input(data, nbatches=1)
 
         self.assertEqual(len(sampled), 1)
         counts = sampled[0]["real_natoms_vec"][:, 2:].sum(axis=0)
         self.assertTrue(np.all(counts > 0))
+        self.assertEqual(sampled[0]["atype"].shape, (2, 2))
         self.assertEqual(sampled[0]["energy"].shape[0], 2)
+        self.assertNotIn("fid", sampled[0])
+
+    def test_make_stat_input_covers_rare_type_in_mixed_systems(self) -> None:
+        """Cover a low-probability system in ``batch_size: mixed`` sampling."""
+        sampled = make_stat_input(_FakeMixedSystemsData(), nbatches=1)
+
+        self.assertEqual(len(sampled), 2)
+        for system in sampled:
+            counts = system["real_natoms_vec"][:, 2:].sum(axis=0)
+            self.assertTrue(np.all(counts > 0))
+            self.assertEqual(system["atype"].shape, (2, 2))
+            self.assertEqual(system["energy"].shape, (2, 1))
+            self.assertNotIn("fid", system)
 
 
 if __name__ == "__main__":
diff --git a/source/tests/pt/test_observed_type.py b/source/tests/pt/test_observed_type.py
index a03a3ee5ee..31f7861eba 100644
--- a/source/tests/pt/test_observed_type.py
+++ b/source/tests/pt/test_observed_type.py
@@ -106,8 +106,8 @@ class _FakeMixedDataSystem:
     prefix_sum: list[int]
 
     def __init__(self) -> None:
-        self.dirs = [_FakeSetDir(np.array([[0, -1], [1, -1]], dtype=np.int32))]
-        self.prefix_sum = [2]
+        self.dirs = [_FakeSetDir(np.array([[0, -1], [1, -1], [1, -1]], dtype=np.int32))]
+        self.prefix_sum = [3]
 
     def get_ntypes(self) -> int:
         """Return the number of real atom types."""
@@ -119,17 +119,23 @@ class _FakeMixedDataset:
 
     data_system: _FakeMixedDataSystem
 
-    def __init__(self) -> None:
+    def __init__(self, min_pair_distances: tuple[float, ...] = (1.0, 1.0, 1.0)) -> None:
         self.data_system = _FakeMixedDataSystem()
+        self.min_pair_distances = min_pair_distances
 
     def __getitem__(self, index: int) -> dict:
         """Return the representative frame containing the missing type."""
-        assert index == 1
+        atom_type = 0 if index == 0 else 1
         return {
-            "coord": np.zeros((6,), dtype=np.float32),
-            "atype": np.array([1, -1], dtype=np.int32),
+            "coord": np.full((6,), index, dtype=np.float32),
+            "atype": np.array([atom_type, -1], dtype=np.int32),
             "box": np.eye(3, dtype=np.float32).reshape(-1),
-            "real_natoms_vec": np.array([2, 2, 0, 1], dtype=np.int32),
+            "real_natoms_vec": np.array(
+                [2, 2, int(atom_type == 0), int(atom_type == 1)], dtype=np.int32
+            ),
+            "min_pair_dist": np.array(
+                [self.min_pair_distances[index]], dtype=np.float64
+            ),
         }
 
 
@@ -154,6 +160,44 @@ def test_append_missing_mixed_type_frame(self) -> None:
             )
             self.assertTrue(torch.all(sampled_counts > 0))
 
+    def test_append_missing_type_uses_later_valid_distance_frame(self) -> None:
+        """Skip a too-close rare-type frame and append a later valid one."""
+        with torch.device("cpu"):
+            sys_stat = {
+                "coord": [torch.zeros((1, 6), dtype=torch.float32)],
+                "atype": [torch.tensor([[0, -1]], dtype=torch.int32)],
+                "box": [torch.eye(3, dtype=torch.float32).reshape(1, 9)],
+                "real_natoms_vec": [torch.tensor([[2, 2, 1, 0]], dtype=torch.int32)],
+            }
+
+            _append_missing_type_frames(
+                sys_stat,
+                _FakeMixedDataset((1.0, 0.1, 0.8)),
+                min_pair_dist=0.5,
+            )
+
+            self.assertEqual(len(sys_stat["real_natoms_vec"]), 2)
+            torch.testing.assert_close(sys_stat["coord"][-1], torch.full((1, 6), 2.0))
+
+    def test_append_missing_type_warns_when_all_frames_are_too_close(self) -> None:
+        """Leave a type uncovered when no representative passes the threshold."""
+        with torch.device("cpu"):
+            sys_stat = {
+                "coord": [torch.zeros((1, 6), dtype=torch.float32)],
+                "atype": [torch.tensor([[0, -1]], dtype=torch.int32)],
+                "box": [torch.eye(3, dtype=torch.float32).reshape(1, 9)],
+                "real_natoms_vec": [torch.tensor([[2, 2, 1, 0]], dtype=torch.int32)],
+            }
+
+            with self.assertLogs("deepmd.pt.utils.stat", level="WARNING"):
+                _append_missing_type_frames(
+                    sys_stat,
+                    _FakeMixedDataset((1.0, 0.1, 0.2)),
+                    min_pair_dist=0.5,
+                )
+
+            self.assertEqual(len(sys_stat["real_natoms_vec"]), 1)
+
 
 class TestObservedTypeStatFile(unittest.TestCase):
     """Test stat file save/load round-trip for observed_type."""

From bf878c5bd1c35fb0458fe5c7aaf05b12f570da4d Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Wed, 28 Oct 2026 16:02:26 +0800
Subject: [PATCH 5/8] fix(stat): search later valid frames for missing atom
 types

Coding-Agent: Codex
Codex-Version: codex-cli 0.164.0
Model: gpt-6-astra
Reasoning-Effort: xhigh
---
 deepmd/utils/model_stat.py             | 42 ++++++++++++++----------
 source/tests/common/test_model_stat.py | 45 ++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 17 deletions(-)

diff --git a/deepmd/utils/model_stat.py b/deepmd/utils/model_stat.py
index f9035684c6..2a7d930010 100644
--- a/deepmd/utils/model_stat.py
+++ b/deepmd/utils/model_stat.py
@@ -181,7 +181,7 @@ def _append_missing_type_frames(
 
     Energy/output bias statistics regress one bias per atom type from the sampled
     frame compositions.  Mixed-type systems can contain types that do not appear
-    in the small random statistics sample.  When that happens, append the first
+    in the small random statistics sample.  When that happens, append the first acceptable
     frame containing each missing type so the regression is constrained for every
     type that exists in the underlying system.  Standard (non-mixed) systems have
     fixed composition and do not need augmentation.
@@ -193,12 +193,11 @@ def _append_missing_type_frames(
         dataset_counts, candidate_frames = _mixed_system_coverage(data)
     else:
         data_system = data.data_systems[sys_idx]
-        dataset_counts, first_frame_for_type = _mixed_type_coverage(data_system)
-        if dataset_counts is None or first_frame_for_type is None:
+        dataset_counts, frames_for_type = _mixed_type_coverage(data_system)
+        if dataset_counts is None or frames_for_type is None:
             return
         candidate_frames = [
-            [] if frame_idx < 0 else [(sys_idx, int(frame_idx))]
-            for frame_idx in first_frame_for_type
+            [(sys_idx, frame_idx) for frame_idx in frames] for frames in frames_for_type
         ]
     if dataset_counts is None or candidate_frames is None:
         return
@@ -317,26 +316,36 @@ def _mixed_system_coverage(
     counts = np.zeros(ntypes, dtype=np.int64)
     candidate_frames: list[list[tuple[int, int]]] = [[] for _ in range(ntypes)]
     for system_idx, data_system in enumerate(data.data_systems):
-        system_counts, first_frame_for_type = _mixed_type_coverage(data_system)
-        if system_counts is None or first_frame_for_type is None:
+        system_counts, frames_for_type = _mixed_type_coverage(data_system)
+        if system_counts is None or frames_for_type is None:
             system_counts = np.asarray(data.natoms_vec[system_idx][2:], dtype=np.int64)
-            first_frame_for_type = np.where(system_counts > 0, 0, -1)
+            # Fixed composition does not imply identical geometry: a later
+            # frame may pass the pair-margin filter when the first one fails.
+            nframes = getattr(data_system, "prefix_sum", [1])[-1]
+            frames_for_type = [
+                list(range(nframes)) if count > 0 else [] for count in system_counts
+            ]
         counts[: len(system_counts)] += system_counts
-        for type_i, frame_idx in enumerate(first_frame_for_type):
-            if frame_idx >= 0:
-                candidate_frames[type_i].append((system_idx, int(frame_idx)))
+        for type_i, frames in enumerate(frames_for_type):
+            candidate_frames[type_i].extend(
+                (system_idx, frame_idx) for frame_idx in frames
+            )
     return counts, candidate_frames
 
 
 def _mixed_type_coverage(
     data_system: Any,
-) -> tuple[np.ndarray | None, np.ndarray | None]:
-    """Return full mixed-type counts and a representative frame per type."""
+) -> tuple[np.ndarray | None, list[list[int]] | None]:
+    """Return type counts and all candidates, in on-disk order.
+
+    A representative can be rejected by the geometry filter or batch schema;
+    retaining later frames lets statistics cover each type that has valid data.
+    """
     if not getattr(data_system, "mixed_type", False):
         return None, None
     ntypes = data_system.get_ntypes()
     counts = np.zeros(ntypes, dtype=np.int64)
-    first_frame_for_type = np.full(ntypes, -1, dtype=np.int64)
+    frames_for_type: list[list[int]] = [[] for _ in range(ntypes)]
     frame_offset = 0
     for set_dir, frame_end in zip(
         data_system.dirs, data_system.prefix_sum, strict=True
@@ -351,10 +360,9 @@ def _mixed_type_coverage(
         for type_i in range(ntypes):
             frame_hits = np.flatnonzero((real_type == type_i).any(axis=1))
             counts[type_i] += int((real_type == type_i).sum())
-            if first_frame_for_type[type_i] < 0 and len(frame_hits) > 0:
-                first_frame_for_type[type_i] = frame_offset + int(frame_hits[0])
+            frames_for_type[type_i].extend((frame_hits + frame_offset).tolist())
         frame_offset = frame_end
-    return counts, first_frame_for_type
+    return counts, frames_for_type
 
 
 def make_stat_input(
diff --git a/source/tests/common/test_model_stat.py b/source/tests/common/test_model_stat.py
index 80d3a41d2d..44bdbfc04f 100644
--- a/source/tests/common/test_model_stat.py
+++ b/source/tests/common/test_model_stat.py
@@ -96,6 +96,39 @@ def get_batch(self, sys_idx: int | None = None) -> dict:
         }
 
 
+class _FilteredMixedDataSystem(_FakeMixedDataSystem):
+    """Two rare-type frames, of which only the later frame is valid."""
+
+    def __init__(self) -> None:
+        self.dirs = [_FakeSetDir(np.array([[0, -1], [1, -1], [1, -1]]))]
+        self.prefix_sum = [3]
+        self.read_frames = []
+
+    def get_single_frame(self, index: int, num_worker: int = 1) -> dict:
+        self.read_frames.append(index)
+        frame = super().get_single_frame(1, num_worker)
+        frame["pair_margin"] = np.array([0.2 if index == 1 else 2.0])
+        frame["find_pair_margin"] = np.float32(1.0)
+        return frame
+
+
+class _FilteredMixedData(_FakeMixedData):
+    """Keep the initial type-zero sample valid while filtering candidates."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.data_systems = [_FilteredMixedDataSystem()]
+
+    def get_nbatches(self) -> list[int]:
+        return [3]
+
+    def get_batch(self, sys_idx: int | None = None) -> dict:
+        batch = super().get_batch(sys_idx)
+        batch["pair_margin"] = np.array([[2.0]])
+        batch["find_pair_margin"] = np.float32(1.0)
+        return batch
+
+
 class _FakeFixedDataSystem:
     """One-frame fixed-composition system used by mixed-batch tests."""
 
@@ -195,6 +228,18 @@ def test_make_stat_input_appends_missing_mixed_type_frame(self) -> None:
         self.assertEqual(sampled[0]["energy"].shape[0], 2)
         self.assertNotIn("fid", sampled[0])
 
+    def test_filtered_statistics_try_later_rare_type_frames(self) -> None:
+        """A rejected first representative must not hide a valid rare type."""
+        for threshold, expected_reads in [(0.0, [1]), (0.5, [1, 2])]:
+            with self.subTest(min_pair_dist=threshold):
+                data = _FilteredMixedData()
+                sampled = make_stat_input(data, nbatches=1, min_pair_dist=threshold)
+                counts = sampled[0]["real_natoms_vec"][:, 2:].sum(axis=0)
+                np.testing.assert_array_equal(counts, [1, 1])
+                self.assertEqual(data.data_systems[0].read_frames, expected_reads)
+                if threshold:
+                    self.assertTrue(np.all(sampled[0]["pair_margin"] >= 1.0))
+
     def test_make_stat_input_covers_rare_type_in_mixed_systems(self) -> None:
         """Cover a low-probability system in ``batch_size: mixed`` sampling."""
         sampled = make_stat_input(_FakeMixedSystemsData(), nbatches=1)

From a96d113b98eebf7906414a8236573ea1946c40d0 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Wed, 28 Oct 2026 16:49:52 +0800
Subject: [PATCH 6/8] test(pt): cover filtered rare types through stat input
 caller

Coding-Agent: Codex
Codex-Version: codex-cli 0.164.0
Model: gpt-6-astra
Reasoning-Effort: xhigh
---
 source/tests/pt/test_observed_type.py | 35 +++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/source/tests/pt/test_observed_type.py b/source/tests/pt/test_observed_type.py
index 4d449c7a82..f9388cd2d4 100644
--- a/source/tests/pt/test_observed_type.py
+++ b/source/tests/pt/test_observed_type.py
@@ -17,11 +17,15 @@
 import numpy as np
 import torch
 
+from deepmd.pt.utils.dataloader import (
+    collate_batch,
+)
 from deepmd.pt.utils.stat import (
     _append_missing_type_frames,
     _restore_observed_type_from_file,
     _save_observed_type_to_file,
     collect_observed_types,
+    make_stat_input,
 )
 from deepmd.utils.path import (
     DPPath,
@@ -143,6 +147,37 @@ def __getitem__(self, index: int) -> dict:
 class TestStatSamplingCoverage(unittest.TestCase):
     """Mixed-type statistics samples should cover all dataset types."""
 
+    def test_make_stat_input_covers_later_valid_rare_type(self) -> None:
+        """The public caller supplements its bounded sample after filtering."""
+        dataset = _FakeMixedDataset((1.0, 0.1, 0.8))
+        with torch.device("cpu"):
+            # The one-batch budget initially observes only type 0. Type 1's
+            # first representative is rejected; its later frame must survive
+            # supplementation, concatenation, and final device conversion.
+            dataloader = [collate_batch([dataset[index]]) for index in range(3)]
+            for options in (
+                {"min_pair_dist": 0.5},
+                {"filter_enabled_by_system": [True]},
+            ):
+                with self.subTest(options=options):
+                    sampled = make_stat_input(
+                        [dataset], [dataloader], nbatches=1, **options
+                    )
+                    self.assertEqual(len(sampled), 1)
+                    result = sampled[0]
+                    torch.testing.assert_close(
+                        result["atype"].cpu(),
+                        torch.tensor([[0, -1], [1, -1]], dtype=torch.int32),
+                    )
+                    torch.testing.assert_close(
+                        result["coord"].cpu(),
+                        torch.stack([torch.zeros(6), torch.full((6,), 2.0)]),
+                    )
+                    self.assertTrue(torch.all(result["pair_margin"] >= 1.0))
+                    self.assertTrue(
+                        torch.all(result["real_natoms_vec"][:, 2:].sum(dim=0) > 0)
+                    )
+
     def test_append_missing_mixed_type_frame(self) -> None:
         """Append a representative frame under an explicit CPU device."""
         with torch.device("cpu"):

From 95e28000a4afffc672a509c8aa30fd8a05c635b9 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Wed, 28 Oct 2026 17:52:43 +0800
Subject: [PATCH 7/8] test(stat): initialize filtered fixtures through their
 base classes

Coding-Agent: Codex
Codex-Version: codex-cli 0.164.0
Model: gpt-6-astra
Reasoning-Effort: xhigh
---
 source/tests/common/test_model_stat.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/source/tests/common/test_model_stat.py b/source/tests/common/test_model_stat.py
index 44bdbfc04f..97612d3307 100644
--- a/source/tests/common/test_model_stat.py
+++ b/source/tests/common/test_model_stat.py
@@ -42,9 +42,11 @@ class _FakeMixedDataSystem:
     dirs: list[_FakeSetDir]
     prefix_sum: list[int]
 
-    def __init__(self) -> None:
-        self.dirs = [_FakeSetDir(np.array([[0, -1], [1, -1]], dtype=np.int32))]
-        self.prefix_sum = [2]
+    def __init__(self, real_types: np.ndarray | None = None) -> None:
+        if real_types is None:
+            real_types = np.array([[0, -1], [1, -1]], dtype=np.int32)
+        self.dirs = [_FakeSetDir(real_types)]
+        self.prefix_sum = [len(real_types)]
 
     def get_ntypes(self) -> int:
         """Return the number of real atom types."""
@@ -72,8 +74,8 @@ class _FakeMixedData:
     natoms_vec: list[np.ndarray]
     default_mesh: list[np.ndarray]
 
-    def __init__(self) -> None:
-        self.data_systems = [_FakeMixedDataSystem()]
+    def __init__(self, system: _FakeMixedDataSystem | None = None) -> None:
+        self.data_systems = [_FakeMixedDataSystem() if system is None else system]
         self.natoms_vec = [np.array([2, 2, 1, 0], dtype=np.int32)]
         self.default_mesh = [np.array([], dtype=np.int32)]
 
@@ -100,8 +102,7 @@ class _FilteredMixedDataSystem(_FakeMixedDataSystem):
     """Two rare-type frames, of which only the later frame is valid."""
 
     def __init__(self) -> None:
-        self.dirs = [_FakeSetDir(np.array([[0, -1], [1, -1], [1, -1]]))]
-        self.prefix_sum = [3]
+        super().__init__(np.array([[0, -1], [1, -1], [1, -1]], dtype=np.int32))
         self.read_frames = []
 
     def get_single_frame(self, index: int, num_worker: int = 1) -> dict:
@@ -116,8 +117,7 @@ class _FilteredMixedData(_FakeMixedData):
     """Keep the initial type-zero sample valid while filtering candidates."""
 
     def __init__(self) -> None:
-        super().__init__()
-        self.data_systems = [_FilteredMixedDataSystem()]
+        super().__init__(_FilteredMixedDataSystem())
 
     def get_nbatches(self) -> list[int]:
         return [3]

From 9d204c8c21ee1a25e37736736e1e9a0293435980 Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Wed, 28 Oct 2026 19:40:24 +0800
Subject: [PATCH 8/8] docs: time out stalled external inventory requests

Coding-Agent: Codex
Codex-Version: codex-cli 0.164.0
Model: gpt-6-astra
Reasoning-Effort: xhigh
---
 doc/conf.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/doc/conf.py b/doc/conf.py
index 5753b2afed..78b3bc2a85 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -142,6 +142,10 @@
 # This pattern also affects html_static_path and html_extra_path.
 exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
+# External inventories are optional cross-reference metadata. Limit connection
+# and read-inactivity waits so Sphinx can warn and continue if a host stalls.
+# This is not a deadline for the full download or documentation build.
+intersphinx_timeout = 30
 intersphinx_mapping = {
     "numpy": ("https://docs.scipy.org/doc/numpy/", None),
     "python": ("https://docs.python.org/", None),