from __future__ import annotations import json from pathlib import Path import numpy as np import trackio from model import RecurrentFeatures, fit_readout, parameter_count, unpack_genome from safetensors.numpy import save_file PROJECT_DIR = Path(__file__).resolve().parent ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "evolino-pocket" DATA_DIR = PROJECT_DIR / "data" HIDDEN_SIZE = 10 HORIZONS = 12 SEEDS = [2311, 2333, 2351] GENERATIONS = 40 POPULATION = 14 def mackey_glass(length: int, delay: int = 17) -> np.ndarray: values = np.full(length + delay + 1, 1.2, dtype=np.float64) values[:delay] += np.linspace(-0.02, 0.02, delay) for index in range(delay, len(values) - 1): delayed = values[index - delay] derivative = 0.2 * delayed / (1 + delayed**10) - 0.1 * values[index] values[index + 1] = values[index] + derivative return values[delay + 1 :] def make_problem() -> tuple[np.ndarray, np.ndarray, dict]: raw = mackey_glass(2_450) train_mean = raw[:1_200].mean() train_std = raw[:1_200].std() series = (raw - train_mean) / train_std targets = np.stack( [ np.roll(series, -horizon) for horizon in range(1, HORIZONS + 1) ], axis=1, ) split = { "train": np.arange(50, 1_200), "validation": np.arange(1_200, 1_600), "test": np.arange(1_600, 2_400), } return series, targets, split def score_genome( genome: np.ndarray, series: np.ndarray, targets: np.ndarray, split: dict, ) -> tuple[float, RecurrentFeatures]: model = unpack_genome(genome, HIDDEN_SIZE) fit_readout(model, series, targets, split["train"]) predictions = model.predict(series)[split["validation"]] rmse = float(np.sqrt(np.mean((predictions - targets[split["validation"]]) ** 2))) return rmse, model def evaluate_model( model: RecurrentFeatures, series: np.ndarray, targets: np.ndarray, indices: np.ndarray, ) -> dict: predictions = model.predict(series)[indices] truth = targets[indices] return { "rmse_all_horizons": float(np.sqrt(np.mean((predictions - truth) ** 2))), "rmse_by_horizon": { str(horizon): float( np.sqrt(np.mean((predictions[:, horizon - 1] - truth[:, horizon - 1]) ** 2)) ) for horizon in [1, 6, 12] }, "examples": len(indices), } def evolve(seed: int, series: np.ndarray, targets: np.ndarray, split: dict) -> dict: rng = np.random.default_rng(seed) genome_size = HIDDEN_SIZE**2 + HIDDEN_SIZE * 2 initial_genome = rng.normal(0, 0.25, genome_size) best_genome = initial_genome.copy() initial_rmse, initial_model = score_genome( initial_genome, series, targets, split, ) best_rmse = initial_rmse sigma = 0.18 history = [] for generation in range(1, GENERATIONS + 1): noise = rng.normal(size=(POPULATION, genome_size)) candidates = best_genome + sigma * noise scored = [ (*score_genome(candidate, series, targets, split), candidate) for candidate in candidates ] candidate_rmse, _, candidate_genome = min(scored, key=lambda item: item[0]) improved = candidate_rmse < best_rmse if improved: best_rmse = candidate_rmse best_genome = candidate_genome.copy() sigma *= 1.015 else: sigma *= 0.97 history.append(best_rmse) if generation % 5 == 0: trackio.log( { "seed": seed, "generation": generation, "validation_rmse": best_rmse, "mutation_sigma": sigma, } ) _, evolved_model = score_genome(best_genome, series, targets, split) fit_readout( evolved_model, series, targets, np.concatenate([split["train"], split["validation"]]), ) fit_readout( initial_model, series, targets, np.concatenate([split["train"], split["validation"]]), ) return { "seed": seed, "initial_genome": initial_genome, "evolved_genome": best_genome, "random_model": initial_model, "evolved_model": evolved_model, "initial_validation_rmse": initial_rmse, "best_validation_rmse": best_rmse, "history": history, } def model_tensors(model: RecurrentFeatures) -> dict[str, np.ndarray]: assert model.readout is not None return { "input_weight": model.input_weight.astype(np.float32), "recurrent_weight": model.recurrent_weight.astype(np.float32), "bias": model.bias.astype(np.float32), "readout": model.readout.astype(np.float32), } def main() -> None: series, targets, split = make_problem() trackio.init( project="evolino-pocket", name="evolved-recurrent-features-v1", config={ "hidden_size": HIDDEN_SIZE, "horizons": HORIZONS, "generations": GENERATIONS, "population": POPULATION, "seeds": SEEDS, }, ) runs = [] saved = None for seed in SEEDS: outcome = evolve(seed, series, targets, split) run = { "seed": seed, "initial_validation_rmse": outcome["initial_validation_rmse"], "best_validation_rmse": outcome["best_validation_rmse"], "random_control": evaluate_model( outcome["random_model"], series, targets, split["test"], ), "evolved": evaluate_model( outcome["evolved_model"], series, targets, split["test"], ), } runs.append(run) if saved is None: saved = outcome assert saved is not None report = { "experiment": "Evolino-inspired recurrent features plus optimal linear readout", "parameters_per_model": parameter_count(HIDDEN_SIZE, HORIZONS), "runs": runs, "aggregate": { "random_test_rmse_mean": float( np.mean( [run["random_control"]["rmse_all_horizons"] for run in runs] ) ), "evolved_test_rmse_mean": float( np.mean([run["evolved"]["rmse_all_horizons"] for run in runs]) ), }, } ARTIFACT_DIR.mkdir(parents=True, exist_ok=True) DATA_DIR.mkdir(parents=True, exist_ok=True) save_file( model_tensors(saved["evolved_model"]), ARTIFACT_DIR / "evolved_features.safetensors", ) save_file( model_tensors(saved["random_model"]), ARTIFACT_DIR / "random_control.safetensors", ) (ARTIFACT_DIR / "evaluation.json").write_text( json.dumps(report, indent=2), encoding="utf-8", ) np.savetxt( DATA_DIR / "mackey_glass_test.csv", np.column_stack( [ np.arange(len(split["test"])), series[split["test"]], ] ), delimiter=",", header="step,normalized_value", comments="", ) trackio.log(report["aggregate"]) trackio.finish() print(json.dumps(report, indent=2)) if __name__ == "__main__": main()