tuning/optimal_ordinal.py

"""Tools for choosing multiple thresholds optimally under various decision criteria via dynamic programming.

The abstract machinery is contained in the classes:
- `OrdinalThresholding`
- `OptimalOrdinalThresholdingViaDynamicProgramming`
- `OptimalCostPerSampleOrdinalThresholding`
- `ClassWeightedOptimalCostPerSampleOrdinalThresholding`
- `OptimalCostPerClassOrdinalThresholding`
These can be subclassed to efficiently implement new decision criteria, depending on their structure.

The main intended user-facing classes are the subclasses implementing different decision criteria:
- `MaxAccuracyOrdinalThresholding`
- `MaxMacroRecallOrdinalThresholding`
- `MinAbsoluteErrorOrdinalThresholding`
- `MaxMacroPrecisionOrdinalThresholding`
- `MaxMacroF1OrdinalThresholding`

"""

from abc import ABC, abstractmethod
from typing import Literal, Optional, Union

import torch


class OrdinalThresholding(torch.nn.Module):
    """Basic 1d thresholding logic."""

    def __init__(self, num_classes: int):
        """Init thresholding module with the specified number of classes (one more than the number of thresholds)."""
        super().__init__()
        self.num_classes = num_classes
        self.register_buffer("thresholds", torch.zeros(num_classes - 1))
        self.thresholds: torch.Tensor

    def is_valid(self) -> bool:
        """Check whether the thresholds are monotone non-decreasing."""
        return all(torch.greater_equal(self.thresholds[1:], self.thresholds[:-1]))

    def forward(self, scores) -> torch.Tensor:
        """Find which thresholds each score lies between."""
        return torch.searchsorted(self.thresholds, scores)

    def tune_thresholds(
        self,
        *,
        scores: torch.Tensor,
        labels: torch.Tensor,
        available_thresholds: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Adapt the thresholds to the given data.

        This is essentially an abstract method, but for testing purposes it's helpful to be able to instantiate the
        class with a no-op version.

        Parameters
        ----------
        scores : a vector of `float` scores for each example in the validation set
        labels : a vector of `int` labels having the same shape as `scores` containing the corresponding labels
        available_thresholds : a vector of `float` score values over which to optimize choice of thresholds;
            `None`, then thresholds between every score in the validation set are allowed. +/- inf are always allowed.

        Returns
        -------
        scalar `float` mean cost on the validation set using optimal thresholds

        """


class OptimalOrdinalThresholdingViaDynamicProgramming(OrdinalThresholding, ABC):
    """Super-class for general dynamic programming implementations of ordinal threshold tuning.

    Subclasses implement different ways of computing the mean cost and corresponding DP step.

    """

    direction: Literal["min", "max"]  # provided by subclasses

    def __init__(self, num_classes: int):
        super().__init__(num_classes=num_classes)
        if self.direction not in ("min", "max"):
            raise ValueError(
                f"Got direction {self.direction!r}, expected 'min' or 'max'."
            )

    @abstractmethod
    def mean_cost(
        self, *, labels: torch.Tensor, preds: Union[int, torch.Tensor]
    ) -> torch.Tensor:
        """Compute the mean cost of assigning label(s) `preds` when the ground truth is `labels`."""

    def best_constant_output_classifier(self, labels: torch.Tensor):
        """Find the optimal mean cost of a constant-output classifier for given `labels` and the associated constant."""
        if self.direction == "min":
            optimize = torch.min
        else:
            optimize = torch.max
        optimum = optimize(
            torch.tensor(
                [
                    self.mean_cost(labels=labels, preds=c)
                    for c in range(self.num_classes)
                ],
                device=labels.device,
            ),
            0,
        )
        return optimum.values, optimum.indices

    @abstractmethod
    def dp_step(
        self,
        c_idx: int,
        *,
        scores: torch.Tensor,
        labels: torch.Tensor,
        available_thresholds: torch.Tensor,
        prev_cost: Optional[torch.Tensor] = None,
    ) -> (torch.Tensor, Optional[torch.Tensor]):
        """Given optimal cost `prev_cost` of classes < `c_idx`, optimize cost of `c_idx` as a function of threshold.

        Arguments
        ---------
            c_idx : current class index
            scores, labels, available_thresholds : see `tune_thresholds`
            prev_cost (optional float tensor) : optimal cost of classes < `c_idx` as a function of upper threshold
                for class `c_idx - 1`; ignored if `c_idx == 0`

        Returns
        -------
            cost: `cost[i]` is for choosing upper threshold of class `c_idx` equal to `available_thresholds[i]`
                when thresholds for lower classes are chosen optimally
            indices : to achieve `cost[i]`, optimal upper threshold for class `c_idx - 1` is
                `available_thresholds[indices[i]]`; `None` if `c_idx == 0`

        """

    def tune_thresholds(
        self,
        *,
        scores: torch.Tensor,
        labels: torch.Tensor,
        available_thresholds: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Set `self.thresholds` to optimize mean cost of given `scores` and `labels`.

        Arguments
        ---------
            scores (1d float tensor) : scores of examples on tuning dataset
            labels (1d int tensor) : labels in {0, ..., self.num_classes - 1} of same shape as scores
            available_thresholds (optional 1d float tensor) : thresholds which will be considered when tuning.
                +/-inf will be added automatically to ensure all examples are classified. If omitted, will
                insert thresholds between each element of sorted(unique(scores)).

        Returns
        -------
            float tensor : optimal mean cost achieved on the provided dataset at the tuned `self.thresholds`

        """
        inf = torch.tensor([torch.inf], device=scores.device)
        if available_thresholds is None:  # use all possible thresholds
            unique_scores = torch.unique(scores)
            available_thresholds = (unique_scores[:-1] + unique_scores[1:]) / 2.0
        # Always allow some classes to be omitted entirely by setting thresholds to +/- inf.
        # This simplifies the algorithm and also guarantees that the baseline constant-output
        # classifiers are feasible choices for tuning, which is needed to assure that the
        # optimum is at least as good as a constant-output classifier.
        available_thresholds = torch.concatenate(
            [
                -inf,
                available_thresholds,
                inf,
            ]
        )
        indices = torch.empty(
            (self.num_classes - 2, len(available_thresholds)),
            dtype=torch.int,
            device=scores.device,
        )

        # cost[j] = optimal total cost of items assigned pred <= c if the
        # threshold between class c and c+1 is available_thresholds[j] (by appropriate choice of lower thresholds).
        cost, _ = self.dp_step(
            c_idx=0,
            scores=scores,
            labels=labels,
            available_thresholds=available_thresholds,
        )
        for c in range(1, self.num_classes - 1):
            cost, indices[c - 1, :] = self.dp_step(
                c_idx=c,
                scores=scores,
                labels=labels,
                available_thresholds=available_thresholds,
                prev_cost=cost,
            )
        cost, best_index = self.dp_step(
            c_idx=self.num_classes - 1,
            scores=scores,
            labels=labels,
            available_thresholds=available_thresholds,
            prev_cost=cost,
        )
        if self.direction == "min":
            cost *= -1

        # Follow DP path backwards to find thresholds which optimized cost
        self.thresholds[self.num_classes - 2] = available_thresholds[
            best_index
        ]  # final threshold
        for c in range(self.num_classes - 3, -1, -1):  # counting down to zero
            best_index = indices[c, best_index.long()]
            self.thresholds[c] = available_thresholds[best_index.long()]

        return cost


def cumsum_with_0(t: torch.Tensor):
    return torch.nn.functional.pad(torch.cumsum(t, dim=0), (1, 0))


class OptimalCostPerSampleOrdinalThresholding(
    OptimalOrdinalThresholdingViaDynamicProgramming, ABC
):
    """Optimal 1d thresholding based on tuning thresholds to optimize the mean of a sample-wise cost function."""

    @abstractmethod
    def cost(self, *, labels: torch.Tensor, preds: Union[int, torch.Tensor]):
        """Compute the sample-wise cost of assigning label(s) `preds` when the ground truth is `labels`."""

    def mean_cost(
        self, *, labels: torch.Tensor, preds: Union[int, torch.Tensor]
    ) -> torch.Tensor:
        """Compute the mean cost of assigning label(s) `preds` when the ground truth is `labels`."""
        return torch.mean(self.cost(labels=labels, preds=preds))

    def dp_step(
        self,
        c_idx: int,
        *,
        scores: torch.Tensor,
        labels: torch.Tensor,
        available_thresholds: torch.Tensor,
        prev_cost: Optional[torch.Tensor] = None,
    ) -> (torch.Tensor, Optional[torch.Tensor]):
        """O(len(scores)) implementation for per-sample cost."""
        # Compute running_cost[i] = sum of costs of elements with score less than available_thresholds[i] if assigned label c
        item_costs = self.cost(labels=labels, preds=c_idx) / len(scores)
        if self.direction == "min":
            item_costs *= -1
        # move tensors to and from CPU because histogram has no CUDA implementation
        cost_new_class_by_thresh, _ = torch.histogram(
            scores.cpu().float(),
            weight=item_costs.cpu().float(),
            bins=available_thresholds.cpu().float(),
        )
        running_cost = cumsum_with_0(cost_new_class_by_thresh.to(labels.device))

        # Combine with running_cost with prev_cost
        if c_idx == 0:
            return running_cost, None
        diff = prev_cost - running_cost
        cummax = torch.cummax(diff, dim=0)
        cost = running_cost + cummax.values
        if c_idx == self.num_classes - 1:
            # -1 to always set the *upper* threshold for class `num_classes - 1` to include the rest of the data
            return cost[-1], cummax.indices[-1]
        return cost, cummax.indices


class MaxAccuracyOrdinalThresholding(OptimalCostPerSampleOrdinalThresholding):
    """Threshold to maximize accuracy."""

    direction = "max"

    def cost(self, *, labels: torch.Tensor, preds: Union[int, torch.Tensor]):
        return torch.eq(labels, preds).float()


class MaxMacroRecallOrdinalThresholding(OptimalCostPerSampleOrdinalThresholding):
    """Threshold to maximize macro-averaged recall."""

    direction = "max"

    def cost(self, *, labels: torch.Tensor, preds: Union[int, torch.Tensor]):
        counts = torch.bincount(labels, minlength=self.num_classes).float()
        ratios = counts.sum() / (self.num_classes * counts)
        return torch.eq(labels, preds).float() * torch.gather(
            ratios, 0, labels.type(torch.int64)
        )


class MinAbsoluteErrorOrdinalThresholding(OptimalCostPerSampleOrdinalThresholding):
    """Threshold to minimize mean absolute error."""

    direction = "min"

    def cost(self, *, labels: torch.Tensor, preds: Union[int, torch.Tensor]):
        return torch.abs(preds - labels).float()


class ClassWeightedOptimalCostPerSampleOrdinalThresholding(
    OptimalCostPerSampleOrdinalThresholding
):
    """Compute cost weighted equally over classes instead of equally over samples.

    This class takes another instance of OptimalCostPerSampleOrdinalThresholding
    which computes its cost independently for each sample and reweights the cost
    based on label frequencies.

    Note: this class depends on an implementation detail of its superclass:
      namely calling `self.cost` with the full tuning or eval set of labels,
      rather than a single label. This is required to do the re-weighting properly.

    """

    def __init__(self, unweighted_instance: OptimalCostPerSampleOrdinalThresholding):
        self.direction = unweighted_instance.direction
        super().__init__(unweighted_instance.num_classes)
        self.unweighted_instance = unweighted_instance

    def cost(self, *, labels: torch.Tensor, preds: Union[int, torch.Tensor]):
        counts = torch.bincount(labels, minlength=self.num_classes)
        (indices,) = torch.where(counts == 0)
        if len(indices) > 0:
            raise ValueError(
                f"Cannot compute class-weighted cost because classes {set(indices.tolist())} are missing."
            )
        unweighted_cost = self.unweighted_instance.cost(labels=labels, preds=preds)
        weights = len(labels) / (self.num_classes * counts[labels].float())
        return weights * unweighted_cost


class OptimalCostPerClassOrdinalThresholding(
    OptimalOrdinalThresholdingViaDynamicProgramming, ABC
):
    """General DP case for when the linear algorithm for per-sample costs is not applicable.

    Complexity depends on the implementation of `cost_matrix`.

    """

    @abstractmethod
    def cost_matrix(
        self,
        c_idx: int,
        *,
        scores: torch.Tensor,
        labels: torch.Tensor,
        available_thresholds: torch.Tensor,
        start: bool,
        end: bool,
    ) -> torch.Tensor:
        """Each output[i, j] = cost for when scores in range `available_thresholds[i:j]` are assigned label `c_idx`."""

    def mean_cost(
        self, *, labels: torch.Tensor, preds: Union[int, torch.Tensor]
    ) -> torch.Tensor:
        """Compute the mean cost of assigning label(s) `preds` when the ground truth is `labels`."""

        if isinstance(preds, int) or preds.numel() == 1:
            preds = preds * torch.ones_like(labels, dtype=torch.int)

        total_cost = torch.tensor(0.0, device=labels.device)
        for c_idx in range(self.num_classes):
            thresholds = torch.tensor([c_idx - 0.5, c_idx + 0.5], device=labels.device)
            total_cost += self.cost_matrix(
                c_idx, preds.float(), labels, thresholds, start=True, end=True
            )[0, 0]
        return total_cost / self.num_classes

    def dp_step(
        self,
        c_idx: int,
        *,
        scores: torch.Tensor,
        labels: torch.Tensor,
        available_thresholds: torch.Tensor,
        prev_cost: Optional[torch.Tensor] = None,
    ) -> (torch.Tensor, Optional[torch.Tensor]):
        cost_matrix = (
            self.cost_matrix(
                c_idx,
                scores=scores,
                labels=labels,
                available_thresholds=available_thresholds,
                start=c_idx == 0,
                end=c_idx == self.num_classes - 1,
            )
            / self.num_classes
        )
        if self.direction == "min":
            cost_matrix *= -1
        if prev_cost is not None:
            cost_matrix += prev_cost[:, None]
        max_ = torch.max(cost_matrix, dim=0)
        return max_.values, max_.indices


def _compute_metrics_matrices(
    scores: torch.Tensor,
    binary_labels: torch.Tensor,
    thresholds: torch.Tensor,
    start: bool = False,
    end: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Each output[i, j] = stats for when scores between thresholds[i] and thresolds[j] are assigned `True`.

    Helper function for `MaxMacroPrecisionOrdinalThresholding` and `MaxMacroF1OrdinalThresholding`

    Computed in O(len(thresholds)**2 + len(scores)*log(len(thresholds))) operations instead of the naive
    O(len(scores)*len(thresholds)**2) operations to compute each element of the output independently.

    Arguments
    ---------
        scores (float Tensor) : scores of labeled examples for which to compute metrics
        binary_labels (bool Tensor) : corresponding binary labels of same shape as `scores`
        thresholds (float Tensor) : thresholds between which to compute metrics
        start : compute only the first row of the output (lower threshold at its minimum value)
        end : compute only the last column of the output (upper threshold at its maximum value)

    Returns
    -------
        tp : tp[i, j] = number of true positives if scores between thresholds[i:j] are classified as positive
        tp_plus_fp: tp_plus_fp[i, j] = number of scores between thresholds[i:j]

    """
    # move tensors to and from CPU because histogram has no CUDA implementation
    scores = scores.float().cpu()
    thresholds = thresholds.float().cpu()
    labeled_true_by_thresh, _ = torch.histogram(
        scores,
        weight=binary_labels.float().cpu(),
        bins=thresholds,
    )
    count_by_thresh, _ = torch.histogram(
        scores,
        bins=thresholds,
    )
    running_labeled_true_by_thresh = cumsum_with_0(
        labeled_true_by_thresh.to(binary_labels.device)
    )
    running_count_by_thresh = cumsum_with_0(
        count_by_thresh.to(binary_labels.device).float()
    )

    def start_slice(t):
        return t[: (1 if start else None), None]

    def end_slice(t):
        return t[None, (-1 if end else None) :]

    tp = end_slice(running_labeled_true_by_thresh) - start_slice(
        running_labeled_true_by_thresh
    )
    tp_plus_fp = end_slice(running_count_by_thresh) - start_slice(
        running_count_by_thresh
    )
    return tp, tp_plus_fp


class MaxMacroPrecisionOrdinalThresholding(OptimalCostPerClassOrdinalThresholding):
    """Threshold to maximize macro-averaged precision."""

    direction = "max"

    def cost_matrix(
        self,
        c_idx: int,
        scores: torch.Tensor,
        labels: torch.Tensor,
        available_thresholds: torch.Tensor,
        start: bool,
        end: bool,
    ) -> torch.Tensor:
        tp, tp_plus_fp = _compute_metrics_matrices(
            scores, torch.eq(labels, c_idx), available_thresholds, start=start, end=end
        )
        safe_tp_plus_fp = torch.maximum(
            tp_plus_fp, torch.ones(1, device=tp_plus_fp.device)
        )
        return torch.where(torch.ge(tp_plus_fp, 0.0), tp / safe_tp_plus_fp, -torch.inf)


class MaxMacroF1OrdinalThresholding(OptimalCostPerClassOrdinalThresholding):
    """Threshold to maximize macro-averaged F1 score."""

    direction = "max"

    def cost_matrix(
        self,
        c_idx: int,
        scores: torch.Tensor,
        labels: torch.Tensor,
        available_thresholds: torch.Tensor,
        start: bool,
        end: bool,
    ) -> torch.Tensor:
        tp, tp_plus_fp = _compute_metrics_matrices(
            scores, torch.eq(labels, c_idx), available_thresholds, start=start, end=end
        )
        tp_plus_fn = torch.eq(labels, c_idx).float().sum()  # scalar
        safe_tp_plus_fp = torch.maximum(
            tp_plus_fp, torch.ones(1, device=tp_plus_fp.device)
        )
        return torch.where(
            torch.ge(tp_plus_fp, 0.0),
            2 * tp / (safe_tp_plus_fp + tp_plus_fn),
            -torch.inf,
        )