scripts/build_wisemoor_training_records.py

#!/usr/bin/env python3
"""Build empirical-format PINO records from resolved WiseMoor candidates."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "scripts"))

from generate_empirical_bootstrap import (  # noqa: E402
    DURATION_SECONDS,
    INTERVAL_SECONDS,
    _build_record,
    apply_commercial_solvent_dilution,
    expand_naturals_to_cas_ratios,
)
from pino.registry import AromaRegistry  # noqa: E402
from pino.verifier import FragrancePipelineVerifier  # noqa: E402


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, default=Path("data/wisemoor_training_candidates.jsonl"))
    parser.add_argument("--output", type=Path, default=Path("data/wisemoor_training_records.jsonl"))
    parser.add_argument("--reject-output", type=Path, default=Path("artifacts/wisemoor_training_record_rejections.json"))
    parser.add_argument("--concentrate-ratio", type=float, default=0.15)
    return parser.parse_args()


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    with path.open(encoding="utf-8") as handle:
        return [json.loads(line) for line in handle if line.strip()]


def enrich_formula_names(formula: list[dict[str, Any]], registry: AromaRegistry) -> list[dict[str, Any]]:
    records = registry.all_records()
    enriched = []
    for item in formula:
        row = records.get(str(item.get("cas")))
        enriched.append(
            {
                "cas": item["cas"],
                "name": item.get("name") or (row or {}).get("name") or item["cas"],
                "smiles": (row or {}).get("smiles", ""),
                "weight_fraction": float(item["weight_fraction"]),
            }
        )
    return enriched


def main() -> None:
    args = parse_args()
    registry = AromaRegistry()
    verifier = FragrancePipelineVerifier()

    records: list[dict[str, Any]] = []
    rejections: list[dict[str, Any]] = []

    for candidate in load_jsonl(args.input):
        formula = enrich_formula_names(candidate["formula"], registry)
        diluted = apply_commercial_solvent_dilution(formula, concentrate_ratio=args.concentrate_ratio)
        expanded_ratios = expand_naturals_to_cas_ratios(diluted)
        if not expanded_ratios:
            rejections.append({"formula_id": candidate["formula_id"], "reason": "empty_expanded_formula"})
            continue

        expanded_formula = [{"cas": cas, "weight_fraction": wf} for cas, wf in expanded_ratios.items()]
        try:
            result = verifier.run_sim(
                expanded_formula,
                duration_seconds=DURATION_SECONDS,
                interval_seconds=INTERVAL_SECONDS,
                skip_ifra=True,
            )
        except Exception as exc:
            rejections.append({"formula_id": candidate["formula_id"], "reason": f"simulation_error: {exc}"})
            continue

        if result.get("status") not in {"passed", "depleted"}:
            rejections.append(
                {
                    "formula_id": candidate["formula_id"],
                    "reason": f"simulation_rejected: {result.get('message')}",
                }
            )
            continue

        trajectory = result.get("trajectory") or []
        if not trajectory:
            rejections.append({"formula_id": candidate["formula_id"], "reason": "empty_trajectory"})
            continue

        composition = [{"cas": cas} for cas in expanded_ratios.keys()]
        record = _build_record(
            formula_id=candidate["formula_id"],
            formula_list=formula,
            trajectory=trajectory,
            composition=composition,
            strategy="wildcard",
            is_control=False,
        )
        record["metadata"].update(
            {
                "generation_strategy": "wildcard",
                "source": candidate["source"],
                "source_url": candidate.get("source_url"),
                "source_name": candidate.get("name"),
                "source_description": candidate.get("description"),
                "source_notes": candidate.get("notes"),
                "source_description_provenance": (
                    (candidate.get("metadata") or {}).get("description_provenance")
                ),
                "wisemoor": candidate.get("metadata", {}),
                "concentrate_ratio": args.concentrate_ratio,
            }
        )
        records.append(record)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")

    args.reject_output.parent.mkdir(parents=True, exist_ok=True)
    args.reject_output.write_text(
        json.dumps(
            {
                "input": str(args.input),
                "output": str(args.output),
                "records_written": len(records),
                "records_rejected": len(rejections),
                "rejections": rejections,
            },
            ensure_ascii=False,
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )
    registry.close()
    print(json.dumps({"records_written": len(records), "records_rejected": len(rejections)}, indent=2))


if __name__ == "__main__":
    main()