stallion.py

# %%
import skrub
import pandas as pd
from pathlib import Path


def load_data(path: Path):
    return pd.read_csv(path)


# %%
train_path = Path("data/train")

historical = skrub.var("historical_volume", train_path / "historical_volume.csv")
price_promo = skrub.var("price_promo", train_path / "price_sales_promotion.csv")
events = skrub.var("events", train_path / "event_calendar.csv")
industry = skrub.var("industry", train_path / "industry_volume.csv")
soda = skrub.var("soda", train_path / "industry_soda_sales.csv")
weather = skrub.var("weather", train_path / "weather.csv")
demographics = skrub.var("demographics", train_path / "demographics.csv")

# %%
historical = historical.skb.apply_func(load_data)
historical

# %%
price_promo = price_promo.skb.apply_func(load_data)
price_promo

# %%
events = events.skb.apply_func(load_data)
events
# %%
industry = industry.skb.apply_func(load_data)
industry
# %%
soda = soda.skb.apply_func(load_data)
soda
# %%
weather = weather.skb.apply_func(load_data)
weather

# %%
demographics = demographics.skb.apply_func(load_data)
demographics


# %%
def normalize_data(df: pd.DataFrame):
    """Parse the date and sort the data."""
    df["YearMonth"] = pd.to_datetime(df["YearMonth"].astype(str), format="%Y%m")
    df = df.sort_values(by=["Agency", "SKU", "YearMonth"]).reset_index(drop=True)
    return df


# %%
normalized_data = historical.skb.apply_func(normalize_data)
normalized_data


# %%
def add_lagged_volume(query: pd.DataFrame, historical_path: Path):
    """Add lagged volume features to query rows."""
    history = normalize_data(load_data(historical_path))

    # consolidate historical data with missing values if date appear in the query but
    # not in the history
    keys = ["YearMonth", "Agency", "SKU"]
    unseen = (
        query[keys]
        .merge(history[keys], on=keys, how="left", indicator=True)
        .loc[lambda d: d["_merge"].eq("left_only"), keys]
        .assign(Volume=pd.NA)  # or np.nan if you prefer float NaNs
    )
    if not unseen.empty:
        history = pd.concat([history, unseen], ignore_index=True).sort_values(
            ["Agency", "SKU", "YearMonth"], ignore_index=True
        )

    # same calendar month span for every (Agency, SKU); fill gaps so shift/rolling are
    # one row per month within each series.
    grid_start, grid_end = history["YearMonth"].min(), history["YearMonth"].max()
    full_span = pd.date_range(grid_start, grid_end, freq="MS")

    # because pandas rolling is not time based, we need to create a grid of all possible
    # dates and add missing values
    gap_parts = []
    for (agency, sku), g in history.groupby(["Agency", "SKU"], sort=False):
        full_months = pd.DataFrame(
            {
                "YearMonth": full_span,
                "Agency": agency,
                "SKU": sku,
            }
        )
        gaps = (
            full_months.merge(g[keys], on=keys, how="left", indicator=True)
            .loc[lambda d: d["_merge"].eq("left_only"), keys]
            .assign(Volume=pd.NA)
        )
        gap_parts.append(gaps)
    missing_months = (
        pd.concat(gap_parts, ignore_index=True)
        if gap_parts
        else pd.DataFrame(columns=history.columns)
    )
    if not missing_months.empty:
        history = pd.concat([history, missing_months], ignore_index=True).sort_values(
            ["Agency", "SKU", "YearMonth"], ignore_index=True
        )

    for lag, window in [(1, 1), (2, 1), (3, 1), (12, 1), (1, 3), (12, 3)]:
        assert lag >= 1
        col_name = f"Volume_lag_{lag}mo_window_{window}mo"

        history_lag = history.copy()
        history_lag[col_name] = history_lag.groupby(["Agency", "SKU"], sort=False)[
            "Volume"
        ].transform(
            lambda x, l=lag, w=window: x.shift(l).rolling(w, min_periods=1).mean()
        )
        lagged_volume = history_lag[keys + [col_name]]

        query = query.merge(lagged_volume, on=keys, how="left")

        agency_agg = (
            lagged_volume.groupby(["YearMonth", "Agency"])[col_name]
            .sum()
            .rename(f"Agency_{col_name}")
            .reset_index()
        )
        query = query.merge(agency_agg, on=["YearMonth", "Agency"], how="left")

        sku_agg = (
            lagged_volume.groupby(["YearMonth", "SKU"])[col_name]
            .sum()
            .rename(f"SKU_{col_name}")
            .reset_index()
        )
        query = query.merge(sku_agg, on=["YearMonth", "SKU"], how="left")

    return query


