core.py

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/src/core/core.ipynb.

# %% auto 0
__all__ = ['StatsForecast']

# %% ../nbs/src/core/core.ipynb 6
import datetime as dt
import errno
import inspect
import logging
import os
import pickle
import re
import reprlib
import time
import warnings
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import numpy as np
import pandas as pd
import utilsforecast.processing as ufp
from fugue.execution.factory import (
    make_execution_engine,
    try_get_context_execution_engine,
)
from threadpoolctl import threadpool_limits
from tqdm.auto import tqdm
from triad import conditional_dispatcher
from utilsforecast.compat import DataFrame, pl_DataFrame, pl_Series
from utilsforecast.grouped_array import GroupedArray as BaseGroupedArray
from utilsforecast.validation import ensure_time_dtype, validate_freq

from .utils import ConformalIntervals

# %% ../nbs/src/core/core.ipynb 7
if __name__ == "__main__":
    logging.basicConfig(
        format="%(asctime)s %(name)s %(levelname)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
logger = logging.getLogger(__name__)

# %% ../nbs/src/core/core.ipynb 10
class GroupedArray(BaseGroupedArray):

    def __eq__(self, other):
        if not hasattr(other, "data") or not hasattr(other, "indptr"):
            return False
        return np.allclose(self.data, other.data) and np.array_equal(
            self.indptr, other.indptr
        )

    def fit(self, models, fallback_model=None):
        fm = np.full((self.n_groups, len(models)), np.nan, dtype=object)
        for i, grp in enumerate(self):
            y = grp[:, 0] if grp.ndim == 2 else grp
            X = grp[:, 1:] if (grp.ndim == 2 and grp.shape[1] > 1) else None
            for i_model, model in enumerate(models):
                try:
                    new_model = model.new()
                    fm[i, i_model] = new_model.fit(y=y, X=X)
                except Exception as error:
                    if fallback_model is not None:
                        new_fallback_model = fallback_model.new()
                        new_fallback_model.alias = model.alias
                        fm[i, i_model] = new_fallback_model.fit(y=y, X=X)
                    else:
                        raise error
        return fm

    def _get_cols(self, models, attr, h, X, level=tuple()):
        n_models = len(models)
        cuts = np.full(n_models + 1, fill_value=0, dtype=np.int32)
        has_level_models = np.full(n_models, fill_value=False, dtype=bool)
        cuts[0] = 0
        for i_model, model in enumerate(models):
            len_cols = 1  # mean
            has_level = (
                "level" in inspect.signature(getattr(model, attr)).parameters
                and len(level) > 0
            )
            has_level_models[i_model] = has_level
            if has_level:
                len_cols += 2 * len(level)  # levels
            cuts[i_model + 1] = len_cols + cuts[i_model]
        return cuts, has_level_models

    def _output_fcst(self, models, attr, h, X, level=tuple()):
        # returns empty output according to method
        cuts, has_level_models = self._get_cols(
            models=models, attr=attr, h=h, X=X, level=level
        )
        out = np.full(
            (self.n_groups * h, cuts[-1]), fill_value=np.nan, dtype=np.float32
        )
        return out, cuts, has_level_models

    def predict(self, fm, h, X=None, level=tuple()):
        # fm stands for fitted_models
        # and fm should have fitted_model
        fcsts, cuts, has_level_models = self._output_fcst(
            models=fm[0], attr="predict", h=h, X=X, level=level
        )
        matches = ["mean", "lo", "hi"]
        cols = []
        for i_model in range(fm.shape[1]):
            has_level = has_level_models[i_model]
            kwargs = {}
            if has_level:
                kwargs["level"] = level
            for i, _ in enumerate(self):
                if X is not None:
                    X_ = X[i]
                else:
                    X_ = None
                res_i = fm[i, i_model].predict(h=h, X=X_, **kwargs)
                cols_m = [
                    key
                    for key in res_i.keys()
                    if any(key.startswith(m) for m in matches)
                ]
                fcsts_i = np.vstack([res_i[key] for key in cols_m]).T
                model_name = repr(fm[i, i_model])
                cols_m = [
                    f"{model_name}" if col == "mean" else f"{model_name}-{col}"
                    for col in cols_m
                ]
                if fcsts_i.ndim == 1:
                    fcsts_i = fcsts_i[:, None]
                fcsts[i * h : (i + 1) * h, cuts[i_model] : cuts[i_model + 1]] = fcsts_i
            cols += cols_m
        return fcsts, cols

    def fit_predict(self, models, h, X=None, level=tuple()):
        # fitted models
        fm = self.fit(models=models)
        # forecasts
        fcsts, cols = self.predict(fm=fm, h=h, X=X, level=level)
        return fm, fcsts, cols

    def forecast(
        self,
        models,
        h,
        fallback_model=None,
        fitted=False,
        X=None,
        level=tuple(),
        verbose=False,
        target_col="y",
    ):
        fcsts, cuts, has_level_models = self._output_fcst(
            models=models, attr="forecast", h=h, X=X, level=level
        )
        matches = ["mean", "lo", "hi"]
        matches_fitted = ["fitted", "fitted-lo", "fitted-hi"]
        if fitted:
            # for the moment we dont return levels for fitted values in
            # forecast mode
            fitted_vals = np.full(
                (self.data.shape[0], 1 + cuts[-1]), np.nan, dtype=np.float32
            )
            if self.data.ndim == 1:
                fitted_vals[:, 0] = self.data
            else:
                fitted_vals[:, 0] = self.data[:, 0]
        iterable = tqdm(
            enumerate(self), disable=(not verbose), total=len(self), desc="Forecast"
        )
        times = {repr(m): 0.0 for m in models}
        for i, grp in iterable:
            y_train = grp[:, 0] if grp.ndim == 2 else grp
            X_train = grp[:, 1:] if (grp.ndim == 2 and grp.shape[1] > 1) else None
            if X is not None:
                X_f = X[i]
            else:
                X_f = None
            cols = []
            cols_fitted = []
            for i_model, model in enumerate(models):
                has_level = has_level_models[i_model]
                kwargs = {}
                if has_level:
                    kwargs["level"] = level
                start = time.perf_counter()
                try:
                    res_i = model.forecast(
                        h=h, y=y_train, X=X_train, X_future=X_f, fitted=fitted, **kwargs
                    )
                except Exception as error:
                    if fallback_model is not None:
                        res_i = fallback_model.forecast(
                            h=h,
                            y=y_train,
                            X=X_train,
                            X_future=X_f,
                            fitted=fitted,
                            **kwargs,
                        )
                    else:
                        raise error
                times[repr(model)] += time.perf_counter() - start
                cols_m = [
                    key
                    for key in res_i.keys()
                    if any(key.startswith(m) for m in matches)
                ]
                fcsts_i = np.vstack([res_i[key] for key in cols_m]).T
                cols_m = [
                    f"{repr(model)}" if col == "mean" else f"{repr(model)}-{col}"
                    for col in cols_m
                ]
                if fcsts_i.ndim == 1:
                    fcsts_i = fcsts_i[:, None]
                fcsts[i * h : (i + 1) * h, cuts[i_model] : cuts[i_model + 1]] = fcsts_i
                cols += cols_m
                if fitted:
                    cols_m_fitted = [
                        key
                        for key in res_i.keys()
                        if any(key.startswith(m) for m in matches_fitted)
                    ]
                    fitted_i = np.vstack([res_i[key] for key in cols_m_fitted]).T
                    cols_m_fitted = [
                        (
                            f"{repr(model)}"
                            if col == "fitted"
                            else f"{repr(model)}-{col.replace('fitted-', '')}"
                        )
                        for col in cols_m_fitted
                    ]
                    fitted_vals[
                        self.indptr[i] : self.indptr[i + 1],
                        (cuts[i_model] + 1) : (cuts[i_model + 1] + 1),
                    ] = fitted_i
                    cols_fitted += cols_m_fitted
        result = {"forecasts": fcsts, "cols": cols, "times": times}
        if fitted:
            result["fitted"] = {"values": fitted_vals}
            result["fitted"]["cols"] = [target_col] + cols_fitted
        return result

    def cross_validation(
        self,
        models,
        h,
        test_size,
        fallback_model=None,
        step_size=1,
        input_size=None,
        fitted=False,
        level=tuple(),
        refit=True,
        verbose=False,
        target_col="y",
    ):
        # output of size: (ts, window, h)
        if (test_size - h) % step_size:
            raise Exception("`test_size - h` should be module `step_size`")
        n_windows = int((test_size - h) / step_size) + 1
        n_models = len(models)
        cuts, has_level_models = self._get_cols(
            models=models, attr="forecast", h=h, X=None, level=level
        )
        # first column of out is the actual y
        out = np.full(
            (self.n_groups, n_windows, h, 1 + cuts[-1]), np.nan, dtype=np.float32
        )
        if fitted:
            fitted_vals = np.full(
                (self.data.shape[0], n_windows, n_models + 1), np.nan, dtype=np.float32
            )
            fitted_idxs = np.full((self.data.shape[0], n_windows), False, dtype=bool)
            last_fitted_idxs = np.full_like(fitted_idxs, False, dtype=bool)
        matches = ["mean", "lo", "hi"]
        steps = list(range(-test_size, -h + 1, step_size))
        for i_ts, grp in enumerate(self):
            iterable = tqdm(
                enumerate(steps, start=0),
                desc=f"Cross Validation Time Series {i_ts + 1}",
                disable=(not verbose),
                total=len(steps),
            )
            fitted_models = [None for _ in range(n_models)]
            for i_window, cutoff in iterable:
                should_fit = i_window == 0 or (refit > 0 and i_window % refit == 0)
                end_cutoff = cutoff + h
                in_size_disp = cutoff if input_size is None else input_size
                y = grp[(cutoff - in_size_disp) : cutoff]
                y_train = y[:, 0] if y.ndim == 2 else y
                X_train = y[:, 1:] if (y.ndim == 2 and y.shape[1] > 1) else None
                y_test = grp[cutoff:] if end_cutoff == 0 else grp[cutoff:end_cutoff]
                X_future = (
                    y_test[:, 1:]
                    if (y_test.ndim == 2 and y_test.shape[1] > 1)
                    else None
                )
                out[i_ts, i_window, :, 0] = y_test[:, 0] if y.ndim == 2 else y_test
                if fitted:
                    fitted_vals[self.indptr[i_ts] : self.indptr[i_ts + 1], i_window, 0][
                        (cutoff - in_size_disp) : cutoff
                    ] = y_train
                    fitted_idxs[self.indptr[i_ts] : self.indptr[i_ts + 1], i_window][
                        (cutoff - in_size_disp) : cutoff
                    ] = True
                    last_fitted_idxs[
                        self.indptr[i_ts] : self.indptr[i_ts + 1], i_window
                    ][cutoff - 1] = True
                cols = [target_col]
                for i_model, model in enumerate(models):
                    has_level = has_level_models[i_model]
                    kwargs = {}
                    if has_level:
                        kwargs["level"] = level
                    # this is implemented like this because not all models have a forward method
                    # so we can't do fit + forward
                    if refit is True:
                        forecast_kwargs = dict(
                            h=h,
                            y=y_train,
                            X=X_train,
                            X_future=X_future,
                            fitted=fitted,
                            **kwargs,
                        )
                        try:
                            res_i = model.forecast(**forecast_kwargs)
                        except Exception as error:
                            if fallback_model is None:
                                raise error
                            res_i = fallback_model.forecast(**forecast_kwargs)
                    else:
                        if should_fit:
                            try:
                                fitted_models[i_model] = model.fit(y=y_train, X=X_train)
                            except Exception as error:
                                if fallback_model is None:
                                    raise error
                                fitted_models[i_model] = fallback_model.new().fit(
                                    y=y_train, X=X_train
                                )
                        res_i = fitted_models[i_model].forward(
                            h=h,
                            y=y_train,
                            X=X_train,
                            X_future=X_future,
                            fitted=fitted,
                            **kwargs,
                        )
                    cols_m = [
                        key
                        for key in res_i.keys()
                        if any(key.startswith(m) for m in matches)
                    ]
                    fcsts_i = np.vstack([res_i[key] for key in cols_m]).T
                    cols_m = [
                        f"{repr(model)}" if col == "mean" else f"{repr(model)}-{col}"
                        for col in cols_m
                    ]
                    out[
                        i_ts, i_window, :, (1 + cuts[i_model]) : (1 + cuts[i_model + 1])
                    ] = fcsts_i
                    if fitted:
                        fitted_vals[
                            self.indptr[i_ts] : self.indptr[i_ts + 1],
                            i_window,
                            i_model + 1,
                        ][(cutoff - in_size_disp) : cutoff] = res_i["fitted"]
                    cols += cols_m
        result = {"forecasts": out.reshape(-1, 1 + cuts[-1]), "cols": cols}
        if fitted:
            result["fitted"] = {
                "values": fitted_vals,
                "idxs": fitted_idxs,
                "last_idxs": last_fitted_idxs,
                "cols": [target_col] + [repr(model) for model in models],
            }
        return result

    def take(self, idxs):
        data, indptr = super().take(idxs)
        return GroupedArray(data, indptr)

    def split(self, n_chunks):
        n_chunks = min(n_chunks, self.n_groups)
        return [
            self.take(idxs) for idxs in np.array_split(range(self.n_groups), n_chunks)
        ]

    def split_fm(self, fm, n_chunks):
        return [
            fm[idxs]
            for idxs in np.array_split(range(self.n_groups), n_chunks)
            if idxs.size
        ]

    def _single_threaded_fit(self, models, fallback_model=None):
        with threadpool_limits(limits=1):
            return self.fit(models=models, fallback_model=fallback_model)

    def _single_threaded_predict(self, fm, h, X=None, level=tuple()):
        with threadpool_limits(limits=1):
            return self.predict(fm=fm, h=h, X=X, level=level)

    def _single_threaded_fit_predict(self, models, h, X=None, level=tuple()):
        with threadpool_limits(limits=1):
            return self.fit_predict(models=models, h=h, X=X, level=level)

    def _single_threaded_forecast(
        self,
        models,
        h,
        fallback_model=None,
        fitted=False,
        X=None,
        level=tuple(),
        verbose=False,
        target_col="y",
    ):
        with threadpool_limits(limits=1):
            return self.forecast(
                models=models,
                h=h,
                fallback_model=fallback_model,
                fitted=fitted,
                X=X,
                level=level,
                verbose=verbose,
                target_col=target_col,
            )

    def _single_threaded_cross_validation(
        self,
        models,
        h,
        test_size,
        fallback_model=None,
        step_size=1,
        input_size=None,
        fitted=False,
        level=tuple(),
        refit=True,
        verbose=False,
        target_col="y",
    ):
        with threadpool_limits(limits=1):
            return self.cross_validation(
                models=models,
                h=h,
                test_size=test_size,
                fallback_model=fallback_model,
                step_size=step_size,
                input_size=input_size,
                fitted=fitted,
                level=level,
                refit=refit,
                verbose=verbose,
                target_col=target_col,
            )

# %% ../nbs/src/core/core.ipynb 24
def _get_n_jobs(n_groups, n_jobs):
    if n_jobs == -1 or (n_jobs is None):
        actual_n_jobs = os.cpu_count()
    else:
        actual_n_jobs = n_jobs
    return min(n_groups, actual_n_jobs)

# %% ../nbs/src/core/core.ipynb 27
def _warn_df_constructor():
    warnings.warn(
        "The `df` argument of the StatsForecast constructor as well as reusing stored "
        "dfs from other methods is deprecated and will raise an error in a future version. "
        "Please provide the `df` argument to the corresponding method instead, e.g. fit/forecast.",
        category=FutureWarning,
    )


def _maybe_warn_sort_df(sort_df):
    if not sort_df:
        warnings.warn(
            "The `sort_df` argument is deprecated and will be removed in a future version. "
            "You can leave it to its default value (True) to supress this warning",
            category=FutureWarning,
        )


def _warn_id_as_idx():
    warnings.warn(
        "In a future version the predictions will have the id as a column. "
        "You can set the `NIXTLA_ID_AS_COL` environment variable "
        "to adopt the new behavior and to suppress this warning.",
        category=FutureWarning,
    )


def _id_as_idx() -> bool:
    return not bool(os.getenv("NIXTLA_ID_AS_COL", ""))

# %% ../nbs/src/core/core.ipynb 28
_param_descriptions = {
    "freq": """freq : str or int
            Frequency of the data. Must be a valid pandas or polars offset alias, or an integer.""",
    "df": """df : pandas or polars DataFrame, optional (default=None)
            DataFrame with ids, times, targets and exogenous.""",
    "sort_df": """sort_df : bool (default=True)
            Sort `df` by ids and times.""",
    "fallback_model": """fallback_model : Any, optional (default=None)
            Any, optional (default=None)
            Model to be used if a model fails.
            Only works with the `forecast` and `cross_validation` methods.""",
    "id_col": """id_col : str (default='unique_id')
            Column that identifies each serie.""",
    "time_col": """time_col : str (default='ds')
            Column that identifies each timestep, its values can be timestamps or integers.""",
    "target_col": """target_col : str (default='y')
            Column that contains the target.""",
    "h": """h : int
            Forecast horizon.""",
    "X_df": """X_df : pandas or polars DataFrame, optional (default=None)
            DataFrame with ids, times and future exogenous.""",
    "level": """level : List[float], optional (default=None)
            Confidence levels between 0 and 100 for prediction intervals.""",
    "prediction_intervals": """prediction_intervals : ConformalIntervals, optional (default=None)
            Configuration to calibrate prediction intervals (Conformal Prediction).""",
    "fitted": """fitted : bool (default=False)
            Store in-sample predictions.""",
    "n_jobs": """n_jobs : int (default=1)
            Number of jobs used in the parallel processing, use -1 for all cores.""",
    "verbose": """verbose : bool (default=True)
            Prints TQDM progress bar when `n_jobs=1`.""",
    "models": """models : List[Any]
            List of instantiated objects models.StatsForecast.""",
    "n_windows": """n_windows : int (default=1)
            Number of windows used for cross validation.""",
    "step_size": """step_size : int (default=1)
            Step size between each window.""",
    "test_size": """test_size : int, optional (default=None)
            Length of test size. If passed, set `n_windows=None`.""",
    "input_size": """input_size : int, optional (default=None)
            Input size for each window, if not none rolled windows.""",
    "refit": """refit : bool or int (default=True)
            Wether or not refit the model for each window.
            If int, train the models every `refit` windows.""",
}

# %% ../nbs/src/core/core.ipynb 29
class _StatsForecast:
    """The `StatsForecast` class allows you to efficiently fit multiple `StatsForecast` models
    for large sets of time series. It operates on a DataFrame `df` with at least three columns
    ids, times and targets.

    The class has memory-efficient `StatsForecast.forecast` method that avoids storing partial
    model outputs. While the `StatsForecast.fit` and `StatsForecast.predict` methods with
    Scikit-learn interface store the fitted models.

    The `StatsForecast` class offers parallelization utilities with Dask, Spark and Ray back-ends.
    See distributed computing example [here](https://github.com/Nixtla/statsforecast/tree/main/experiments/ray).
    """

    def __init__(
        self,
        models: List[Any],
        freq: Union[str, int],
        n_jobs: int = 1,
        df: Optional[DataFrame] = None,
        sort_df: bool = True,
        fallback_model: Optional[Any] = None,
        verbose: bool = False,
    ):
        """Train statistical models.

        Parameters
        ----------
        {models}
        {freq}
        {n_jobs}
        {df}
        {sort_df}
        {fallback_model}
        {verbose}
        """
        # TODO @fede: needed for residuals, think about it later
        self.models = models
        self._validate_model_names()
        self.freq = freq
        self.n_jobs = n_jobs
        self.fallback_model = fallback_model
        self.verbose = verbose
        if df is not None:
            _warn_df_constructor()
            self._prepare_fit(df=df, sort_df=sort_df)
        else:
            _maybe_warn_sort_df(sort_df)

    __init__.__doc__ = __init__.__doc__.format(**_param_descriptions)  # type: ignore[union-attr]

    def _validate_model_names(self):
        # Some test models don't have alias
        names = [getattr(model, "alias", lambda: None) for model in self.models]
        names = [x for x in names if x is not None]
        if len(names) != len(set(names)):
            raise ValueError(
                "Model names must be unique. You can use `alias` to set a unique name for each model."
            )

    def _prepare_fit(
        self,
        df: Optional[DataFrame],
        sort_df: bool = True,
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
    ) -> None:
        if df is None:
            if not hasattr(self, "ga"):
                raise ValueError("You must provide the `df` argument.")
            _warn_df_constructor()
            return
        df = ensure_time_dtype(df, time_col)
        validate_freq(df[time_col], self.freq)
        if isinstance(df, pd.DataFrame) and df.index.name == id_col:
            warnings.warn(
                "Passing unique_id as the index is deprecated. "
                "Please provide it as a column instead.",
                category=FutureWarning,
            )
            df = df.reset_index()
        _maybe_warn_sort_df(sort_df)
        self.uids, last_times, data, indptr, sort_idxs = ufp.process_df(
            df, id_col, time_col, target_col
        )
        if isinstance(df, pd.DataFrame):
            self.last_dates = pd.Index(last_times, name=time_col)
        else:
            self.last_dates = pl_Series(last_times)
        self.ga = GroupedArray(data, indptr)
        self.og_dates = df[time_col].to_numpy()
        if sort_idxs is not None:
            self.og_dates = self.og_dates[sort_idxs]
        self.n_jobs = _get_n_jobs(len(self.ga), self.n_jobs)
        self.df_constructor = type(df)
        self.id_col = id_col
        self.time_col = time_col
        self.target_col = target_col
        self._exog = [c for c in df.columns if c not in (id_col, time_col, target_col)]

    def _validate_sizes_for_prediction_intervals(
        self,
        prediction_intervals: Optional[ConformalIntervals],
        offset: int = 0,
    ) -> None:
        if prediction_intervals is None:
            return
        sizes = np.diff(self.ga.indptr) - offset
        # the absolute minimum requires two windows
        min_samples = 2 * prediction_intervals.h + 1
        if np.any(sizes < min_samples):
            raise ValueError(
                f"Minimum samples for computing prediction intervals are {min_samples + offset:,}, "
                "some series have less. Please remove them or adjust the horizon."
            )
        # required samples for current configuration
        required_samples = prediction_intervals.n_windows * prediction_intervals.h + 1
        if np.any(sizes < required_samples):
            warnings.warn(
                f"Prediction intervals settings require at least {required_samples + offset:,} samples, "
                "some series have less and will use less windows."
            )

    def _set_prediction_intervals(
        self, prediction_intervals: Optional[ConformalIntervals]
    ) -> None:
        for model in self.models:
            interval = getattr(model, "prediction_intervals", None)
            if interval is None:
                setattr(model, "prediction_intervals", prediction_intervals)

    def fit(
        self,
        df: Optional[DataFrame] = None,
        sort_df: bool = True,
        prediction_intervals: Optional[ConformalIntervals] = None,
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
    ):
        """Fit statistical models.

        Fit `models` to a large set of time series from DataFrame `df`
        and store fitted models for later inspection.

        Parameters
        ----------
        {df}
            If None, the `StatsForecast` class should have been instantiated using `df`.
        {sort_df}
        {prediction_intervals}
        {id_col}
        {time_col}
        {target_col}

        Returns
        -------
        self : StatsForecast
            Returns with stored `StatsForecast` fitted `models`.
        """
        self._prepare_fit(
            df=df,
            sort_df=sort_df,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )
        self._validate_sizes_for_prediction_intervals(prediction_intervals)
        self._set_prediction_intervals(prediction_intervals=prediction_intervals)
        if self.n_jobs == 1:
            self.fitted_ = self.ga.fit(
                models=self.models, fallback_model=self.fallback_model
            )
        else:
            self.fitted_ = self._fit_parallel()
        return self

    fit.__doc__ = fit.__doc__.format(**_param_descriptions)  # type: ignore[union-attr]

    def _make_future_df(self, h: int):
        start_dates = ufp.offset_times(self.last_dates, freq=self.freq, n=1)
        dates = ufp.time_ranges(start_dates, freq=self.freq, periods=h)
        uids = ufp.repeat(self.uids, n=h)
        df = self.df_constructor({self.id_col: uids, self.time_col: dates})
        if isinstance(df, pd.DataFrame):
            if _id_as_idx():
                _warn_id_as_idx()
                df = df.set_index(self.id_col)
            else:
                df = df.reset_index(drop=True)
        return df

    def _parse_X_level(
        self, h: int, X: Optional[DataFrame], level: Optional[List[int]]
    ):
        if level is None:
            level = []
        if X is None:
            return X, level
        expected_shape = (h * len(self.ga), self.ga.data.shape[1] + 1)
        if X.shape != expected_shape:
            raise ValueError(
                f"Expected X to have shape {expected_shape}, but got {X.shape}"
            )
        _, _, data, indptr, _ = ufp.process_df(X, self.id_col, self.time_col, None)
        return GroupedArray(data, indptr), level

    def _validate_exog(self, X_df: Optional[DataFrame] = None) -> None:
        if not any(m.uses_exog for m in self.models) or not self._exog:
            return
        err_msg = (
            f"Models require the following exogenous features {self._exog} "
            "for the forecasting step. Please provide them through `X_df`."
        )
        if X_df is None:
            raise ValueError(err_msg)
        missing_exog = [c for c in self._exog if c not in X_df.columns]
        if missing_exog:
            raise ValueError(err_msg)

    def predict(
        self,
        h: int,
        X_df: Optional[DataFrame] = None,
        level: Optional[List[int]] = None,
    ):
        """Predict statistical models.

        Use stored fitted `models` to predict large set of time series from DataFrame `df`.

        Parameters
        ----------
        {h}
        {X_df}
        {level}

        Returns
        -------
        fcsts_df : pandas or polars DataFrame
            DataFrame with `models` columns for point predictions and probabilistic
            predictions for all fitted `models`.
        """
        if not hasattr(self, "fitted_"):
            raise ValueError("You must call the fit method before calling predict.")
        if (
            any(
                getattr(m, "prediction_intervals", None) is not None
                for m in self.models
            )
            and level is None
        ):
            warnings.warn(
                "Prediction intervals are set but `level` was not provided. "
                "Predictions won't have intervals."
            )
        self._validate_exog(X_df)
        X, level = self._parse_X_level(h=h, X=X_df, level=level)
        if self.n_jobs == 1:
            fcsts, cols = self.ga.predict(fm=self.fitted_, h=h, X=X, level=level)
        else:
            fcsts, cols = self._predict_parallel(h=h, X=X, level=level)
        fcsts_df = self._make_future_df(h=h)
        fcsts_df[cols] = fcsts
        return fcsts_df

    predict.__doc__ = predict.__doc__.format(**_param_descriptions)  # type: ignore[union-attr]

    def fit_predict(
        self,
        h: int,
        df: Optional[DataFrame] = None,
        X_df: Optional[DataFrame] = None,
        level: Optional[List[int]] = None,
        sort_df: bool = True,
        prediction_intervals: Optional[ConformalIntervals] = None,
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
    ) -> DataFrame:
        """Fit and Predict with statistical models.

        This method avoids memory burden due from object storage.
        It is analogous to Scikit-Learn `fit_predict` without storing information.
        It requires the forecast horizon `h` in advance.

        In contrast to `StatsForecast.forecast` this method stores partial models outputs.

        Parameters
        ----------
        {h}
        {df}
            If None, the `StatsForecast` class should have been instantiated using `df`.
        {X_df}
        {level}
        {sort_df}
        {prediction_intervals}
        {id_col}
        {time_col}
        {target_col}

        Returns
        -------
        fcsts_df : pandas or polars DataFrame
            DataFrame with `models` columns for point predictions and probabilistic
            predictions for all fitted `models`.
        """
        self._prepare_fit(
            df=df,
            sort_df=sort_df,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )
        self._validate_exog(X_df)
        if prediction_intervals is not None and level is None:
            raise ValueError(
                "You must specify `level` when using `prediction_intervals`"
            )
        self._validate_sizes_for_prediction_intervals(prediction_intervals)
        self._set_prediction_intervals(prediction_intervals=prediction_intervals)
        X, level = self._parse_X_level(h=h, X=X_df, level=level)
        if self.n_jobs == 1:
            self.fitted_, fcsts, cols = self.ga.fit_predict(
                models=self.models, h=h, X=X, level=level
            )
        else:
            self.fitted_, fcsts, cols = self._fit_predict_parallel(
                h=h, X=X, level=level
            )
        fcsts_df = self._make_future_df(h=h)
        fcsts_df[cols] = fcsts
        return fcsts_df

    fit_predict.__doc__ = fit_predict.__doc__.format(**_param_descriptions)  # type: ignore[union-attr]

    def forecast(
        self,
        h: int,
        df: Optional[DataFrame] = None,
        X_df: Optional[DataFrame] = None,
        level: Optional[List[int]] = None,
        fitted: bool = False,
        sort_df: bool = True,
        prediction_intervals: Optional[ConformalIntervals] = None,
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
    ) -> DataFrame:
        """Memory Efficient predictions.

        This method avoids memory burden due from object storage.
        It is analogous to Scikit-Learn `fit_predict` without storing information.
        It requires the forecast horizon `h` in advance.

        Parameters
        ----------
        {h}
        {df}
        {X_df}
        {level}
        {fitted}
        {sort_df}
        {prediction_intervals}
        {id_col}
        {time_col}
        {target_col}

        Returns
        -------
        fcsts_df : pandas or polars DataFrame
            DataFrame with `models` columns for point predictions and probabilistic
            predictions for all fitted `models`.
        """
        self.__dict__.pop("fcst_fitted_values_", None)
        self._prepare_fit(
            df=df,
            sort_df=sort_df,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )
        self._validate_exog(X_df)
        self._validate_sizes_for_prediction_intervals(prediction_intervals)
        self._set_prediction_intervals(prediction_intervals=prediction_intervals)
        X, level = self._parse_X_level(h=h, X=X_df, level=level)
        if self.n_jobs == 1:
            res_fcsts = self.ga.forecast(
                models=self.models,
                h=h,
                fallback_model=self.fallback_model,
                fitted=fitted,
                X=X,
                level=level,
                verbose=self.verbose,
                target_col=target_col,
            )
        else:
            res_fcsts = self._forecast_parallel(
                h=h,
                fitted=fitted,
                X=X,
                level=level,
                target_col=target_col,
            )
        if fitted:
            self.fcst_fitted_values_ = res_fcsts["fitted"]
        fcsts = res_fcsts["forecasts"]
        cols = res_fcsts["cols"]
        fcsts_df = self._make_future_df(h=h)
        fcsts_df[cols] = fcsts
        self.forecast_times_ = res_fcsts["times"]
        return fcsts_df

    forecast.__doc__ = forecast.__doc__.format(**_param_descriptions)  # type: ignore[union-attr]

    def forecast_fitted_values(self):
        """Access insample predictions.

        After executing `StatsForecast.forecast`, you can access the insample
        prediction values for each model. To get them, you need to pass `fitted=True`
        to the `StatsForecast.forecast` method and then use the
        `StatsForecast.forecast_fitted_values` method.

        Returns
        -------
        fcsts_df : pandas.DataFrame | polars.DataFrame
            DataFrame with insample `models` columns for point predictions and probabilistic
            predictions for all fitted `models`.
        """
        if not hasattr(self, "fcst_fitted_values_"):
            raise Exception("Please run `forecast` method using `fitted=True`")
        cols = self.fcst_fitted_values_["cols"]
        df = self.df_constructor(
            {
                self.id_col: ufp.repeat(self.uids, np.diff(self.ga.indptr)),
                self.time_col: self.og_dates,
            }
        )
        df[cols] = self.fcst_fitted_values_["values"]
        if isinstance(df, pd.DataFrame):
            if _id_as_idx():
                _warn_id_as_idx()
                df = df.set_index(self.id_col)
            else:
                df = df.reset_index(drop=True)
        return df

    def cross_validation(
        self,
        h: int,
        df: Optional[DataFrame] = None,
        n_windows: int = 1,
        step_size: int = 1,
        test_size: Optional[int] = None,
        input_size: Optional[int] = None,
        level: Optional[List[int]] = None,
        fitted: bool = False,
        refit: Union[bool, int] = True,
        sort_df: bool = True,
        prediction_intervals: Optional[ConformalIntervals] = None,
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
    ) -> DataFrame:
        """Temporal Cross-Validation.

        Efficiently fits a list of `StatsForecast`
        models through multiple training windows, in either chained or rolled manner.

        `StatsForecast.models`' speed allows to overcome this evaluation technique
        high computational costs. Temporal cross-validation provides better model's
        generalization measurements by increasing the test's length and diversity.

        Parameters
        ----------
        {h}
        {df}
            If None, the `StatsForecast` class should have been instantiated using `df`.
        {n_windows}
        {step_size}
        {test_size}
        {input_size}
        {level}
        {fitted}
        {refit}
        {sort_df}
        {prediction_intervals}
        {id_col}
        {time_col}
        {target_col}

        Returns
        -------
        fcsts_df : pandas or polars DataFrame
            DataFrame with insample `models` columns for point predictions and probabilistic
            predictions for all fitted `models`.
        """
        if n_windows is None and test_size is None:
            raise ValueError("you must define `n_windows` or `test_size`")
        if test_size is None:
            test_size = h + step_size * (n_windows - 1)
        if prediction_intervals is not None and level is None:
            raise ValueError(
                "You must specify `level` when using `prediction_intervals`"
            )
        if refit != True:
            no_forward = [m for m in self.models if not hasattr(m, "forward")]
            if no_forward:
                raise ValueError(
                    "Can only use integer refit or refit=False with models that implement the forward method. "
                    f"The following models do not implement the forward method: {no_forward}."
                )
            if self.fallback_model is not None and not hasattr(
                self.fallback_model, "forward"
            ):
                raise ValueError(
                    "Can only use integer refit or refit=False with a fallback model that implements the forward method."
                )
        self.__dict__.pop("cv_fitted_values_", None)
        self._prepare_fit(
            df=df,
            sort_df=sort_df,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )
        series_sizes = np.diff(self.ga.indptr)
        short_series = series_sizes <= test_size
        if short_series.any():
            short_ids = self.uids[short_series].to_numpy().tolist()
            raise ValueError(
                f"The following series are too short for the cross validation settings: {reprlib.repr(short_ids)}\n"
                "Please remove these series or change the settings, e.g. reducing the horizon or the number of windows."
            )
        self._validate_sizes_for_prediction_intervals(
            prediction_intervals=prediction_intervals, offset=test_size
        )
        self._set_prediction_intervals(prediction_intervals=prediction_intervals)
        _, level = self._parse_X_level(h=h, X=None, level=level)
        if self.n_jobs == 1:
            res_fcsts = self.ga.cross_validation(
                models=self.models,
                h=h,
                test_size=test_size,
                fallback_model=self.fallback_model,
                step_size=step_size,
                input_size=input_size,
                fitted=fitted,
                level=level,
                verbose=self.verbose,
                refit=refit,
                target_col=target_col,
            )
        else:
            res_fcsts = self._cross_validation_parallel(
                h=h,
                test_size=test_size,
                step_size=step_size,
                input_size=input_size,
                fitted=fitted,
                level=level,
                refit=refit,
                target_col=target_col,
            )
        if fitted:
            self.cv_fitted_values_ = res_fcsts["fitted"]
            self.n_cv_ = n_windows
        fcsts_df = ufp.cv_times(
            times=self.og_dates,
            uids=self.uids,
            indptr=self.ga.indptr,
            h=h,
            test_size=test_size,
            step_size=step_size,
            id_col=id_col,
            time_col=time_col,
        )
        # the cv_times is sorted by window and then id
        fcsts_df = ufp.sort(fcsts_df, [id_col, "cutoff", time_col])
        fcsts_df = ufp.assign_columns(
            fcsts_df, res_fcsts["cols"], res_fcsts["forecasts"]
        )
        if isinstance(fcsts_df, pd.DataFrame) and _id_as_idx():
            _warn_id_as_idx()
            fcsts_df = fcsts_df.set_index(id_col)
        return fcsts_df

    cross_validation.__doc__ = cross_validation.__doc__.format(**_param_descriptions)  # type: ignore[union-attr]

    def cross_validation_fitted_values(self) -> DataFrame:
        """Access insample cross validated predictions.

        After executing `StatsForecast.cross_validation`, you can access the insample
        prediction values for each model and window. To get them, you need to pass `fitted=True`
        to the `StatsForecast.cross_validation` method and then use the
        `StatsForecast.cross_validation_fitted_values` method.

        Returns
        -------
        fcsts_df : pandas or polars DataFrame
            DataFrame with insample `models` columns for point predictions
            and probabilistic predictions for all fitted `models`.
        """
        if not hasattr(self, "cv_fitted_values_"):
            raise Exception("Please run `cross_validation` method using `fitted=True`")
        idxs = self.cv_fitted_values_["idxs"].flatten(order="F")
        train_uids = ufp.repeat(self.uids, np.diff(self.ga.indptr))
        cv_uids = ufp.vertical_concat([train_uids for _ in range(self.n_cv_)])
        used_uids = ufp.take_rows(cv_uids, idxs)
        dates = np.tile(self.og_dates, self.n_cv_)[idxs]
        cutoffs_mask = self.cv_fitted_values_["last_idxs"].flatten(order="F")[idxs]
        cutoffs_sizes = np.diff(np.append(0, np.where(cutoffs_mask)[0] + 1))
        cutoffs = np.repeat(dates[cutoffs_mask], cutoffs_sizes)
        df = self.df_constructor(
            {
                self.id_col: used_uids,
                self.time_col: dates,
                "cutoff": cutoffs,
            }
        )
        fitted_vals = np.reshape(
            self.cv_fitted_values_["values"],
            (-1, len(self.models) + 1),
            order="F",
        )
        df = ufp.assign_columns(df, self.cv_fitted_values_["cols"], fitted_vals[idxs])
        df = ufp.drop_index_if_pandas(df)
        if isinstance(df, pd.DataFrame):
            if _id_as_idx():
                _warn_id_as_idx()
                df = df.set_index(self.id_col)
            else:
                df = df.reset_index(drop=True)
        return df

    def _get_pool(self):
        from multiprocessing import Pool

        pool_kwargs = dict()
        return Pool, pool_kwargs

    def _fit_parallel(self):
        gas = self.ga.split(self.n_jobs)
        Pool, pool_kwargs = self._get_pool()
        with Pool(self.n_jobs, **pool_kwargs) as executor:
            futures = []
            for ga in gas:
                future = executor.apply_async(
                    ga._single_threaded_fit, (self.models, self.fallback_model)
                )
                futures.append(future)
            fm = np.vstack([f.get() for f in futures])
        return fm

    def _get_gas_Xs(self, X):
        gas = self.ga.split(self.n_jobs)
        if X is not None:
            Xs = X.split(self.n_jobs)
        else:
            from itertools import repeat

            Xs = repeat(None)
        return gas, Xs

    def _predict_parallel(self, h, X, level):
        # create elements for each core
        gas, Xs = self._get_gas_Xs(X=X)
        fms = self.ga.split_fm(self.fitted_, self.n_jobs)
        Pool, pool_kwargs = self._get_pool()
        # compute parallel forecasts
        with Pool(self.n_jobs, **pool_kwargs) as executor:
            futures = []
            for ga, fm, X_ in zip(gas, fms, Xs):
                future = executor.apply_async(
                    ga._single_threaded_predict,
                    (fm, h, X_, level),
                )
                futures.append(future)
            out = [f.get() for f in futures]
            fcsts, cols = list(zip(*out))
            fcsts = np.vstack(fcsts)
            cols = cols[0]
        return fcsts, cols

    def _fit_predict_parallel(self, h, X, level):
        # create elements for each core
        gas, Xs = self._get_gas_Xs(X=X)
        Pool, pool_kwargs = self._get_pool()
        # compute parallel forecasts
        with Pool(self.n_jobs, **pool_kwargs) as executor:
            futures = []
            for ga, X_ in zip(gas, Xs):
                future = executor.apply_async(
                    ga._single_threaded_fit_predict,
                    (self.models, h, X_, level),
                )
                futures.append(future)
            out = [f.get() for f in futures]
            fm, fcsts, cols = list(zip(*out))
            fm = np.vstack(fm)
            fcsts = np.vstack(fcsts)
            cols = cols[0]
        return fm, fcsts, cols

    def _forecast_serie(self, level, **forecast_kwargs):
        forecast_res = {}
        fitted_res = {}
        times = {}
        with threadpool_limits(limits=1):
            for model in self.models:
                if "level" in inspect.signature(model.forecast).parameters and level:
                    model_kwargs = {**forecast_kwargs, "level": level}
                else:
                    model_kwargs = forecast_kwargs
                start = time.perf_counter()
                try:
                    model_res = model.forecast(**model_kwargs)
                except Exception as e:
                    if self.fallback_model is None:
                        raise e
                    model_res = self.fallback_model.forecast(**model_kwargs)
                model_name = repr(model)
                times[model_name] = time.perf_counter() - start
                for k, v in model_res.items():
                    if k == "mean":
                        forecast_res[model_name] = v
                    elif k.startswith(("lo", "hi")):
                        col_name = f"{model_name}-{k}"
                        forecast_res[col_name] = v
                    elif k == "fitted":
                        fitted_res[model_name] = v
                    elif k.startswith(("fitted-lo", "fitted-hi")):
                        col_name = f'{model_name}-{k.replace("fitted-", "")}'
                        fitted_res[col_name] = v
        return forecast_res, fitted_res, times

    def _forecast_parallel(self, h, fitted, X, level, target_col):
        n_series = self.ga.n_groups
        forecast_res = defaultdict(lambda: np.empty(n_series * h, dtype=np.float32))
        fitted_res = defaultdict(
            lambda: np.empty(self.ga.data.shape[0], dtype=self.ga.data.dtype)
        )
        fitted_res[target_col] = self.ga.data[:, 0]
        future2pos = {}
        times = {repr(m): 0.0 for m in self.models}
        with ProcessPoolExecutor(self.n_jobs) as executor:
            for i, serie in enumerate(self.ga):
                y_train = serie[:, 0]
                X_train = serie[:, 1:] if serie.shape[1] > 1 else None
                if X is None:
                    X_future = None
                else:
                    X_future = X[i]
                future = executor.submit(
                    self._forecast_serie,
                    h=h,
                    y=y_train,
                    X=X_train,
                    X_future=X_future,
                    fitted=fitted,
                    level=level,
                )
                future2pos[future] = i
            iterable = tqdm(
                as_completed(future2pos),
                disable=not self.verbose,
                total=n_series,
                desc="Forecast",
            )
            for future in iterable:
                i = future2pos[future]
                fcst_idxs = slice(i * h, (i + 1) * h)
                serie_idxs = slice(self.ga.indptr[i], self.ga.indptr[i + 1])
                serie_fcst, serie_fitted, serie_times = future.result()
                for k, v in serie_fcst.items():
                    forecast_res[k][fcst_idxs] = v
                for k, v in serie_fitted.items():
                    fitted_res[k][serie_idxs] = v
                for model_name, model_time in serie_times.items():
                    times[model_name] += model_time
        return {
            "cols": list(forecast_res.keys()),
            "forecasts": np.hstack([v[:, None] for v in forecast_res.values()]),
            "fitted": {
                "cols": list(fitted_res.keys()),
                "values": np.hstack([v[:, None] for v in fitted_res.values()]),
            },
            "times": times,
        }

    def _cross_validation_parallel(
        self, h, test_size, step_size, input_size, fitted, level, refit, target_col
    ):
        # create elements for each core
        gas = self.ga.split(self.n_jobs)
        Pool, pool_kwargs = self._get_pool()
        # compute parallel forecasts
        result = {}
        with Pool(self.n_jobs, **pool_kwargs) as executor:
            futures = []
            for ga in gas:
                future = executor.apply_async(
                    ga._single_threaded_cross_validation,
                    tuple(),
                    dict(
                        models=self.models,
                        h=h,
                        test_size=test_size,
                        fallback_model=self.fallback_model,
                        step_size=step_size,
                        input_size=input_size,
                        fitted=fitted,
                        level=level,
                        refit=refit,
                        verbose=self.verbose,
                        target_col=target_col,
                    ),
                )
                futures.append(future)
            out = [f.get() for f in futures]
            fcsts = [d["forecasts"] for d in out]
            fcsts = np.vstack(fcsts)
            cols = out[0]["cols"]
            result["forecasts"] = fcsts
            result["cols"] = cols
            if fitted:
                result["fitted"] = {}
                result["fitted"]["values"] = np.concatenate(
                    [d["fitted"]["values"] for d in out]
                )
                for key in ["last_idxs", "idxs"]:
                    result["fitted"][key] = np.concatenate(
                        [d["fitted"][key] for d in out]
                    )
                result["fitted"]["cols"] = out[0]["fitted"]["cols"]
        return result

    @staticmethod
    def plot(
        df: DataFrame,
        forecasts_df: Optional[DataFrame] = None,
        unique_ids: Union[Optional[List[str]], np.ndarray] = None,
        plot_random: bool = True,
        models: Optional[List[str]] = None,
        level: Optional[List[float]] = None,
        max_insample_length: Optional[int] = None,
        plot_anomalies: bool = False,
        engine: str = "matplotlib",
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
        resampler_kwargs: Optional[Dict] = None,
    ):
        """Plot forecasts and insample values.

        Parameters
        ----------
        {df}
        forecasts_df : pandas or polars DataFrame, optional (default=None)
            DataFrame ids, times and models.
        unique_ids : list of str, optional (default=None)
            ids to plot. If None, they're selected randomly.
        plot_random : bool (default=True)
            Select time series to plot randomly.
        models : List[str], optional (default=None)
            List of models to plot.
        level : List[float], optional (default=None)
            List of prediction intervals to plot if paseed.
        max_insample_length : int, optional (default=None)
            Max number of train/insample observations to be plotted.
        plot_anomalies : bool (default=False)
            Plot anomalies for each prediction interval.
        engine : str (default='matplotlib')
            Library used to plot. 'plotly', 'plotly-resampler' or 'matplotlib'.
        {id_col}
        {time_col}
        {target_col}
        resampler_kwargs : dict
            Kwargs to be passed to plotly-resampler constructor.
            For further custumization ("show_dash") call the method,
            store the plotting object and add the extra arguments to
            its `show_dash` method.
        """
        from utilsforecast.plotting import plot_series

        df = ensure_time_dtype(df, time_col)
        if isinstance(df, pd.DataFrame) and df.index.name == id_col:
            warnings.warn(
                "Passing the ids as the index is deprecated. "
                "Please provide them as a column instead.",
                category=FutureWarning,
            )
            df = df.reset_index()
        if forecasts_df is not None:
            forecasts_df = ensure_time_dtype(forecasts_df, time_col)
        if isinstance(forecasts_df, pd.DataFrame) and forecasts_df.index.name == id_col:
            warnings.warn(
                "Passing the ids as the index is deprecated. "
                "Please provide them as a column instead.",
                category=FutureWarning,
            )
            forecasts_df = forecasts_df.reset_index()
        return plot_series(
            df=df,
            forecasts_df=forecasts_df,
            ids=unique_ids,
            plot_random=plot_random,
            models=models,
            level=level,
            max_insample_length=max_insample_length,
            plot_anomalies=plot_anomalies,
            engine=engine,
            resampler_kwargs=resampler_kwargs,
            palette="tab20b",
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )

    def save(
        self,
        path: Optional[Union[Path, str]] = None,
        max_size: Optional[str] = None,
        trim: bool = False,
    ):
        """Function that will save StatsForecast class with certain settings to make it
        reproducible.

        Parameters
        ----------
        path : str or pathlib.Path, optional (default=None)
            Path of the file to be saved. If `None` will create one in the current
            directory using the current UTC timestamp.
        max_size : str, optional (default = None)
            StatsForecast object should not exceed this size.
            Available byte naming: ['B', 'KB', 'MB', 'GB']
        trim : bool (default = False)
            Delete any attributes not needed for inference.
        """
        # Will be used to find the size of the fitted models
        # Never expecting anything higher than GB (even that's a lot')
        bytes_hmap = {
            "B": 1,
            "KB": 2**10,
            "MB": 2**20,
            "GB": 2**30,
        }

        # Removing unnecessary attributes
        # @jmoralez decide future implementation
        trim_attr: list = ["fcst_fitted_values_", "cv_fitted_values_"]
        if trim:
            for attr in trim_attr:
                # remove unnecessary attributes here
                self.__dict__.pop(attr, None)

        sf_size = len(pickle.dumps(self))

        if max_size is not None:
            cap_size = self._get_cap_size(max_size, bytes_hmap)
            if sf_size >= cap_size:
                err_messg = "StatsForecast is larger than the specified max_size"
                raise OSError(errno.EFBIG, err_messg)

        converted_size, sf_byte = None, None
        for key in reversed(list(bytes_hmap.keys())):
            x_byte = bytes_hmap[key]
            if sf_size >= x_byte:
                converted_size = sf_size / x_byte
                sf_byte = key
                break

        if converted_size is None or sf_byte is None:
            err_messg = "Internal Error, this shouldn't happen, please open an issue"
            raise RuntimeError(err_messg)

        print(f"Saving StatsForecast object of size {converted_size:.2f}{sf_byte}.")

        if path is None:
            datetime_record = dt.datetime.utcnow().strftime("%Y-%m-%d_%H-%M-%S")
            path = f"StatsForecast_{datetime_record}.pkl"

        with open(path, "wb") as m_file:
            pickle.dump(self, m_file)
        print("StatsForecast object saved")

    def _get_cap_size(self, max_size, bytes_hmap):
        max_size = max_size.upper().replace(" ", "")
        match = re.match(r"(\d+\.\d+|\d+)(\w+)", max_size)
        if (
            match is None
            or len(match.groups()) < 2
            or match[2] not in bytes_hmap.keys()
        ):
            parsing_error = (
                "Couldn't parse `max_size`, it should be `None`",
                " or a number followed by one of the following units: ['B', 'KB', 'MB', 'GB']",
            )
            raise ValueError(parsing_error)
        else:
            m_size = float(match[1])
            key_ = match[2]
            cap_size = m_size * bytes_hmap[key_]
        return cap_size

    @staticmethod
    def load(path: Union[Path, str]):
        """
        Automatically loads the model into ready StatsForecast.

        Parameters
        ----------
        path : str or pathlib.Path
            Path to saved StatsForecast file.

        Returns
        -------
        sf: StatsForecast
            Previously saved StatsForecast
        """
        if not Path(path).exists():
            raise ValueError("Specified path does not exist, check again and retry.")
        with open(path, "rb") as f:
            return pickle.load(f)

    def __repr__(self):
        return f"StatsForecast(models=[{','.join(map(repr, self.models))}])"


_StatsForecast.plot.__doc__ = _StatsForecast.plot.__doc__.format(**_param_descriptions)  # type: ignore[union-attr]

# %% ../nbs/src/core/core.ipynb 30
class ParallelBackend:
    def forecast(
        self,
        *,
        models,
        fallback_model,
        freq,
        h,
        df,
        X_df,
        level,
        fitted,
        prediction_intervals,
        id_col,
        time_col,
        target_col,
    ) -> Any:
        model = _StatsForecast(
            models=models,
            freq=freq,
            fallback_model=fallback_model,
        )
        return model.forecast(
            df=df,
            h=h,
            X_df=X_df,
            level=level,
            fitted=fitted,
            prediction_intervals=prediction_intervals,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )

    def cross_validation(
        self,
        *,
        df,
        models,
        freq,
        fallback_model,
        h,
        n_windows,
        step_size,
        test_size,
        input_size,
        level,
        refit,
        fitted,
        prediction_intervals,
        id_col,
        time_col,
        target_col,
    ) -> Any:
        model = _StatsForecast(
            models=models,
            freq=freq,
            fallback_model=fallback_model,
        )
        return model.cross_validation(
            df=df,
            h=h,
            n_windows=n_windows,
            step_size=step_size,
            test_size=test_size,
            input_size=input_size,
            level=level,
            refit=refit,
            fitted=fitted,
            prediction_intervals=prediction_intervals,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )


@conditional_dispatcher
def make_backend(obj: Any, *args: Any, **kwargs: Any) -> ParallelBackend:
    return ParallelBackend()

# %% ../nbs/src/core/core.ipynb 31
class StatsForecast(_StatsForecast):
    def forecast(
        self,
        h: int,
        df: Any = None,
        X_df: Optional[DataFrame] = None,
        level: Optional[List[int]] = None,
        fitted: bool = False,
        sort_df: bool = True,
        prediction_intervals: Optional[ConformalIntervals] = None,
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
    ):
        if prediction_intervals is not None and level is None:
            raise ValueError(
                "You must specify `level` when using `prediction_intervals`"
            )
        if self._is_native(df=df):
            return super().forecast(
                df=df,
                h=h,
                X_df=X_df,
                level=level,
                fitted=fitted,
                sort_df=sort_df,
                prediction_intervals=prediction_intervals,
                id_col=id_col,
                time_col=time_col,
                target_col=target_col,
            )
        assert df is not None
        engine = make_execution_engine(infer_by=[df])
        self._backend = make_backend(engine)
        return self._backend.forecast(
            models=self.models,
            fallback_model=self.fallback_model,
            freq=self.freq,
            df=df,
            h=h,
            X_df=X_df,
            level=level,
            fitted=fitted,
            prediction_intervals=prediction_intervals,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )

    def forecast_fitted_values(self):
        if hasattr(self, "_backend"):
            res = self._backend.forecast_fitted_values()
        else:
            res = super().forecast_fitted_values()
        return res

    def cross_validation(
        self,
        h: int,
        df: Any = None,
        n_windows: int = 1,
        step_size: int = 1,
        test_size: Optional[int] = None,
        input_size: Optional[int] = None,
        level: Optional[List[int]] = None,
        fitted: bool = False,
        refit: Union[bool, int] = True,
        sort_df: bool = True,
        prediction_intervals: Optional[ConformalIntervals] = None,
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
    ):
        if self._is_native(df=df):
            return super().cross_validation(
                h=h,
                df=df,
                n_windows=n_windows,
                step_size=step_size,
                test_size=test_size,
                input_size=input_size,
                level=level,
                fitted=fitted,
                refit=refit,
                sort_df=sort_df,
                prediction_intervals=prediction_intervals,
                id_col=id_col,
                time_col=time_col,
                target_col=target_col,
            )
        assert df is not None
        engine = make_execution_engine(infer_by=[df])
        backend = make_backend(engine)
        return backend.cross_validation(
            df=df,
            models=self.models,
            freq=self.freq,
            fallback_model=self.fallback_model,
            h=h,
            n_windows=n_windows,
            step_size=step_size,
            test_size=test_size,
            input_size=input_size,
            level=level,
            refit=refit,
            fitted=fitted,
            prediction_intervals=prediction_intervals,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )

    def _is_native(self, df) -> bool:
        engine = try_get_context_execution_engine()
        return engine is None and (
            df is None or isinstance(df, pd.DataFrame) or isinstance(df, pl_DataFrame)
        )