Source code for baybe.recommenders.pure.nonpredictive.sampling

"""Recommenders based on sampling."""

from typing import ClassVar

import numpy as np
import pandas as pd
from attrs import define, field, fields
from attrs.validators import instance_of
from typing_extensions import override

from baybe.exceptions import InfeasibilityError
from baybe.recommenders.pure.nonpredictive.base import NonPredictiveRecommender
from baybe.searchspace import SearchSpace, SearchSpaceType, SubspaceDiscrete
from baybe.settings import Settings, active_settings
from baybe.utils.conversion import to_string
from baybe.utils.sampling_algorithms import FPSInitialization, farthest_point_sampling




class RandomRecommender(NonPredictiveRecommender):
    """Recommends experiments randomly."""

    # Class variables
    compatibility: ClassVar[SearchSpaceType] = SearchSpaceType.HYBRID
    # See base class.

    supports_discrete_subset_generating_constraints: ClassVar[bool] = True
    # See base class.

    @override
    def _recommend_hybrid(
        self,
        searchspace: SearchSpace,
        candidates_exp: pd.DataFrame,
        batch_size: int,
    ) -> pd.DataFrame:
        if searchspace.type is SearchSpaceType.CONTINUOUS:
            return searchspace.continuous.sample_uniform(batch_size=batch_size)

        # Restrict to a random subset if subset-generating constraints are present
        is_hybrid = searchspace.type is SearchSpaceType.HYBRID
        if searchspace.discrete.n_subsets > 0:
            masks = searchspace.discrete.sample_subset_masks(
                candidates_exp,
                n=1,
                min_candidates=None if is_hybrid else batch_size,
            )
            if not masks:
                raise InfeasibilityError(
                    "No feasible subset found for the given "
                    "subset-generating constraints. All subsets have fewer "
                    f"candidates than the requested {batch_size=}."
                )
            candidates_exp = candidates_exp.loc[masks[0]]

        disc_random = candidates_exp.sample(
            n=batch_size,
            replace=is_hybrid or len(candidates_exp) < batch_size,
        )

        if not is_hybrid:
            return disc_random

        cont_random = searchspace.continuous.sample_uniform(batch_size=batch_size)
        cont_random.index = disc_random.index
        return pd.concat([disc_random, cont_random], axis=1)

    @override
    def __str__(self) -> str:
        fields = [to_string("Compatibility", self.compatibility, single_line=True)]
        return to_string(self.__class__.__name__, *fields)





@define
class FPSRecommender(NonPredictiveRecommender):
    """An initial recommender that selects candidates via Farthest Point Sampling.

    If the optional `fpsample <https://github.com/leonardodalinky/fpsample>`_ package is
    installed, a more efficient implementation is available that can be (de-)activated
    via the :attr:`~baybe.settings.Settings.use_fpsample` setting. Otherwise, a custom
    fallback implementation is used.
    """

    # Class variables
    compatibility: ClassVar[SearchSpaceType] = SearchSpaceType.DISCRETE
    # See base class.

    initialization: FPSInitialization = field(
        default=FPSInitialization.FARTHEST, converter=FPSInitialization
    )
    """See :func:`~baybe.utils.sampling_algorithms.farthest_point_sampling`.

    If the optional package 'fpsample' is used, only
    :attr:`~baybe.utils.sampling_algorithms.FPSInitialization.FARTHEST`
    is supported.
    """

    random_tie_break: bool = field(validator=instance_of(bool), kw_only=True)
    """See :func:`~baybe.utils.sampling_algorithms.farthest_point_sampling`.

    If the optional package 'fpsample' is used, only ``False`` is supported.
    """

    @initialization.validator
    def _validate_initialization(self, _, value):
        if active_settings.use_fpsample and value is not FPSInitialization.FARTHEST:
            raise ValueError(
                f"'{self.__class__.__name__}' is currently using the optional "
                f"'fpsample' package, which does not support the "
                f"'{self.initialization}' mode. "
                f"Please choose a supported initialization mode or deactivate "
                f"`fpsample` usage via the '{fields(Settings)._use_fpsample.alias}' "
                f"option in BayBE's settings."
            )

    @random_tie_break.default
    def _default_random_tie_break(self) -> bool:
        return self.initialization is not FPSInitialization.FARTHEST

    @random_tie_break.validator
    def _validate_random_tie_break(self, _, value):
        if active_settings.use_fpsample and value:
            raise ValueError(
                f"'{self.__class__.__name__}' is currently using the optional "
                f"'fpsample' package, which does not support random tie-breaking. "
                f"Either disable the mechanism by passing "
                f"'{fields(self.__class__).random_tie_break.name}=False' or deactivate "
                f"`fpsample` usage via the '{fields(Settings)._use_fpsample.alias}' "
                f"option in BayBE's settings."
            )

    @override
    def _recommend_discrete(
        self,
        subspace_discrete: SubspaceDiscrete,
        candidates_exp: pd.DataFrame,
        batch_size: int,
    ) -> pd.Index:
        # Fit scaler on entire search space
        from sklearn.preprocessing import StandardScaler

        # TODO [Scaling]: scaling should be handled by search space object
        scaler = StandardScaler()
        scaler.fit(subspace_discrete.comp_rep)

        # Scale and sample
        candidates_comp = subspace_discrete.transform(candidates_exp)
        candidates_scaled = np.ascontiguousarray(scaler.transform(candidates_comp))

        if active_settings.use_fpsample:
            from baybe._optional.fpsample import fps_sampling

            ilocs = fps_sampling(
                candidates_scaled,
                n_samples=batch_size,
            )
        else:
            ilocs = farthest_point_sampling(
                candidates_scaled,
                batch_size,
                initialization=self.initialization,
                random_tie_break=self.random_tie_break,
            )
        return candidates_comp.index[ilocs]

    @override
    def __str__(self) -> str:
        fields = [to_string("Compatibility", self.compatibility, single_line=True)]
        return to_string(self.__class__.__name__, *fields)