# %%
class Splitter:
    """Time-based cross-validation strategy.

    Create a time-based walk-forward validation strategy. The training data will
    contain at least `min_train_months` months. Each test set will be `interval`
    months apart.

    Parameters
    ----------
    min_train_months : int, default=6
        Minimum number of months to use for training.
    interval : int, default=1
        Interval in months between test sets.

    Notes
    -----
    A scikit-learn compatible splitter should expose the two methods `split` and
    `get_n_splits`.
    """

    def __init__(self, min_train_months=6, interval=1):
        self.min_train_months = min_train_months
        self.interval = interval

    def split(self, X, y=None, groups=None):
        """Split the data into training and test sets.

        Parameters
        ----------
        X : pd.DataFrame
            The data to split.
        y : pd.Series, optional
            No-op for compatibility with general splitter interface.
        groups : pd.Series, optional
            No-op for compatibility with general splitter interface.

        Yields
        ------
        train_indices, test_indices : tuple of numpy.ndarrays
            Train and test indices.
        """
        start = X["YearMonth"].min() + pd.DateOffset(months=self.min_train_months)
        end = X["YearMonth"].max()
        test_months = pd.date_range(start, end, freq=f"{self.interval}MS")
        for test_month in test_months:
            train_indices = X.index[X["YearMonth"] < test_month].to_numpy()
            test_indices = X.index[X["YearMonth"] == test_month].to_numpy()
            if len(train_indices) and len(test_indices):
                yield train_indices, test_indices

    def get_n_splits(self, X, y=None, groups=None):
        """Get the number of splits."""
        return sum(1 for _ in self.split(X, y=y, groups=groups))


# %%
X = normalized_data.skb.drop("Volume").skb.mark_as_X(cv=Splitter(), split_kwargs={})
y = normalized_data.skb.apply_func(lambda df: df["Volume"]).skb.mark_as_y()

# %%
features = X.skb.apply_func(
    add_lagged_volume, historical_path=train_path / "historical_volume.csv"
)
features

# %%
vectorized_features = features.skb.apply(
    skrub.TableVectorizer(
        datetime=skrub.DatetimeEncoder(resolution="month"),
        low_cardinality=skrub.ToCategorical(),
        cardinality_threshold=100,
    )
)
vectorized_features

# %%
from sklearn.ensemble import HistGradientBoostingRegressor

pred = vectorized_features.skb.apply(
    HistGradientBoostingRegressor(categorical_features="from_dtype"), y=y
)
pred

# %%
from sklearn.dummy import DummyRegressor

dummy_learner = X.skb.apply(DummyRegressor(), y=y).skb.make_learner()
hgbdt_learner = vectorized_features.skb.apply(
    HistGradientBoostingRegressor(categorical_features="from_dtype"), y=y
).skb.make_learner()

# %%
import skore

report = skore.evaluate(
    [dummy_learner, hgbdt_learner],
    data={"historical_volume": train_path / "historical_volume.csv"},
    splitter=Splitter(min_train_months=24, interval=12),  # limit to only 3 splits
)
report

# %%
hgbdt_learner.fit({"historical_volume": train_path / "historical_volume.csv"})

# %%
test_path = Path("data/test")
test_input = load_data(test_path / "volume_forecast.csv")
test_input["YearMonth"] = int(pd.Timestamp("2018-01-01").strftime("%Y%m"))
test_input = test_input[["Agency", "SKU", "YearMonth", "Volume"]]
test_input.to_csv(test_path / "volume_forecast_to_predict.csv", index=False)


# %%
pred = hgbdt_learner.predict(
    {"historical_volume": test_path / "volume_forecast_to_predict.csv"}
)
pred

# %%
vectorized_features.skb.eval(
    {"historical_volume": test_path / "volume_forecast_to_predict.csv"}
)

# %%

With the latest version, we can do the following:

# %%
from sklearn.metrics import mean_absolute_error

report.metrics.add(mean_absolute_error)
report

to get the absolute error