Source code for alf_core.optimizer.search

# Copyright 2026 InstaDeep Ltd. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import abc

import numpy as np
from alf_core.dataclasses import Candidate, LabelledCandidates, State
from alf_core.model.base_model import BaseModel
from alf_core.utils.metrics.acquisition_batch import compute_recall, compute_regret




class BaseSearch(abc.ABC):
    """Base class for all types of search methods.

    Search methods provide a way of searching the design space in order to propose new candidates
    not yet observed in the training dataset. This can be achieved in multiple ways depending on
    the search protocol, or availability of ground truth labels.
    """

    @abc.abstractmethod
    def __call__(self, state: State, **kwargs) -> list[Candidate]:
        """Generate a pool of candidates to evaluate.

        Args:
            state: Current task state.
            **kwargs: Additional keyword arguments for search.

        Returns:
            List of candidate objects to consider for acquisition.
        """
        pass



    def get_metrics(self, state: State) -> dict[str, float]:
        """Get metrics for the search function.

        Args:
            state: Current task state.

        Returns:
            Dictionary of metric names to values. Returns empty dict by default;
            subclasses should override.
        """
        return {}






class SearchProtocol:
    """Search protocol is used to define the search process."""

    @abc.abstractmethod
    def __call__(self, state: State) -> list[Candidate]:
        """Apply the search protocol to generate a candidate pool.

        Args:
            state: Current task state.

        Returns:
            List of candidate objects generated by the protocol.
        """
        pass





class GeneratorSearch(BaseSearch):
    """Search method based on a model generating samples."""

    def __init__(self, model: BaseModel) -> None:
        """Initialize generator search with a model.

        Args:
            model: BaseModel instance to use for generating candidates.
        """
        self.model: BaseModel = model

    def __call__(self, state: State, **kwargs) -> list[Candidate]:
        """Sample candidates from the model to define the search pool.

        Args:
            state: Current task state (not used for model-based generation).
            **kwargs: Additional keyword arguments (not used).

        Returns:
            List of candidates sampled from the model.
        """
        return self.model.sample()





class DatasetSearch(BaseSearch):
    """Offline search method based on a dataset defining the search pool."""

    def __call__(self, state: State, **kwargs) -> list[Candidate]:
        """Get the candidate pool from the dataset.
        If the acquisition batch size is larger than the remaining candidate pool,
        raise a ValueError.

        Args:
            state: Current task state containing the dataset.
            **kwargs: Additional keyword arguments (not used).

        Returns:
            List of candidates from the dataset's candidate pool.

        Raises:
            ValueError: If the acquisition batch size is larger than the remaining candidate pool.
        """
        if len(state.dataset.candidate_pool) <= state.acq_batch_size:
            raise ValueError(
                f"The acquisition batch size ({state.acq_batch_size}) is larger than the "
                f"remaining candidate pool ({len(state.dataset.candidate_pool)}), breaking ..."
            )
        return state.dataset.candidate_pool.candidates



    def get_metrics(self, state: State) -> dict[str, float]:
        """Return recall and regret metrics for the dataset search method.

        Both metrics compare the initial candidate pool against the candidates
        acquired during the loop. The acquired-only set is reconstructed from
        `state.history` so the initial labelled seed (which is not part of the
        candidate pool) does not enter the comparison.

        Args:
            state: Current task state.

        Returns:
            Dictionary containing recall and regret metrics comparing the initial
            candidate pool to the acquired candidates. The metrics are omitted
            until at least one candidate has been acquired.
        """
        if not state.history:
            return {}
        init_candidate_pool = state.dataset.init_candidate_pool
        acquired_candidates = LabelledCandidates(
            candidates=[c for acquired in state.history for c in acquired.candidates],
            labels=np.concatenate([acquired.labels for acquired in state.history]),
        )
        recall_metrics = compute_recall(init_candidate_pool, acquired_candidates)
        regret_metrics = compute_regret(init_candidate_pool, acquired_candidates)
        return {**recall_metrics, **regret_metrics}






class ProtocolSearch(BaseSearch):
    """Search method based on a function/procedure to define the search pool."""

    def __init__(self, protocol: SearchProtocol) -> None:
        """Initialize protocol search with a search protocol.

        Args:
            protocol: SearchProtocol instance to use for generating candidates.
        """
        self.protocol: SearchProtocol = protocol

    def __call__(self, state: State, **kwargs) -> list[Candidate]:
        """Apply the search protocol to return a pool of candidates.

        Args:
            state: Current task state.
            **kwargs: Additional keyword arguments passed to the protocol.

        Returns:
            List of candidates generated by the protocol.
        """
        return self.protocol(state, **kwargs)





class ModelProtocolSearch(BaseSearch):
    """Search method based on a model and a protocol to define the search pool."""

    def __init__(self, model: BaseModel, protocol: SearchProtocol) -> None:
        """Initialize model-protocol search with a model and protocol.

        Args:
            model: BaseModel instance to use.
            protocol: SearchProtocol instance to use.
        """
        self.model: BaseModel = model
        self.protocol: SearchProtocol = protocol

    @abc.abstractmethod
    def __call__(self, state: State, **kwargs) -> list[Candidate]:
        """Apply the model and search protocol to return a pool of candidates.

        Args:
            state: Current task state.
            **kwargs: Additional keyword arguments.

        Returns:
            List of candidates generated by combining the model and protocol.
        """
        pass