code/utils/dispdataset1d.py

import random
from typing import Optional, Dict, Any, List

import h5py
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader


class NCFDispersionDataset(Dataset):
    """
    基于 HDF5 读取：
    - waveform
    - dispersion
    - mask

    自动跳过以下无效样本：
    1. 没有 ncf group
    2. ncf group 下面没有任何 item
    3. 缺少 disp_periods / disp_velocity / disp_mask
    4. （可选）disp_mask 全 0
    """

    def __init__(
        self,
        h5_path: str,
        split: str = "train",
        waveform_length: int = 1536,
        random_ncf: bool = True,
        pad_value: float = 0.0,
        return_time: bool = True,
        seed: Optional[int] = None,
        skip_zero_mask: bool = True,   # 新增：mask 全 0 的样本也跳过
        verbose: bool = True,
    ):
        super().__init__()

        if split not in ("train", "test"):
            raise ValueError(f"split must be 'train' or 'test', got {split}")

        self.h5_path = h5_path
        self.split = split
        self.waveform_length = waveform_length
        self.random_ncf = random_ncf
        self.pad_value = pad_value
        self.return_time = return_time
        self.skip_zero_mask = skip_zero_mask
        self.verbose = verbose

        self.rng = random.Random(seed)
        self._h5_file = None

        with h5py.File(self.h5_path, "r") as f:
            key_name = f"{split}_keys"
            if key_name not in f:
                raise KeyError(f"HDF5 中不存在 {key_name}")

            raw_keys = f[key_name][()]
            raw_keys = [self._decode_if_needed(x) for x in raw_keys]

            valid_keys = []
            dropped = []

            for key in raw_keys:
                ok, reason = self._is_valid_key(f, key)
                if ok:
                    valid_keys.append(key)
                else:
                    dropped.append((key, reason))

        self.keys = valid_keys

        if self.verbose:
            print(f"[NCFDispersionDataset] split={split}")
            print(f"[NCFDispersionDataset] 原始样本数: {len(raw_keys)}")
            print(f"[NCFDispersionDataset] 有效样本数: {len(self.keys)}")
            print(f"[NCFDispersionDataset] 跳过样本数: {len(dropped)}")
            if len(dropped) > 0:
                print("[NCFDispersionDataset] 前10个被跳过样本示例:")
                for k, r in dropped[:10]:
                    print(f"  - {k}: {r}")

        if len(self.keys) == 0:
            raise RuntimeError("过滤后没有可用样本，请检查 HDF5 数据集。")

    @staticmethod
    def _decode_if_needed(x):
        if isinstance(x, bytes):
            return x.decode("utf-8")
        return str(x)

    def _get_h5(self):
        if self._h5_file is None:
            self._h5_file = h5py.File(self.h5_path, "r")
        return self._h5_file

    def _is_valid_key(self, f: h5py.File, key: str):
        """
        检查一个 key 是否有效。
        """
        if "paths" not in f:
            return False, "HDF5 缺少 /paths"

        if key not in f["paths"]:
            return False, "key 不存在于 /paths"

        grp = f["paths"][key]

        # dispersion 必须存在
        for name in ["disp_periods", "disp_velocity", "disp_mask"]:
            if name not in grp:
                return False, f"缺少 {name}"

        try:
            mask = np.asarray(grp["disp_mask"][()])
        except Exception as e:
            return False, f"读取 disp_mask 失败: {e}"

        if self.skip_zero_mask and mask.sum() <= 0:
            return False, "disp_mask 全 0"

        # ncf 必须存在且至少一个 item
        if "ncf" not in grp:
            return False, "缺少 ncf group"

        ncf_group = grp["ncf"]
        item_names = sorted(list(ncf_group.keys()))
        if len(item_names) == 0:
            return False, "ncf group 为空"

        # 至少有一个 item 同时有 time 和 waveform
        found_valid_ncf = False
        for item_name in item_names:
            item = ncf_group[item_name]
            if ("time" in item) and ("waveform" in item):
                found_valid_ncf = True
                break

        if not found_valid_ncf:
            return False, "ncf 中没有合法 item（缺少 time/waveform）"

        return True, "ok"

    def __len__(self):
        return len(self.keys)

    def _pad_or_truncate_1d(self, arr: np.ndarray, target_len: int, pad_value: float = 0.0):
        arr = np.asarray(arr, dtype=np.float32).reshape(-1)
        n = arr.shape[0]

        if n == target_len:
            return arr
        elif n > target_len:
            return arr[:target_len]
        else:
            out = np.full((target_len,), pad_value, dtype=np.float32)
            out[:n] = arr
            return out

    def _get_valid_ncf_item_names(self, ncf_group) -> List[str]:
        """
        返回当前 ncf_group 下所有合法 item 名称。
        """
        valid_names = []
        for item_name in sorted(list(ncf_group.keys())):
            item = ncf_group[item_name]
            if ("time" in item) and ("waveform" in item):
                valid_names.append(item_name)
        return valid_names

    def _choose_ncf_item_name(self, ncf_group) -> str:
        valid_names = self._get_valid_ncf_item_names(ncf_group)

        if len(valid_names) == 0:
            raise RuntimeError("当前样本没有可用的 NCF 子项")

        if self.random_ncf and len(valid_names) > 1:
            return self.rng.choice(valid_names)
        return valid_names[0]

    def __getitem__(self, index: int) -> Dict[str, Any]:
        f = self._get_h5()
        key = self.keys[index]

        try:
            path_group = f["paths"][key]

            periods = np.asarray(path_group["disp_periods"][()], dtype=np.float32)
            disp = np.asarray(path_group["disp_velocity"][()], dtype=np.float32)
            mask = np.asarray(path_group["disp_mask"][()], dtype=np.float32)

            ncf_group = path_group["ncf"]
            chosen_item_name = self._choose_ncf_item_name(ncf_group)
            chosen_ncf = ncf_group[chosen_item_name]

            time = np.asarray(chosen_ncf["time"][()], dtype=np.float32)
            waveform = np.asarray(chosen_ncf["waveform"][()], dtype=np.float32)

            waveform = self._pad_or_truncate_1d(
                waveform,
                target_len=self.waveform_length,
                pad_value=self.pad_value,
            )
            time = self._pad_or_truncate_1d(
                time,
                target_len=self.waveform_length,
                pad_value=0.0,
            )

            sample = {
                "key": key,
                "waveform": torch.from_numpy(waveform),   # [T]
                "disp": torch.from_numpy(disp),           # [49]
                "mask": torch.from_numpy(mask),           # [49]
                "periods": torch.from_numpy(periods),     # [49]
            }

            if self.return_time:
                sample["ncf_time"] = torch.from_numpy(time)

            return sample

        except Exception as e:
            # 极端情况下，单个样本还是坏了，就随机回退到另一个样本
            new_index = self.rng.randint(0, len(self.keys) - 1)
            if new_index == index and len(self.keys) > 1:
                new_index = (index + 1) % len(self.keys)

            if self.verbose:
                print(f"[WARN] __getitem__ 跳过坏样本 key={key}, error={e}, 改为读取 index={new_index}")

            return self.__getitem__(new_index)

    def close(self):
        if self._h5_file is not None:
            try:
                self._h5_file.close()
            except Exception:
                pass
            self._h5_file = None

    def __del__(self):
        self.close()


