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() and (out / "manifest.json").exists():64 with open(out / "manifest.json") as f:65 man = json.load(f)66 if man.get("config_hash") == cfg.config_hash(): # resume only an identical configuration67 with open(out / "outcomes.json") as f:68 return json.load(f)69 import shutil70 shutil.rmtree(out)71 res = run_config(cfg, out_dir=out, cache_path=cache_path)72 oc = run_outcomes(res)73 oc["run_dir"] = str(out)74 with open(out / "outcomes.json", "w") as f:75 json.dump(oc, f, indent=1, default=float)76 return oc777879def run_design(design_path: str | Path, results_root: str | Path = "results/runs", workers: int = 4,80 cache_path: str | None = None, only_cells: list[str] | None = None, max_seeds: int | None = None,81 progress: bool = True) -> pd.DataFrame:82 design_path = Path(design_path)83 design = load_design(design_path)84 cells = expand_cells(design)85 if only_cells:86 cells = [c for c in cells if c["name"] in only_cells]87 seeds = seeds_of(design)88 if max_seeds:89 seeds = seeds[:max_seeds]90 root = Path(results_root) / design["name"]91 root.mkdir(parents=True, exist_ok=True)92 jobs = []93 for c in cells:94 for s in seeds:95 cfg = cell_config(design, c, s, design_path.parent)96 out = root / c["name"] / f"seed{s:04d}"97 jobs.append((cfg.model_dump(mode="json"), str(out), cache_path))98 t0 = time.perf_counter()99 rows = []100 if workers <= 1:101 for i, j in enumerate(jobs):102 rows.append(_run_one(j))103 if progress:104 print(f" [{i + 1}/{len(jobs)}] {j[1]} ({time.perf_counter() - t0:.0f}s)", flush=True)105 else:106 with ProcessPoolExecutor(max_workers=workers) as ex:107 futs = {ex.submit(_run_one, j): j for j in jobs}108 for i, f in enumerate(as_completed(futs)):109 rows.append(f.result())110 if progress and (i + 1) % 10 == 0:111 print(f" [{i + 1}/{len(jobs)}] done ({time.perf_counter() - t0:.0f}s)", flush=True)112 df = pd.DataFrame(rows)113 df.to_parquet(root / "outcomes.parquet", index=False)114 with open(root / "design.json", "w") as f:115 json.dump({"design": design, "cells": cells, "seeds": seeds, "n_runs": len(jobs), "wall_s": time.perf_counter() - t0}, f, indent=1, default=str)116 return df117