1"""Experimental designs: a design = list of cells (config overrides + treatment labels) × seeds.2Runs are executed with multiprocessing (classical) or sequentially with a shared LLM cache (API cells),3each writing its own run directory + manifest; outcomes are appended to a parquet table."""4from __future__ import annotations56import json7import time8from concurrent.futures import ProcessPoolExecutor, as_completed9from pathlib import Path1011import pandas as pd12import yaml1314from ..metrics.outcomes import run_outcomes15from .config import RunConfig16from .runner import run_config171819def load_design(path: str | Path) -> dict:20 with open(path) as f:21 return yaml.safe_load(f)222324def expand_cells(design: dict) -> list[dict]:25 """Design file: {base: <config path or dict>, seeds: [..] | {start, n}, cells: [{name, overrides, treatment}], factors: {...}}.26 `factors` (optional) is a dict factor -> list of {label, overrides}; the full factorial product is expanded."""27 cells = list(design.get("cells", []))28 factors = design.get("factors")29 if factors:30 import itertools31 names = list(factors)32 for combo in itertools.product(*[factors[n] for n in names]):33 ov, tr, label = {}, {}, []34 for n, lev in zip(names, combo):35 ov.update(lev.get("overrides", {}))36 tr[n] = lev["label"]37 label.append(f"{n}={lev['label']}")38 cells.append({"name": "__".join(label), "overrides": ov, "treatment": tr})39 return cells404142def seeds_of(design: dict) -> list[int]:43 s = design.get("seeds", {"start": 1, "n": 30})44 if isinstance(s, list):45 return [int(x) for x in s]46 return list(range(int(s.get("start", 1)), int(s.get("start", 1)) + int(s["n"])))474849def cell_config(design: dict, cell: dict, seed: int, base_dir: Path) -> RunConfig:50 base = design["base"]51 cfg = RunConfig.from_yaml(base_dir / base) if isinstance(base, str) else RunConfig.model_validate(base)52 ov = dict(design.get("overrides", {}))53 ov.update(cell.get("overrides", {}))54 cfg = cfg.with_overrides(**ov)55 return cfg.with_overrides(name=f"{design['name']}/{cell['name']}", seed=seed,56 treatment={**design.get("treatment", {}), **cell.get("treatment", {})})575859def _run_one(args):60 cfg_json, out_dir, cache_path = args61 cfg = RunConfig.model_validate(cfg_json)62 out = Path(out_dir)63 if (out / "outcomes.json").exists():64 with open(out / "outcomes.json") as f:65 return json.load(f)66 res = run_config(cfg, out_dir=out, cache_path=cache_path)67 oc = run_outcomes(res)68 oc["run_dir"] = str(out)69 with open(out / "outcomes.json", "w") as f:70 json.dump(oc, f, indent=1, default=float)71 return oc727374def run_design(design_path: str | Path, results_root: str | Path = "results/runs", workers: int = 4,75 cache_path: str | None = None, only_cells: list[str] | None = None, max_seeds: int | None = None,76 progress: bool = True) -> pd.DataFrame:77 design_path = Path(design_path)78 design = load_design(design_path)79 cells = expand_cells(design)80 if only_cells:81 cells = [c for c in cells if c["name"] in only_cells]82 seeds = seeds_of(design)83 if max_seeds:84 seeds = seeds[:max_seeds]85 root = Path(results_root) / design["name"]86 root.mkdir(parents=True, exist_ok=True)87 jobs = []88 for c in cells:89 for s in seeds:90 cfg = cell_config(design, c, s, design_path.parent)91 out = root / c["name"] / f"seed{s:04d}"92 jobs.append((cfg.model_dump(mode="json"), str(out), cache_path))93 t0 = time.perf_counter()94 rows = []95 if workers <= 1:96 for i, j in enumerate(jobs):97 rows.append(_run_one(j))98 if progress:99 print(f" [{i + 1}/{len(jobs)}] {j[1]} ({time.perf_counter() - t0:.0f}s)", flush=True)100 else:101 with ProcessPoolExecutor(max_workers=workers) as ex:102 futs = {ex.submit(_run_one, j): j for j in jobs}103 for i, f in enumerate(as_completed(futs)):104 rows.append(f.result())105 if progress and (i + 1) % 10 == 0:106 print(f" [{i + 1}/{len(jobs)}] done ({time.perf_counter() - t0:.0f}s)", flush=True)107 df = pd.DataFrame(rows)108 df.to_parquet(root / "outcomes.parquet", index=False)109 with open(root / "design.json", "w") as f:110 json.dump({"design": design, "cells": cells, "seeds": seeds, "n_runs": len(jobs), "wall_s": time.perf_counter() - t0}, f, indent=1, default=str)111 return df112