Source code for alf_tools.optimizer.acquisition_functions.core_set

# 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 numpy as np
import torch
from alf_core import AcquisitionFunction, Candidate, LabelledCandidates, State
from scipy.spatial.distance import cdist

_CDIST_CHUNK_SIZE = 512


def _to_numpy(features: np.ndarray | torch.Tensor) -> np.ndarray:
    """Convert model features to a numpy array.

    Args:
        features: Feature output from a model's featurise method.

    Returns:
        Numpy array representation of the features.

    Raises:
        ValueError: If features is None or cannot be converted to a numpy array.
    """
    if features is None:
        raise ValueError("featurise returned None; expected a numpy array or torch.Tensor")
    if isinstance(features, torch.Tensor):
        return features.detach().cpu().numpy()
    try:
        return np.asarray(features)
    except (TypeError, ValueError) as e:
        raise ValueError(
            f"featurise returned a value that cannot be converted to a numpy array: {e}"
        ) from e




class CoreSet(AcquisitionFunction):
    """Core-set acquisition function using greedy k-centres.

    Greedily selects candidates that maximise the minimum distance to
    the training set and previously selected candidates (greedy k-centres).

    Candidates are scored by their selection rank (n_select - step), so the
    first selected candidate receives the highest score and the last receives 1.
    Unselected candidates receive a score of 0.
    This is a maximising acquisition function that uses features as a
    2-D array of shape (n_inputs, d).
    """

    def __call__(
        self,
        search_candidates: list[Candidate],
        state: State,
    ) -> LabelledCandidates:
        """Compute CoreSet acquisition values for unlabelled candidates.
        The model's featurise() must return a 2-D array of shape (n_inputs, d).

        Args:
            search_candidates: List of unlabelled candidates to score.
            state: The task state containing the current datasets and surrogate model.

        Raises:
            ValueError: If featurise returns None or a non-array-like object.

        Returns:
            LabelledCandidates with CoreSet acquisition values.
        """
        if not search_candidates:
            return LabelledCandidates(candidates=[], labels=np.zeros(0))

        training_candidates = state.dataset.train_dataset.candidates
        features = state.surrogate.featurise(training_candidates + search_candidates)
        embeddings = _to_numpy(features)

        if embeddings.ndim != 2:
            raise ValueError(
                "featurise must return a 2-D array of shape (n_inputs, d), "
                f"got shape {embeddings.shape}"
            )

        n_train = len(training_candidates)
        n_cands = len(search_candidates)

        if len(embeddings) != n_train + n_cands:
            raise ValueError(
                f"featurise returned {len(embeddings)} rows for "
                f"{n_train} training + {n_cands} candidates (expected {n_train + n_cands})"
            )

        training_embs = embeddings[:n_train]
        candidate_embs = embeddings[n_train:]

        n_select = min(state.acq_batch_size, n_cands)

        if n_train == 0:
            min_dists = np.full(n_cands, np.inf)
        else:
            min_dists = np.empty(n_cands)
            for _i in range(0, n_cands, _CDIST_CHUNK_SIZE):
                _sl = candidate_embs[_i : _i + _CDIST_CHUNK_SIZE]
                min_dists[_i : _i + _CDIST_CHUNK_SIZE] = cdist(_sl, training_embs).min(axis=1)
        # Guard against re-selection: np.minimum zeroes selected entries in the
        # common case, but if remaining candidates share identical embeddings
        # (all pairwise distances = 0), selected_mask prevents argmax from
        # re-picking an already-selected index.
        selected_mask = np.zeros(n_cands, dtype=bool)
        acquisition_values = np.zeros(n_cands)

        for step in range(n_select):
            masked = np.where(~selected_mask, min_dists, -np.inf)
            best_idx = int(np.argmax(masked))
            acquisition_values[best_idx] = float(n_select - step)
            selected_mask[best_idx] = True
            dists_to_new = cdist(candidate_embs, candidate_embs[[best_idx]])[:, 0]
            min_dists = np.minimum(min_dists, dists_to_new)

        return LabelledCandidates(candidates=search_candidates, labels=acquisition_values)