stonks-oracle/services/intelligence_pipeline_v3/sentiment/calibrator.py

"""Sentiment probability calibration.

Applies isotonic or Platt calibration to raw FinBERT probabilities
to produce better-calibrated confidence estimates. The calibrator
preserves probability ordering (monotonicity for isotonic) while
improving expected calibration error (ECE).
"""

from __future__ import annotations

import logging
from typing import Literal

import numpy as np

logger = logging.getLogger(__name__)

# Default calibration version when no artifact is loaded
DEFAULT_CALIBRATION_VERSION = "uncalibrated"


class SentimentCalibrator:
    """Calibrates raw sentiment probabilities using isotonic or Platt scaling.

    The calibrator fits on a held-out calibration set from the Gold_Corpus
    and transforms raw model probabilities to better-calibrated values.

    Parameters
    ----------
    method
        Calibration method: "isotonic" or "platt".
    """

    def __init__(self, method: Literal["isotonic", "platt"] = "isotonic") -> None:
        self._method = method
        self._calibration_version = DEFAULT_CALIBRATION_VERSION
        self._fitted = False
        self._calibrators: list | None = None  # One per class

    @property
    def calibration_version(self) -> str:
        """Return the current calibration artifact version."""
        return self._calibration_version

    @property
    def is_fitted(self) -> bool:
        """Return whether the calibrator has been fitted."""
        return self._fitted

    @property
    def method(self) -> str:
        """Return the calibration method."""
        return self._method

    def fit(
        self,
        raw_probs: list[list[float]],
        true_labels: list[int],
        version: str = "v1.0",
    ) -> None:
        """Fit the calibrator on a calibration dataset.

        Parameters
        ----------
        raw_probs
            List of [positive, negative, neutral] probability vectors.
        true_labels
            True class labels: 0=positive, 1=negative, 2=neutral.
        version
            Version string for this calibration artifact.
        """
        if not raw_probs or not true_labels:
            raise ValueError("raw_probs and true_labels must not be empty")

        if len(raw_probs) != len(true_labels):
            raise ValueError("raw_probs and true_labels must have the same length")

        raw_array = np.array(raw_probs, dtype=np.float64)
        labels_array = np.array(true_labels, dtype=np.int32)

        n_classes = raw_array.shape[1] if raw_array.ndim > 1 else 3

        if self._method == "isotonic":
            self._fit_isotonic(raw_array, labels_array, n_classes)
        else:
            self._fit_platt(raw_array, labels_array, n_classes)

        self._calibration_version = version
        self._fitted = True
        logger.info(
            "Calibrator fitted: method=%s, samples=%d, version=%s",
            self._method,
            len(true_labels),
            version,
        )

    def calibrate(self, raw_probs: list[float]) -> list[float]:
        """Calibrate a single probability vector.

        Parameters
        ----------
        raw_probs
            Raw [positive, negative, neutral] probabilities.

        Returns
        -------
        list[float]
            Calibrated probabilities that sum to 1.0 and preserve
            relative ordering within each class.
        """
        if not self._fitted:
            # Pass through uncalibrated
            return list(raw_probs)

        calibrated = []
        for i, prob in enumerate(raw_probs):
            if self._calibrators and i < len(self._calibrators):
                cal = self._calibrators[i]
                cal_prob = float(cal.predict(np.array([[prob]]))[0])
                # Clamp to [0, 1]
                cal_prob = max(0.0, min(1.0, cal_prob))
                calibrated.append(cal_prob)
            else:
                calibrated.append(prob)

        # Normalize to sum to 1.0
        total = sum(calibrated)
        if total > 0:
            calibrated = [p / total for p in calibrated]
        else:
            calibrated = [1.0 / len(calibrated)] * len(calibrated)

        return calibrated

    def calibrate_batch(self, raw_probs_batch: list[list[float]]) -> list[list[float]]:
        """Calibrate a batch of probability vectors.

        Parameters
        ----------
        raw_probs_batch
            List of raw [positive, negative, neutral] probability vectors.

        Returns
        -------
        list[list[float]]
            Calibrated probability vectors.
        """
        return [self.calibrate(probs) for probs in raw_probs_batch]

    def _fit_isotonic(
        self,
        raw_array: np.ndarray,
        labels_array: np.ndarray,
        n_classes: int,
    ) -> None:
        """Fit isotonic regression calibrators per class."""
        from sklearn.isotonic import IsotonicRegression

        self._calibrators = []
        for cls_idx in range(n_classes):
            # Binary indicator: is this the true class?
            binary_labels = (labels_array == cls_idx).astype(np.float64)
            class_probs = raw_array[:, cls_idx]

            iso = IsotonicRegression(y_min=0.0, y_max=1.0, out_of_bounds="clip")
            iso.fit(class_probs, binary_labels)
            self._calibrators.append(iso)

    def _fit_platt(
        self,
        raw_array: np.ndarray,
        labels_array: np.ndarray,
        n_classes: int,
    ) -> None:
        """Fit Platt (logistic) scaling calibrators per class."""
        from sklearn.linear_model import LogisticRegression

        self._calibrators = []
        for cls_idx in range(n_classes):
            binary_labels = (labels_array == cls_idx).astype(np.int32)
            class_probs = raw_array[:, cls_idx].reshape(-1, 1)

            lr = LogisticRegression(solver="lbfgs", max_iter=1000)
            # Need at least 2 classes in binary labels
            if len(np.unique(binary_labels)) < 2:
                # If only one class present, use identity
                self._calibrators.append(_IdentityCalibrator())
            else:
                lr.fit(class_probs, binary_labels)
                self._calibrators.append(_PlattWrapper(lr))


class _IdentityCalibrator:
    """Pass-through calibrator when insufficient data for fitting."""

    def predict(self, x: np.ndarray) -> np.ndarray:
        return x.ravel()


class _PlattWrapper:
    """Wrapper that extracts probability of the positive class."""

    def __init__(self, lr) -> None:  # noqa: ANN001
        self._lr = lr

    def predict(self, x: np.ndarray) -> np.ndarray:
        return self._lr.predict_proba(x)[:, 1]