fix: cover mixed types in stat sampling

deepmd/pt/utils/stat.py

@@ -217,25 +217,12 @@ def make_stat_input(
                 if frame_mask is not None:
                     stat_data = select_batch_frames(stat_data, frame_mask)
                 accepted_batches += 1
-                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[system_index],
+                filter_enabled=filter_enabled,
+            )
 
         if not sys_stat:
             if min_pair_dist > 0.0:
@@ -263,6 +250,129 @@ def make_stat_input(
     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, filter_enabled: bool = False
+) -> 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, 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))
+    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:
+        extra_batch: dict[str, Any] | None = None
+        for type_i in missing_types:
+            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 = pair_margin_frame_mask(candidate_batch, filter_enabled)
+                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 extra_batch is None:
+            log.warning(
+                "No frame containing atom types %s satisfies "
+                "filter_enabled=%s; statistics will not cover these types.",
+                missing_types.tolist(),
+                filter_enabled,
+            )
+            break
+        _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, 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)
+    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
+    ):
+        type_path = set_dir / "real_atom_types.npy"
+        real_type = type_path.load_numpy()
+        if getattr(data_system, "enforce_type_map", False):
+            real_type = real_type.copy()
+            real_mask = real_type >= 0
+            # Negative types denote padding, not indices into the type map.
+            real_type[real_mask] = data_system.type_idx_map[real_type[real_mask]]
+        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())
+            candidate_frames[type_i].extend(
+                frame_offset + int(frame_idx) for frame_idx in frame_hits
+            )
+        frame_offset = frame_end
+    return counts, candidate_frames
+
+
 def _full_pass_loader(dataloader: Any) -> Any:
     """Return a loader that covers the dataset once, whatever sampler it uses.
 

deepmd/utils/model_stat.py

@@ -159,6 +159,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]:
@@ -168,6 +169,202 @@ def collect_batches(
     return all_stat
 
 
+def _append_missing_type_frames(
+    data: Any,
+    sys_idx: int,
+    sys_stat: dict[str, list[Any]],
+    *,
+    normalized: bool = False,
+    filter_enabled: bool = False,
+) -> 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 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.
+    """
+    if "real_natoms_vec" not in sys_stat or not hasattr(data, "data_systems"):
+        return
+    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, frames_for_type = _mixed_type_coverage(data_system)
+        if dataset_counts is None or frames_for_type is None:
+            return
+        candidate_frames = [
+            [(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
+    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[tuple[int, int]] = set()
+    while len(missing_types) > 0:
+        appended = False
+        for type_i in missing_types:
+            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,
+                )
+                # The generic statistics path now consumes canonical batches;
+                # the TF collect_batches entry point still needs legacy keys.
+                if normalized:
+                    extra_batch = normalize_batch(extra_batch)
+                frame_mask = pair_margin_frame_mask(extra_batch, filter_enabled)
+                if frame_mask is not None and not frame_mask.all():
+                    continue
+                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 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
+        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, 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)
+            # 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, 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, 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)
+    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
+    ):
+        real_type = (set_dir / "real_atom_types.npy").load_numpy()
+        if getattr(data_system, "enforce_type_map", False):
+            real_type = real_type.copy()
+            real_mask = real_type >= 0
+            # Negative types denote padding, not indices into the type map.
+            real_type[real_mask] = data_system.type_idx_map[real_type[real_mask]]
+        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())
+            frames_for_type[type_i].extend((frame_hits + frame_offset).tolist())
+        frame_offset = frame_end
+    return counts, frames_for_type
+
+
 def make_stat_input(
     data: Any,
     nbatches: int,
@@ -218,11 +415,19 @@ def make_stat_input(
             )
             log.info(f"Skipping statistical system {ii}: {reason}.")
             continue
+        sys_stat = {key: [batch[key] for batch in batches] for key in batches[0]}
+        _append_missing_type_frames(
+            data,
+            ii,
+            sys_stat,
+            normalized=True,
+            filter_enabled=min_pair_dist > 0.0,
+        )
         # Arrays of two or more dimensions lead with the frame axis; every
         # other entry (find_* flags, an absent box) is constant per system.
         lst.append(
             {
-                key: np.concatenate([batch[key] for batch in batches], axis=0)
+                key: np.concatenate(sys_stat[key], axis=0)
                 if isinstance(value, np.ndarray) and value.ndim >= 2
                 else value
                 for key, value in batches[0].items()