YohouPyfuncModel¶

class YohouPyfuncModel:
    """The object ``mlflow.pyfunc.load_model`` wraps for a yohou model.

    Each call works on a copy of the loaded forecaster: it observes ``y`` (and the
    other frames) when given, then predicts once from the updated state. Nothing a
    call changes is visible to the next call.

    Parameters
    ----------
    forecaster : yohou forecaster
        The loaded, fitted forecaster.
    default_type : str
        ``prediction_type`` used when a call does not name one.
    """

    def __init__(self, forecaster: Any, default_type: str) -> None:
        self._forecaster = forecaster
        self._default_type = default_type
        self._supported = supported_prediction_types(forecaster)

    def get_raw_model(self) -> Any:
        """Return the loaded forecaster itself, as ``PyFuncModel.get_raw_model`` expects.

        Returns
        -------
        yohou forecaster
            The forecaster this model wraps. Changing it changes later predictions.
        """
        return self._forecaster

    def predict(self, model_input: Any, params: Mapping[str, Any] | None = None) -> pl.DataFrame:
        """Observe the given frames on a copy of the forecaster, then predict once.

        Parameters
        ----------
        model_input : dict of str to polars.DataFrame
            A subset of ``y``, ``X_actual``, ``X_future`` and ``X_forecast``. With
            ``y``, the copy observes the frames before predicting. Without ``y``, it
            predicts from the saved state.
        params : mapping or None, default=None
            ``prediction_type``, ``forecasting_horizon``, ``coverage_rates`` and
            ``groups``. An empty ``coverage_rates`` or ``groups`` means the
            forecaster's default.

        Returns
        -------
        polars.DataFrame
            The frame the selected yohou predict method returns.

        Raises
        ------
        ValueError
            If ``prediction_type`` is unknown or unsupported by the forecaster, or
            ``coverage_rates`` is given for a prediction type other than ``interval``.
        """
        frames = validate_input(model_input)
        params = dict(params or {})
        kind = params.get("prediction_type") or self._default_type
        if kind not in self._supported:
            msg = f"prediction_type={kind!r} is not supported by this forecaster. Supported: {sorted(self._supported)}."
            raise ValueError(msg)
        horizon = params.get("forecasting_horizon")
        # yohou applies its defaults only for None, and MLflow sends [] for an omitted
        # list param: [] would raise on a non-panel forecaster for `groups`, select no
        # series on a panel one, and produce no intervals for `coverage_rates`.
        coverage_rates = list(params.get("coverage_rates") or []) or None
        groups = list(params.get("groups") or []) or None
        if coverage_rates is not None and kind != "interval":
            msg = f"coverage_rates only applies to prediction_type='interval', not {kind!r}."
            raise ValueError(msg)

        forecaster = copy.deepcopy(self._forecaster)
        if "y" in frames:
            forecaster.observe(
                frames["y"],
                X_actual=frames.get("X_actual"),
                groups=groups,
                X_future=frames.get("X_future"),
                X_forecast=frames.get("X_forecast"),
            )
        kwargs: dict[str, Any] = {
            "X_future": frames.get("X_future"),
            "X_forecast": frames.get("X_forecast"),
            "forecasting_horizon": None if horizon is None else int(horizon),
            "groups": groups,
        }
        if kind == "interval":
            kwargs["coverage_rates"] = coverage_rates
        return getattr(forecaster, _PREDICT_METHODS[kind])(**kwargs)