tests/test_algorithms.py - master - Handdas/FirstOrderAlgorythms

tests/test_algorithms.py
128 строк4 KB
Your Name
initial commit
09 июн 2026, 10:13
09 июн 2026, 10:13f88c506
"""Tests for first-order optimizers and WIND integration."""

from __future__ import annotations

import numpy as np
import pytest

from blackbox_optimizer import (
    AdamWOptimizer,
    BayesianSearch,
    CosineAnnealing,
    GridSearch,
    LionOptimizer,
    MomentumOptimizer,
    NesterovOptimizer,
    OptimizationPipeline,
    RMSpropOptimizer,
    RandomSearch,
    SGDOptimizer,
    SophiaOptimizer,
    StepLR,
    WINDAdapter,
    Warmup,
)


def sphere_loss_grad(params: np.ndarray) -> tuple[float, np.ndarray]:
    loss = float(np.sum(params**2))
    grad = 2.0 * params
    return loss, grad


@pytest.fixture
def initial_params() -> np.ndarray:
    return np.array([2.0, -1.0, 0.5], dtype=float)


class FakeWINDProblem:
    dim = 3
    lb = np.array([-5.0, -5.0, -5.0])
    ub = np.array([5.0, 5.0, 5.0])

    def loss_and_grad(self, params: np.ndarray) -> tuple[float, np.ndarray]:
        return sphere_loss_grad(params)


@pytest.mark.parametrize(
    "optimizer",
    [
        SGDOptimizer(learning_rate=0.05),
        MomentumOptimizer(learning_rate=0.05, momentum=0.9),
        NesterovOptimizer(learning_rate=0.05, momentum=0.9),
        AdamWOptimizer(learning_rate=0.05),
        RMSpropOptimizer(learning_rate=0.05),
        LionOptimizer(learning_rate=0.03),
        SophiaOptimizer(learning_rate=0.05),
    ],
)
def test_optimizer_step_vectorized(optimizer: SGDOptimizer, initial_params: np.ndarray) -> None:
    _, grad = sphere_loss_grad(initial_params)
    updated = optimizer.step(initial_params, grad)
    assert updated.shape == initial_params.shape
    assert np.isfinite(updated).all()


def test_pipeline_collects_metrics(initial_params: np.ndarray) -> None:
    pipeline = OptimizationPipeline()
    optimizer = SGDOptimizer(learning_rate=0.08)
    result = pipeline.run(optimizer, sphere_loss_grad, initial_params, iterations=40, log_every=20)
    assert result.trace.iterations == 40
    assert len(result.trace.losses) == 40
    assert len(result.trace.gradient_norms) == 40
    assert len(result.trace.wall_times) == 40
    assert result.trace.losses[-1] < result.trace.losses[0]


def test_wind_adapter() -> None:
    adapter = WINDAdapter.from_wind_object(FakeWINDProblem())
    loss, grad = adapter.loss_and_grad(np.array([10.0, -10.0, 1.0]))
    assert loss >= 0.0
    assert grad.shape == (3,)


def test_lr_schedulers() -> None:
    step = StepLR(base_lr=0.1, step_size=2, gamma=0.5)
    cosine = CosineAnnealing(base_lr=0.1, max_steps=10, min_lr=0.01)
    warmup = Warmup(base_lr=0.1, warmup_steps=4)
    assert step.get_lr(0) == 0.1
    assert step.get_lr(2) == 0.05
    assert 0.01 <= cosine.get_lr(5) <= 0.1
    assert warmup.get_lr(0) < warmup.get_lr(3) <= 0.1


def _objective_for_search(params: dict[str, float]) -> float:
    x = params["x"]
    y = params["y"]
    return (x - 1.2) ** 2 + (y + 0.5) ** 2


def test_grid_search() -> None:
    result = GridSearch().run(
        grid={"x": [0.0, 1.2, 2.0], "y": [-1.0, -0.5, 0.0]},
        objective=_objective_for_search,
        early_stop_threshold=None,
    )
    assert result.best_score == pytest.approx(0.0)


def test_random_search_with_early_stop() -> None:
    result = RandomSearch(seed=7).run(
        space={"x": (-2.0, 2.0), "y": (-2.0, 2.0)},
        objective=_objective_for_search,
        trials=40,
        early_stop_threshold=0.1,
    )
    assert len(result.history) <= 40
    assert result.best_score < 0.5


def test_bayesian_search_runs() -> None:
    result = BayesianSearch(seed=42).run(
        space={"x": (-2.0, 2.0), "y": (-2.0, 2.0)},
        objective=_objective_for_search,
        iterations=12,
        init_points=4,
        candidates=64,
    )
    assert result.best_score < 0.6