"""Strict experiment configuration loading and validation.""" from __future__ import annotations import math from dataclasses import dataclass from pathlib import Path from typing import Any, Literal import yaml from mcss.model.system import ModelConfig DatasetName = Literal["synthetic", "manifest", "replica", "hypersim"] SpatialProtocol = Literal["legacy_manifest_bounds", "context_local_metric"] Bounds3D = tuple[tuple[float, float, float], tuple[float, float, float]] @dataclass(frozen=True) class DatasetConfig: name: DatasetName root: str | None image_size: tuple[int, int] context_views: int target_views: int length: int = 64 sample_stride: int = 1 spatial_protocol: SpatialProtocol = "legacy_manifest_bounds" local_bounds_m: Bounds3D | None = None scene_ids: tuple[str, ...] | None = None @dataclass(frozen=True) class TrainingConfig: output_dir: str batch_size: int learning_rate: float max_steps: int num_workers: int amp: bool amp_dtype: Literal["float16", "bfloat16"] log_every: int checkpoint_every: int measurements: tuple[str, ...] gradient_accumulation: int = 1 weight_decay: float = 0.0 resume: str | None = None lr_schedule: Literal["constant", "warmup_cosine"] = "constant" warmup_steps: int = 0 min_lr_ratio: float = 0.05 ema_decay: float = 0.98 evidence_residual_weight: float = 0.0 context_subset_geometry_weight: float = 0.0 @dataclass(frozen=True) class ExperimentConfig: seed: int device: str dataset: DatasetConfig model: ModelConfig training: TrainingConfig def __post_init__(self) -> None: if self.training.context_subset_geometry_weight == 0.0: return if self.dataset.context_views != 4: raise ValueError( "context_subset_geometry_weight requires exactly four context views" ) if self.model.state_architecture != "evidence_residual": raise ValueError( "context_subset_geometry_weight requires state_architecture: evidence_residual" ) if self.model.mode != "fixed": raise ValueError("context_subset_geometry_weight requires model.mode: fixed") def load_config(path: str | Path) -> ExperimentConfig: config_path = Path(path) try: raw = yaml.safe_load(config_path.read_text(encoding="utf-8")) except yaml.YAMLError as error: raise ValueError(f"invalid YAML in {config_path}") from error if not isinstance(raw, dict): raise ValueError("config root must be a mapping") _reject_unknown(raw, {"seed", "device", "dataset", "model", "training"}, "config") _require_keys(raw, {"seed", "device", "dataset", "model", "training"}, "config") if not isinstance(raw["seed"], int): raise ValueError("seed must be an integer") if not isinstance(raw["device"], str): raise ValueError("device must be a string") return ExperimentConfig( seed=raw["seed"], device=raw["device"], dataset=_parse_dataset(raw["dataset"]), model=_parse_model(raw["model"]), training=_parse_training(raw["training"]), ) def _parse_dataset(value: Any) -> DatasetConfig: if not isinstance(value, dict): raise ValueError("dataset must be a mapping") allowed = { "name", "root", "image_size", "context_views", "target_views", "length", "sample_stride", "spatial_protocol", "local_bounds_m", "scene_ids", } _reject_unknown(value, allowed, "dataset") _require_keys( value, {"name", "root", "image_size", "context_views", "target_views", "length"}, "dataset" ) name = value["name"] if name not in {"synthetic", "manifest", "replica", "hypersim"}: raise ValueError("dataset.name must be synthetic, manifest, replica, or hypersim") root = value["root"] if root is not None and not isinstance(root, str): raise ValueError("dataset.root must be a string or null") if name != "synthetic" and not root: raise ValueError("prepared real datasets require dataset.root") image_size = _parse_image_size(value["image_size"]) context_views = _positive_int(value["context_views"], "dataset.context_views") target_views = _positive_int(value["target_views"], "dataset.target_views") length = _positive_int(value["length"], "dataset.length") sample_stride = _positive_int(value.get("sample_stride", 1), "dataset.sample_stride") spatial_protocol = value.get("spatial_protocol", "legacy_manifest_bounds") if spatial_protocol not in {"legacy_manifest_bounds", "context_local_metric"}: raise ValueError( "dataset.spatial_protocol must be legacy_manifest_bounds or context_local_metric" ) raw_local_bounds = value.get("local_bounds_m") local_bounds_m = None if raw_local_bounds is None else _parse_bounds(raw_local_bounds) if spatial_protocol == "context_local_metric" and local_bounds_m is None: raise ValueError("dataset.local_bounds_m is required for context_local_metric") raw_scene_ids = value.get("scene_ids") scene_ids = None if raw_scene_ids is not None: if not isinstance(raw_scene_ids, list) or not raw_scene_ids: raise ValueError("dataset.scene_ids must be a non-empty list") if not all(isinstance(scene_id, str) and scene_id for scene_id in raw_scene_ids): raise ValueError("dataset.scene_ids must contain non-empty strings") if len(set(raw_scene_ids)) != len(raw_scene_ids): raise ValueError("dataset.scene_ids must contain unique values") scene_ids = tuple(raw_scene_ids) return DatasetConfig( name=name, root=root, image_size=image_size, context_views=context_views, target_views=target_views, length=length, sample_stride=sample_stride, spatial_protocol=spatial_protocol, local_bounds_m=local_bounds_m, scene_ids=scene_ids, ) def _parse_model(value: Any) -> ModelConfig: if not isinstance(value, dict): raise ValueError("model must be a mapping") allowed = { "mode", "state_architecture", "voxel_resolution", "image_feature_dim", "state_feature_dim", "refinement_blocks", "evidence_temperature", "observed_residual_floor", "completion_residual_scale", "surface_peak_temperature", "appearance_resolution_scale", "n_samples", "ray_chunk_size", } required = { "mode", "voxel_resolution", "image_feature_dim", "state_feature_dim", "refinement_blocks", "n_samples", "ray_chunk_size", } _reject_unknown(value, allowed, "model") _require_keys(value, required, "model") parsed = dict(value) parsed["voxel_resolution"] = _parse_voxel_resolution(value["voxel_resolution"]) try: return ModelConfig(**parsed) except (TypeError, ValueError) as error: raise ValueError(f"invalid model config: {error}") from error def _parse_training(value: Any) -> TrainingConfig: if not isinstance(value, dict): raise ValueError("training must be a mapping") allowed = { "output_dir", "batch_size", "learning_rate", "max_steps", "num_workers", "amp", "amp_dtype", "log_every", "checkpoint_every", "measurements", "gradient_accumulation", "weight_decay", "resume", "lr_schedule", "warmup_steps", "min_lr_ratio", "ema_decay", "evidence_residual_weight", "context_subset_geometry_weight", } required = { "output_dir", "batch_size", "learning_rate", "max_steps", "num_workers", "amp", "log_every", "checkpoint_every", "measurements", } _reject_unknown(value, allowed, "training") _require_keys(value, required, "training") if not isinstance(value["output_dir"], str) or not value["output_dir"]: raise ValueError("training.output_dir must be a non-empty string") if not isinstance(value["learning_rate"], (int, float)) or value["learning_rate"] <= 0: raise ValueError("training.learning_rate must be positive") if not isinstance(value["amp"], bool): raise ValueError("training.amp must be boolean") amp_dtype = value.get("amp_dtype", "float16") if amp_dtype not in {"float16", "bfloat16"}: raise ValueError("training.amp_dtype must be float16 or bfloat16") if not isinstance(value["measurements"], list) or not all( isinstance(name, str) for name in value["measurements"] ): raise ValueError("training.measurements must be a list of strings") resume = value.get("resume") if resume is not None and not isinstance(resume, str): raise ValueError("training.resume must be a string or null") max_steps = _positive_int(value["max_steps"], "training.max_steps") gradient_accumulation = _positive_int( value.get("gradient_accumulation", 1), "training.gradient_accumulation" ) if max_steps % gradient_accumulation != 0: raise ValueError("training.max_steps must be divisible by gradient_accumulation") lr_schedule = value.get("lr_schedule", "constant") if lr_schedule not in {"constant", "warmup_cosine"}: raise ValueError("training.lr_schedule must be constant or warmup_cosine") warmup_steps = _nonnegative_int(value.get("warmup_steps", 0), "training.warmup_steps") total_updates = max_steps // gradient_accumulation if warmup_steps > total_updates: raise ValueError("training.warmup_steps cannot exceed optimizer updates") if lr_schedule == "constant" and warmup_steps: raise ValueError("training.warmup_steps requires lr_schedule: warmup_cosine") min_lr_ratio = value.get("min_lr_ratio", 0.05) if not isinstance(min_lr_ratio, (int, float)) or not 0.0 <= min_lr_ratio <= 1.0: raise ValueError("training.min_lr_ratio must be between zero and one") ema_decay = value.get("ema_decay", 0.98) if not isinstance(ema_decay, (int, float)) or not 0.0 <= ema_decay < 1.0: raise ValueError("training.ema_decay must be in [0, 1)") evidence_residual_weight = value.get("evidence_residual_weight", 0.0) if ( not isinstance(evidence_residual_weight, (int, float)) or evidence_residual_weight < 0.0 ): raise ValueError("training.evidence_residual_weight must be non-negative") context_subset_geometry_weight = value.get("context_subset_geometry_weight", 0.0) if ( isinstance(context_subset_geometry_weight, bool) or not isinstance(context_subset_geometry_weight, (int, float)) or not math.isfinite(context_subset_geometry_weight) ): raise ValueError( "training.context_subset_geometry_weight must be a finite non-negative number" ) if context_subset_geometry_weight < 0.0: raise ValueError("training.context_subset_geometry_weight must be non-negative") checkpoint_every = _positive_int( value["checkpoint_every"], "training.checkpoint_every" ) if ( context_subset_geometry_weight > 0.0 and checkpoint_every % gradient_accumulation != 0 ): raise ValueError( "training.checkpoint_every must be divisible by gradient_accumulation " "when context_subset_geometry_weight is enabled" ) return TrainingConfig( output_dir=value["output_dir"], batch_size=_positive_int(value["batch_size"], "training.batch_size"), learning_rate=float(value["learning_rate"]), max_steps=max_steps, num_workers=_nonnegative_int(value["num_workers"], "training.num_workers"), amp=value["amp"], amp_dtype=amp_dtype, log_every=_positive_int(value["log_every"], "training.log_every"), checkpoint_every=checkpoint_every, measurements=tuple(value["measurements"]), gradient_accumulation=gradient_accumulation, weight_decay=float(value.get("weight_decay", 0.0)), resume=resume, lr_schedule=lr_schedule, warmup_steps=warmup_steps, min_lr_ratio=float(min_lr_ratio), ema_decay=float(ema_decay), evidence_residual_weight=float(evidence_residual_weight), context_subset_geometry_weight=float(context_subset_geometry_weight), ) def _parse_image_size(value: Any) -> tuple[int, int]: if not isinstance(value, list) or len(value) != 2: raise ValueError("dataset.image_size must be [height, width]") return ( _positive_int(value[0], "dataset.image_size[0]"), _positive_int(value[1], "dataset.image_size[1]"), ) def _parse_voxel_resolution(value: Any) -> tuple[int, int, int]: if isinstance(value, int): resolution = (value, value, value) elif isinstance(value, list) and len(value) == 3: resolution = tuple( _positive_int(axis, f"model.voxel_resolution[{index}]") for index, axis in enumerate(value) ) else: raise ValueError("model.voxel_resolution must be an integer or [D, H, W]") if min(resolution) < 4: raise ValueError("model.voxel_resolution values must be at least four") return resolution def _parse_bounds(value: Any) -> Bounds3D: if not isinstance(value, list) or len(value) != 2: raise ValueError("dataset.local_bounds_m must be [[minx, miny, minz], [maxx, maxy, maxz]]") rows: list[tuple[float, float, float]] = [] for row_index, row in enumerate(value): if not isinstance(row, list) or len(row) != 3: raise ValueError(f"dataset.local_bounds_m[{row_index}] must contain three values") if not all(isinstance(axis, (int, float)) and math.isfinite(axis) for axis in row): raise ValueError("dataset.local_bounds_m values must be finite numbers") rows.append(tuple(float(axis) for axis in row)) # type: ignore[arg-type] if any(upper <= lower for lower, upper in zip(rows[0], rows[1], strict=True)): raise ValueError("dataset.local_bounds_m maximum must exceed minimum on every axis") return rows[0], rows[1] def _positive_int(value: Any, name: str) -> int: if not isinstance(value, int) or value < 1: raise ValueError(f"{name} must be a positive integer") return value def _nonnegative_int(value: Any, name: str) -> int: if not isinstance(value, int) or value < 0: raise ValueError(f"{name} must be a non-negative integer") return value def _reject_unknown(value: dict[str, Any], allowed: set[str], name: str) -> None: unknown = set(value) - allowed if unknown: raise ValueError(f"unknown {name} keys: {sorted(unknown)}") def _require_keys(value: dict[str, Any], required: set[str], name: str) -> None: missing = required - set(value) if missing: raise ValueError(f"missing {name} keys: {sorted(missing)}")