213 lines
8.3 KiB
Python
213 lines
8.3 KiB
Python
|
"""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")
|