code/scripts/export_training_manifests.py

#!/usr/bin/env python3
"""Export exact training-key manifests for the manuscript experiments.

The manifests contain dataset keys, labels, SNR thresholds, and random seeds
needed to reconstruct the training pools from the public datasets. They do not
contain raw waveform or dispersion arrays.
"""

from __future__ import annotations

import argparse
import csv
import gzip
import hashlib
import json
import math
import sys
from pathlib import Path
from typing import Any

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from scripts.disp_snr_transfer_experiment import matched_threshold_subsets
from scripts.reproduce_paper_stats import PHASE_TO_GROUP, PhasePick, Record
from scripts.snr_transfer_phase_balanced_experiment import (
    PHASE_COMPOSITION_ORDER,
    collect_phase_snr,
    filter_records_by_any_thresholds,
    filter_records_by_complete_thresholds,
    filter_records_by_phase_thresholds,
    matched_records_by_phase_composition,
    matched_records,
    phase_composition_counts,
    phase_counts,
    phase_snr_key,
    record_key,
)


def read_json(path: Path) -> Any:
    opener = gzip.open if path.suffix == ".gz" else open
    with opener(path, "rt", encoding="utf-8") as handle:
        return json.load(handle)


def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    opener = gzip.open if path.suffix == ".gz" else open
    with opener(path, "wt", encoding="utf-8") as handle:
        json.dump(payload, handle, ensure_ascii=False, indent=2)


