Compressor/backend/model_assistant_generalized/split_datasets.py

#!/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()