TableWeighter¶

class TableWeighter(BaseWeighter):
    """DataFrame-driven weights resolved by joining on a key column.

    Replaces the former ``pl.DataFrame`` weight input. Weights are looked up by
    left-joining the key series to ``frame`` on column ``on``. For panel data, a
    ``{group_name}__weight`` column is used when present, otherwise ``weight``.
    A key with no match raises ``ValueError``.

    Parameters
    ----------
    frame : pl.DataFrame
        Weight table with the join column ``on`` and a ``"weight"`` column
        (and/or ``"{group}__weight"`` columns for panel data). Per-group weight
        columns use the library-wide ``group__column`` double-underscore panel
        convention (e.g. ``"store_1__weight"``).
    on : str, default="time"
        Join column name (``"time"``, ``"vintage_time"``, or
        ``"forecasting_step"``).

    See Also
    --------
    - [`LookupWeighter`][yohou.weighting.weighters.LookupWeighter] : Explicit per-key weights.
    - [`CompositeWeighter`][yohou.weighting.weighters.CompositeWeighter] : Combine weighters by product or mean.

    """

    _parameter_constraints: dict = {
        "frame": [pl.DataFrame, None],
        "on": [str],
    }

    def __init__(self, frame: pl.DataFrame | None = None, on: str = "time") -> None:
        self.frame = frame
        self.on = on

    def compute_weights(self, key: pl.Series, group_name: str | None = None) -> pl.Series:
        """Compute join-based weights for ``key``."""
        self._validate_params()
        if self.frame is None:
            raise ValueError("TableWeighter requires a non-None `frame` to compute weights")
        key_df = pl.DataFrame({self.on: key})
        joined = key_df.join(self.frame, on=self.on, how="left")

        weight_col = None
        if group_name is not None:
            group_col = f"{group_name}__weight"
            if group_col in joined.columns:
                weight_col = group_col
        if weight_col is None and "weight" in joined.columns:
            weight_col = "weight"
        if weight_col is None:
            if group_name is not None:
                raise ValueError(
                    f"Weight DataFrame missing both '{group_name}__weight' and "
                    f"'weight' columns for panel group '{group_name}'"
                )
            raise ValueError("Weight DataFrame must have 'weight' column")

        weights_np = joined[weight_col].to_numpy().astype(np.float64)

        nan_mask = np.isnan(weights_np)
        if nan_mask.any():
            bad_keys = key.filter(pl.Series(nan_mask)).unique()
            present_keys = set(self.frame[self.on].to_list())
            absent = [k for k in bad_keys.to_list() if k not in present_keys]
            null_valued = [k for k in bad_keys.to_list() if k in present_keys]
            if absent:
                raise ValueError(f"Weight DataFrame has no values for {self.on}s: {absent}")
            raise ValueError(f"Weight DataFrame has explicit null '{weight_col}' values for {self.on}s: {null_valued}")

        return pl.Series(weights_np, dtype=pl.Float64).alias("weight")