def write_jsonl_gz(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with gzip.open(path, "wt", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")


def sha256_text(values: list[str]) -> str:
    h = hashlib.sha256()
    for value in values:
        h.update(value.encode("utf-8"))
        h.update(b"\n")
    return h.hexdigest()


def record_from_dict(row: dict[str, Any]) -> Record:
    phases = tuple(
        PhasePick(phase=str(p["phase"]), index=int(p["index"]), source=str(p.get("source", "")))
        for p in row.get("phases", [])
    )
    return Record(
        event=str(row["event"]),
        station=str(row["station"]),
        length=int(row["length"]),
        delta=float(row["delta"]),
        distance_km=float(row["distance_km"]),
        phases=phases,
    )


def record_to_manifest_row(record: Record, *, seed: int, subset_slug: str) -> dict[str, Any]:
    phases = [
        {"phase": pick.phase, "index": pick.index, "source": pick.source}
        for pick in record.phases
    ]
    phase_groups = [PHASE_TO_GROUP[pick.phase] for pick in record.phases]
    return {
        "seed": seed,
        "subset_slug": subset_slug,
        "record_key": record_key(record),
        "event": record.event,
        "station": record.station,
        "length": record.length,
        "delta": record.delta,
        "distance_km": record.distance_km,
        "phase_count": len(record.phases),
        "phase_groups": "".join(phase_groups),
        "phases": phases,
    }


def export_phase(args: argparse.Namespace) -> None:
    raw = read_json(args.phase_records_json)
    records = [record_from_dict(row) for row in raw["records"]]
    phase_snr = {str(k): float(v) for k, v in read_json(args.phase_snr_json).items()}
    by_group = collect_phase_snr(records, phase_snr)

    if args.phase_s_threshold_mode == "same-as-p":
        s5 = 5.0
        s10 = 10.0
        p5_label = "P/S>=5 dB record-any"
        p10_label = "P/S>=10 dB record-any"
    else:
        # This branch is preserved for auditability, but the manuscript archive
        # uses same-as-P thresholds for the training-manifest export.
        from scripts.snr_transfer_phase_balanced_experiment import choose_s_threshold

        target_p5 = sum(v >= 5.0 for v in by_group["P"])
        target_p10 = sum(v >= 10.0 for v in by_group["P"])
        s5, kept_s5 = choose_s_threshold(by_group["S"], target_p5)
        s10, kept_s10 = choose_s_threshold(by_group["S"], target_p10)
        p5_label = f"P>=5 dB, S>={s5:.2f} dB phase-balanced"
        p10_label = f"P>=10 dB, S>={s10:.2f} dB phase-balanced"

    subset_defs = [
        ("full", "Full", None, None),
        ("p5_s_bal", p5_label, 5.0, s5),
        ("p10_s_bal", p10_label, 10.0, s10),
    ]
    filter_fn = {
        "phase-label": filter_records_by_phase_thresholds,
        "record-complete": filter_records_by_complete_thresholds,
        "record-any": filter_records_by_any_thresholds,
    }[args.phase_filter_mode]
    candidate_records = {
        slug: filter_fn(records, phase_snr, p_thr, s_thr)
        for slug, _label, p_thr, s_thr in subset_defs
    }
    candidate_phase_composition = {
        slug: phase_composition_counts(pool) for slug, pool in candidate_records.items()
    }
    candidate_phase_counts = {slug: phase_counts(pool) for slug, pool in candidate_records.items()}

    if args.phase_match_mode == "phase-composition":
        matched_composition = {
            key: min(counts[key] for counts in candidate_phase_composition.values())
            for key in PHASE_COMPOSITION_ORDER
        }
        train_pools_by_seed = {}
        for seed in args.seeds:
            train_pools_by_seed[seed] = {
                slug: matched_records_by_phase_composition(candidate_records[slug], matched_composition, seed + idx * 8191)
                for idx, (slug, _label, _p, _s) in enumerate(subset_defs)
            }
    else:
        matched_count = min(len(pool) for pool in candidate_records.values())
        matched_composition = None
        train_pools_by_seed = {}
        for seed in args.seeds:
            train_pools_by_seed[seed] = {
                slug: matched_records(candidate_records[slug], matched_count, seed + idx * 8191)
                for idx, (slug, _label, _p, _s) in enumerate(subset_defs)
            }

    out_dir = args.out_dir / "phase_picker"
    out_dir.mkdir(parents=True, exist_ok=True)
    summary_rows: list[dict[str, Any]] = []
    for seed, pools in train_pools_by_seed.items():
        for slug, records_selected in pools.items():
            rows = [record_to_manifest_row(record, seed=seed, subset_slug=slug) for record in records_selected]
            out_path = out_dir / f"seed{seed}_{slug}_train_records.jsonl.gz"
            write_jsonl_gz(out_path, rows)
            keys = [row["record_key"] for row in rows]
            phase_counts_selected = phase_counts(records_selected)
            comp_counts_selected = phase_composition_counts(records_selected)
            summary_rows.append(
                {
                    "task": "phase_picker",
                    "seed": seed,
                    "subset_slug": slug,
                    "manifest": str(out_path.relative_to(args.out_dir)),
                    "records": len(rows),
                    "sha256_record_keys": sha256_text(keys),
                    "P_labels": phase_counts_selected["P"],
                    "S_labels": phase_counts_selected["S"],
                    "P_only_records": comp_counts_selected["P_only"],
                    "PS_records": comp_counts_selected["PS"],
                    "S_only_records": comp_counts_selected["S_only"],
                    "none_records": comp_counts_selected["none"],
                }
            )

    with (out_dir / "phase_training_manifest_summary.csv").open("w", newline="", encoding="utf-8") as handle:
        fieldnames = [
            "task",
            "seed",
            "subset_slug",
            "manifest",
            "records",
            "sha256_record_keys",
            "P_labels",
            "S_labels",
            "P_only_records",
            "PS_records",
            "S_only_records",
            "none_records",
        ]
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(summary_rows)

    config = {
        "task": "phase_picker",
        "public_dataset": "CREDIT-X1local",
        "source_records": str(args.phase_records_json),
        "source_phase_snr": str(args.phase_snr_json),
        "seeds": args.seeds,
        "filter_mode": args.phase_filter_mode,
        "s_threshold_mode": args.phase_s_threshold_mode,
        "match_mode": args.phase_match_mode,
        "subset_definitions": [
            {"slug": slug, "label": label, "p_threshold_db": p_thr, "s_threshold_db": s_thr}
            for slug, label, p_thr, s_thr in subset_defs
        ],
        "candidate_records": {slug: len(pool) for slug, pool in candidate_records.items()},
        "candidate_phase_counts": candidate_phase_counts,
        "candidate_phase_composition_counts": candidate_phase_composition,
        "matched_phase_composition": matched_composition,
        "snr_summary_db": {
            group: {
                "count": len(values),
                "median": float(np.median(values)) if values else math.nan,
                "p25": float(np.percentile(values, 25)) if values else math.nan,
                "p75": float(np.percentile(values, 75)) if values else math.nan,
            }
            for group, values in by_group.items()
        },
    }
    write_json(out_dir / "phase_training_manifest_config.json", config)
    print(f"Wrote phase training manifests to {out_dir}")


def export_dispersion(args: argparse.Namespace) -> None:
    rows = read_json(args.dispersion_snr_json)
    out_dir = args.out_dir / "dispersion"
    out_dir.mkdir(parents=True, exist_ok=True)
    summary_rows: list[dict[str, Any]] = []
    config = {
        "task": "dispersion",
        "public_dataset": "SeisDispFusion-NCF",
        "source_snr_cache": str(args.dispersion_snr_json),
        "seeds": args.seeds,
        "subsets": {},
    }
    for seed in args.seeds:
        subsets, thresholds = matched_threshold_subsets(rows, seed)
        config["subsets"][str(seed)] = {"thresholds": thresholds}
        for slug, keys in subsets.items():
            out_path = out_dir / f"seed{seed}_{slug}_train_keys.txt"
            out_path.write_text("\n".join(keys) + "\n", encoding="utf-8")
            snrs = [float(rows[key]["snr_db"]) for key in keys]
            summary_rows.append(
                {
                    "task": "dispersion",
                    "seed": seed,
                    "subset_slug": slug,
                    "manifest": str(out_path.relative_to(args.out_dir)),
                    "records": len(keys),
                    "sha256_keys": sha256_text(keys),
                    "snr_min": min(snrs),
                    "snr_median": float(np.median(snrs)),
                    "snr_max": max(snrs),
                    "q1_threshold": thresholds["q1"],
                    "q2_threshold": thresholds["q2"],
                }
            )
    with (out_dir / "dispersion_training_manifest_summary.csv").open("w", newline="", encoding="utf-8") as handle:
        fieldnames = [
            "task",
            "seed",
            "subset_slug",
            "manifest",
            "records",
            "sha256_keys",
            "snr_min",
            "snr_median",
            "snr_max",
            "q1_threshold",
            "q2_threshold",
        ]
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(summary_rows)
    write_json(out_dir / "dispersion_training_manifest_config.json", config)
    print(f"Wrote dispersion training manifests to {out_dir}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out-dir", type=Path, default=Path("training_manifests"))
    parser.add_argument("--seeds", nargs="+", type=int, default=[20260609, 20260610, 20260611])
    parser.add_argument("--phase-records-json", type=Path, default=None)
    parser.add_argument("--phase-snr-json", type=Path, default=None)
    parser.add_argument("--phase-filter-mode", choices=["phase-label", "record-complete", "record-any"], default="record-any")
    parser.add_argument("--phase-s-threshold-mode", choices=["balanced", "same-as-p"], default="same-as-p")
    parser.add_argument("--phase-match-mode", choices=["record-count", "phase-composition"], default="phase-composition")
    parser.add_argument("--dispersion-snr-json", type=Path, default=None)
    args = parser.parse_args()

    args.out_dir.mkdir(parents=True, exist_ok=True)
    if args.phase_records_json and args.phase_snr_json:
        export_phase(args)
    if args.dispersion_snr_json:
        export_dispersion(args)
    if not ((args.phase_records_json and args.phase_snr_json) or args.dispersion_snr_json):
        raise SystemExit("No manifests requested. Provide phase and/or dispersion source cache paths.")


if __name__ == "__main__":
    main()