Source code for caliber.regression.conformal_regression.cqr.base
from typing import Literal
import numpy as np
from caliber.regression.conformal_regression.base import (
ConformalizedScoreRegressionModel,
)
from caliber.utils.quantile_checks import (
both_quantile_check,
lower_quantile_check,
upper_quantile_check,
)
from caliber.utils.quantile_error import which_quantile_error
[docs]
class ConformalizedQuantileRegressionModel(ConformalizedScoreRegressionModel):
def __init__(
self,
confidence: float,
which_quantile: Literal["both", "lower", "upper"] = "both",
):
super().__init__(confidence=confidence)
self.which_quantile = which_quantile
[docs]
def fit(self, quantiles: np.ndarray, targets: np.ndarray) -> None:
if self.which_quantile == "both":
both_quantile_check(quantiles)
scores = np.maximum(quantiles[:, 0] - targets, targets - quantiles[:, 1])
elif self.which_quantile == "lower":
lower_quantile_check(quantiles)
scores = quantiles - targets
elif self.which_quantile == "upper":
upper_quantile_check(quantiles)
scores = targets - quantiles
else:
which_quantile_error(self.which_quantile)
super().fit(scores, targets)
[docs]
def predict(self, quantiles: np.ndarray) -> np.ndarray:
if self.which_quantile == "both":
both_quantile_check(quantiles)
lowers = quantiles[:, 0] - self._params
uppers = quantiles[:, 1] + self._params
bounds = np.stack((lowers, uppers), axis=1)
elif self.which_quantile == "lower":
lower_quantile_check(quantiles)
bounds = quantiles - self._params
elif self.which_quantile == "upper":
upper_quantile_check(quantiles)
bounds = quantiles + self._params
else:
which_quantile_error(self.which_quantile)
return bounds