Spaces:
Sleeping
Sleeping
| #!/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()) | |