#!/usr/bin/env python3 """Evolve a dataset-conditioned BERNN policy across whole biological datasets.""" from __future__ import annotations import argparse import csv import hashlib import importlib.metadata import json import os import sys import uuid from pathlib import Path import numpy as np ROOT = Path(__file__).resolve().parent.parent if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from scripts import hp_search from src.dataset_splits import load_dataset_partitions from src.evolutionary_meta import ( EvolutionConfig, PolicyShape, aggregate_dataset_scores, breed_next_population, config_dict, decode_config, genome_digest, initialize_population, normalize_meta_features, recommended_batch_size, ) from src.zero_shot_recommender.meta_features import ( META_FEATURE_NAMES, extract_meta_features, ) from src.meta_policy import META_POLICY_REPO, load_policy_metadata, publish_policy_if_improved SCORE_THRESHOLDS = { "massbench_adenocarcinoma": {"reference": 0.24, "acceptable": 0.19}, "massbench_alzheimer": {"reference": 0.44, "acceptable": 0.39}, "massbench_benchmark": {"reference": 0.75, "acceptable": 0.70}, } def parse_args(argv=None): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--split-manifest", type=Path, default=ROOT / "config" / "evolution_development_datasets.json") parser.add_argument("--output-dir", type=Path, default=ROOT / "results" / "meta_evolution") parser.add_argument("--population-size", type=int, default=12) parser.add_argument("--generations", type=int, default=25) parser.add_argument( "--candidate-budget", type=int, default=None, help="Exact total number of source-only candidate policies to evaluate; " "defaults to population-size * generations and may end mid-generation.", ) parser.add_argument("--hidden-size", type=int, default=16) parser.add_argument("--elite-count", type=int, default=2) parser.add_argument("--tournament-size", type=int, default=3) parser.add_argument("--crossover-rate", type=float, default=0.9) parser.add_argument("--mutation-rate", type=float, default=0.05) parser.add_argument("--mutation-scale", type=float, default=0.1) parser.add_argument("--worst-dataset-weight", type=float, default=0.25) parser.add_argument("--validation-patience", type=int, default=8) parser.add_argument("--n-epochs", type=int, default=100) parser.add_argument("--n-repeats", type=int, default=3) parser.add_argument("--num-workers", type=int, default=4) parser.add_argument( "--validation-epoch-interval", type=int, default=10, help="Log held-out validation epoch telemetry at epoch 0 and this interval.", ) parser.add_argument("--device", default="cuda", choices=["cpu", "cuda"]) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--resume", action="store_true") parser.add_argument("--smoke-evaluator", action="store_true", help="Exercise evolution without training BERNN") parser.add_argument("--no-wandb", action="store_true", help="Disable live Weights & Biases telemetry") parser.add_argument("--wandb-project", default="BE_leaderboard_meta_evolution") parser.add_argument("--wandb-run-name", default=None) parser.add_argument("--hf-meta-policy-repo", default=META_POLICY_REPO) parser.add_argument( "--no-hf-meta-policy-push", action="store_true", help="Do not compare/promote the validation champion on Hugging Face.", ) return parser.parse_args(argv) def _sha256(path: Path) -> str: return hashlib.sha256(path.read_bytes()).hexdigest() def _atomic_json(path: Path, value: dict) -> None: temporary = path.with_suffix(path.suffix + ".tmp") temporary.write_text(json.dumps(value, indent=2, default=str) + "\n") os.replace(temporary, path) def _load_datasets(dataset_ids): datasets = {} meta = {} for dataset_id in dataset_ids: X, y, batches = hp_search.load_dataset(dataset_id) datasets[dataset_id] = (X, y, batches) extracted = extract_meta_features(X, y, batches) meta[dataset_id] = np.asarray([extracted[name] for name in META_FEATURE_NAMES], dtype=np.float32) return datasets, meta def predict_config_for_dataset( X, y, batches, *, model_path=None, meta_model_path=None, checkpoint_path=None, max_warmup=50, **kwargs, ): """Predict exactly one BERNN config for one unseen dataset. The frozen evolutionary meta-policy is loaded from ``best_policy.npz``. Meta-features for the unseen dataset are normalized with the mean/scale learned from meta-training datasets and stored in that artifact. """ policy_path = model_path or meta_model_path or checkpoint_path if policy_path is None: policy_path = ROOT / "results" / "meta_evolution" / "best_policy.npz" policy_path = Path(policy_path) if policy_path.is_dir(): policy_path = policy_path / "best_policy.npz" if not policy_path.exists(): raise FileNotFoundError( f"Meta-policy not found: {policy_path}. " "Pass --meta-model-path /path/to/best_policy.npz." ) saved = np.load(policy_path) required = {"genome", "meta_mean", "meta_scale", "n_inputs", "hidden_size"} missing = required.difference(saved.files) if missing: raise ValueError( f"Meta-policy {policy_path} is missing required arrays: {sorted(missing)}" ) genome = saved["genome"] meta_mean = np.asarray(saved["meta_mean"], dtype=np.float32) meta_scale = np.asarray(saved["meta_scale"], dtype=np.float32) n_inputs = int(np.asarray(saved["n_inputs"]).reshape(-1)[0]) hidden_size = int(np.asarray(saved["hidden_size"]).reshape(-1)[0]) extracted = extract_meta_features(X, y, batches) raw_meta = np.asarray( [extracted[name] for name in META_FEATURE_NAMES], dtype=np.float32, ) if raw_meta.shape[0] != n_inputs: raise ValueError( f"Meta-feature count mismatch: policy expects {n_inputs}, " f"current extractor produced {raw_meta.shape[0]}." ) if meta_mean.shape != raw_meta.shape or meta_scale.shape != raw_meta.shape: raise ValueError( "Saved meta normalization shape does not match current meta-features: " f"mean={meta_mean.shape}, scale={meta_scale.shape}, current={raw_meta.shape}." ) safe_scale = np.where(np.abs(meta_scale) < 1e-12, 1.0, meta_scale) normalized_meta = (raw_meta - meta_mean) / safe_scale shape = PolicyShape(n_inputs=n_inputs, hidden_size=hidden_size) return decode_config( genome, normalized_meta, shape, max_warmup=int(max_warmup), ) def _append_solution_record(output_dir: Path, record: dict, train_datasets: list[str]) -> None: jsonl_path = output_dir / "solutions.jsonl" with jsonl_path.open("a", encoding="utf-8") as stream: stream.write(json.dumps(record, default=str) + "\n") csv_path = output_dir / "solutions.csv" columns = [ "solution_step", "generation", "population_index", "genome_digest", "train_fitness", "mean_mcc", "worst_mcc", *[f"mcc_{name}" for name in train_datasets], *[f"test_mcc_{name}" for name in train_datasets], ] row = {name: record.get(name) for name in columns} for name in train_datasets: row[f"mcc_{name}"] = record["train_scores"][name] row[f"test_mcc_{name}"] = record["test_scores"].get(name) write_header = not csv_path.exists() or csv_path.stat().st_size == 0 with csv_path.open("a", newline="", encoding="utf-8") as stream: writer = csv.DictWriter(stream, fieldnames=columns) if write_header: writer.writeheader() writer.writerow(row) def _mean_epoch_metric(metrics: dict, metric_name: str) -> dict[int, float]: """Average one BERNN epoch metric over grouped outer-CV folds/repeats.""" values: dict[int, list[float]] = {} for rows in metrics.get("_epoch_traces_by_fold", {}).values(): for row in rows: epoch = int(row["epoch"]) matching = [ float(value) for key, value in row.items() if key == metric_name or key.endswith(f"/{metric_name}") ] if matching: values.setdefault(epoch, []).extend(matching) return {epoch: float(np.mean(epoch_values)) for epoch, epoch_values in values.items()} def _classifier_widths(config: dict) -> list[int]: widths = [max(16, int(config["layer1"]))] for _ in range(1, max(1, int(config["n_layers"]))): widths.append(max(widths[-1] // 2, 16)) return widths def _numeric_config_telemetry(dataset_id: str, config: dict) -> dict[str, float]: telemetry = {} for name, value in config.items(): if isinstance(value, (bool, np.bool_)): telemetry[f"solutions/config/{dataset_id}/{name}"] = float(value) elif isinstance(value, (int, float, np.integer, np.floating)): telemetry[f"solutions/config/{dataset_id}/{name}"] = float(value) widths = _classifier_widths(config) telemetry[f"solutions/config/{dataset_id}/n_neurons"] = float(sum(widths)) for layer_index, width in enumerate(widths, 1): telemetry[f"solutions/config/{dataset_id}/layer{layer_index}_width"] = float(width) return telemetry def _score_figure( dataset_id: str, rows: list[tuple[int, float, float]], thresholds: dict[str, float] | None = None, ): """Plot validation/test MCC plus BERNN-best and acceptable-goal traces.""" import plotly.graph_objects as go x = [row[0] for row in rows] valid = [row[1] for row in rows] test = [row[2] for row in rows] figure = go.Figure() figure.add_trace(go.Scatter(x=x, y=valid, mode="lines+markers", name="valid_mcc")) figure.add_trace(go.Scatter(x=x, y=test, mode="lines+markers", name="test_mcc")) thresholds = dict(SCORE_THRESHOLDS.get(dataset_id, {}) if thresholds is None else thresholds) if "reference" in thresholds: figure.add_trace(go.Scatter( x=x, y=[thresholds["reference"]] * len(x), mode="lines", line={"dash": "solid", "width": 2}, name="Best valid MCC", )) if "acceptable" in thresholds: figure.add_trace(go.Scatter( x=x, y=[thresholds["acceptable"]] * len(x), mode="lines", line={"dash": "dash", "width": 2}, name="goal", )) figure.update_layout( title=f"{dataset_id}: MCC versus BERNN thresholds", xaxis_title="solution step", yaxis_title="grouped-CV MCC", yaxis_range=[-1.0, 1.0], ) return figure def _solution_records(output_dir: Path) -> list[dict]: path = output_dir / "solutions.jsonl" if not path.exists(): return [] return [json.loads(line) for line in path.read_text().splitlines() if line.strip()] def _save_source_champion( output_dir: Path, genome: np.ndarray, *, mean: np.ndarray, scale: np.ndarray, shape: PolicyShape, record: dict, ) -> None: np.savez_compressed( output_dir / "best_policy.npz", genome=np.asarray(genome, dtype=np.float32), meta_mean=mean, meta_scale=scale, n_inputs=np.asarray([shape.n_inputs]), hidden_size=np.asarray([shape.hidden_size]), ) _atomic_json(output_dir / "best_policy.json", record) def _replay_source_champion( output_dir: Path, shape: PolicyShape, evolution: EvolutionConfig, seed: int, ) -> tuple[np.ndarray, dict] | None: """Reconstruct evaluated genomes without rerunning any BERNN fit. Evolution and breeding use source fitness only. Replaying the deterministic initializer and breeder therefore recovers the exact genome corresponding to every durable solution record, including legacy runs whose saved policy was selected with target validation telemetry. """ records = _solution_records(output_dir) if not records: return None replay_rng = np.random.default_rng(seed) population = initialize_population(shape, evolution, replay_rng) best_genome = None best_record = None generations = sorted({int(row["generation"]) for row in records}) for generation in generations: rows = sorted( (row for row in records if int(row["generation"]) == generation), key=lambda row: int(row["population_index"]), ) fitness = np.full(evolution.population_size, np.nan, dtype=float) for row in rows: index = int(row["population_index"]) if not 0 <= index < evolution.population_size: raise ValueError(f"Invalid population index in generation {generation}: {index}") observed_digest = str(row["genome_digest"]) recovered_digest = genome_digest(population[index]) if observed_digest != recovered_digest: raise ValueError( f"Cannot replay generation {generation} index {index}: " f"digest {recovered_digest} != {observed_digest}" ) value = float(row["train_fitness"]) fitness[index] = value if best_record is None or value > float(best_record["train_fitness"]): best_genome = population[index].copy() best_record = { "selection_protocol": "source_only_best_fitness", "generation": generation, "population_index": index, "solution_step": int(row["solution_step"]), "train_fitness": value, "train_scores": row["train_scores"], "genome_digest": observed_digest, "target_scores_observed_for_selection": False, } if len(rows) == evolution.population_size: if not np.isfinite(fitness).all(): raise ValueError(f"Incomplete fitness vector for generation {generation}") population = breed_next_population(population, fitness, evolution, replay_rng) elif generation != generations[-1]: raise ValueError(f"Only the final generation may be partial; generation {generation} is incomplete") return best_genome, best_record def main(argv=None) -> int: args = parse_args(argv) if args.n_repeats != 3: raise ValueError("Meta-evolution requires grouped CV=3") if importlib.metadata.version("bernn") != "1.0.6": raise RuntimeError("This experiment requires bernn==1.0.6") if args.generations < 1: raise ValueError("generations must be positive") candidate_budget = int( args.candidate_budget if args.candidate_budget is not None else args.population_size * args.generations ) if candidate_budget < 1: raise ValueError("candidate_budget must be positive") partitions = load_dataset_partitions(args.split_manifest) # Target datasets provide descriptors for a frozen recommendation only. # Their BERNN scores and fixed tests are never opened by source search. heldout_ids = tuple(getattr(partitions, "test", ())) datasets, raw_meta = _load_datasets(partitions.train) fixed_tests = { dataset_id: hp_search.load_fixed_test_dataset(dataset_id) for dataset_id in partitions.train } train_matrix = np.stack([raw_meta[name] for name in partitions.train]) normalized_train, mean, scale = normalize_meta_features(train_matrix) normalized_meta = { name: row for name, row in zip(partitions.train, normalized_train) } evolution = EvolutionConfig( population_size=args.population_size, elite_count=args.elite_count, tournament_size=args.tournament_size, crossover_rate=args.crossover_rate, mutation_rate=args.mutation_rate, mutation_scale=args.mutation_scale, worst_dataset_weight=args.worst_dataset_weight, ) evolution.validate() shape = PolicyShape(len(META_FEATURE_NAMES), args.hidden_size) args.output_dir.mkdir(parents=True, exist_ok=True) manifest_hash = _sha256(args.split_manifest) checkpoint_path = args.output_dir / "checkpoint.npz" state_path = args.output_dir / "state.json" rng = np.random.default_rng(args.seed) generation_start = 0 best_train_fitness = -np.inf solution_step = 0 if args.resume: if not checkpoint_path.exists() or not state_path.exists(): raise FileNotFoundError("--resume requires checkpoint.npz and state.json") state = json.loads(state_path.read_text()) if state["development_manifest_sha256"] != manifest_hash: raise ValueError("Development dataset manifest changed since checkpoint creation") expected_shape = {"n_inputs": shape.n_inputs, "hidden_size": shape.hidden_size} if state.get("policy_shape") != expected_shape: raise ValueError("Policy shape differs from the saved checkpoint") if state.get("evolution_config") != config_dict(evolution): raise ValueError("Evolution configuration differs from the saved checkpoint") saved = np.load(checkpoint_path) population = saved["population"] generation_start = int(state["next_generation"]) best_train_fitness = float(state.get("best_train_fitness", -np.inf)) solution_step = int(state.get("solution_step", generation_start * args.population_size)) rng.bit_generator.state = state["rng_state"] else: population = initialize_population(shape, evolution, rng) hp_args = hp_search.parse_args([]) hp_args.n_epochs = args.n_epochs hp_args.n_repeats = args.n_repeats hp_args.num_workers = args.num_workers hp_args.log1p = True hp_args.device = args.device hp_args.no_wandb = True hp_args.combine_test = False hp_args.max_warmup = max(1, min(50, args.n_epochs)) score_cache = {} run_metadata_path = args.output_dir / "run_metadata.json" if run_metadata_path.exists(): run_metadata = json.loads(run_metadata_path.read_text()) else: run_metadata = {"wandb_run_id": uuid.uuid4().hex[:8]} run_metadata.update({ "arm": "meta_evolution", "candidate_budget": candidate_budget, "budget_unit": "candidate_policy", "population_size": int(args.population_size), "generations": int(args.generations), "bernn_version": importlib.metadata.version("bernn"), "log1p": True, "num_workers": int(args.num_workers), "train_datasets": list(partitions.train), "validation_datasets": list(partitions.validation), "selection_protocol": "source_only_best_fitness", "target_scores_observed": False, "cross_test": 1, "fixed_test_labels": "monitoring_only_excluded_from_fitness_and_selection", }) _atomic_json(run_metadata_path, run_metadata) wandb_run = None if not args.no_wandb: import wandb wandb_run = wandb.init( project=args.wandb_project, name=args.wandb_run_name, id=run_metadata["wandb_run_id"], resume="allow", config={ **vars(args), "split_manifest": str(args.split_manifest), "output_dir": str(args.output_dir), "train_datasets": list(partitions.train), "validation_datasets": list(partitions.validation), "heldout_datasets": list(heldout_ids), }, ) wandb.define_metric("solution_step") wandb.define_metric("solutions/*", step_metric="solution_step") wandb.define_metric("generation") wandb.define_metric("champion/*", step_metric="generation") for dataset_id, thresholds in SCORE_THRESHOLDS.items(): wandb_run.summary[f"reference_mcc/{dataset_id}"] = thresholds["reference"] wandb_run.summary[f"acceptable_mcc/{dataset_id}"] = thresholds["acceptable"] wandb_run.summary["selection_protocol"] = "source_only_best_fitness" wandb_run.summary["target_scores_observed_after_contract_fix"] = False def evaluate(genome, dataset_id): key = f"{genome_digest(genome)}:{dataset_id}" if key in score_cache: return score_cache[key] metrics = {} try: if args.smoke_evaluator: target = raw_meta[dataset_id] score = float(np.tanh(np.mean(genome[: min(len(genome), len(target))]) + np.mean(target) * 1e-4)) else: X, y, batches = datasets[dataset_id] config = decode_config(genome, normalized_meta[dataset_id], shape, max_warmup=hp_args.max_warmup) run_args = argparse.Namespace(**vars(hp_args)) run_args.dataset = dataset_id run_args.seed = args.seed run_args.bs = recommended_batch_size(batches, cap=hp_args.bs) run_args.results_dir = str(args.output_dir / dataset_id) run_args.cv_split_cache = str(args.output_dir / "cv_splits" / f"{dataset_id}.npz") run_args.resolved_n_repeats = hp_search.resolve_n_repeats(run_args.n_repeats, batches) exp_id = f"meta_evolution_{dataset_id}_{genome_digest(genome)[:12]}" score, metrics = hp_search.run_trial( config, run_args, (X, y, batches), exp_id, fixed_test_data=fixed_tests[dataset_id], ) except Exception as exc: # noqa: BLE001 - invalid candidates receive sentinel fitness score = -1.0 print( f"[evolution] failed genome={genome_digest(genome)[:12]} " f"dataset={dataset_id}: {type(exc).__name__}: {exc}", flush=True, ) score_cache[key] = (float(score), metrics) return score_cache[key] replayed = _replay_source_champion(args.output_dir, shape, evolution, args.seed) if replayed is not None: replayed_genome, replayed_record = replayed best_train_fitness = float(replayed_record["train_fitness"]) _save_source_champion( args.output_dir, replayed_genome, mean=mean, scale=scale, shape=shape, record=replayed_record, ) ledger = args.output_dir / "source_generations.jsonl" existing_solutions = _solution_records(args.output_dir) score_history = {dataset_id: [] for dataset_id in partitions.train} for row in existing_solutions: for dataset_id in partitions.train: test_score = float(row.get("test_scores", {}).get(dataset_id, np.nan)) score_history[dataset_id].append(( int(row["solution_step"]), float(row["train_scores"][dataset_id]), test_score, )) def save_state(next_generation: int, checkpoint_population: np.ndarray) -> None: np.savez_compressed( checkpoint_path, population=checkpoint_population, meta_mean=mean, meta_scale=scale, ) _atomic_json(state_path, { "next_generation": next_generation, "best_train_fitness": best_train_fitness, "solution_step": solution_step, "rng_state": rng.bit_generator.state, "selection_protocol": "source_only_best_fitness", "development_manifest_sha256": manifest_hash, "policy_shape": {"n_inputs": shape.n_inputs, "hidden_size": shape.hidden_size}, "evolution_config": config_dict(evolution), "train_datasets": list(partitions.train), "validation_datasets": list(partitions.validation), "heldout_datasets": list(heldout_ids), }) for generation in range(generation_start, args.generations): prior_rows = sorted( (row for row in _solution_records(args.output_dir) if int(row["generation"]) == generation), key=lambda row: int(row["population_index"]), ) expected_indices = list(range(len(prior_rows))) observed_indices = [int(row["population_index"]) for row in prior_rows] if observed_indices != expected_indices: raise ValueError( f"Generation {generation} durable records are not a contiguous prefix: " f"{observed_indices}" ) train_scores = np.full((len(population), len(partitions.train)), np.nan, dtype=np.float64) train_fitness = np.full(len(population), np.nan, dtype=np.float64) for row in prior_rows: index = int(row["population_index"]) if genome_digest(population[index]) != row["genome_digest"]: raise ValueError(f"Checkpoint/ledger digest mismatch at generation {generation} index {index}") train_fitness[index] = float(row["train_fitness"]) train_scores[index] = [float(row["train_scores"][name]) for name in partitions.train] for population_index in range(len(prior_rows), len(population)): if solution_step >= candidate_budget: break genome = population[population_index] scores = np.asarray([evaluate(genome, name)[0] for name in partitions.fitness_datasets()]) fitness = aggregate_dataset_scores(scores, evolution.worst_dataset_weight) train_scores[population_index] = scores train_fitness[population_index] = fitness decoded = { name: decode_config(genome, normalized_meta[name], shape, max_warmup=hp_args.max_warmup) for name in partitions.train } solution_record = { "solution_step": solution_step, "generation": generation, "population_index": population_index, "genome_digest": genome_digest(genome), "train_fitness": float(fitness), "mean_mcc": float(np.mean(scores)), "worst_mcc": float(np.min(scores)), "train_scores": dict(zip(partitions.train, scores.tolist())), "test_scores": { name: float(evaluate(genome, name)[1].get("test_mcc", np.nan)) for name in partitions.train }, "config_by_dataset": decoded, "target_scores_observed": False, } solution_record["total_valid_mcc"] = float(fitness) solution_record["total_test_mcc"] = float(aggregate_dataset_scores( solution_record["test_scores"].values(), evolution.worst_dataset_weight, )) for dataset_id, score in zip(partitions.train, scores): score_history[dataset_id].append(( solution_step, float(score), solution_record["test_scores"][dataset_id], )) _append_solution_record(args.output_dir, solution_record, list(partitions.train)) if fitness > best_train_fitness: best_train_fitness = float(fitness) champion_record = { "selection_protocol": "source_only_best_fitness", "generation": generation, "population_index": population_index, "solution_step": solution_step, "train_fitness": float(fitness), "train_scores": solution_record["train_scores"], "genome_digest": solution_record["genome_digest"], "target_scores_observed_for_selection": False, } _save_source_champion( args.output_dir, genome, mean=mean, scale=scale, shape=shape, record=champion_record, ) if wandb_run is not None: config_telemetry = {} for dataset_id, dataset_config in decoded.items(): config_telemetry.update(_numeric_config_telemetry(dataset_id, dataset_config)) figure_telemetry = { f"solutions/mcc/{dataset_id}": _score_figure( dataset_id, score_history[dataset_id] ) for dataset_id in partitions.train } wandb_run.log({ "solution_step": solution_step, "solutions/train_fitness": float(fitness), "solutions/mean_mcc": float(np.mean(scores)), "solutions/worst_mcc": float(np.min(scores)), "solutions/total_valid_mcc": solution_record["total_valid_mcc"], **{ f"solutions/metrics/{name}/valid_mcc": float(score) for name, score in zip(partitions.train, scores) }, **{ f"solutions/metrics/{name}/test_mcc": solution_record["test_scores"][name] for name in partitions.train }, "solutions/total_test_mcc": solution_record["total_test_mcc"], **config_telemetry, **figure_telemetry, }) print(f"[evolution] source solution {solution_step} complete", flush=True) solution_step += 1 evaluated = int(np.isfinite(train_fitness).sum()) if evaluated < len(population): save_state(generation, population) print( f"[evolution] exact candidate budget reached at {solution_step}; " f"generation {generation} is {evaluated}/{len(population)} complete", flush=True, ) break champion_index = int(np.argmax(train_fitness)) source_record = { "generation": generation, "selection_protocol": "source_only_best_fitness", "train_champion_fitness": float(train_fitness[champion_index]), "train_scores": dict(zip(partitions.train, train_scores[champion_index].tolist())), "champion_digest": genome_digest(population[champion_index]), "target_scores_observed": False, } with ledger.open("a", encoding="utf-8") as stream: stream.write(json.dumps(source_record, default=str) + "\n") if wandb_run is not None: wandb_run.log({ "generation": generation, "champion/train_fitness": float(train_fitness[champion_index]), }) next_population = breed_next_population(population, train_fitness, evolution, rng) save_state(generation + 1, next_population) population = next_population best_policy_path = args.output_dir / "best_policy.npz" if not best_policy_path.exists(): raise FileNotFoundError("Source-only evolution did not produce best_policy.npz") best_policy = np.load(best_policy_path) best_genome = best_policy["genome"] _, target_raw_meta = _load_datasets(partitions.validation) target_matrix = np.stack([target_raw_meta[name] for name in partitions.validation]) normalized_targets, _, _ = normalize_meta_features(train_matrix, target_matrix) target_meta = { name: row for name, row in zip(partitions.validation, normalized_targets) } target_configs = { name: decode_config( best_genome, target_meta[name], shape, max_warmup=hp_args.max_warmup, ) for name in partitions.validation } _atomic_json(args.output_dir / "final_recommendations.json", { "selection_protocol": "source_only_best_fitness", "source_candidates": int(solution_step), "source_best_train_fitness": float(best_train_fitness), "target_scores_observed": False, "policy": best_policy_path.name, "policy_digest": genome_digest(best_genome), "target_configs": target_configs, }) if wandb_run is not None: wandb_run.summary["source_best_train_fitness"] = float(best_train_fitness) wandb_run.summary["source_candidates"] = int(solution_step) if wandb_run is not None: wandb_run.finish() return 0 if __name__ == "__main__": raise SystemExit(main())