#!/usr/bin/env python3 """Split PKS candidate samples by whole event groups without leakage.""" from __future__ import annotations import argparse import csv import json import random import shutil from collections import Counter, defaultdict from pathlib import Path SPLIT_NAMES = ("train", "validation", "test") def normalize_splits(values: list[float]) -> tuple[tuple[str, float], ...]: total = sum(values) if total <= 0: raise ValueError("切分比例之和必须大于 0") return tuple((name, value / total) for name, value in zip(SPLIT_NAMES, values)) def load_manifest(path: Path) -> list[dict]: with path.open(encoding="utf-8-sig", newline="") as handle: return list(csv.DictReader(handle)) def allocate(groups: list[tuple], splits: tuple[tuple[str, float], ...], assignments: dict, weight=None) -> None: """Assign whole groups so one stratum approaches the target split ratios. `weight` maps a group's members to its contribution to the target counts. Defaults to member count; pass a constant to balance by group count instead. """ if weight is None: weight = len weights = [weight(members) for _, members, _, _ in groups] total = sum(weights) targets = {name: total * ratio for name, ratio in splits} counts = Counter() for (group_id, members, _, _), group_weight in zip(groups, weights): # Put each group in the currently most under-filled partition. split = min( (name for name, _ in splits), key=lambda name: (counts[name] - targets[name], counts[name], name), ) assignments[group_id] = split counts[split] += group_weight def assign_groups( rows: list[dict], splits: tuple[tuple[str, float], ...], seed: int, positive_basis: str = "samples" ) -> dict[str, str]: """Distribute whole event groups by seeded random sampling, stratified by label and unit. Groups are the indivisible unit. A one-sample event remains one group, even when that means exact file-level target ratios are impossible. Positive and per-unit normal strata are sampled independently so every split keeps a proportional share, and every split is guaranteed at least one independent positive event group. `positive_basis` controls how positive events are balanced: "samples" weights each event by its window count, "groups" weights every event equally (useful when a few long events would otherwise dominate a split). """ grouped = defaultdict(list) for row in rows: grouped[row["event_group_id"]].append(row) positives = [] normal_by_unit = defaultdict(list) for group_id, members in grouped.items(): first = min(members, key=lambda row: row["start_hour"]) item = (group_id, members, first["start_hour"], first["unit"]) if first["label"] == "1": positives.append(item) else: normal_by_unit[first["unit"]].append(item) if len(positives) < 3: raise RuntimeError("独立异常事件少于3个，无法同时建立训练、验证和测试正例") assignments = {} rng = random.Random(seed) # Positive event groups are sampled globally: a unit with only a few events # cannot independently supply all three partitions. rng.shuffle(positives) positive_weight = (lambda members: 1) if positive_basis == "groups" else None allocate(positives, splits, assignments, weight=positive_weight) by_split = defaultdict(list) for group_id, members, _, _ in positives: by_split[assignments[group_id]].append((group_id, members)) for name, _ in splits: if by_split[name]: continue donor = max((other for other, _ in splits), key=lambda other: len(by_split[other])) group_id, members = by_split[donor].pop() assignments[group_id] = name by_split[name].append((group_id, members)) for unit in sorted(normal_by_unit): groups = normal_by_unit[unit] rng.shuffle(groups) allocate(groups, splits, assignments) if not any(assignments[group_id] == "train" for group_id, _, _, _ in groups): raise RuntimeError(f"{unit} 的正常事件组没有进入训练集，无法计算该机组基准") return assignments def write_split(dataset: Path, splits: tuple[tuple[str, float], ...], seed: int, positive_basis: str) -> dict: manifest = load_manifest(dataset / "manifest.csv") candidates = [row for row in manifest if row["label"] in {"0", "1"}] assignments = assign_groups(candidates, splits, seed, positive_basis) for split, _ in splits: target = dataset / split if target.exists(): shutil.rmtree(target) target.mkdir() for row in candidates: split = assignments[row["event_group_id"]] source = dataset / "candidates" / f"{row['sample_id']}.json" target = dataset / split / source.name target.hardlink_to(source) row["split"] = split fields = list(manifest[0]) with (dataset / "manifest.csv").open("w", encoding="utf-8-sig", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=fields) writer.writeheader() writer.writerows(manifest) summary = defaultdict(Counter) group_summary = defaultdict(Counter) group_first = {} for row in candidates: split = row["split"] summary[split][f"samples_{row['label_name']}"] += 1 summary[split]["samples_total"] += 1 group_first.setdefault(row["event_group_id"], row) for row in group_first.values(): group_summary[row["split"]][f"groups_{row['label_name']}"] += 1 group_summary[row["split"]]["groups_total"] += 1 return { "dataset": dataset.name, "sample_counts": {split: dict(summary[split]) for split, _ in splits}, "group_counts": {split: dict(group_summary[split]) for split, _ in splits}, "excluded_samples": sum(row["label"] == "" for row in manifest), "target_ratios": {name: round(ratio, 4) for name, ratio in splits}, "seed": seed, "positive_event_rule": "positive event groups are sampled globally with a fixed seed; every split keeps at least one independent positive event", "rule": "event_group_id is never split across train, validation, and test", } def main() -> None: parser = argparse.ArgumentParser(description="按完整事件组切分 PKS 模型样本") parser.add_argument("--datasets-dir", default=str(Path(__file__).resolve().parent / "datasets")) parser.add_argument("--model-types", nargs="+", choices=("pressure", "vibration"), default=["pressure", "vibration"]) parser.add_argument( "--split-ratios", nargs=3, type=float, default=[75.0, 15.0, 10.0], metavar=("TRAIN", "VALIDATION", "TEST"), help="训练/验证/测试目标比例，任意正数，内部按和归一化（默认 75 15 10）", ) parser.add_argument("--seed", type=int, default=20260921, help="事件组随机分配种子") parser.add_argument( "--positive-basis", choices=("samples", "groups"), default="samples", help="异常事件分配依据：samples 按窗口数（默认），groups 每个事件等权", ) args = parser.parse_args() if any(value <= 0 for value in args.split_ratios): parser.error("切分比例必须为正数") try: splits = normalize_splits(args.split_ratios) except ValueError as error: parser.error(str(error)) datasets = Path(args.datasets_dir) result = {name: write_split(datasets / name, splits, args.seed, args.positive_basis) for name in args.model_types} report = datasets / "split_report.json" report.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8") print(json.dumps(result, ensure_ascii=False, indent=2)) print(f"切分报告：{report}") if __name__ == "__main__": main()