Example Optim Aggregated

"""Auto-generated example: Aggregated Optimisation."""

import math
import random

import bencher as bn

class AlgorithmBench(bn.ParametrizedSweep):
    """Finds best learning rate across algorithms (aggregated)."""

    algorithm = bn.StringSweep(
        ["gradient_descent", "adam", "rmsprop"],
        doc="Optimization algorithm",
        optimize=False,  # sweep but don't optimize — aggregate results
    )
    learning_rate = bn.FloatSweep(default=0.01, bounds=[0.001, 1.0], doc="Learning rate")

    loss = bn.ResultFloat("loss", bn.OptDir.minimize, doc="Training loss (minimize)")

    def benchmark(self):
        algo_sensitivity = {"gradient_descent": 1.0, "adam": 0.6, "rmsprop": 0.8}
        optimal_lr = 0.01 * algo_sensitivity[self.algorithm]
        self.loss = (math.log10(self.learning_rate) - math.log10(optimal_lr)) ** 2
        self.loss += random.gauss(0, 0.02)


def example_optim_aggregated(run_cfg: bn.BenchRunCfg | None = None) -> bn.Bench:
    """Aggregated Optimisation."""
    bench = AlgorithmBench().to_bench(run_cfg)
    bench.plot_sweep(input_vars=["algorithm", "learning_rate"], result_vars=["loss"], description='Finds the best learning rate averaged across algorithms. algorithm has optimize=False so Optuna only suggests learning_rate and averages loss over all algorithm choices.', post_description="The importance plot shows learning_rate and repeat — algorithm is aggregated away. Compare 'With Repeats' vs 'Without Repeats' to see if measurement noise affects the result.")

    return bench


if __name__ == "__main__":
    bn.run(example_optim_aggregated, subsampling_divisions=3, repeats=3, optimise=30)