test_balancing_learner.py

from __future__ import annotations

import pytest

from adaptive.learner import BalancingLearner, Learner1D, SequenceLearner
from adaptive.runner import simple

strategies = ["loss", "loss_improvements", "npoints", "cycle"]


def test_balancing_learner_loss_cache():
    learner = Learner1D(lambda x: x, bounds=(-1, 1))
    learner.tell(-1, -1)
    learner.tell(1, 1)
    learner.tell_pending(0)

    real_loss = learner.loss(real=True)
    pending_loss = learner.loss(real=False)

    # Test if the real and pending loss are cached correctly
    bl = BalancingLearner([learner])
    assert bl.loss(real=True) == real_loss
    assert bl.loss(real=False) == pending_loss

    # Test if everything is still fine when executed in the reverse order
    bl = BalancingLearner([learner])
    assert bl.loss(real=False) == pending_loss
    assert bl.loss(real=True) == real_loss


@pytest.mark.parametrize("strategy", strategies)
def test_distribute_first_points_over_learners(strategy):
    for initial_points in [0, 3]:
        learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
        learner = BalancingLearner(learners, strategy=strategy)

        points = learner.ask(initial_points)[0]
        learner.tell_many(points, points)

        points, _ = learner.ask(100)
        i_learner, xs = zip(*points)
        # assert that are all learners in the suggested points
        assert len(set(i_learner)) == len(learners)


@pytest.mark.parametrize("strategy", strategies)
def test_ask_0(strategy):
    learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
    learner = BalancingLearner(learners, strategy=strategy)
    points, _ = learner.ask(0)
    assert len(points) == 0


@pytest.mark.parametrize(
    "strategy, goal_type, goal",
    [
        ("loss", "loss_goal", 0.1),
        ("loss_improvements", "loss_goal", 0.1),
        ("npoints", "goal", lambda bl: all(lrn.npoints > 10 for lrn in bl.learners)),
        ("cycle", "loss_goal", 0.1),
    ],
)
def test_strategies(strategy, goal_type, goal):
    learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
    learner = BalancingLearner(learners, strategy=strategy)
    simple(learner, **{goal_type: goal})


def test_sequential_strategy() -> None:
    learners = [SequenceLearner(lambda x: x, sequence=[0, 1, 2, 3]) for i in range(10)]
    learner = BalancingLearner(learners, strategy="sequential")  # type: ignore[arg-type]
    simple(learner, goal=lambda lrn: sum(x.npoints for x in lrn.learners) >= 4 * 5)
    assert all(lrn.done() for lrn in learners[:5])
    assert all(not lrn.done() for lrn in learners[5:])