Compressor/models/generalized/split_datasets.py

#!/usr/bin/env python3
"""Split pressure samples by whole event groups without leakage."""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from models.common import paths, splitting


def main() -> None:
    parser = argparse.ArgumentParser(description="按完整事件组切分压力模型样本")
    parser.add_argument("--datasets-dir", default=str(paths.PRESSURE_DATAS))
    parser.add_argument("--split-ratios", nargs=3, type=float, default=[65.0, 20.0, 15.0], metavar=("TRAIN", "VALIDATION", "TEST"))
    parser.add_argument("--seed", type=int, default=20260921)
    parser.add_argument("--positive-basis", choices=("samples", "groups"), default="samples")
    args = parser.parse_args()
    if any(value <= 0 for value in args.split_ratios):
        parser.error("切分比例必须为正数")
    splits = splitting.normalize_splits(args.split_ratios)
    dataset = Path(args.datasets_dir)
    result = splitting.write_split(dataset, splits, args.seed, args.positive_basis)
    report = dataset / "split_report.json"
    report.write_text(json.dumps({"pressure": result}, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps({"pressure": result}, ensure_ascii=False, indent=2))
    print(f"切分报告：{report}")


if __name__ == "__main__":
    main()