source/train.py

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()