def ncf_disp_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    out = {
        "key": [item["key"] for item in batch],
        "waveform": torch.stack([item["waveform"] for item in batch], dim=0),  # [B, T]
        "disp": torch.stack([item["disp"] for item in batch], dim=0),          # [B, 49]
        "mask": torch.stack([item["mask"] for item in batch], dim=0),          # [B, 49]
        "periods": torch.stack([item["periods"] for item in batch], dim=0),    # [B, 49]
    }

    if "ncf_time" in batch[0]:
        out["ncf_time"] = torch.stack([item["ncf_time"] for item in batch], dim=0)

    return out


def build_dataloader(
    h5_path: str,
    split: str = "train",
    batch_size: int = 16,
    num_workers: int = 0,
    shuffle: Optional[bool] = None,
    waveform_length: int = 1536,
    random_ncf: bool = True,
    pin_memory: bool = False,
    drop_last: bool = False,
    seed: Optional[int] = None,
    skip_zero_mask: bool = True,
    verbose: bool = True,
) -> DataLoader:
    if shuffle is None:
        shuffle = (split == "train")

    dataset = NCFDispersionDataset(
        h5_path=h5_path,
        split=split,
        waveform_length=waveform_length,
        random_ncf=random_ncf,
        return_time=True,
        seed=seed,
        skip_zero_mask=skip_zero_mask,
        verbose=verbose,
    )

    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=pin_memory,
        drop_last=drop_last,
        collate_fn=ncf_disp_collate_fn,
    )
    return loader