"""Deterministic paired-trial plan expansion.""" from __future__ import annotations import itertools import json from pathlib import Path from typing import Any, Mapping from .hashing import stable_hash, to_jsonable from .rng import named_seed_record from .schema import ( SCHEMA_VERSION, STORAGE_FORMATS, ContractError, require_nonempty_string, require_schema, require_sequence, ) def load_document(path: Path) -> dict[str, Any]: """Load JSON, with optional YAML support when PyYAML is available.""" path = Path(path) suffix = path.suffix.lower() with path.open("r", encoding="utf-8") as stream: if suffix == ".json": document = json.load(stream) elif suffix in {".yaml", ".yml"}: try: import yaml except ImportError as exc: raise RuntimeError( "YAML input requires PyYAML; use JSON in minimal environments" ) from exc document = yaml.safe_load(stream) else: raise ValueError("experiment documents must use .json, .yaml, or .yml") if not isinstance(document, dict): raise ContractError(f"{path} must contain a mapping") return document def _normalize_methods(methods: Any) -> list[dict[str, Any]]: normalized: list[dict[str, Any]] = [] for entry in require_sequence(methods, "methods"): if isinstance(entry, str): method = {"method_id": require_nonempty_string(entry, "method")} elif isinstance(entry, Mapping): method = dict(to_jsonable(entry)) source_id = method.pop("id", method.get("method_id")) method["method_id"] = require_nonempty_string( source_id, "method.method_id" ) else: raise ContractError("each method must be a string or mapping") normalized.append(method) identifiers = [method["method_id"] for method in normalized] if len(set(identifiers)) != len(identifiers): raise ContractError("method identifiers must be unique") return normalized def _normalize_trajectories(trajectories: Any) -> list[dict[str, Any]]: normalized: list[dict[str, Any]] = [] for entry in require_sequence(trajectories, "trajectories"): if isinstance(entry, str): trajectory = { "trajectory_id": require_nonempty_string(entry, "trajectory") } elif isinstance(entry, Mapping): trajectory = dict(to_jsonable(entry)) source_id = trajectory.pop("id", trajectory.get("trajectory_id")) trajectory["trajectory_id"] = require_nonempty_string( source_id, "trajectory.trajectory_id" ) else: raise ContractError("each trajectory must be a string or mapping") normalized.append(trajectory) identifiers = [trajectory["trajectory_id"] for trajectory in normalized] if len(set(identifiers)) != len(identifiers): raise ContractError("trajectory identifiers must be unique") return normalized def _factor_cells(factors: Any) -> list[dict[str, Any]]: if factors is None: return [{}] if not isinstance(factors, Mapping): raise ContractError("factors must be a mapping from name to levels") names = sorted(factors) levels = [list(require_sequence(factors[name], f"factors.{name}")) for name in names] return [ dict(zip(names, to_jsonable(combination))) for combination in itertools.product(*levels) ] def build_trial_plan(specification: Mapping[str, Any]) -> dict[str, Any]: """Expand a study specification into immutable paired trial records.""" if not isinstance(specification, Mapping): raise ContractError("study specification must be a mapping") spec = dict(to_jsonable(specification)) study_id = require_nonempty_string(spec.get("study_id"), "study_id") split = require_nonempty_string(spec.get("split"), "split") if split not in {"calibration", "pilot", "locked"}: raise ContractError("split must be calibration, pilot, or locked") root_seed = int(spec.get("root_seed", 0)) if root_seed < 0: raise ContractError("root_seed must be non-negative") replicates = int(spec.get("replicates", 1)) if replicates <= 0: raise ContractError("replicates must be positive") methods = _normalize_methods(spec.get("methods")) trajectories = _normalize_trajectories(spec.get("trajectories")) factor_cells = _factor_cells(spec.get("factors", {})) trials: list[dict[str, Any]] = [] pair_ids: set[str] = set() for trajectory, factors, replicate in itertools.product( trajectories, factor_cells, range(replicates) ): pair_basis = { "study_id": study_id, "split": split, "trajectory": trajectory, "factors": factors, "replicate": replicate, "root_seed": root_seed, } pair_id = f"pair-{stable_hash(pair_basis, prefix='pair')[:16]}" seeds = named_seed_record(root_seed, pair_basis) pair_ids.add(pair_id) for method in methods: trial_basis = { "pair_id": pair_id, "method": method, "seeds": seeds, } trial_id = f"trial-{stable_hash(trial_basis, prefix='trial')[:16]}" trial = { "trial_index": len(trials), "trial_id": trial_id, "pair_id": pair_id, "study_id": study_id, "split": split, "method": method, "trajectory": trajectory, "factors": factors, "replicate": replicate, "seeds": seeds, } trial["trial_spec_hash"] = stable_hash(trial, prefix="trial-spec") trials.append(trial) spec_hash = stable_hash(spec, prefix="study-spec") plan = { "kind": "trial_plan", "schema_version": SCHEMA_VERSION, "study_id": study_id, "split": split, "specification": spec, "spec_hash": spec_hash, "storage_formats": dict(STORAGE_FORMATS), "pair_count": len(pair_ids), "trial_count": len(trials), "trials": trials, } plan["plan_hash"] = stable_hash(plan, prefix="trial-plan") validate_trial_plan(plan) return plan def validate_trial_plan(plan: Mapping[str, Any]) -> None: require_schema(plan, "trial_plan") trials = require_sequence(plan.get("trials"), "trials") if int(plan.get("trial_count", -1)) != len(trials): raise ContractError("trial_count does not match trials") trial_ids: set[str] = set() pairs: dict[str, dict[str, Any]] = {} for trial in trials: if not isinstance(trial, Mapping): raise ContractError("each trial must be a mapping") trial_id = require_nonempty_string(trial.get("trial_id"), "trial_id") pair_id = require_nonempty_string(trial.get("pair_id"), "pair_id") if trial_id in trial_ids: raise ContractError(f"duplicate trial_id {trial_id}") trial_ids.add(trial_id) trial_without_hash = dict(trial) recorded_trial_hash = trial_without_hash.pop("trial_spec_hash", None) expected_trial_hash = stable_hash(trial_without_hash, prefix="trial-spec") if recorded_trial_hash != expected_trial_hash: raise ContractError(f"trial_spec_hash mismatch for {trial_id}") pair_signature = { "trajectory": trial.get("trajectory"), "factors": trial.get("factors"), "replicate": trial.get("replicate"), "seeds": trial.get("seeds"), "study_id": trial.get("study_id"), "split": trial.get("split"), } previous = pairs.setdefault(pair_id, pair_signature) if previous != pair_signature: raise ContractError(f"paired inputs differ inside {pair_id}") if int(plan.get("pair_count", -1)) != len(pairs): raise ContractError("pair_count does not match unique pair identifiers") plan_without_hash = dict(plan) recorded_hash = plan_without_hash.pop("plan_hash", None) expected_hash = stable_hash(plan_without_hash, prefix="trial-plan") if recorded_hash != expected_hash: raise ContractError("plan_hash mismatch")