tabpackage 0.0.1__py3-none-any.whl

tabpackage/ensemble.py

@@ -0,0 +1,263 @@
+"""PyTorch ensembling utilities."""
+
+from collections.abc import Callable
+from typing import TypeAlias
+
+import torch
+from torch import Tensor
+
+_ScoreFn: TypeAlias = Callable[
+    [Tensor],  # Shape: (n_predictions, *)
+    Tensor,  # Shape: (n_predictions,)
+]
+
+
+def _validate_predictions(predictions: Tensor) -> None:
+    if predictions.ndim < 2:
+        raise ValueError(
+            'Predictions must have a dimension for ensemble members and at least one'
+            " dimension for each member's predictions"
+        )
+    if predictions.numel() == 0:
+        raise ValueError('Predictions must not be empty')
+
+
+def compute_unsanitized_ensemble_predictions(
+    member_predictions: Tensor, member_weights: Tensor | None
+) -> Tensor:
+    """Aggregate predictions of ensemble members into one _RAW_ prediction.
+
+    Warning:
+        Due to the floating-point arithmetic limitations:
+
+        - The computed raw predictions can contain values slightly beyond the range of
+          `member_predictions`, i.e. below/above their minimum/maximum values.
+          Transform the result as needed, e.g. with `.clamp(max=1.0)` for
+          probabilities.
+        - If `member_predictions` are probabilities, the computed distributions may
+          not sum to `1.0`.
+
+    Args:
+        member_predictions: the predictions of the ensemble members.
+            Shape: `(n_members, *)`.
+        member_weights: weights to aggregate the predictions with, normalized by their
+            sum. Pass `None` for equal weights. Shape: `(n_members,)`.
+    Returns:
+        The aggregated prediction. Shape: `(*)`.
+    """
+    _validate_predictions(member_predictions)
+
+    if member_weights is None:
+        return member_predictions.mean(0)
+    else:
+        if member_weights.ndim != 1:
+            raise ValueError('`member_weights` must be a one-dimensional tensor')
+        if len(member_weights) != len(member_predictions):
+            raise ValueError(
+                'The number of member weights must match the number of ensemble members'
+            )
+        return (
+            (member_weights / member_weights.sum()) @ member_predictions.flatten(1)
+        ).view(member_predictions.shape[1:])
+
+
+def topk_ensemble(predictions: Tensor, *, score_fn: _ScoreFn, k: int) -> Tensor:
+    """Build an equally weighted ensemble of the top-k individual performers.
+
+    Args:
+        predictions: the predictions to build the ensemble from.
+        score_fn: the scoring function taking multiple stacked predictions as input and
+            returning their performance scores (one per prediction). Higher scores are
+            better. Ties are resolved by selecting the earliest predictions.
+        k: the number of predictions to select, capped by the number available.
+    Returns:
+        The selected prediction indices.
+    """
+    _validate_predictions(predictions)
+    if k <= 0:
+        raise ValueError('`k` must be positive to select at least one ensemble member')
+
+    prediction_scores = score_fn(predictions)
+    if prediction_scores.ndim != 1:
+        raise ValueError(
+            'The scoring function must return a one-dimensional tensor of scores'
+        )
+    if len(prediction_scores) != len(predictions):
+        raise ValueError('The scoring function must return one score per prediction')
+
+    return (
+        torch.argmax(prediction_scores)[None]
+        if k == 1
+        else torch.argsort(prediction_scores, descending=True, stable=True)[:k]
+    )
+
+
+def greedy_ensemble(
+    predictions: Tensor,
+    *,
+    score_fn: _ScoreFn,
+    init_top_k: int = 1,
+    max_ensemble_size: int | None = None,
+    with_replacement: bool = False,
+) -> tuple[Tensor, Tensor | None]:
+    """Build an ensemble by greedily adding the best candidate one at a time.
+
+    The algorithm initializes the ensemble by unconditionally selecting the top-k
+    predictions based on their individual performance, and then greedily adds more
+    predictions one at a time, stopping when reaching the allowed ensemble size or when
+    adding more predictions does not improve performance.
+
+    Args:
+        predictions: the predictions to build the ensemble from.
+        score_fn: the scoring function taking multiple stacked predictions as input and
+            returning their performance scores (one per prediction). Higher scores are
+            better. Usually, it is based on the main task metric (e.g., accuracy,
+            _negative_ RMSE, etc.). The algorithm stops when the ensemble performance
+            stops improving according to this score.
+        init_top_k: the number of the best individual performers to unconditionally add
+            to the ensemble on initialization.
+        max_ensemble_size: the maximum allowed ensemble size. When
+            `with_replacement=False`, it is the maximum number of distinct ensemble
+            members. Otherwise, it is the maximum number of member selections that the
+            algorithm performs, including the unconditional selections in the beginning.
+            Must be specified when `with_replacement=True`.
+        with_replacement: if `True`, each prediction can be added to the ensemble any
+            number of times.
+    Returns:
+        The selected indices in original prediction order, and their selection counts
+        when `with_replacement=True`; otherwise, the second item is `None`.
+    """
+    _validate_predictions(predictions)
+    if not torch.is_floating_point(predictions):
+        raise TypeError('Predictions must be floating-point values')
+    if init_top_k <= 0:
+        raise ValueError(
+            '`init_top_k` must be positive to initialize the ensemble with at least one'
+            ' member'
+        )
+    if with_replacement and max_ensemble_size is None:
+        raise ValueError(
+            'Set `max_ensemble_size` to limit the number of ensemble members when'
+            ' `with_replacement=True`'
+        )
+    if max_ensemble_size is not None and init_top_k > max_ensemble_size:
+        raise ValueError('`init_top_k` must not exceed `max_ensemble_size`')
+
+    device = predictions.device
+
+    # Individual candidate scores.
+    prediction_scores = score_fn(predictions)
+
+    # Initial ensemble members.
+    init_idx = topk_ensemble(
+        predictions, score_fn=lambda _: prediction_scores, k=init_top_k
+    )
+
+    # Prediction-related constants.
+    n_predictions = len(predictions)
+    prediction_range_idx = torch.arange(n_predictions, device=device)
+
+    # The mask indicating if a given prediction is selected as an ensemble member.
+    ensemble_mask = torch.zeros(n_predictions, dtype=torch.bool, device=device)
+    ensemble_mask[init_idx] = True
+
+    if with_replacement:
+        # Each prediction's weight is the number of times it was selected.
+        ensemble_weights = torch.zeros(
+            n_predictions, dtype=predictions.dtype, device=device
+        )
+        ensemble_weights[init_idx] = 1.0
+    else:
+        ensemble_weights = None
+
+    # Current ensemble prediction.
+    ensemble_prediction = predictions[init_idx].mean(0)
+
+    # Current ensemble score.
+    ensemble_score = score_fn(ensemble_prediction[None])[0]
+
+    ensemble_size = len(init_idx)  # The number of distinct ensemble members.
+    ensemble_weight_sum = ensemble_size
+
+    # Without replacement, the weight sum equals the number of distinct members.
+    while ensemble_weight_sum < (
+        n_predictions
+        if max_ensemble_size is None  # True only when `with_replacement=False`.
+        else max_ensemble_size
+        if with_replacement
+        else min(n_predictions, max_ensemble_size)
+    ):
+        # Get the indices of predictions available for selection.
+        candidate_idx = (
+            prediction_range_idx
+            if with_replacement
+            else torch.nonzero_static(
+                ~ensemble_mask, size=n_predictions - ensemble_size
+            )[:, 0]
+        )
+
+        # Compute the ensemble predictions for all possible ensembles formed by adding
+        # one of the candidate predictions.
+        candidate_predictions = (
+            predictions.clone()
+            if with_replacement
+            else predictions.index_select(0, candidate_idx)
+        )
+        candidate_predictions.lerp_(
+            ensemble_prediction[None], ensemble_weight_sum / (ensemble_weight_sum + 1)
+        )
+
+        # Compute ensemble scores of all possible ensembles.
+        candidates_scores = score_fn(candidate_predictions)
+
+        # Get the best possible ensemble score.
+        best_candidate_score = candidates_scores.max()
+        # If the obtained score is no better than that of the current ensemble, stop the
+        # algorithm.
+        if best_candidate_score <= ensemble_score:
+            break
+
+        # Find all candidates with the best score.
+        candidate_local_mask = candidates_scores == best_candidate_score
+        best_candidate_local_idx = torch.nonzero(candidate_local_mask)[:, 0]
+        best_candidate_idx = candidate_idx[candidate_local_mask]
+        if len(best_candidate_idx) == 0:
+            raise ValueError(
+                'Cannot select the best ensemble candidate because the scoring function'
+                ' returned a NaN score'
+            )
+
+        if len(best_candidate_idx) == 1:
+            # Select the only candidate with the best score.
+            best_candidate_local_index = best_candidate_local_idx[0]
+            best_candidate_index = best_candidate_idx[0]
+        else:
+            # Break ties by individual score, then by the original prediction order.
+            best_candidate_local_index = best_candidate_local_idx[
+                torch.argmax(prediction_scores[best_candidate_idx])
+            ]
+            best_candidate_index = candidate_idx[best_candidate_local_index]
+
+        # Update the ensemble state.
+        ensemble_mask[best_candidate_index] = True
+        if ensemble_weights is not None:
+            ensemble_weights[best_candidate_index] += 1.0
+        ensemble_score = candidates_scores[best_candidate_local_index]
+        # Clone the only needed prediction and delete the rest to leave more free memory
+        # for the next iteration.
+        ensemble_prediction = candidate_predictions[best_candidate_local_index].clone()
+        del candidate_predictions
+        if (
+            ensemble_weights is None
+            or ensemble_weights[best_candidate_index].item() == 1.0
+        ):
+            ensemble_size += 1
+        ensemble_weight_sum += 1
+
+        if not with_replacement:
+            assert ensemble_size == ensemble_weight_sum
+
+    ensemble_idx = torch.nonzero_static(ensemble_mask, size=ensemble_size)[:, 0]
+    return ensemble_idx, (
+        None if ensemble_weights is None else ensemble_weights[ensemble_idx]
+    )

