"""Generate versioned trial metrics and paper source-data tables.""" from __future__ import annotations import argparse import csv import json import os import tempfile from pathlib import Path from typing import Any, Mapping, Sequence import numpy as np from experiments.hashing import file_sha256, stable_hash, to_jsonable from experiments.io import atomic_write_json, atomic_write_jsonl, utc_now from experiments.plan import load_document from experiments.schema import SCHEMA_VERSION, STORAGE_FORMATS from experiments.validate import require_valid_batch from .metrics import derive_trial_metrics PROVENANCE_IDENTITY_FIELDS = ( "h2_data_group_id", "h2_data_seed_record_hash", "bilateral_data_group_id", "bilateral_data_seed_record_hash", "network_pair_group_id", "network_profile_id", ) def _method_selected( family: str, family_configuration: Mapping[str, Any], method_id: str, ) -> bool: """Apply explicit per-family method filters without silent typos.""" include = family_configuration.get("include_methods") exclude = family_configuration.get("exclude_methods") for name, value in ( ("include_methods", include), ("exclude_methods", exclude), ): if value is not None and ( not isinstance(value, list) or any(not isinstance(item, str) or not item for item in value) ): raise ValueError(f"{family}.{name} must be a list of method IDs") if include is not None and method_id not in include: return False if exclude is not None and method_id in exclude: return False return True def _trial_metric_configuration( configuration: Mapping[str, Any], method_id: str, ) -> dict[str, Any] | None: enabled = configuration.get("enabled") if not isinstance(enabled, list) or not enabled: raise ValueError("metric configuration needs a non-empty enabled list") selected = [] trial_configuration: dict[str, Any] = {} for family in enabled: if not isinstance(family, str) or not family: raise ValueError("enabled metric family names must be non-empty strings") family_configuration = configuration.get(family, {}) if not isinstance(family_configuration, Mapping): raise ValueError(f"{family} metric configuration must be a mapping") if _method_selected(family, family_configuration, method_id): selected.append(family) trial_configuration[family] = dict(family_configuration) if not selected: return None trial_configuration["enabled"] = selected return trial_configuration def _scalar_cell(value: Any) -> Any: if value is None or isinstance(value, (str, int, float, bool)): return value return json.dumps( to_jsonable(value), sort_keys=True, separators=(",", ":"), ensure_ascii=False, ) def _identity_row( trial: Mapping[str, Any], trial_metadata: Mapping[str, Any] | None = None, ) -> dict[str, Any]: trajectory = trial["trajectory"] method = trial["method"] row = { "trial_id": trial["trial_id"], "pair_id": trial["pair_id"], "study_id": trial["study_id"], "split": trial["split"], "method_id": method["method_id"], "trajectory_id": trajectory["trajectory_id"], "trajectory_family": trajectory.get("family", ""), "replicate": trial["replicate"], } for name, value in sorted(trial.get("factors", {}).items()): row[f"factor_{name}"] = _scalar_cell(value) metadata = {} if trial_metadata is None else trial_metadata for name in PROVENANCE_IDENTITY_FIELDS: value = metadata.get(name) if value is not None: row[name] = _scalar_cell(value) if ( "bilateral_data_seed_record_hash" not in row and metadata.get("bilateral_data_seed_record") is not None ): row["bilateral_data_seed_record_hash"] = stable_hash( metadata["bilateral_data_seed_record"], prefix="bilateral-data-seed-record", ) if "network_profile_id" not in row: profile = trial.get("factors", {}).get("network_profile") if isinstance(profile, Mapping) and profile.get("profile_id") is not None: row["network_profile_id"] = _scalar_cell(profile["profile_id"]) return row def _atomic_write_csv(path: Path, rows: Sequence[Mapping[str, Any]]) -> None: path.parent.mkdir(parents=True, exist_ok=True) identity_order = [ "trial_id", "pair_id", "study_id", "split", "method_id", "trajectory_id", "trajectory_family", "replicate", *PROVENANCE_IDENTITY_FIELDS, ] all_fields = {key for row in rows for key in row} fieldnames = [name for name in identity_order if name in all_fields] fieldnames.extend(sorted(all_fields - set(fieldnames))) fd, temporary_name = tempfile.mkstemp( prefix=f".{path.name}.", suffix=".tmp", dir=path.parent ) temporary = Path(temporary_name) try: with os.fdopen(fd, "w", encoding="utf-8", newline="") as stream: writer = csv.DictWriter(stream, fieldnames=fieldnames) writer.writeheader() for row in rows: writer.writerow( {name: _scalar_cell(row.get(name)) for name in fieldnames} ) stream.flush() os.fsync(stream.fileno()) os.replace(temporary, path) except Exception: temporary.unlink(missing_ok=True) raise def _finite_metric(row: Mapping[str, Any], name: str) -> float: if name not in row: raise ValueError( f"network paired gates require metric {name!r}" ) try: value = float(row[name]) except (TypeError, ValueError) as error: raise ValueError( f"network paired metric {name!r} must be numeric" ) from error if not np.isfinite(value): raise ValueError( f"network paired metric {name!r} must be finite" ) return value def _boolean_metric(row: Mapping[str, Any], name: str) -> bool: if name not in row: raise ValueError( f"network paired gates require metric {name!r}" ) value = row[name] if isinstance(value, (bool, np.bool_)): return bool(value) if isinstance(value, (int, float, np.integer, np.floating)): numeric = float(value) if np.isfinite(numeric) and numeric in (0.0, 1.0): return bool(numeric) raise ValueError( f"network paired metric {name!r} must be boolean or 0/1" ) def _apply_network_paired_gates( rows: list[dict[str, Any]], metric_configuration: Mapping[str, Any], ) -> None: network_configuration = metric_configuration.get("network") if not isinstance(network_configuration, Mapping): return paired = network_configuration.get("paired_gates") if paired is None: return if not isinstance(paired, Mapping): raise ValueError("network.paired_gates must be a mapping") nominal_profile_id = paired.get("nominal_profile_id", "nominal") if not isinstance(nominal_profile_id, str) or not nominal_profile_id: raise ValueError( "network.paired_gates.nominal_profile_id must be non-empty" ) maximum_tracking_delta = paired.get( "maximum_tracking_rmse_delta_vs_nominal_rad", paired.get("max_tracking_delta"), ) maximum_contact_relative_change = paired.get( "maximum_abs_contact_rms_relative_change_vs_nominal", paired.get("max_abs_contact_relative_change"), ) if maximum_tracking_delta is None: raise ValueError( "network paired gates require " "maximum_tracking_rmse_delta_vs_nominal_rad" ) if maximum_contact_relative_change is None: raise ValueError( "network paired gates require " "maximum_abs_contact_rms_relative_change_vs_nominal" ) maximum_tracking_delta = float(maximum_tracking_delta) maximum_contact_relative_change = float( maximum_contact_relative_change ) contact_denominator_epsilon_N = float( paired.get("contact_rms_denominator_epsilon_N", 1e-12) ) if ( not np.isfinite(maximum_tracking_delta) or maximum_tracking_delta < 0.0 or not np.isfinite(maximum_contact_relative_change) or maximum_contact_relative_change < 0.0 or not np.isfinite(contact_denominator_epsilon_N) or contact_denominator_epsilon_N <= 0.0 ): raise ValueError( "network paired-gate thresholds must be finite/non-negative " "and the contact denominator epsilon must be positive" ) network_rows = [ row for row in rows if "network_local_full_gate_pass" in row ] groups: dict[tuple[Any, ...], list[dict[str, Any]]] = {} for row in network_rows: missing_identity = [ name for name in ("network_pair_group_id", "network_profile_id") if row.get(name) in (None, "") ] if missing_identity: raise ValueError( "network paired gates require identity fields " f"{missing_identity} for trial {row.get('trial_id')}" ) group_key = ( row["network_pair_group_id"], row["method_id"], row["trajectory_id"], row["replicate"], ) groups.setdefault(group_key, []).append(row) for group_key, group_rows in groups.items(): nominal_rows = [ row for row in group_rows if row["network_profile_id"] == nominal_profile_id ] if len(nominal_rows) != 1: reason = "missing" if not nominal_rows else "duplicate" raise ValueError( f"{reason} nominal network profile for group {group_key}; " f"expected exactly one {nominal_profile_id!r}" ) nominal = nominal_rows[0] nominal_tracking = _finite_metric( nominal, "network_slave_zero_delay_tracking_rmse_rad", ) nominal_feedback_lag = _finite_metric( nominal, "network_return_feedback_lag_rmse_Nm", ) contact_expected = _boolean_metric( nominal, "network_contact_expected", ) nominal_contact_rms: float | None = None contact_denominator: float | None = None if contact_expected: nominal_contact_rms = _finite_metric( nominal, "network_contact_force_rms_N", ) contact_denominator = max( abs(nominal_contact_rms), contact_denominator_epsilon_N, ) for row in group_rows: if ( _boolean_metric(row, "network_contact_expected") != contact_expected ): raise ValueError( "network paired gates require a consistent " f"contact condition for group {group_key}" ) if row is nominal: tracking_delta = 0.0 feedback_lag_delta = 0.0 contact_relative_change = ( 0.0 if contact_expected else None ) else: tracking_delta = ( _finite_metric( row, "network_slave_zero_delay_tracking_rmse_rad", ) - nominal_tracking ) feedback_lag_delta = ( _finite_metric( row, "network_return_feedback_lag_rmse_Nm", ) - nominal_feedback_lag ) contact_relative_change = ( ( _finite_metric( row, "network_contact_force_rms_N", ) - nominal_contact_rms ) / contact_denominator if contact_expected else None ) tracking_gate = bool( tracking_delta <= maximum_tracking_delta ) contact_gate = ( bool( abs(contact_relative_change) <= maximum_contact_relative_change ) if contact_expected else None ) paired_gate = bool( tracking_gate and (contact_gate if contact_expected else True) ) row.update( { ( "network_slave_zero_delay_tracking_rmse_" "delta_vs_nominal_rad" ): tracking_delta, ( "network_zero_delay_tracking_rmse_" "delta_vs_nominal_rad" ): tracking_delta, ( "network_return_feedback_lag_rmse_" "delta_vs_nominal_Nm" ): feedback_lag_delta, ( "network_contact_rms_relative_change_" "vs_nominal" ): contact_relative_change, ( "network_contact_force_rms_relative_change_" "vs_nominal" ): contact_relative_change, "network_paired_tracking_gate_pass": tracking_gate, "network_paired_contact_gate_pass": contact_gate, "network_paired_gate_pass": paired_gate, "network_full_gate_pass": bool( row["network_local_full_gate_pass"] and paired_gate ), } ) def generate_paper_source_data( batch_dir: Path, metric_configuration: Mapping[str, Any], *, output_root: Path | None = None, ) -> dict[str, Any]: """Recompute all configured metrics and emit traceable source-data tables.""" batch_dir = Path(batch_dir) require_valid_batch(batch_dir, require_complete=True) plan = load_document(batch_dir / "plan.json") output_root = batch_dir if output_root is None else Path(output_root) derived_dir = output_root / "derived" source_dir = output_root / "paper" / "source_data" derived_dir.mkdir(parents=True, exist_ok=True) source_dir.mkdir(parents=True, exist_ok=True) rows: list[dict[str, Any]] = [] input_files: dict[str, str] = {} for trial in plan["trials"]: trial_dir = batch_dir / "raw" / trial["trial_id"] sample_path = trial_dir / "samples.npz" trial_manifest_path = trial_dir / "trial_manifest.json" input_files[str(sample_path.relative_to(batch_dir))] = file_sha256(sample_path) input_files[ str(trial_manifest_path.relative_to(batch_dir)) ] = file_sha256(trial_manifest_path) trial_manifest = load_document(trial_manifest_path) trial_metadata = trial_manifest.get("metadata", {}) if not isinstance(trial_metadata, Mapping): raise ValueError( f"{trial_manifest_path} metadata must be a mapping" ) with np.load(sample_path, allow_pickle=False) as archive: samples = {name: archive[name] for name in archive.files} method_id = trial["method"]["method_id"] trial_configuration = _trial_metric_configuration( metric_configuration, method_id, ) metrics = ( {} if trial_configuration is None else derive_trial_metrics(samples, trial_configuration) ) rows.append( { **_identity_row(trial, trial_metadata), **metrics, } ) _apply_network_paired_gates(rows, metric_configuration) metric_path = derived_dir / "trial_metrics.jsonl" atomic_write_jsonl(metric_path, rows) table_files: dict[str, str] = {} family_row_counts: dict[str, int] = {} for family in metric_configuration["enabled"]: family_rows: list[dict[str, Any]] = [] prefix = f"{family}_" for row in rows: provenance_identity_fields = set(PROVENANCE_IDENTITY_FIELDS) if not any( key.startswith(prefix) and key not in provenance_identity_fields for key in row ): continue selected = { key: value for key, value in row.items() if key in { "trial_id", "pair_id", "study_id", "split", "method_id", "trajectory_id", "trajectory_family", "replicate", *PROVENANCE_IDENTITY_FIELDS, } or key.startswith("factor_") or key.startswith(prefix) } family_rows.append(selected) if not family_rows: raise ValueError( f"metric family {family!r} selected no methods in this plan" ) table_path = source_dir / f"{family}.csv" _atomic_write_csv(table_path, family_rows) table_files[str(table_path.relative_to(output_root))] = file_sha256(table_path) family_row_counts[family] = len(family_rows) manifest = { "kind": "paper_artifact_manifest", "schema_version": SCHEMA_VERSION, "created_utc": utc_now(), "source_batch": str(batch_dir), "source_plan_hash": plan["plan_hash"], "metric_configuration": to_jsonable(metric_configuration), "metric_configuration_hash": stable_hash( metric_configuration, prefix="metric-configuration" ), "row_count": len(rows), "family_row_counts": family_row_counts, "row_hash": stable_hash(rows, prefix="trial-metric-rows"), "storage_formats": { "trial_metrics": STORAGE_FORMATS["trial_metrics"], "paper_source_data": STORAGE_FORMATS["paper_source_data"], }, "inputs": input_files, "outputs": { str(metric_path.relative_to(output_root)): file_sha256(metric_path), **table_files, }, } manifest_path = output_root / "paper" / "artifact_manifest.json" atomic_write_json(manifest_path, manifest) return manifest def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="Recompute trial metrics and paper source-data tables" ) parser.add_argument("--batch-dir", type=Path, required=True) parser.add_argument("--metric-config", type=Path, required=True) parser.add_argument("--output-root", type=Path, default=None) return parser def main(argv: list[str] | None = None) -> int: args = build_parser().parse_args(argv) configuration = load_document(args.metric_config) manifest = generate_paper_source_data( args.batch_dir, configuration, output_root=args.output_root, ) print(json.dumps(manifest, indent=2, ensure_ascii=False, sort_keys=True)) return 0 if __name__ == "__main__": raise SystemExit(main())