tabpackage/metrics.py

@@ -0,0 +1,253 @@
+"""Metric computation using PyTorch.
+
+Common conventions:
+
+- The pack dimension must be `tabpackage.nn.PACK_DIM`.
+- The `*_pack` functions (`accuracy_pack`, `mse_pack`, etc.) compute common metrics for
+  prediction packs; their distinct property is that the pack dimension is preserved
+  during the metric aggregation.
+- The argument naming follows that of Scikit-learn 1.9+. In particular, predicted
+  labels are accepted as `y_pred`, predicted probabilities as `y_proba` and decision
+  scores (the higher - the better) as `y_score`.
+- Prediction tensors are expected to have exactly one pack dimension.
+- Ground-truth labels with omitted pack and batch dimensions are broadcast along these
+  dimensions to match predictions. This in particular means that labels without the pack
+  dimension are shared between all pack members.
+"""
+
+import torch
+import torch.nn as nn  #  noqa: PLR0402
+from torch import Tensor
+
+import tabpackage.utils.torch
+from tabpackage.nn import PACK_DIM
+
+
+def _check_y_shapes(*, y_true: Tensor, y_pred: Tensor, n_trailing_dims: int) -> None:
+    assert n_trailing_dims >= 0
+
+    if y_pred.ndim < 1 + n_trailing_dims:
+        raise ValueError(
+            'The prediction tensor has too few dimensions to represent a pack of'
+            ' outputs'
+        )
+    if y_true.ndim > y_pred.ndim - n_trailing_dims:
+        raise ValueError('Labels have too many dimensions for the provided predictions')
+
+    if y_true.ndim > 0:
+        # Only missing leading dimensions may be broadcast; existing ones must match.
+        end = y_pred.ndim - n_trailing_dims
+        if y_true.shape != y_pred.shape[end - y_true.ndim : end]:
+            raise ValueError(
+                'The label and prediction shapes are incompatible; only missing leading'
+                ' label dimensions may be broadcast'
+            )
+
+    if y_pred.numel() == 0:
+        raise ValueError('Predictions must not be empty')
+
+
+def _mean_nonpack(x: Tensor) -> Tensor:
+    assert x.ndim >= 1  # Must have at least the pack dimension.
+
+    if x.ndim == 1:
+        # Relying on the `else` branch would be incorrect, because it would result in
+        # `x.mean(())`, which means aggregating across all dimensions in PyTorch.
+        # The only remaining dimension must be the pack dimension.
+        assert PACK_DIM == 0
+        return x  # No dimensions to aggregate across.
+    else:
+        nonpack_dims = tuple(dim for dim in range(x.ndim) if dim != PACK_DIM)
+        return x.mean(nonpack_dims)
+
+
+# The implementation follows `_mean_nonpack`.
+def _sum_nonpack(x: Tensor) -> Tensor:
+    assert x.ndim >= 1
+
+    if x.ndim == 1:
+        assert PACK_DIM == 0
+        return x
+    else:
+        nonpack_dims = tuple(dim for dim in range(x.ndim) if dim != PACK_DIM)
+        return x.sum(nonpack_dims)
+
+
+def accuracy_pack(*, y_true: Tensor, y_pred: Tensor) -> Tensor:
+    """Compute the accuracy for each pack member."""
+    if not tabpackage.utils.torch.is_integer(y_true):
+        raise TypeError('Labels must have an integer data type')
+    if y_true.dtype != y_pred.dtype:
+        raise TypeError('Labels and predictions must have the same data type')
+    _check_y_shapes(y_true=y_true, y_pred=y_pred, n_trailing_dims=0)
+
+    return _mean_nonpack((y_true == y_pred).float())
+
+
+def _midranks_pack(x: Tensor) -> Tensor:
+    """Compute the one-based ranks of sorted values, averaged within tied groups.
+
+    For example, the midranks of `[10, 20, 20, 20, 30]` are `[1, 3, 3, 3, 5]`: the
+    three tied values occupy the ranks 2, 3 and 4, hence the average rank 3.
+
+    The input must be sorted along the last (batch) dimension in the ascending order.
+    """
+    assert x.ndim == 2
+
+    n_objects = x.shape[-1]
+    indices = torch.arange(n_objects, device=x.device)
+    ones = torch.ones((x.shape[PACK_DIM], 1), dtype=torch.bool, device=x.device)
+
+    changes = x[:, 1:] != x[:, :-1]
+    # Whether a position starts a tied group.
+    starts = torch.cat([ones, changes], dim=-1)
+    # Whether a position ends a tied group.
+    ends = torch.cat([changes, ones], dim=-1)
+
+    # For each position, the first and the last position of its tied group. The running
+    # extrema propagate the group boundaries to all the positions within the group.
+    first = torch.cummax(torch.where(starts, indices, -1), dim=-1).values
+    last = torch.cummin(
+        torch.where(ends, indices, n_objects).flip(-1), dim=-1
+    ).values.flip(-1)
+    # `+ 1` turns the zero-based positions into one-based ranks.
+    return (first + last).to(x.dtype) / 2 + 1
+
+
+def binary_roc_auc_pack(*, y_true: Tensor, y_score: Tensor) -> Tensor:
+    """Compute the ROC AUC for each pack member."""
+    # Because of the implementation complexity, the supported input data types and
+    # shapes are limited to those that were actually tested.
+    if y_true.dtype != torch.int64:
+        raise TypeError('Labels must have the data type `torch.int64`')
+    if y_score.dtype not in (torch.float32, torch.float64):
+        raise TypeError(
+            'Scores must have the data type `torch.float32` or `torch.float64`'
+        )
+
+    if y_true.ndim != 1:
+        raise ValueError('The label tensor must have one dimension')
+    if y_score.ndim != 2:
+        raise ValueError('The score tensor must have two dimensions')
+    if y_true.shape[-1] != y_score.shape[-1]:
+        raise ValueError('Labels and scores must have the same batch size')
+    if y_score.numel() == 0:
+        raise ValueError('Scores must not be empty')
+
+    sorted_indices = torch.argsort(y_score, dim=-1, stable=True)
+    y_score_sorted = torch.take_along_dim(y_score, sorted_indices, dim=-1)
+    y_true_sorted = y_true[sorted_indices]
+
+    n_positives = y_true.sum()
+    n_negatives = len(y_true) - n_positives
+    n_pairs = n_positives * n_negatives
+
+    # The U statistic is the number of the (positive, negative) pairs ranked correctly,
+    # with the ties counted as one half.
+    positive_rank_sum = (_midranks_pack(y_score_sorted) * y_true_sorted).sum(-1)
+    u_statistic = positive_rank_sum - n_positives * (n_positives + 1) / 2
+
+    return torch.where(n_pairs == 0, torch.nan, u_statistic / n_pairs).to(y_score.dtype)
+
+
+def binary_cross_entropy_pack(*, y_true: Tensor, y_proba: Tensor) -> Tensor:
+    """Compute the binary cross-entropy for each pack member."""
+    if not tabpackage.utils.torch.is_integer(y_true):
+        raise TypeError('Labels must have an integer data type')
+    if not torch.is_floating_point(y_proba):
+        raise TypeError('Probabilities must be floating-point values')
+    _check_y_shapes(y_true=y_true, y_pred=y_proba, n_trailing_dims=0)
+
+    return _mean_nonpack(
+        nn.functional.binary_cross_entropy(
+            input=y_proba,
+            target=y_true.to(y_proba.dtype).expand_as(y_proba),
+            reduction='none',
+        )
+    )
+
+
+def multiclass_cross_entropy_pack(*, y_true: Tensor, y_proba: Tensor) -> Tensor:
+    """Compute the multiclass cross-entropy for each pack member."""
+    # Only class labels are supported as `y_true`; plus, this tensor is used to index
+    # another tensor. Hence the `torch.int64` requirement.
+    if y_true.dtype != torch.int64:
+        raise TypeError('Labels must have the int64 data type')
+    if not torch.is_floating_point(y_proba):
+        raise TypeError('Probabilities must be floating-point values')
+    _check_y_shapes(y_true=y_true, y_pred=y_proba, n_trailing_dims=1)
+
+    correct_class_probs = y_proba.gather(
+        -1, y_true.expand(y_proba.shape[:-1]).unsqueeze(-1)
+    ).squeeze(-1)
+    return _mean_nonpack(
+        torch.log(correct_class_probs)
+        .clamp_min(-100.0)  # The clamping follows `torch.nn.BCELoss`.
+        .neg()
+    )
+
+
+def mse_pack(*, y_true: Tensor, y_pred: Tensor) -> Tensor:
+    """Compute the mean squared error for each pack member."""
+    if not torch.is_floating_point(y_true):
+        raise TypeError('Labels must be floating-point values')
+    if y_true.dtype != y_pred.dtype:
+        raise TypeError('Labels and predictions must have the same data type')
+    _check_y_shapes(y_true=y_true, y_pred=y_pred, n_trailing_dims=0)
+
+    return _mean_nonpack(
+        nn.functional.mse_loss(y_pred, y_true.expand_as(y_pred), reduction='none')
+    )
+
+
+def mae_pack(*, y_true: Tensor, y_pred: Tensor) -> Tensor:
+    """Compute the mean absolute error for each pack member."""
+    if not torch.is_floating_point(y_true):
+        raise TypeError('Labels must be floating-point values')
+    if y_true.dtype != y_pred.dtype:
+        raise TypeError('Labels and predictions must have the same data type')
+    _check_y_shapes(y_true=y_true, y_pred=y_pred, n_trailing_dims=0)
+
+    return _mean_nonpack(
+        nn.functional.l1_loss(y_pred, y_true.expand_as(y_pred), reduction='none')
+    )
+
+
+def r2_pack(*, y_true: Tensor, y_pred: Tensor, force_finite: bool = True) -> Tensor:
+    """Compute the R2 score for each pack member."""
+    if not torch.is_floating_point(y_true):
+        raise TypeError('Labels must be floating-point values')
+    if y_true.dtype != y_pred.dtype:
+        raise TypeError('Labels and predictions must have the same data type')
+    _check_y_shapes(y_true=y_true, y_pred=y_pred, n_trailing_dims=0)
+
+    numerator = _sum_nonpack((y_pred - y_true).square())
+    if y_true.ndim < y_pred.ndim:
+        shape_without_pack_dim = torch.Size(
+            [x for i, x in enumerate(y_pred.shape) if i != PACK_DIM]
+        )
+        denominator = torch.sum(
+            torch.broadcast_to(
+                (y_true - y_true.mean()).square(), shape_without_pack_dim
+            )
+        )
+    else:
+        denominator = _sum_nonpack(
+            (
+                y_true
+                - _mean_nonpack(y_true).reshape(
+                    [-1 if dim == PACK_DIM else 1 for dim in range(y_true.ndim)]
+                )
+            ).square()
+        )
+    ratio = numerator / denominator
+
+    if force_finite:
+        # Map constant-target cases to scores of 1 (perfect) or 0 before checking
+        # the sample count below.
+        ratio = torch.where(
+            numerator == 0.0, 0.0, torch.where(denominator == 0.0, 1.0, ratio)
+        )
+
+    n_objects = y_pred.numel() // y_pred.shape[PACK_DIM]
+    return torch.full_like(ratio, torch.nan) if n_objects < 2 else 1.0 - ratio