diff --git a/mlforecast/core.py b/mlforecast/core.py index 93e40eba..c956a297 100644 --- a/mlforecast/core.py +++ b/mlforecast/core.py @@ -371,26 +371,29 @@ def _get_local_tfms( local[name] = tfm return local - def _trim_pooled_states(self) -> None: - """Trim each pooled state's history under ``keep_last_n``. - - Parity with the ``self.ga`` trim: a pooled state whose transforms are - *all* finite-window drops its unused history prefix, while a state - containing any Expanding*/EWM transform keeps full history (pooled has - no carried accumulator -- it recomputes over the full aggregate vectors - at predict, so trimming those would move predictions). - - Retention is ``max(keep_last_n, W_state)`` ordinals, where ``W_state`` - is the state's largest finite window. The floor is required and is where - pooled legitimately diverges from the local coreforecast path: local - rolling survives an undersized explicit ``keep_last_n`` because - coreforecast carries a per-transform window buffer; pooled has none (the - aggregates *are* the buffer), so trimming below ``W_state`` would compute - windows off a truncated prefix. When ``keep_last_n`` is inferred it - already equals the global max window, so the floor is then a no-op. + def _trim_pooled_states(self, keep_last_n: Optional[int]) -> None: + """Trim each pooled state to only the history its own transforms need. + + A pooled state whose transforms are *all* finite-window drops its unused + history prefix, while a state containing any Expanding*/EWM transform + keeps full history (its predict-time accumulator is initialized from the + whole aggregate prefix before being carried across steps, so trimming + those would move predictions). + + Retention is ``max(keep_last_n or 0, W_state)`` ordinals, where + ``W_state`` is the state's largest finite window and ``keep_last_n`` is + the value the *user* passed (``None`` when it was inferred). Each state + is therefore trimmed to its own window regardless of the global inferred + ``keep_last_n`` — a small-window state next to a wide-window one keeps + only what it needs. The ``W_state`` floor is required and is where pooled + legitimately diverges from the local coreforecast path: local rolling + survives an undersized explicit ``keep_last_n`` because coreforecast + carries a per-transform window buffer; pooled has none (the aggregates + *are* the buffer), so trimming below ``W_state`` would compute windows + off a truncated prefix. An *explicit* ``keep_last_n`` is honored as a + floor (parity with the ``self.ga`` trim / extra buffer for ``update``). """ - if self.keep_last_n is None: - return + floor = keep_last_n if keep_last_n is not None else 0 for key, tfms in self._get_pooled_tfms().items(): state = self._pooled_states.get(key) if state is None: @@ -399,7 +402,7 @@ def _trim_pooled_states(self) -> None: if not all(tfm._is_finite_window for tfm in tfm_list): continue w_state = max(tfm.update_samples for tfm in tfm_list) - state.trim_to_last(max(self.keep_last_n, w_state)) + state.trim_to_last(max(floor, w_state)) def _apply_keep_last_n(self) -> None: """Resolve ``keep_last_n`` and trim the stored history accordingly. @@ -410,6 +413,10 @@ def _apply_keep_last_n(self) -> None: the lag transforms have computed their state, since local stateful transforms (Expanding*/EWM) warm their buffers from the full history. """ + # the user-provided value (``None`` when inferred below) is the only + # retention floor the pooled states honor — each is otherwise trimmed to + # its own window, not the global inferred ``keep_last_n``. + user_keep_last_n = self.keep_last_n update_samples = [ getattr(tfm, "update_samples", -1) for tfm in self.transforms.values() ] @@ -422,7 +429,7 @@ def _apply_keep_last_n(self) -> None: self.keep_last_n = max(update_samples) if self.keep_last_n is not None: self.ga = self.ga.take_from_groups(slice(-self.keep_last_n, None)) - self._trim_pooled_states() + self._trim_pooled_states(user_keep_last_n) def _initialize_lag_transform_states(self) -> None: """Materialize lag transform state for subsequent update-based prediction. @@ -689,6 +696,10 @@ def _fit( df_for_pooled = df for key, tfms in pooled_tfms.items(): mode, group_cols_t, part_cols_t = key + # A state stores raw per-ordinal values only if one of its + # transforms needs them (a raw-row quantile); the flag is fixed + # here and threaded through every later rebuild via the state. + store_values = any(tfm._needs_value_store for tfm in tfms.values()) if mode == "global" and not part_cols_t: self._pooled_states[key] = PooledState.from_global( sorted_df, @@ -697,6 +708,7 @@ def _fit( target_col=target_col, ga_data_dtype=ga.data.dtype, n_series=len(ga.indptr) - 1, + store_values=store_values, ) elif mode == "groupby" and not part_cols_t: for col in group_cols_t: @@ -723,6 +735,7 @@ def _fit( target_col=target_col, ga_data_dtype=ga.data.dtype, static_features=self.static_features_, + store_values=store_values, ) else: all_cols = list(group_cols_t) + list(part_cols_t) @@ -743,6 +756,7 @@ def _fit( ga_data_dtype=ga.data.dtype, static_features=self.static_features_, n_series=len(ga.indptr) - 1, + store_values=store_values, ) for key, state in self._pooled_states.items(): if state.groups is not None: @@ -1340,10 +1354,17 @@ def _compute_features_df(self) -> DataFrame: state = self._pooled_states[key] n_series = len(self.uids) slow_tfms: Dict[str, _BaseLagTransform] = {} + cache_root = getattr(self, "_latest_agg_cache", None) for name, tfm in tfms.items(): + sub_cache = ( + cache_root.setdefault((key, name), {}) + if cache_root is not None + else None + ) latest = tfm._compute_latest_from_aggs( state._ts_aggs, state.next_time_index_by_bucket, + sub_cache, ) if latest is not None: if state.groups is None: @@ -1590,6 +1611,12 @@ def _backup(self) -> Iterator[None]: # them faithfully without deep-copying every aggregate array per model. pooled_states = getattr(self, "_pooled_states", {}) pooled_snaps = {key: state.snapshot() for key, state in pooled_states.items()} + # Per-model cache for the pooled predict fast paths, keyed by + # (state_key, feature_name) -> per-bucket running accumulators. Created + # on enter and dropped on exit, so it is scoped to a single model's + # recursive walk and can never carry a stale prefix across restores + # (a later model predicts different values for the same appended dates). + self._latest_agg_cache: Dict[Tuple, dict] = {} try: yield finally: @@ -1597,6 +1624,7 @@ def _backup(self) -> Iterator[None]: self.transforms = lag_tfms for key, snap in pooled_snaps.items(): self._pooled_states[key].restore(snap) + self._latest_agg_cache = {} def _predict_setup(self) -> None: # TODO: move to utils diff --git a/mlforecast/lag_transforms.py b/mlforecast/lag_transforms.py index 10a72f12..fd388a4c 100644 --- a/mlforecast/lag_transforms.py +++ b/mlforecast/lag_transforms.py @@ -187,7 +187,7 @@ def _bucket_feature_rows_impl( return None def _latest_from_aggs_impl( - self, _ts_aggs, _target_ords + self, _ts_aggs, _target_ords, _cache=None ) -> Optional[Dict[int, float]]: return None @@ -223,16 +223,25 @@ def _compute_latest_from_aggs( self, ts_aggs, target_ords: Dict[int, int], + cache: Optional[dict] = None, ) -> Optional[Dict[int, float]]: """Compute feature value at the target timestamp per bucket from cached aggregates. ``target_ords`` maps bucket_id to the time-index ordinal at which to evaluate the statistic (typically ``next_time_index_by_bucket``). - Returns None if this transform doesn't support the fast path. + ``cache`` is an optional per-(state, transform) dict, owned by + ``TimeSeries._backup`` and thus scoped to a single model's recursive + predict walk; prefix-dependent transforms (Expanding*/EWM) use it to + carry a running accumulator so they avoid rescanning the full aggregate + vectors every step. ``None`` (the default) forces a full recompute -- + the behaviour every other ``_impl`` ignores it and keeps. Returns None + if this transform doesn't support the fast path. """ if not ts_aggs: return None - return self._latest_from_aggs_impl(self._maybe_reagg(ts_aggs), target_ords) + return self._latest_from_aggs_impl( + self._maybe_reagg(ts_aggs), target_ords, cache + ) def _compute_ts_level_from_aggs(self, ts_aggs): if not ts_aggs: @@ -308,14 +317,29 @@ def _is_finite_window(self) -> bool: A pooled state may be trimmed under ``keep_last_n`` only if *every* one of its transforms is finite-window: the dropped prefix can then never enter a window, so trimming is prediction-neutral. Unbounded transforms - (Expanding*/EWM) recompute over the full aggregate vectors at predict — - pooled has no carried accumulator — so they must keep full history. + (Expanding*/EWM) depend on the whole history prefix — their predict-time + accumulator is initialized from it before being carried across steps — + so they must keep full history. Defaults to ``False`` so an unknown/custom transform is never silently trimmed (it keeps full history; correctness over the perf win). """ return False + @property + def _needs_value_store(self) -> bool: + """Whether the pooled aggregate cache must retain the raw per-ordinal + observations (the CSR value store on ``_TimestampAggregates``). + + Only raw-row quantiles need it: a quantile is not derivable from the + scalar aggregates, so its fit/predict fast paths read the stored values + directly. A ``time_agg`` quantile instead works off the single collapsed + scalar per timestamp (derived from the scalar aggregates by + ``_ReaggregatedAggregates``), so it needs no O(rows) store. Every other + transform defaults to ``False``. + """ + return False + class Lag(_BaseLagTransform): def __init__(self, lag: int): @@ -443,6 +467,7 @@ def _compute_latest_from_aggs( self, ts_aggs, _target_ords, + _cache=None, ) -> Optional[Dict[int, float]]: # Fast predict path: the looked-up value is the target `lag` occurrences # back. In local partition mode each bucket is a single series (one @@ -637,25 +662,32 @@ def _latest_from_aggs_impl( self, ts_aggs, target_ords: Dict[int, int], + _cache=None, ) -> Dict[int, float]: lag = self._core_tfm.lag w = self.window_size min_samples = _resolve_min_samples(self) result: Dict[int, float] = {} for bid, agg in ts_aggs.items(): - if len(agg.unique_times) == 0: + ut = agg.unique_times + if len(ut) == 0: result[bid] = float("nan") continue - cum_sum = np.cumsum(agg.sums) - cum_cnt = np.cumsum(agg.counts) t = target_ords[bid] upper = t - lag lower = t - lag - w - ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1 - li = int(np.searchsorted(agg.unique_times, lower, side="right")) - 1 - s = (cum_sum[ui] if ui >= 0 else 0.0) - (cum_sum[li] if li >= 0 else 0.0) - c = (cum_cnt[ui] if ui >= 0 else 0.0) - (cum_cnt[li] if li >= 0 else 0.0) - result[bid] = s / c if c >= min_samples and c > 0 else float("nan") + ui = int(np.searchsorted(ut, upper, side="right")) - 1 + li = int(np.searchsorted(ut, lower, side="right")) - 1 + # The window is the half-open agg range ``(li, ui]``. Sum that + # ``O(window)`` slice directly rather than differencing full-vector + # cumsums, which would cost ``O(len(ut))`` at every recursive + # predict step (``ut`` grows by one each step). + lo, hi = li + 1, ui + 1 + c = agg.counts[lo:hi].sum() + if c >= min_samples and c > 0: + result[bid] = float(agg.sums[lo:hi].sum() / c) + else: + result[bid] = float("nan") return result def _ts_level_from_aggs_impl(self, ts_aggs): @@ -764,29 +796,29 @@ def _latest_from_aggs_impl( self, ts_aggs, target_ords: Dict[int, int], + _cache=None, ) -> Dict[int, float]: lag = self._core_tfm.lag w = self.window_size min_samples = _resolve_min_samples(self) result: Dict[int, float] = {} for bid, agg in ts_aggs.items(): - if len(agg.unique_times) == 0: + ut = agg.unique_times + if len(ut) == 0: result[bid] = float("nan") continue - cum_sum = np.cumsum(agg.sums) - cum_cnt = np.cumsum(agg.counts) - cum_sum_sq = np.cumsum(agg.sum_sq) t = target_ords[bid] upper = t - lag lower = t - lag - w - ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1 - li = int(np.searchsorted(agg.unique_times, lower, side="right")) - 1 - s = (cum_sum[ui] if ui >= 0 else 0.0) - (cum_sum[li] if li >= 0 else 0.0) - sq = (cum_sum_sq[ui] if ui >= 0 else 0.0) - ( - cum_sum_sq[li] if li >= 0 else 0.0 - ) - c = (cum_cnt[ui] if ui >= 0 else 0.0) - (cum_cnt[li] if li >= 0 else 0.0) + ui = int(np.searchsorted(ut, upper, side="right")) - 1 + li = int(np.searchsorted(ut, lower, side="right")) - 1 + # Half-open agg window ``(li, ui]``; sum the ``O(window)`` slice + # instead of differencing full-vector cumsums (see RollingMean). + lo, hi = li + 1, ui + 1 + c = agg.counts[lo:hi].sum() if c >= min_samples and c > 1: + s = agg.sums[lo:hi].sum() + sq = agg.sum_sq[lo:hi].sum() var = (sq - s**2 / c) / (c - 1) var = max(var, 0.0) result[bid] = float(np.sqrt(var)) @@ -853,28 +885,30 @@ def _latest_from_aggs_impl( self, ts_aggs, target_ords: Dict[int, int], + _cache=None, ) -> Dict[int, float]: lag = self._core_tfm.lag w = self.window_size min_samples = _resolve_min_samples(self) result: Dict[int, float] = {} for bid, agg in ts_aggs.items(): - if len(agg.unique_times) == 0: + ut = agg.unique_times + if len(ut) == 0: result[bid] = float("nan") continue - sparse = _build_sparse_table(agg.mins, np.fmin) - cum_cnt = np.cumsum(agg.counts) t = target_ords[bid] upper = t - lag lower = t - lag - w + 1 - ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1 - li = int(np.searchsorted(agg.unique_times, lower, side="left")) - c = (cum_cnt[ui] if ui >= 0 else 0.0) - (cum_cnt[li - 1] if li > 0 else 0.0) + ui = int(np.searchsorted(ut, upper, side="right")) - 1 + li = int(np.searchsorted(ut, lower, side="left")) + # Inclusive agg window ``[li, ui]``. Reduce over that ``O(window)`` + # slice with NaN-skipping ``fmin`` (matching the sparse table) + # instead of rebuilding an ``O(n log n)`` sparse table at every + # recursive predict step. ``c > 0`` guarantees a non-empty slice + # with at least one non-NaN entry. + c = agg.counts[li : ui + 1].sum() if c >= min_samples and c > 0: - val = _query_sparse_table( - sparse, np.array([li]), np.array([ui]), np.fmin - ) - result[bid] = float(val[0]) + result[bid] = float(np.fmin.reduce(agg.mins[li : ui + 1])) else: result[bid] = float("nan") return result @@ -910,28 +944,28 @@ def _latest_from_aggs_impl( self, ts_aggs, target_ords: Dict[int, int], + _cache=None, ) -> Dict[int, float]: lag = self._core_tfm.lag w = self.window_size min_samples = _resolve_min_samples(self) result: Dict[int, float] = {} for bid, agg in ts_aggs.items(): - if len(agg.unique_times) == 0: + ut = agg.unique_times + if len(ut) == 0: result[bid] = float("nan") continue - sparse = _build_sparse_table(agg.maxs, np.fmax) - cum_cnt = np.cumsum(agg.counts) t = target_ords[bid] upper = t - lag lower = t - lag - w + 1 - ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1 - li = int(np.searchsorted(agg.unique_times, lower, side="left")) - c = (cum_cnt[ui] if ui >= 0 else 0.0) - (cum_cnt[li - 1] if li > 0 else 0.0) + ui = int(np.searchsorted(ut, upper, side="right")) - 1 + li = int(np.searchsorted(ut, lower, side="left")) + # Inclusive agg window ``[li, ui]``; reduce over the ``O(window)`` + # slice with NaN-skipping ``fmax`` instead of rebuilding a sparse + # table each recursive predict step (see RollingMin). + c = agg.counts[li : ui + 1].sum() if c >= min_samples and c > 0: - val = _query_sparse_table( - sparse, np.array([li]), np.array([ui]), np.fmax - ) - result[bid] = float(val[0]) + result[bid] = float(np.fmax.reduce(agg.maxs[li : ui + 1])) else: result[bid] = float("nan") return result @@ -946,15 +980,48 @@ def _ts_level_from_aggs_impl(self, ts_aggs): } +def _rolling_quantile_from_agg(agg, query_ords, lag, window_size, min_samples, p): + """Rolling quantile at each query ordinal, from the CSR value store. + + ``query_ords`` are ordinal *values* (fit passes ``agg.unique_times``; predict + a single target ordinal). The window for target ``t`` spans the ordinals in + ``(t - lag - window_size, t - lag]`` -- the same searchsorted bounds as + :func:`_rolling_mean_from_agg` -- gathered as one contiguous ``values`` + slice. Returns one quantile per query ordinal, or ``None`` when the bucket + carries no value store (caller falls back to the row slow path). + """ + if agg.values is None or agg.row_offsets is None: + return None + ut = agg.unique_times + out = np.full(len(query_ords), np.nan) + if len(ut) == 0: + return out + off = agg.row_offsets + vals = agg.values + ui = np.searchsorted(ut, query_ords - lag, side="right") - 1 + li = np.searchsorted(ut, query_ords - lag - window_size, side="right") - 1 + for j in range(len(query_ords)): + u = int(ui[j]) + if u < 0: + continue + start = int(off[int(li[j]) + 1]) + end = int(off[u + 1]) + if end - start >= min_samples and end > start: + out[j] = np.quantile(vals[start:end], p) + return out + + class RollingQuantile(_RollingBase): """Rolling quantile. Note: - In pooled modes (``global_``/``groupby``/``partition_by``) this - transform has no aggregate-cache fast path: it falls back to a - row-level pass whose cost grows with ``unique timestamps x bucket - rows`` at fit, and aggregates are rebuilt at every recursive - prediction step. Can be slow on large panels. + In pooled modes (``global_``/``groupby``/``partition_by``) the fit and + predict fast paths read each bucket's per-ordinal value store (see + ``_TimestampAggregates``): a quantile is not derivable from the scalar + aggregates, so with ``time_agg=None`` the state retains every + observation (its memory grows O(rows)). With ``time_agg`` set the + quantile runs over the single collapsed scalar per timestamp and needs + no raw store. """ def __init__( @@ -988,20 +1055,66 @@ def _set_core_tfm(self, lag: int): ) return self + @property + def _needs_value_store(self) -> bool: + # Raw-row quantile: store observations. A time_agg quantile derives its + # one-scalar-per-timestamp view from the scalar aggregates instead. + return self.time_agg is None + def _window_stat(self, vals: np.ndarray) -> float: return float(np.quantile(vals, self.p)) + def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs): + lag = self._core_tfm.lag + w = self.window_size + min_samples = _resolve_min_samples(self) + result = np.full(len(bid_arr), np.nan) + for bid, agg in ts_aggs.items(): + feat_u = _rolling_quantile_from_agg( + agg, agg.unique_times, lag, w, min_samples, self.p + ) + if feat_u is None: + return None + idxs = np.where(bid_arr == bid)[0] + inv = np.searchsorted(agg.unique_times, ord_arr[idxs]) + result[idxs] = feat_u[inv] + return result -class _Seasonal_RollingBase(_BaseLagTransform): - """Rolling statistic over seasonal periods + def _latest_from_aggs_impl( + self, + ts_aggs, + target_ords: Dict[int, int], + _cache=None, + ) -> Optional[Dict[int, float]]: + lag = self._core_tfm.lag + w = self.window_size + min_samples = _resolve_min_samples(self) + result: Dict[int, float] = {} + for bid, agg in ts_aggs.items(): + t = np.array([target_ords[bid]], dtype=np.int64) + feat = _rolling_quantile_from_agg(agg, t, lag, w, min_samples, self.p) + if feat is None: + return None + result[bid] = float(feat[0]) + return result - Note: - In pooled modes (``global_``/``groupby``/``partition_by``) seasonal - rolling transforms have no aggregate-cache fast path: they fall back - to a row-level pass whose cost grows with ``unique timestamps x - bucket rows`` at fit, and aggregates are rebuilt at every recursive - prediction step. Can be slow on large panels. - """ + def _ts_level_from_aggs_impl(self, ts_aggs): + lag = self._core_tfm.lag + w = self.window_size + min_samples = _resolve_min_samples(self) + out: Dict[int, np.ndarray] = {} + for bid, agg in ts_aggs.items(): + feat_u = _rolling_quantile_from_agg( + agg, agg.unique_times, lag, w, min_samples, self.p + ) + if feat_u is None: + return None + out[bid] = feat_u + return out + + +class _Seasonal_RollingBase(_BaseLagTransform): + """Rolling statistic over seasonal periods""" def __init__( self, @@ -1132,25 +1245,358 @@ def _seasonal_stat(self, vals: np.ndarray) -> float: ) +def _seasonal_step_positions( + unique_times, target_ords, lag, season_length, window_size +): + """Locate every seasonal step's target ordinal in ``unique_times``. + + Seasonal windows stride in ordinal *value* space (``t - lag - + i*season_length`` for ``i in range(window_size)``); the calendar may have + holes (partition mode), so targets are matched by value via searchsorted, + never by position. Returns one ``(positions, valid)`` pair per seasonal + step: ``positions`` are clip-safe indices into ``unique_times`` and + ``valid`` flags which targets exist (matching the slow path's ``t >= 0`` + filter). ``unique_times`` must be non-empty. + """ + last = len(unique_times) - 1 + steps = [] + for i in range(window_size): + tgt = target_ords - lag - i * season_length + pos = np.minimum(np.searchsorted(unique_times, tgt), last) + valid = (tgt >= 0) & (unique_times[pos] == tgt) + steps.append((pos, valid)) + return steps + + +def _seasonal_mean_from_agg( + agg, target_ords, lag, season_length, window_size, min_samples +): + if len(agg.unique_times) == 0: + return np.full(len(target_ords), np.nan) + sums = agg.sums + counts = agg.counts + win_sum = np.zeros(len(target_ords)) + # Neumaier compensation: seasonal steps stride newest-to-oldest, so summing + # per-timestamp sums naively can drop a small residual added after a large + # partial (e.g. [1, -1e16, 1e16]). Counts are small exact integers, so their + # plain running sum needs no compensation. + comp = np.zeros(len(target_ords)) + win_cnt = np.zeros(len(target_ords)) + for pos, valid in _seasonal_step_positions( + agg.unique_times, target_ords, lag, season_length, window_size + ): + v = np.where(valid, sums[pos], 0.0) + t = win_sum + v + comp += np.where( + np.abs(win_sum) >= np.abs(v), (win_sum - t) + v, (v - t) + win_sum + ) + win_sum = t + win_cnt += np.where(valid, counts[pos], 0.0) + win_sum = win_sum + comp + safe_cnt = np.where(win_cnt > 0, win_cnt, 1.0) + return np.where( + (win_cnt >= min_samples) & (win_cnt > 0), win_sum / safe_cnt, np.nan + ) + + +def _seasonal_std_from_agg( + agg, target_ords, lag, season_length, window_size, min_samples +): + if len(agg.unique_times) == 0: + return np.full(len(target_ords), np.nan) + sums = agg.sums + counts = agg.counts + sum_sq = agg.sum_sq + # Combine the seasonal steps' per-timestamp moments with the numerically + # stable parallel (Chan et al.) algorithm rather than the textbook + # ``sum_sq - sum**2 / n``: for large-magnitude targets with small spread the + # latter cancels catastrophically (both terms ~n*mean**2) and collapses the + # variance to 0. The combine only ever forms differences of nearby means, so + # it stays accurate. ``M2_i`` is each timestamp's own sum of squared + # deviations (0 for single-row timestamps, exactly ``x**2 - x**2``). + n = len(target_ords) + count_n = np.zeros(n) + mean = np.zeros(n) + m2 = np.zeros(n) + for pos, valid in _seasonal_step_positions( + agg.unique_times, target_ords, lag, season_length, window_size + ): + ni = np.where(valid, counts[pos], 0.0) + has = ni > 0 + safe_ni = np.where(has, ni, 1.0) + si = np.where(has, sums[pos], 0.0) + mean_i = si / safe_ni + m2_i = np.maximum(np.where(has, sum_sq[pos] - si * si / safe_ni, 0.0), 0.0) + new_n = count_n + ni + safe_new_n = np.where(new_n > 0, new_n, 1.0) + delta = mean_i - mean + mean = np.where(has, mean + delta * ni / safe_new_n, mean) + m2 = np.where(has, m2 + m2_i + delta * delta * count_n * ni / safe_new_n, m2) + count_n = new_n + den = np.where(count_n > 1, count_n - 1, 1.0) + var = np.maximum(m2 / den, 0.0) + return np.where((count_n >= min_samples) & (count_n > 1), np.sqrt(var), np.nan) + + +def _seasonal_min_from_agg( + agg, target_ords, lag, season_length, window_size, min_samples +): + if len(agg.unique_times) == 0: + return np.full(len(target_ords), np.nan) + mins = agg.mins + counts = agg.counts + acc = np.full(len(target_ords), np.nan) + win_cnt = np.zeros(len(target_ords)) + for pos, valid in _seasonal_step_positions( + agg.unique_times, target_ords, lag, season_length, window_size + ): + # per-timestamp mins are NaN where the timestamp has no valid rows, so + # invalid/unobserved steps contribute NaN and fmin skips them + acc = np.fmin(acc, np.where(valid, mins[pos], np.nan)) + win_cnt += np.where(valid, counts[pos], 0.0) + return np.where((win_cnt >= min_samples) & (win_cnt > 0), acc, np.nan) + + +def _seasonal_max_from_agg( + agg, target_ords, lag, season_length, window_size, min_samples +): + if len(agg.unique_times) == 0: + return np.full(len(target_ords), np.nan) + maxs = agg.maxs + counts = agg.counts + acc = np.full(len(target_ords), np.nan) + win_cnt = np.zeros(len(target_ords)) + for pos, valid in _seasonal_step_positions( + agg.unique_times, target_ords, lag, season_length, window_size + ): + acc = np.fmax(acc, np.where(valid, maxs[pos], np.nan)) + win_cnt += np.where(valid, counts[pos], 0.0) + return np.where((win_cnt >= min_samples) & (win_cnt > 0), acc, np.nan) + + class SeasonalRollingMean(_Seasonal_RollingBase): def _seasonal_stat(self, vals: np.ndarray) -> float: return float(np.mean(vals)) + def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs): + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + result = np.full(len(bid_arr), np.nan) + for bid, agg in ts_aggs.items(): + idxs = np.where(bid_arr == bid)[0] + feat_u = _seasonal_mean_from_agg( + agg, agg.unique_times, lag, sl, w, min_samples + ) + inv = np.searchsorted(agg.unique_times, ord_arr[idxs]) + result[idxs] = feat_u[inv] + return result + + def _latest_from_aggs_impl( + self, + ts_aggs, + target_ords: Dict[int, int], + _cache=None, + ) -> Dict[int, float]: + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + result: Dict[int, float] = {} + for bid, agg in ts_aggs.items(): + t = np.array([target_ords[bid]], dtype=np.int64) + result[bid] = float( + _seasonal_mean_from_agg(agg, t, lag, sl, w, min_samples)[0] + ) + return result + + def _ts_level_from_aggs_impl(self, ts_aggs): + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + return { + bid: _seasonal_mean_from_agg(agg, agg.unique_times, lag, sl, w, min_samples) + for bid, agg in ts_aggs.items() + } + class SeasonalRollingStd(_Seasonal_RollingBase): def _seasonal_stat(self, vals: np.ndarray) -> float: return float(np.std(vals, ddof=1)) if len(vals) > 1 else np.nan + def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs): + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + result = np.full(len(bid_arr), np.nan) + for bid, agg in ts_aggs.items(): + idxs = np.where(bid_arr == bid)[0] + feat_u = _seasonal_std_from_agg( + agg, agg.unique_times, lag, sl, w, min_samples + ) + inv = np.searchsorted(agg.unique_times, ord_arr[idxs]) + result[idxs] = feat_u[inv] + return result + + def _latest_from_aggs_impl( + self, + ts_aggs, + target_ords: Dict[int, int], + _cache=None, + ) -> Dict[int, float]: + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + result: Dict[int, float] = {} + for bid, agg in ts_aggs.items(): + t = np.array([target_ords[bid]], dtype=np.int64) + result[bid] = float( + _seasonal_std_from_agg(agg, t, lag, sl, w, min_samples)[0] + ) + return result + + def _ts_level_from_aggs_impl(self, ts_aggs): + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + return { + bid: _seasonal_std_from_agg(agg, agg.unique_times, lag, sl, w, min_samples) + for bid, agg in ts_aggs.items() + } + class SeasonalRollingMin(_Seasonal_RollingBase): def _seasonal_stat(self, vals: np.ndarray) -> float: return float(np.min(vals)) + def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs): + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + result = np.full(len(bid_arr), np.nan) + for bid, agg in ts_aggs.items(): + idxs = np.where(bid_arr == bid)[0] + feat_u = _seasonal_min_from_agg( + agg, agg.unique_times, lag, sl, w, min_samples + ) + inv = np.searchsorted(agg.unique_times, ord_arr[idxs]) + result[idxs] = feat_u[inv] + return result + + def _latest_from_aggs_impl( + self, + ts_aggs, + target_ords: Dict[int, int], + _cache=None, + ) -> Dict[int, float]: + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + result: Dict[int, float] = {} + for bid, agg in ts_aggs.items(): + t = np.array([target_ords[bid]], dtype=np.int64) + result[bid] = float( + _seasonal_min_from_agg(agg, t, lag, sl, w, min_samples)[0] + ) + return result + + def _ts_level_from_aggs_impl(self, ts_aggs): + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + return { + bid: _seasonal_min_from_agg(agg, agg.unique_times, lag, sl, w, min_samples) + for bid, agg in ts_aggs.items() + } + class SeasonalRollingMax(_Seasonal_RollingBase): def _seasonal_stat(self, vals: np.ndarray) -> float: return float(np.max(vals)) + def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs): + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + result = np.full(len(bid_arr), np.nan) + for bid, agg in ts_aggs.items(): + idxs = np.where(bid_arr == bid)[0] + feat_u = _seasonal_max_from_agg( + agg, agg.unique_times, lag, sl, w, min_samples + ) + inv = np.searchsorted(agg.unique_times, ord_arr[idxs]) + result[idxs] = feat_u[inv] + return result + + def _latest_from_aggs_impl( + self, + ts_aggs, + target_ords: Dict[int, int], + _cache=None, + ) -> Dict[int, float]: + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + result: Dict[int, float] = {} + for bid, agg in ts_aggs.items(): + t = np.array([target_ords[bid]], dtype=np.int64) + result[bid] = float( + _seasonal_max_from_agg(agg, t, lag, sl, w, min_samples)[0] + ) + return result + + def _ts_level_from_aggs_impl(self, ts_aggs): + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + return { + bid: _seasonal_max_from_agg(agg, agg.unique_times, lag, sl, w, min_samples) + for bid, agg in ts_aggs.items() + } + + +def _seasonal_quantile_from_agg( + agg, query_ords, lag, season_length, window_size, min_samples, p +): + """Seasonal rolling quantile at each query ordinal, from the CSR store. + + Uses :func:`_seasonal_step_positions` (clip-safe, matched by value) to locate + each seasonal step's ordinal, then concatenates those ordinals' value slices + and takes the quantile -- mirroring the row slow path's ``np.isin`` gather + over ``t - lag - i*season_length``. Returns ``None`` with no value store. + """ + if agg.values is None or agg.row_offsets is None: + return None + ut = agg.unique_times + out = np.full(len(query_ords), np.nan) + if len(ut) == 0: + return out + off = agg.row_offsets + vals = agg.values + steps = _seasonal_step_positions(ut, query_ords, lag, season_length, window_size) + for j in range(len(query_ords)): + parts = [] + for pos, valid in steps: + if valid[j]: + s = int(pos[j]) + parts.append(vals[off[s] : off[s + 1]]) + if parts: + gathered = np.concatenate(parts) + if len(gathered) >= min_samples and len(gathered) > 0: + out[j] = np.quantile(gathered, p) + return out + class SeasonalRollingQuantile(_Seasonal_RollingBase): def __init__( @@ -1177,9 +1623,64 @@ def __init__( ) self.p = p + @property + def _needs_value_store(self) -> bool: + return self.time_agg is None + def _seasonal_stat(self, vals: np.ndarray) -> float: return float(np.quantile(vals, self.p)) + def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs): + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + result = np.full(len(bid_arr), np.nan) + for bid, agg in ts_aggs.items(): + feat_u = _seasonal_quantile_from_agg( + agg, agg.unique_times, lag, sl, w, min_samples, self.p + ) + if feat_u is None: + return None + idxs = np.where(bid_arr == bid)[0] + inv = np.searchsorted(agg.unique_times, ord_arr[idxs]) + result[idxs] = feat_u[inv] + return result + + def _latest_from_aggs_impl( + self, + ts_aggs, + target_ords: Dict[int, int], + _cache=None, + ) -> Optional[Dict[int, float]]: + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + result: Dict[int, float] = {} + for bid, agg in ts_aggs.items(): + t = np.array([target_ords[bid]], dtype=np.int64) + feat = _seasonal_quantile_from_agg(agg, t, lag, sl, w, min_samples, self.p) + if feat is None: + return None + result[bid] = float(feat[0]) + return result + + def _ts_level_from_aggs_impl(self, ts_aggs): + lag = self._core_tfm.lag + sl = self.season_length + w = self.window_size + min_samples = _resolve_min_samples(self) + out: Dict[int, np.ndarray] = {} + for bid, agg in ts_aggs.items(): + feat_u = _seasonal_quantile_from_agg( + agg, agg.unique_times, lag, sl, w, min_samples, self.p + ) + if feat_u is None: + return None + out[bid] = feat_u + return out + class _ExpandingBase(_BaseLagTransform): """Expanding statistic @@ -1232,9 +1733,10 @@ def update_samples(self) -> int: @property def _is_finite_window(self) -> bool: - # Pooled Expanding* recomputes cumsum over the FULL aggregate vectors at - # predict (no carried accumulator, unlike the local coreforecast path), - # so its window is effectively unbounded -- its state is never trimmed. + # Pooled Expanding* depends on the whole history prefix: its predict-time + # accumulator is initialized from the FULL aggregate vectors before being + # carried across steps, so its window is effectively unbounded and its + # state is never trimmed. return False def _bucket_feature_rows_impl( @@ -1298,21 +1800,33 @@ def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs): result[idxs] = feat_u[inv] return result - def _latest_from_aggs_impl(self, ts_aggs, target_ords): + def _latest_from_aggs_impl(self, ts_aggs, target_ords, cache=None): lag = self._core_tfm.lag result: Dict[int, float] = {} for bid, agg in ts_aggs.items(): - if len(agg.unique_times) == 0: + ut = agg.unique_times + if len(ut) == 0: result[bid] = float("nan") continue - cum_sum = np.cumsum(agg.sums) - cum_cnt = np.cumsum(agg.counts) t = target_ords[bid] upper = t - lag - ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1 - s = cum_sum[ui] if ui >= 0 else 0.0 - c = cum_cnt[ui] if ui >= 0 else 0.0 - result[bid] = s / c if c > 0 else float("nan") + ui = int(np.searchsorted(ut, upper, side="right")) - 1 + # Carry the running prefix sum/count across recursive predict steps: + # only the newly consumable ordinals ``(prev_ui, ui]`` are summed, + # so a horizon walk costs ``O(history + horizon)`` total instead of + # an ``O(n)`` cumsum per step. ``cache=None`` (non-predict callers) + # recomputes the whole prefix. + prev_ui, run_sum, run_cnt = ( + cache.get(bid, (-1, 0.0, 0.0)) if cache is not None else (-1, 0.0, 0.0) + ) + if prev_ui > ui: # non-monotone target ordinal: drop the stale prefix + prev_ui, run_sum, run_cnt = -1, 0.0, 0.0 + if ui > prev_ui: + run_sum += agg.sums[prev_ui + 1 : ui + 1].sum() + run_cnt += agg.counts[prev_ui + 1 : ui + 1].sum() + if cache is not None: + cache[bid] = (ui, run_sum, run_cnt) + result[bid] = float(run_sum / run_cnt) if run_cnt > 0 else float("nan") return result def _ts_level_from_aggs_impl(self, ts_aggs): @@ -1351,24 +1865,34 @@ def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs): result[idxs] = feat_u[inv] return result - def _latest_from_aggs_impl(self, ts_aggs, target_ords): + def _latest_from_aggs_impl(self, ts_aggs, target_ords, cache=None): lag = self._core_tfm.lag result: Dict[int, float] = {} for bid, agg in ts_aggs.items(): - if len(agg.unique_times) == 0: + ut = agg.unique_times + if len(ut) == 0: result[bid] = float("nan") continue - cum_sum = np.cumsum(agg.sums) - cum_cnt = np.cumsum(agg.counts) - cum_sum_sq = np.cumsum(agg.sum_sq) t = target_ords[bid] upper = t - lag - ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1 - s = cum_sum[ui] if ui >= 0 else 0.0 - c = cum_cnt[ui] if ui >= 0 else 0.0 - sq = cum_sum_sq[ui] if ui >= 0 else 0.0 - if c > 1: - var = (sq - s**2 / c) / (c - 1) + ui = int(np.searchsorted(ut, upper, side="right")) - 1 + # Carry running prefix sum/count/sum_sq across steps (see + # ExpandingMean); only ``(prev_ui, ui]`` is summed per step. + prev_ui, run_sum, run_cnt, run_sq = ( + cache.get(bid, (-1, 0.0, 0.0, 0.0)) + if cache is not None + else (-1, 0.0, 0.0, 0.0) + ) + if prev_ui > ui: # non-monotone target ordinal: drop the stale prefix + prev_ui, run_sum, run_cnt, run_sq = -1, 0.0, 0.0, 0.0 + if ui > prev_ui: + run_sum += agg.sums[prev_ui + 1 : ui + 1].sum() + run_cnt += agg.counts[prev_ui + 1 : ui + 1].sum() + run_sq += agg.sum_sq[prev_ui + 1 : ui + 1].sum() + if cache is not None: + cache[bid] = (ui, run_sum, run_cnt, run_sq) + if run_cnt > 1: + var = (run_sq - run_sum**2 / run_cnt) / (run_cnt - 1) var = max(var, 0.0) result[bid] = float(np.sqrt(var)) else: @@ -1409,18 +1933,31 @@ def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs): result[idxs] = feat_u[inv] return result - def _latest_from_aggs_impl(self, ts_aggs, target_ords): + def _latest_from_aggs_impl(self, ts_aggs, target_ords, cache=None): lag = self._core_tfm.lag result: Dict[int, float] = {} for bid, agg in ts_aggs.items(): - if len(agg.unique_times) == 0: + ut = agg.unique_times + if len(ut) == 0: result[bid] = float("nan") continue - prefix_min = np.fmin.accumulate(agg.mins) t = target_ords[bid] upper = t - lag - ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1 - result[bid] = float(prefix_min[ui]) if ui >= 0 else float("nan") + ui = int(np.searchsorted(ut, upper, side="right")) - 1 + # Carry the running prefix min across steps: fold only the newly + # consumable ordinals ``(prev_ui, ui]`` with NaN-skipping fmin, so a + # horizon walk costs ``O(history + horizon)`` instead of an ``O(n)`` + # fmin.accumulate per step. + prev_ui, run = ( + cache.get(bid, (-1, np.nan)) if cache is not None else (-1, np.nan) + ) + if prev_ui > ui: # non-monotone target ordinal: drop the stale prefix + prev_ui, run = -1, np.nan + if ui > prev_ui: + run = np.fmin(run, np.fmin.reduce(agg.mins[prev_ui + 1 : ui + 1])) + if cache is not None: + cache[bid] = (ui, run) + result[bid] = float(run) if ui >= 0 and not np.isnan(run) else float("nan") return result def _ts_level_from_aggs_impl(self, ts_aggs): @@ -1443,18 +1980,29 @@ def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs): result[idxs] = feat_u[inv] return result - def _latest_from_aggs_impl(self, ts_aggs, target_ords): + def _latest_from_aggs_impl(self, ts_aggs, target_ords, cache=None): lag = self._core_tfm.lag result: Dict[int, float] = {} for bid, agg in ts_aggs.items(): - if len(agg.unique_times) == 0: + ut = agg.unique_times + if len(ut) == 0: result[bid] = float("nan") continue - prefix_max = np.fmax.accumulate(agg.maxs) t = target_ords[bid] upper = t - lag - ui = int(np.searchsorted(agg.unique_times, upper, side="right")) - 1 - result[bid] = float(prefix_max[ui]) if ui >= 0 else float("nan") + ui = int(np.searchsorted(ut, upper, side="right")) - 1 + # Carry the running prefix max across steps (see ExpandingMin); + # fold only ``(prev_ui, ui]`` with NaN-skipping fmax per step. + prev_ui, run = ( + cache.get(bid, (-1, np.nan)) if cache is not None else (-1, np.nan) + ) + if prev_ui > ui: # non-monotone target ordinal: drop the stale prefix + prev_ui, run = -1, np.nan + if ui > prev_ui: + run = np.fmax(run, np.fmax.reduce(agg.maxs[prev_ui + 1 : ui + 1])) + if cache is not None: + cache[bid] = (ui, run) + result[bid] = float(run) if ui >= 0 and not np.isnan(run) else float("nan") return result def _ts_level_from_aggs_impl(self, ts_aggs): @@ -1462,15 +2010,45 @@ def _ts_level_from_aggs_impl(self, ts_aggs): return {bid: _expanding_max_from_agg(agg, lag) for bid, agg in ts_aggs.items()} +def _expanding_quantile_from_agg(agg, query_ords, lag, p): + """Expanding quantile at each query ordinal, from the CSR value store. + + The window for target ``t`` is every ordinal ``<= t - lag`` (the same + ``upper`` bound as :func:`_expanding_mean_from_agg`), i.e. the value-store + prefix ``values[:row_offsets[ui + 1]]``. No ``min_samples`` (the expanding + row path emits a value whenever the prefix has any observation). Returns + ``None`` with no value store. + """ + if agg.values is None or agg.row_offsets is None: + return None + ut = agg.unique_times + out = np.full(len(query_ords), np.nan) + if len(ut) == 0: + return out + off = agg.row_offsets + vals = agg.values + ui = np.searchsorted(ut, query_ords - lag, side="right") - 1 + for j in range(len(query_ords)): + u = int(ui[j]) + if u < 0: + continue + end = int(off[u + 1]) + if end > 0: + out[j] = np.quantile(vals[:end], p) + return out + + class ExpandingQuantile(_ExpandingBase): """Expanding quantile. Note: - In pooled modes (``global_``/``groupby``/``partition_by``) this - transform has no aggregate-cache fast path: it falls back to a - row-level pass whose cost grows with ``unique timestamps x bucket - rows`` at fit, and aggregates are rebuilt at every recursive - prediction step. Can be slow on large panels. + In pooled modes (``global_``/``groupby``/``partition_by``) the fit and + predict fast paths read each bucket's per-ordinal value store (see + ``_TimestampAggregates``): a quantile is not derivable from the scalar + aggregates, so with ``time_agg=None`` the state retains every + observation (its memory grows O(rows)). With ``time_agg`` set the + quantile runs over the single collapsed scalar per timestamp and needs + no raw store. """ def __init__( @@ -1495,9 +2073,46 @@ def __init__( def update_samples(self) -> int: return -1 + @property + def _needs_value_store(self) -> bool: + return self.time_agg is None + def _expanding_stat(self, vals: np.ndarray) -> float: return float(np.quantile(vals, self.p)) + def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs): + lag = self._core_tfm.lag + result = np.full(len(bid_arr), np.nan) + for bid, agg in ts_aggs.items(): + feat_u = _expanding_quantile_from_agg(agg, agg.unique_times, lag, self.p) + if feat_u is None: + return None + idxs = np.where(bid_arr == bid)[0] + inv = np.searchsorted(agg.unique_times, ord_arr[idxs]) + result[idxs] = feat_u[inv] + return result + + def _latest_from_aggs_impl(self, ts_aggs, target_ords, _cache=None): + lag = self._core_tfm.lag + result: Dict[int, float] = {} + for bid, agg in ts_aggs.items(): + t = np.array([target_ords[bid]], dtype=np.int64) + feat = _expanding_quantile_from_agg(agg, t, lag, self.p) + if feat is None: + return None + result[bid] = float(feat[0]) + return result + + def _ts_level_from_aggs_impl(self, ts_aggs): + lag = self._core_tfm.lag + out: Dict[int, np.ndarray] = {} + for bid, agg in ts_aggs.items(): + feat_u = _expanding_quantile_from_agg(agg, agg.unique_times, lag, self.p) + if feat_u is None: + return None + out[bid] = feat_u + return out + def _ewm_from_agg(agg, lag, alpha): # ``agg`` may be a lazy ``_ReaggregatedAggregates`` view whose ``counts`` / @@ -1607,9 +2222,10 @@ def _pooled_time_agg(self) -> Optional[str]: @property def _is_finite_window(self) -> bool: - # Pooled EWM consumes every observed bucket-aggregate mean up to the lag - # at predict (no carried running state), so it depends on the full - # history -- its state is never trimmed. + # Pooled EWM folds every observed bucket-aggregate mean up to the lag at + # predict: its per-model accumulator resumes from an exact prefix + # midpoint across steps, but that midpoint is seeded from the full + # history, so it depends on all of it -- its state is never trimmed. return False def _bucket_feature_rows_impl( @@ -1667,30 +2283,40 @@ def _bucket_feature_from_aggs_impl(self, bid_arr, ord_arr, ts_aggs): result[idxs] = feat_u[inv] return result - def _latest_from_aggs_impl(self, ts_aggs, target_ords): + def _latest_from_aggs_impl(self, ts_aggs, target_ords, cache=None): lag = self._core_tfm.lag alpha = self.alpha result: Dict[int, float] = {} for bid, agg in ts_aggs.items(): - if len(agg.unique_times) == 0: + ut = agg.unique_times + if len(ut) == 0: result[bid] = float("nan") continue t = target_ords[bid] - counts = agg.counts - mean_per_ord = np.full(len(agg.unique_times), np.nan) - np.divide(agg.sums, counts, out=mean_per_ord, where=counts > 0) - ewm = np.nan upper = t - lag - for k in range(len(agg.unique_times)): - if agg.unique_times[k] > upper: - break - if not np.isnan(mean_per_ord[k]): - ewm = ( - mean_per_ord[k] - if np.isnan(ewm) - else alpha * mean_per_ord[k] + (1 - alpha) * ewm - ) - result[bid] = float(ewm) if not np.isnan(ewm) else float("nan") + # Resume the EWM recurrence from the last consumed ordinal. The fold + # is sequential, so the saved ``(k, ewm)`` is an exact midpoint and + # the result is identical to replaying from ``k=0`` -- but a horizon + # walk costs ``O(history + horizon)`` instead of refolding the whole + # prefix each step. ``cache=None`` (non-predict callers) refolds from + # scratch. Each observed timestamp contributes its bucket mean once. + k, ewm, has = ( + cache.get(bid, (0, np.nan, False)) + if cache is not None + else (0, np.nan, False) + ) + sums = agg.sums + counts = agg.counts + while k < len(ut) and ut[k] <= upper: + c = counts[k] + if c > 0: + m = sums[k] / c + ewm = m if not has else alpha * m + (1 - alpha) * ewm + has = True + k += 1 + if cache is not None: + cache[bid] = (k, ewm, has) + result[bid] = float(ewm) if has else float("nan") return result def _ts_level_from_aggs_impl(self, ts_aggs): @@ -1742,11 +2368,15 @@ def update_samples(self) -> int: def _is_finite_window(self) -> bool: return self.tfm._is_finite_window + @property + def _needs_value_store(self) -> bool: + return self.tfm._needs_value_store + def _compute_ts_level_from_aggs(self, ts_aggs): return self.tfm._compute_ts_level_from_aggs(ts_aggs) - def _compute_latest_from_aggs(self, ts_aggs, target_ords): - return self.tfm._compute_latest_from_aggs(ts_aggs, target_ords) + def _compute_latest_from_aggs(self, ts_aggs, target_ords, cache=None): + return self.tfm._compute_latest_from_aggs(ts_aggs, target_ords, cache) def _compute_bucket_feature( self, @@ -1827,6 +2457,10 @@ def update_samples(self): def _is_finite_window(self) -> bool: return self.tfm1._is_finite_window and self.tfm2._is_finite_window + @property + def _needs_value_store(self) -> bool: + return self.tfm1._needs_value_store or self.tfm2._needs_value_store + def _compute_ts_level_from_aggs(self, ts_aggs): r1 = self.tfm1._compute_ts_level_from_aggs(ts_aggs) r2 = self.tfm2._compute_ts_level_from_aggs(ts_aggs) @@ -1834,9 +2468,15 @@ def _compute_ts_level_from_aggs(self, ts_aggs): return {bid: self.operator(r1[bid], r2[bid]) for bid in r1 if bid in r2} return None - def _compute_latest_from_aggs(self, ts_aggs, target_ords): - r1 = self.tfm1._compute_latest_from_aggs(ts_aggs, target_ords) - r2 = self.tfm2._compute_latest_from_aggs(ts_aggs, target_ords) + def _compute_latest_from_aggs(self, ts_aggs, target_ords, cache=None): + # Each operand keys its cache by bucket id, so give them separate + # sub-dicts (string keys can't collide with the int bucket ids inside). + c1 = c2 = None + if cache is not None: + c1 = cache.setdefault("_tfm1", {}) + c2 = cache.setdefault("_tfm2", {}) + r1 = self.tfm1._compute_latest_from_aggs(ts_aggs, target_ords, c1) + r2 = self.tfm2._compute_latest_from_aggs(ts_aggs, target_ords, c2) if r1 is not None and r2 is not None: return {bid: self.operator(r1[bid], r2[bid]) for bid in r1 if bid in r2} return None diff --git a/mlforecast/pooled.py b/mlforecast/pooled.py index a1291c4f..5c928e16 100644 --- a/mlforecast/pooled.py +++ b/mlforecast/pooled.py @@ -178,12 +178,22 @@ class _TimestampAggregates: sum_sq: np.ndarray mins: np.ndarray maxs: np.ndarray + # Opt-in CSR raw-value store, built only when a state's transforms include a + # raw-row quantile (see ``PooledState._store_values``). ``values`` holds + # every non-NaN observation grouped by ordinal — slot ``k``'s observations + # are ``values[row_offsets[k]:row_offsets[k + 1]]`` — and ``row_offsets`` is + # always integer (``np.intp``) so it can index arrays. ``None`` on both means + # "no store": quantile fast paths return ``None`` and fall back to the row + # slow path (also the state for old pickles, which lack these fields). + values: Optional[np.ndarray] = None + row_offsets: Optional[np.ndarray] = None def _build_ts_aggs( bid_arr: np.ndarray, ord_arr: np.ndarray, y_arr: np.ndarray, + store_values: bool = False, ) -> Dict[int, _TimestampAggregates]: aggs: Dict[int, _TimestampAggregates] = {} for bid in np.unique(bid_arr): @@ -207,6 +217,22 @@ def _build_ts_aggs( no_valid = mins == np.inf mins[no_valid] = np.nan maxs[no_valid] = np.nan + values = row_offsets = None + if store_values: + # Group every non-NaN observation by its ordinal slot. The input + # rows are NOT assumed ordinal-sorted (flat arrays may be in any row + # order, and updates recompute ordinals without reordering them), so + # stable-sort the valid rows by their slot ``inv[valid]`` to lay the + # values out contiguously per ordinal. Offsets come from an *integer* + # per-slot count, never ``cumsum(counts)``: ``counts`` is float (it is + # a weighted ``np.bincount`` above) and float offsets cannot index + # arrays. Empty / all-NaN ordinals get a zero count -> repeated offset. + order = np.argsort(valid_inv, kind="stable") + values = valid_y[order].astype(float, copy=False) + counts_int = np.bincount(valid_inv, minlength=m).astype(np.intp) + row_offsets = np.empty(m + 1, dtype=np.intp) + row_offsets[0] = 0 + np.cumsum(counts_int, out=row_offsets[1:]) aggs[int(bid)] = _TimestampAggregates( unique_times=unique_ord, sums=sums, @@ -214,6 +240,8 @@ def _build_ts_aggs( sum_sq=sum_sq, mins=mins, maxs=maxs, + values=values, + row_offsets=row_offsets, ) return aggs @@ -306,6 +334,24 @@ def mins(self) -> np.ndarray: def maxs(self) -> np.ndarray: return self._values + @property + def values(self) -> np.ndarray: + # CSR value store for the quantile fast paths, derived rather than + # stored: under ``time_agg`` each ordinal collapses to a single scalar + # (``_values``), so a quantile over the collapsed series needs only the + # one observed value per timestamp — never the O(rows) raw store a + # raw-row quantile builds. Observed timestamps only (NaN => unobserved). + return self._values[self._obs] + + @property + def row_offsets(self) -> np.ndarray: + # Each observed timestamp contributes exactly one value; unobserved + # timestamps contribute none (a repeated offset). Integer so it indexes. + offsets = np.empty(len(self._values) + 1, dtype=np.intp) + offsets[0] = 0 + np.cumsum(self._obs.astype(np.intp), out=offsets[1:]) + return offsets + def _reaggregate_ts_aggs( ts_aggs: Dict[int, _TimestampAggregates], time_agg: str @@ -516,6 +562,13 @@ class PooledState: _scope_key_to_parent_id: Optional[Dict[tuple, int]] = None _ts_aggs: Dict[int, _TimestampAggregates] = field(default_factory=dict) _idsorted_to_bucket_pos: Optional[np.ndarray] = None + # Immutable build policy: whether ``_ts_aggs`` carries the CSR raw-value + # store the quantile fast paths need (set once by the constructors from the + # state's transforms). Every build/rebuild/create path reads this so a + # rebuilt aggregate keeps the store the original build chose. Not + # snapshotted (never mutated during prediction); ``getattr(state, + # "_store_values", False)`` keeps old pickles on the slow path. + _store_values: bool = False @property def group_uids(self): @@ -587,6 +640,7 @@ def from_global( target_col: str, ga_data_dtype, n_series: int, + store_values: bool = False, ): keep_cols = _dedupe_preserve_order([id_col, time_col, target_col]) global_df = sorted_df[keep_cols] @@ -600,6 +654,11 @@ def from_global( ord_raw = np.searchsorted(unique_ts, ts_raw).astype(np.int64) bid_arr = np.zeros(len(global_df), dtype=np.int64) y_float = y_raw.astype(float) + # The stored bucket_df is only ever keyed on (id_col, time_col) after + # construction (slow-path join + idsorted permutation); the target lives + # in ``y``. Drop everything else so we don't retain a full-history copy + # of columns nothing reads. + global_df = global_df[_dedupe_preserve_order([id_col, time_col])] return cls( bucket_df=global_df, groups=None, @@ -611,7 +670,10 @@ def from_global( y=y_float, next_time_index_by_bucket={0: len(unique_ts)}, join_cols=[id_col, time_col], - _ts_aggs=_build_ts_aggs(bid_arr, ord_raw, y_float), + _ts_aggs=_build_ts_aggs( + bid_arr, ord_raw, y_float, store_values=store_values + ), + _store_values=store_values, ) @classmethod @@ -624,6 +686,7 @@ def from_groupby( target_col: str, ga_data_dtype, static_features, + store_values: bool = False, ): keep_cols = _dedupe_preserve_order( [id_col] + group_cols_list + [time_col, target_col] @@ -639,6 +702,12 @@ def from_groupby( bid_raw = bucket_df["_bucket_id"].to_numpy() bucket_df = ufp.drop_index_if_pandas(bucket_df) bid_arr = bid_raw.astype(np.int64) + # The stored bucket_df is only ever keyed on (id_col, time_col) after + # construction (slow-path join + idsorted permutation); the group values + # are recoverable from ``groups`` + ``bucket_id`` and the target lives in + # ``y``. Drop everything else so we don't retain a full-history copy of + # the (static) group columns for every series. + bucket_df = bucket_df[_dedupe_preserve_order([id_col, time_col])] ord_arr, next_by_bucket = _compute_time_index(bid_arr, ts_raw) series_bucket_id = lookup_bucket_ids( static_features, groups, group_cols_list @@ -655,7 +724,10 @@ def from_groupby( y=y_float, next_time_index_by_bucket=next_by_bucket, join_cols=[id_col, time_col], - _ts_aggs=_build_ts_aggs(bid_arr, ord_arr, y_float), + _ts_aggs=_build_ts_aggs( + bid_arr, ord_arr, y_float, store_values=store_values + ), + _store_values=store_values, ) @classmethod @@ -671,6 +743,7 @@ def from_partition( ga_data_dtype, static_features, n_series: int, + store_values: bool = False, ): """Build a PooledState for partition_by transforms. @@ -811,7 +884,10 @@ def from_partition( _bucket_to_parent_id=bucket_to_parent, _parent_to_buckets=parent_to_buckets, _scope_key_to_parent_id=scope_key_to_parent, - _ts_aggs=_build_ts_aggs(bid_arr, ord_arr, y_float), + _ts_aggs=_build_ts_aggs( + bid_arr, ord_arr, y_float, store_values=store_values + ), + _store_values=store_values, ) def update_series_bucket_id(self, context_df, _id_col: str): @@ -855,6 +931,12 @@ def update_series_bucket_id(self, context_df, _id_col: str): else: self.next_time_index_by_bucket[bid_int] = 0 if bid_int not in self._ts_aggs: + if self._store_values: + empty_values: Optional[np.ndarray] = np.empty(0) + empty_offsets: Optional[np.ndarray] = np.array([0], dtype=np.intp) + else: + empty_values = None + empty_offsets = None self._ts_aggs[bid_int] = _TimestampAggregates( unique_times=np.array([], dtype=np.intp), sums=np.array([], dtype=np.float64), @@ -862,6 +944,8 @@ def update_series_bucket_id(self, context_df, _id_col: str): sum_sq=np.array([], dtype=np.float64), mins=np.array([], dtype=np.float64), maxs=np.array([], dtype=np.float64), + values=empty_values, + row_offsets=empty_offsets, ) def _resolve_parent_for_bucket(self, bid: int, groups_nw=None) -> Optional[int]: @@ -970,6 +1054,13 @@ def append_predictions(self, curr_dates, predictions, n_series): agg.maxs = np.append( agg.maxs, np.max(valid_vals) if len(valid_vals) > 0 else np.nan ) + # CSR store (present iff the state builds it): the step adds one + # new ordinal whose non-NaN values are all of ``valid_vals``. + if agg.row_offsets is not None: + agg.values = np.append(agg.values, valid_vals) + agg.row_offsets = np.append( + agg.row_offsets, agg.row_offsets[-1] + valid_vals.size + ).astype(np.intp, copy=False) else: sort_order = np.argsort(self.series_bucket_id, kind="stable") sorted_bids = self.series_bucket_id[sort_order] @@ -1003,6 +1094,13 @@ def append_predictions(self, curr_dates, predictions, n_series): agg.maxs = np.append( agg.maxs, np.max(valid_vals) if len(valid_vals) > 0 else np.nan ) + # CSR store (present iff the state builds it): one new + # ordinal for this bucket holding its non-NaN step values. + if agg.row_offsets is not None: + agg.values = np.append(agg.values, valid_vals) + agg.row_offsets = np.append( + agg.row_offsets, agg.row_offsets[-1] + valid_vals.size + ).astype(np.intp, copy=False) if self._parent_time_grids is not None: self._advance_parent_calendars(new_ts_val) else: @@ -1035,7 +1133,9 @@ def append_observations( self.bucket_id = np.concatenate([self.bucket_id, new_bid]) self.time_index = np.concatenate([old_idx, new_idx]) self.next_time_index_by_bucket[0] = len(unique_all) - self._ts_aggs = _build_ts_aggs(self.bucket_id, self.time_index, self.y) + self._ts_aggs = _build_ts_aggs( + self.bucket_id, self.time_index, self.y, store_values=self._store_values + ) old_len = len(self.bucket_df) new_df_nw = nw.from_native(new_df) new_rows_nw = new_df_nw.with_row_index(name="_bucket_pos").with_columns( @@ -1111,7 +1211,9 @@ def append_observations( new_ord_arr, new_next = _compute_time_index(all_bid, all_ts) self.time_index = new_ord_arr self.next_time_index_by_bucket = new_next - self._ts_aggs = _build_ts_aggs(self.bucket_id, self.time_index, self.y) + self._ts_aggs = _build_ts_aggs( + self.bucket_id, self.time_index, self.y, store_values=self._store_values + ) old_len = len(self.bucket_df) bucket_df_nw = nw.from_native(bucket_df) new_rows_nw = bucket_df_nw.with_row_index(name="_bucket_pos").with_columns( @@ -1183,7 +1285,8 @@ def trim_to_last(self, n_ordinals: int) -> None: ``_idsorted_to_bucket_pos``) is regenerated through the same primitives the constructors/append paths use, so all representations stay mutually consistent. The caller must pass an ``n_ordinals`` covering every - transform's window (see the retention rule in ``TimeSeries._transform``), + transform's window (see the retention rule in + ``TimeSeries._trim_pooled_states``), so the dropped prefix can never enter a window and the trim is prediction-neutral. @@ -1250,6 +1353,19 @@ def trim_to_last(self, n_ordinals: int) -> None: m = agg.unique_times >= cut if not m.any(): continue + # The CSR store trims to the same surviving suffix: ``m`` is a + # contiguous tail of ordinals (unique_times is sorted), so the + # surviving observations are one contiguous value slice. + if agg.values is not None and agg.row_offsets is not None: + first = int(np.argmax(m)) + first_off = int(agg.row_offsets[first]) + new_values: Optional[np.ndarray] = agg.values[first_off:] + new_row_offsets: Optional[np.ndarray] = ( + agg.row_offsets[first:] - first_off + ) + else: + new_values = None + new_row_offsets = None new_aggs[bid_int] = _TimestampAggregates( unique_times=agg.unique_times[m] - cut, sums=agg.sums[m], @@ -1257,10 +1373,14 @@ def trim_to_last(self, n_ordinals: int) -> None: sum_sq=agg.sum_sq[m], mins=agg.mins[m], maxs=agg.maxs[m], + values=new_values, + row_offsets=new_row_offsets, ) self._ts_aggs = new_aggs else: - self._ts_aggs = _build_ts_aggs(self.bucket_id, self.time_index, self.y) + self._ts_aggs = _build_ts_aggs( + self.bucket_id, self.time_index, self.y, store_values=self._store_values + ) for bid_int, cal_len in self.next_time_index_by_bucket.items(): self.next_time_index_by_bucket[bid_int] = min(cal_len, n_ordinals) if self._idsorted_to_bucket_pos is not None: @@ -1323,7 +1443,9 @@ def compute_pooled_features( ) -> Dict[str, np.ndarray]: if query_arrays is not None: bid_arr, idx_arr, y_arr = query_arrays - ts_aggs = _build_ts_aggs(bid_arr, idx_arr, y_arr) + ts_aggs = _build_ts_aggs( + bid_arr, idx_arr, y_arr, store_values=state._store_values + ) else: bid_arr = state.bucket_id idx_arr = state.time_index diff --git a/nbs/docs/how-to-guides/pooled_lag_transforms.ipynb b/nbs/docs/how-to-guides/pooled_lag_transforms.ipynb index b2e3b3f9..a0ba09c9 100644 --- a/nbs/docs/how-to-guides/pooled_lag_transforms.ipynb +++ b/nbs/docs/how-to-guides/pooled_lag_transforms.ipynb @@ -589,7 +589,7 @@ "`keep_last_n` controls how much history is carried into the recursive prediction loop. It applies to pooled transforms too, **per state** (one state per `(mode, groupby, partition_by)` combination):\n", "\n", "* A state whose transforms are **all finite-window** — `Lag`, any `Rolling*` or `SeasonalRolling*`, and `Offset`/`Combine` built from them — is trimmed to its last `max(keep_last_n, W_state)` parent-calendar ordinals, where `W_state` is the state's widest window. The dropped prefix can never enter a window, so predictions are unchanged.\n", - "* A state containing any **unbounded** transform — `Expanding*` or `ExponentiallyWeightedMean` — keeps its **full** history. Unlike the local (coreforecast) path, pooled expanding/EWM carry no running accumulator: they recompute over the entire aggregate vectors at every step, so trimming them would change predictions.\n", + "* A state containing any **unbounded** transform — `Expanding*` or `ExponentiallyWeightedMean` — keeps its **full** history. Unlike the local (coreforecast) path, pooled expanding/EWM depend on the entire history prefix (their predict-time running accumulator is rebuilt per model from the full aggregate vectors, then carried across steps), so trimming them would change predictions.\n", "\n", "This applies whether `keep_last_n` is set explicitly or inferred — when you leave it unset it is inferred as the largest window across all transforms — mirroring how the per-series arrays used by local transforms are trimmed.\n", "\n", @@ -647,9 +647,9 @@ "* All pooled modes — `global_`, `groupby`, and `partition_by` (including local, alone) — require **series to end at the same timestamp**.\n", "* Subset prediction (`predict(ids=[...])`) is rejected when nonlocal pooled transforms are in play: omitted series would still contribute to the bucket aggregates.\n", "* All supported rolling, expanding, seasonal-rolling, and exponentially weighted transforms accept `global_`, `groupby`, and `partition_by`. `Offset` and `Combine` delegate to their wrapped transforms.\n", - "* **Performance**: `RollingQuantile`, `ExpandingQuantile` and the `SeasonalRolling*` transforms have no aggregate-cache fast path in pooled modes. They fall back to a row-level pass whose cost grows with `unique timestamps × bucket rows` at fit, and aggregates are rebuilt at every recursive prediction step. For large panels prefer the mean/std/min/max/EWM transforms, which use cached per-timestamp aggregates.\n", - "* **Memory**: finite-window pooled states (`Lag`/`Rolling*`/`SeasonalRolling*`) are trimmed under `keep_last_n` just like the per-series arrays (see the [`keep_last_n` and pooled history](#keep_last_n-and-pooled-history) section above); a state containing an `Expanding*`/`ExponentiallyWeightedMean` transform keeps the full training history because it recomputes over all of it at predict. Each `predict` call also backs up the pooled state per model to isolate recursive mutations. Expect higher peak memory with unbounded pooled transforms on large panels.\n", - "* `time_agg` (`'sum'`/`'count'`/`'mean'`/`'min'`/`'max'`) pre-aggregates rows sharing a timestamp within each bucket before the transform runs; it requires `global_` or `groupby` and counts observed timestamps for `min_samples`. All pooled-capable transforms support it, including the quantile and seasonal ones (via the row-level path).\n" + "* **Performance**: every supported pooled transform computes from the cached per-timestamp aggregates. `SeasonalRolling*` strides its window in ordinal value space over those aggregates, and the quantiles (`RollingQuantile`/`ExpandingQuantile`/`SeasonalRollingQuantile`) read a per-bucket value index built alongside the aggregates. This replaces the earlier row-level pass at fit and the per-step aggregate rebuild at predict, so pooled transforms now scale with the number of unique timestamps rather than bucket rows.\n", + "* **Memory**: finite-window pooled states (`Lag`/`Rolling*`/`SeasonalRolling*`) are trimmed under `keep_last_n` just like the per-series arrays (see the [`keep_last_n` and pooled history](#keep_last_n-and-pooled-history) section above); a state containing an `Expanding*`/`ExponentiallyWeightedMean` transform keeps the full training history because its statistic depends on all of it at predict. A state containing a raw-row quantile (`RollingQuantile`/`ExpandingQuantile`/`SeasonalRollingQuantile` without `time_agg`) additionally keeps every observation — O(rows) — since a quantile can't be recovered from the scalar aggregates; a `time_agg` quantile needs only the one collapsed value per timestamp and adds no such store. Each `predict` call also backs up the pooled state per model to isolate recursive mutations. Expect higher peak memory with unbounded or raw-row-quantile pooled transforms on large panels.\n", + "* `time_agg` (`'sum'`/`'count'`/`'mean'`/`'min'`/`'max'`) pre-aggregates rows sharing a timestamp within each bucket before the transform runs; it requires `global_` or `groupby` and counts observed timestamps for `min_samples`. All pooled-capable transforms support it, including the quantile and seasonal ones, which run their aggregate fast path over the collapsed per-timestamp series.\n" ] }, { diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 9b76cf1e..640a08b8 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -5,7 +5,13 @@ from utilsforecast.losses import smape from mlforecast import MLForecast -from mlforecast.lag_transforms import RollingMax, RollingMean, RollingMin +from mlforecast.lag_transforms import ( + ExpandingMean, + ExponentiallyWeightedMean, + RollingMax, + RollingMean, + RollingMin, +) from mlforecast.target_transforms import Differences, LocalStandardScaler from mlforecast.utils import generate_daily_series @@ -102,6 +108,51 @@ def test_predict(benchmark, fcst, series, use_exog, series_with_exog, exogs, sta assert summary["lr"] < summary["seas_naive"] +@pytest.fixture(scope="module") +def pooled_series(series): + """``series`` plus a low-cardinality categorical static so the groupby + pooled transform has real groups (the generated statics are all floats).""" + s = series.copy() + uids = s["unique_id"].unique() + grp_map = {u: f"g{i % 10}" for i, u in enumerate(uids)} + s["grp"] = s["unique_id"].map(grp_map).astype("category") + return s + + +@pytest.fixture +def pooled_fcst(): + # Exercises every pooled predict fast path PR C touches: the finite-window + # tail-window path (RollingMean/Max) and the carried per-model prefix caches + # (ExpandingMean, EWM). Non-local pooled transforms need equal ends, which + # the ``series`` fixture provides. + return MLForecast( + models={"lr": LinearRegression()}, + freq="D", + lags=[1, 7], + lag_transforms={ + 1: [ + RollingMean(7, global_=True), + ExpandingMean(global_=True), + ExponentiallyWeightedMean(alpha=0.9, global_=True), + ], + 7: [RollingMax(7, groupby=["grp"])], + }, + date_features=["dayofweek"], + ) + + +@pytest.mark.parametrize("num_threads", [1, 2]) +def test_predict_pooled(benchmark, pooled_fcst, pooled_series, statics, num_threads): + horizon = 14 + valid = pooled_series.groupby("unique_id").tail(horizon) + train = pooled_series.drop(valid.index) + pooled_fcst.ts.num_threads = num_threads + pooled_fcst.fit(train, static_features=statics + ["grp"]) + preds = benchmark(pooled_fcst.predict, horizon) + assert len(preds) == pooled_series["unique_id"].nunique() * horizon + assert np.isfinite(preds["lr"].to_numpy()).all() + + def test_drop_auxiliary_columns_default_drops_groupby(series): """Default True should auto-drop all groupby columns but keep computed lag features.""" statics = series.columns.drop(["unique_id", "ds", "y"]).tolist() diff --git a/tests/test_pooled.py b/tests/test_pooled.py index c83a6885..438a44db 100644 --- a/tests/test_pooled.py +++ b/tests/test_pooled.py @@ -10,13 +10,20 @@ ExpandingMax, ExpandingMean, ExpandingMin, + ExpandingQuantile, ExpandingStd, ExponentiallyWeightedMean, LookupLag, RollingMax, RollingMean, RollingMin, + RollingQuantile, RollingStd, + SeasonalRollingMax, + SeasonalRollingMean, + SeasonalRollingMin, + SeasonalRollingQuantile, + SeasonalRollingStd, ) _LAGS = [1, 3] @@ -181,6 +188,42 @@ def test_group_update_preserves_bucket_df(engine, lag): assert new_len == orig_len + 2 +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_bucket_df_stores_only_id_and_time(engine): + """global/groupby states retain only ``[id_col, time_col]`` in bucket_df: + the (static) group columns and the target are not kept per row -- they live + in ``groups``/``bucket_id`` and ``y``. Asserts the exact schema (robust even + when a group column coincides with the id column).""" + df = _make_df( + engine, + { + "unique_id": ["a", "a", "b", "b"], + "ds": [1, 2, 1, 2], + "y": [1.0, 2.0, 10.0, 20.0], + "brand": ["x", "x", "y", "y"], + }, + ) + ts = TimeSeries( + freq=1, + lag_transforms={ + 1: [RollingMean(2, global_=True), RollingMean(2, groupby=["brand"])] + }, + ) + ts.fit_transform( + df, + id_col="unique_id", + time_col="ds", + target_col="y", + dropna=False, + static_features=["brand"], + keep_last_n=10_000, + ) + for key in (("global", (), ()), ("groupby", ("brand",), ())): + state = ts._pooled_states[key] + assert list(state.bucket_df.columns) == ["unique_id", "ds"], key + assert len(state.bucket_df) == len(df), key # rows preserved + + @pytest.mark.parametrize("engine", ["pandas", "polars"]) @pytest.mark.parametrize("lag", _LAGS) def test_global_sequential_updates(engine, lag): @@ -2211,6 +2254,13 @@ def test_pooled_transforms_lag2_groupby(engine): lambda m: ExpandingMin(**m), lambda m: ExpandingMax(**m), lambda m: ExponentiallyWeightedMean(alpha=0.3, **m), + lambda m: SeasonalRollingMean(season_length=3, window_size=2, **m), + lambda m: SeasonalRollingStd(season_length=3, window_size=2, **m), + lambda m: SeasonalRollingMin(season_length=3, window_size=2, **m), + lambda m: SeasonalRollingMax(season_length=3, window_size=2, **m), + lambda m: RollingQuantile(p=0.5, window_size=4, **m), + lambda m: ExpandingQuantile(p=0.5, **m), + lambda m: SeasonalRollingQuantile(p=0.5, season_length=3, window_size=2, **m), ], ids=[ "RollingMean", @@ -2222,6 +2272,13 @@ def test_pooled_transforms_lag2_groupby(engine): "ExpandingMin", "ExpandingMax", "EWM", + "SeasonalRollingMean", + "SeasonalRollingStd", + "SeasonalRollingMin", + "SeasonalRollingMax", + "RollingQuantile", + "ExpandingQuantile", + "SeasonalRollingQuantile", ], ) @pytest.mark.parametrize("lag", _LAGS) @@ -2371,6 +2428,13 @@ def test_fast_vs_slow_equivalence(tfm_factory, lag): lambda m: ExpandingMin(**m), lambda m: ExpandingMax(**m), lambda m: ExponentiallyWeightedMean(alpha=0.3, **m), + lambda m: SeasonalRollingMean(season_length=3, window_size=2, **m), + lambda m: SeasonalRollingStd(season_length=3, window_size=2, **m), + lambda m: SeasonalRollingMin(season_length=3, window_size=2, **m), + lambda m: SeasonalRollingMax(season_length=3, window_size=2, **m), + lambda m: RollingQuantile(p=0.5, window_size=4, **m), + lambda m: ExpandingQuantile(p=0.5, **m), + lambda m: SeasonalRollingQuantile(p=0.5, season_length=3, window_size=2, **m), ], ids=[ "RollingMean", @@ -2382,6 +2446,13 @@ def test_fast_vs_slow_equivalence(tfm_factory, lag): "ExpandingMin", "ExpandingMax", "EWM", + "SeasonalRollingMean", + "SeasonalRollingStd", + "SeasonalRollingMin", + "SeasonalRollingMax", + "RollingQuantile", + "ExpandingQuantile", + "SeasonalRollingQuantile", ], ) @pytest.mark.parametrize("lag", _LAGS) @@ -2505,9 +2576,20 @@ def test_fast_vs_slow_partition(tfm_factory, lag): ) +@pytest.mark.parametrize( + "tfm_factory", + [ + lambda m: RollingMean(window_size=2, min_samples=1, **m), + lambda m: SeasonalRollingMean( + season_length=2, window_size=2, min_samples=1, **m + ), + lambda m: RollingQuantile(p=0.5, window_size=2, min_samples=1, **m), + ], + ids=["RollingMean", "SeasonalRollingMean", "RollingQuantile"], +) @pytest.mark.parametrize("engine", ["pandas", "polars"]) @pytest.mark.parametrize("lag", _LAGS) -def test_fast_vs_slow_local_partition_with_nan(engine, lag): +def test_fast_vs_slow_local_partition_with_nan(engine, lag, tfm_factory): """Local partition_by with a missing partition value: the slow-path join (forced by clearing the aggregate cache and the idsorted permutation) keys on (id, time), so missing-partition rows are matched, not dropped — matching the @@ -2534,7 +2616,7 @@ def test_fast_vs_slow_local_partition_with_nan(engine, lag): }, ) - tfm = RollingMean(window_size=2, min_samples=1, partition_by=["promo"]) + tfm = tfm_factory({"partition_by": ["promo"]}) col = tfm._get_name(lag) ts = TimeSeries(freq=1, lag_transforms={lag: [tfm]}) fast = np.asarray( @@ -2550,9 +2632,7 @@ def test_fast_vs_slow_local_partition_with_nan(engine, lag): ts_slow = TimeSeries( freq=1, - lag_transforms={ - lag: [RollingMean(window_size=2, min_samples=1, partition_by=["promo"])] - }, + lag_transforms={lag: [tfm_factory({"partition_by": ["promo"]})]}, ) ts_slow._fit( df, @@ -2578,6 +2658,412 @@ def test_fast_vs_slow_local_partition_with_nan(engine, lag): assert not np.all(np.isnan(slow)) +def _force_slow_path(monkeypatch, *tfm_classes): + """Route the given transform classes through the row-level slow path by + restoring the base-class (None-returning) aggregate hooks.""" + from mlforecast.lag_transforms import _BaseLagTransform + + for cls in tfm_classes: + for impl in ( + "_ts_level_from_aggs_impl", + "_latest_from_aggs_impl", + "_bucket_feature_from_aggs_impl", + ): + if impl in cls.__dict__: + monkeypatch.setattr(cls, impl, getattr(_BaseLagTransform, impl)) + + +_PREDICT_EQUIV_MODES = [ + ("global", {"global_": True}), + ("groupby", {"groupby": ["grp"]}), + ("global+partition", {"global_": True, "partition_by": ["promo"]}), + ("local+partition", {"partition_by": ["promo"]}), +] + + +@pytest.mark.parametrize( + "tfm_factory,tfm_cls", + [ + ( + lambda m: SeasonalRollingMean( + season_length=3, window_size=3, min_samples=1, **m + ), + SeasonalRollingMean, + ), + ( + lambda m: SeasonalRollingStd( + season_length=3, window_size=3, min_samples=2, **m + ), + SeasonalRollingStd, + ), + ( + lambda m: SeasonalRollingMin( + season_length=3, window_size=3, min_samples=1, **m + ), + SeasonalRollingMin, + ), + ( + lambda m: SeasonalRollingMax( + season_length=3, window_size=3, min_samples=1, **m + ), + SeasonalRollingMax, + ), + ( + lambda m: RollingQuantile(p=0.5, window_size=6, min_samples=1, **m), + RollingQuantile, + ), + ( + lambda m: ExpandingQuantile(p=0.5, **m), + ExpandingQuantile, + ), + ( + lambda m: SeasonalRollingQuantile( + p=0.5, season_length=3, window_size=3, min_samples=1, **m + ), + SeasonalRollingQuantile, + ), + ], + ids=[ + "SeasonalRollingMean", + "SeasonalRollingStd", + "SeasonalRollingMin", + "SeasonalRollingMax", + "RollingQuantile", + "ExpandingQuantile", + "SeasonalRollingQuantile", + ], +) +@pytest.mark.parametrize( + "mode_kwargs", [m[1] for m in _PREDICT_EQUIV_MODES], ids=[m[0] for m in _PREDICT_EQUIV_MODES] +) +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_fast_predict_matches_forced_slow(engine, mode_kwargs, tfm_factory, tfm_cls, monkeypatch): + """Multi-step recursive predict through the aggregate fast path matches the + same model predicted with the class's fast-path hooks disabled (row-level + slow path), including the per-step append_predictions aggregate extension. + + ``season_length`` matches the promo period (both 3, with the transform on + lag 3) so every seasonal target falls in the queried partition bucket's own + residue class: partition buckets keep calendar holes (value-space matching + is exercised) while windows stay populated, keeping features finite for the + sklearn model at every recursive step.""" + from sklearn.linear_model import LinearRegression + from mlforecast.forecast import MLForecast + + rng = np.random.default_rng(8) + n_series, n_times, h = 4, 24, 6 + ids = np.repeat([f"s{i}" for i in range(n_series)], n_times) + times = np.tile(range(n_times), n_series) + y = rng.standard_normal(n_series * n_times) + 5.0 + data = {"unique_id": ids.tolist(), "ds": times.tolist(), "y": y.tolist()} + static_features = [] + if "groupby" in mode_kwargs: + data["grp"] = np.repeat(["A"] * 2 + ["B"] * 2, n_times).tolist() + static_features = ["grp"] + needs_xdf = "partition_by" in mode_kwargs + if needs_xdf: + data["promo"] = np.tile((np.arange(n_times) % 3 == 0).astype(float), n_series).tolist() + df = _make_df(engine, data) + X_df = None + if needs_xdf: + fut_ds = np.tile(np.arange(n_times, n_times + h), n_series) + X_df = _make_df( + engine, + { + "unique_id": np.repeat([f"s{i}" for i in range(n_series)], h).tolist(), + "ds": fut_ds.tolist(), + "promo": (fut_ds % 3 == 0).astype(float).tolist(), + }, + ) + + def fit_predict(): + fcst = MLForecast( + models=[LinearRegression()], + freq=1, + lags=[1], + lag_transforms={3: [tfm_factory(dict(mode_kwargs))]}, + ) + fcst.fit(df, static_features=static_features) + preds = fcst.predict(h, X_df=X_df) + return np.asarray(preds["LinearRegression"]) + + preds_fast = fit_predict() + assert np.isfinite(preds_fast).all() + _force_slow_path(monkeypatch, tfm_cls) + preds_slow = fit_predict() + np.testing.assert_allclose(preds_fast, preds_slow, rtol=1e-9, atol=1e-9) + + +def _seasonal_fit_transform(engine, tfm, lag, y): + """fit_transform a single-series global_ pooled seasonal transform and return + its feature column (the aggregate fast path).""" + n = len(y) + df = _make_df( + engine, + {"unique_id": ["s0"] * n, "ds": list(range(n)), "y": [float(v) for v in y]}, + ) + col = tfm._get_name(lag) + ts = TimeSeries(freq=1, lag_transforms={lag: [tfm]}) + out = ts.fit_transform( + df, + id_col="unique_id", + time_col="ds", + target_col="y", + dropna=False, + static_features=[], + ) + return np.asarray(out[col]) + + +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_seasonal_mean_stable_summation(engine): + """Regression: the seasonal fast path strides its steps newest-to-oldest, so + naive accumulation drops a small residual added after a large partial. The + window for the last row is ordinals {2, 1, 0} = [1, -1e16, 1e16], whose mean + is exactly 1/3 -- a plain running sum collapses it to 0.""" + tfm = SeasonalRollingMean( + season_length=1, window_size=3, min_samples=1, global_=True + ) + vals = _seasonal_fit_transform(engine, tfm, 1, [1e16, -1e16, 1.0, 0.0]) + np.testing.assert_allclose(vals[-1], 1.0 / 3.0, rtol=1e-12) + + +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_seasonal_std_stable_moments(engine, monkeypatch): + """Regression: for large-magnitude targets with small spread the textbook + ``sum_sq - sum**2 / n`` variance cancels catastrophically and collapses to 0. + The stable moment combine must match the row-level slow path.""" + y = (1e8 + np.arange(12) % 3).astype(float).tolist() + fast = _seasonal_fit_transform( + engine, + SeasonalRollingStd(season_length=2, window_size=4, min_samples=2, global_=True), + 1, + y, + ) + slow_tfm = SeasonalRollingStd( + season_length=2, window_size=4, min_samples=2, global_=True + ) + _force_slow_path(monkeypatch, SeasonalRollingStd) + slow = _seasonal_fit_transform(engine, slow_tfm, 1, y) + np.testing.assert_allclose(fast, slow, rtol=1e-9, atol=1e-9, equal_nan=True) + # the fixed fast path recovers a real (non-zero) std where cancellation used + # to zero it out + finite = fast[~np.isnan(fast)] + assert finite.size and np.all(finite > 0.5) + + +def _pooled_predict_fixture(engine, mode_kwargs, h, n_series=4, n_times=24, seed=8): + """Build (df, X_df, static_features) for a pooled recursive-predict case in + any of the ``_PREDICT_EQUIV_MODES``. Buckets keep enough history that the + expanding/EWM statistics stay finite from the first predicted step.""" + rng = np.random.default_rng(seed) + ids = np.repeat([f"s{i}" for i in range(n_series)], n_times) + times = np.tile(range(n_times), n_series) + y = rng.standard_normal(n_series * n_times) + 5.0 + data = {"unique_id": ids.tolist(), "ds": times.tolist(), "y": y.tolist()} + static_features: list = [] + if "groupby" in mode_kwargs: + half = n_series // 2 + data["grp"] = np.repeat(["A"] * half + ["B"] * (n_series - half), n_times).tolist() + static_features = ["grp"] + X_df = None + if "partition_by" in mode_kwargs: + data["promo"] = np.tile((np.arange(n_times) % 3 == 0).astype(float), n_series).tolist() + df = _make_df(engine, data) + if "partition_by" in mode_kwargs: + fut_ds = np.tile(np.arange(n_times, n_times + h), n_series) + X_df = _make_df( + engine, + { + "unique_id": np.repeat([f"s{i}" for i in range(n_series)], h).tolist(), + "ds": fut_ds.tolist(), + "promo": (fut_ds % 3 == 0).astype(float).tolist(), + }, + ) + return df, X_df, static_features + + +_EXPANDING_EWM_FACTORIES = [ + (lambda m: ExpandingMean(**m), ExpandingMean), + (lambda m: ExpandingStd(**m), ExpandingStd), + (lambda m: ExpandingMin(**m), ExpandingMin), + (lambda m: ExpandingMax(**m), ExpandingMax), + (lambda m: ExponentiallyWeightedMean(alpha=0.5, **m), ExponentiallyWeightedMean), +] +_EXPANDING_EWM_IDS = [ + "ExpandingMean", + "ExpandingStd", + "ExpandingMin", + "ExpandingMax", + "EWM", +] + + +@pytest.mark.parametrize( + "tfm_factory,tfm_cls", _EXPANDING_EWM_FACTORIES, ids=_EXPANDING_EWM_IDS +) +@pytest.mark.parametrize( + "mode_kwargs", + [m[1] for m in _PREDICT_EQUIV_MODES], + ids=[m[0] for m in _PREDICT_EQUIV_MODES], +) +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_expanding_ewm_fast_predict_matches_forced_slow( + engine, mode_kwargs, tfm_factory, tfm_cls, monkeypatch +): + """The prefix-dependent transforms carry a per-model running accumulator + across recursive predict steps (commit 2). The cached fast path must match + the same model predicted through the row-level slow path (hooks disabled), + including the per-step ``append_predictions`` extension.""" + from sklearn.linear_model import LinearRegression + from mlforecast.forecast import MLForecast + + h = 6 + df, X_df, static_features = _pooled_predict_fixture(engine, mode_kwargs, h) + + def fit_predict(): + fcst = MLForecast( + models=[LinearRegression()], + freq=1, + lags=[1], + lag_transforms={1: [tfm_factory(dict(mode_kwargs))]}, + ) + fcst.fit(df, static_features=static_features) + return np.asarray(fcst.predict(h, X_df=X_df)["LinearRegression"]) + + preds_fast = fit_predict() + assert np.isfinite(preds_fast).all() + _force_slow_path(monkeypatch, tfm_cls) + preds_slow = fit_predict() + np.testing.assert_allclose(preds_fast, preds_slow, rtol=1e-9, atol=1e-9) + + +@pytest.mark.parametrize( + "mode_kwargs", + [m[1] for m in _PREDICT_EQUIV_MODES], + ids=[m[0] for m in _PREDICT_EQUIV_MODES], +) +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_latest_cache_matches_no_cache(engine, mode_kwargs, monkeypatch): + """The per-model predict cache produces the same forecasts as recomputing + the full aggregate prefix at every step. Forcing ``cache=None`` into each + prefix-dependent ``_latest_from_aggs_impl`` exercises the recompute branch.""" + from sklearn.linear_model import LinearRegression + from mlforecast.forecast import MLForecast + + h = 8 + df, X_df, static_features = _pooled_predict_fixture(engine, mode_kwargs, h) + + def fit_predict(): + tfms = [ + ExpandingMean(**mode_kwargs), + ExpandingStd(**mode_kwargs), + ExpandingMin(**mode_kwargs), + ExpandingMax(**mode_kwargs), + ExponentiallyWeightedMean(alpha=0.4, **mode_kwargs), + ] + fcst = MLForecast( + models=[LinearRegression()], + freq=1, + lags=[1], + lag_transforms={1: tfms}, + ) + fcst.fit(df, static_features=list(static_features)) + return np.asarray(fcst.predict(h, X_df=X_df)["LinearRegression"]) + + preds_cached = fit_predict() + for cls in ( + ExpandingMean, + ExpandingStd, + ExpandingMin, + ExpandingMax, + ExponentiallyWeightedMean, + ): + orig = cls._latest_from_aggs_impl + + def no_cache(self, ts_aggs, target_ords, _cache=None, _orig=orig): + return _orig(self, ts_aggs, target_ords, None) + + monkeypatch.setattr(cls, "_latest_from_aggs_impl", no_cache) + preds_nocache = fit_predict() + np.testing.assert_allclose(preds_cached, preds_nocache, rtol=1e-9, atol=1e-9) + + +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_multi_model_latest_cache_isolation(engine): + """The predict cache is created and dropped inside ``_backup``, so it is + scoped to one model's recursive walk. Two models that predict different + values (hence append different aggregates) must forecast identically whether + run together or alone -- a leaked prefix would corrupt the second model.""" + from sklearn.linear_model import LinearRegression, Ridge + from mlforecast.forecast import MLForecast + + h = 8 + df, X_df, static_features = _pooled_predict_fixture(engine, {"global_": True}, h) + + def make_tfms(): + return { + 1: [ + ExpandingMean(global_=True), + ExponentiallyWeightedMean(alpha=0.5, global_=True), + RollingMean(window_size=4, min_samples=1, global_=True), + ] + } + + fcst = MLForecast( + models={"lr": LinearRegression(), "ridge": Ridge(alpha=1.0)}, + freq=1, + lags=[1], + lag_transforms=make_tfms(), + ) + fcst.fit(df, static_features=static_features) + both = fcst.predict(h, X_df=X_df) + + for name, model in [("lr", LinearRegression()), ("ridge", Ridge(alpha=1.0))]: + solo = MLForecast( + models={name: model}, freq=1, lags=[1], lag_transforms=make_tfms() + ) + solo.fit(df, static_features=static_features) + solo_preds = np.asarray(solo.predict(h, X_df=X_df)[name]) + np.testing.assert_allclose( + np.asarray(both[name]), solo_preds, rtol=1e-12, atol=1e-12 + ) + + +@pytest.mark.parametrize( + "tfm_factory,tfm_cls", _EXPANDING_EWM_FACTORIES, ids=_EXPANDING_EWM_IDS +) +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_expanding_ewm_multi_horizon_matches_forced_slow( + engine, tfm_factory, tfm_cls, monkeypatch +): + """Direct (``max_horizon``) forecasting drives the prefix caches through the + ``_predict_multi`` path, which advances the target ordinal each step without + feeding predictions back. The cached fast path must still match the same + model predicted through the row-level slow path.""" + from sklearn.linear_model import LinearRegression + from mlforecast.forecast import MLForecast + + h = 6 + df, _, static_features = _pooled_predict_fixture(engine, {"global_": True}, h) + + def fit_predict(): + fcst = MLForecast( + models=[LinearRegression()], + freq=1, + lags=[1], + lag_transforms={1: [tfm_factory({"global_": True})]}, + ) + fcst.fit(df, static_features=static_features, max_horizon=h) + return np.asarray(fcst.predict(h)["LinearRegression"]) + + preds_fast = fit_predict() + assert np.isfinite(preds_fast).all() + _force_slow_path(monkeypatch, tfm_cls) + preds_slow = fit_predict() + np.testing.assert_allclose(preds_fast, preds_slow, rtol=1e-9, atol=1e-9) + + @pytest.mark.parametrize("engine", ["pandas", "polars"]) def test_partition_predict_x_df_partition_column_has_nan(engine): """predict with an X_df whose partition column is NaN routes to the existing @@ -3526,9 +4012,11 @@ def _range_quantile_oracle( @pytest.mark.parametrize("engine", ["pandas", "polars"]) @pytest.mark.parametrize("mode", ["local", "global"]) def test_slow_path_quantile_with_partition_by(engine, mode): - """RollingQuantile has no aggregate fast path, so partition_by routes it - through the row-level slow path at fit and through build_query_arrays at - predict. Pin both against a RANGE-window oracle.""" + """RollingQuantile with partition_by, pinned against a RANGE-window oracle at + fit and across recursive predict steps. It now runs the aggregate fast path + (per-ordinal value store); ``test_slow_path_quantile_with_partition_by_forced`` + below monkeypatches the hooks off to keep the row/build_query_arrays slow path + covered too.""" from mlforecast.lag_transforms import RollingQuantile n_series, n_times = 3, 10 @@ -3621,6 +4109,179 @@ def test_slow_path_quantile_with_partition_by(engine, mode): ) +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +@pytest.mark.parametrize("mode", ["local", "global"]) +def test_slow_path_quantile_with_partition_by_forced(engine, mode, monkeypatch): + """Forced-slow twin of the test above: with RollingQuantile's aggregate + hooks disabled the row-level fit path must still match the RANGE-window + oracle, keeping the slow path covered now that the fast path is default.""" + from mlforecast.lag_transforms import RollingQuantile + + _force_slow_path(monkeypatch, RollingQuantile) + n_series, n_times = 3, 10 + rng = np.random.default_rng(11) + ids = np.repeat([f"s{i}" for i in range(n_series)], n_times) + times = np.tile(np.arange(n_times), n_series) + y = rng.standard_normal(n_series * n_times) + promos = np.tile([0, 1, 0, 0, 1, 1, 0, 1, 0, 1], n_series) + df = _make_df( + engine, + { + "unique_id": ids.tolist(), + "ds": times.tolist(), + "y": y.tolist(), + "promo": promos.tolist(), + }, + ) + tfm = RollingQuantile( + 0.5, 3, min_samples=1, global_=(mode == "global"), partition_by=["promo"] + ) + col = tfm._get_name(1) + ts = TimeSeries(freq=1, lag_transforms={1: [tfm]}) + res = ts.fit_transform( + df, + id_col="unique_id", + time_col="ds", + target_col="y", + dropna=False, + static_features=[], + ) + hist = list(zip(ids, times, y, promos)) + expected_fit = np.array( + [_range_quantile_oracle(hist, i, t, pr, mode) for i, t, _, pr in hist] + ) + np.testing.assert_allclose( + np.asarray(res[col], dtype=float), expected_fit, atol=1e-12, equal_nan=True + ) + + +def _assert_csr_ok(agg): + """CSR value-store invariants for one bucket's aggregates.""" + assert agg.values is not None and agg.row_offsets is not None + assert agg.row_offsets.dtype == np.intp + assert len(agg.row_offsets) == len(agg.unique_times) + 1 + assert agg.row_offsets[0] == 0 + assert np.all(np.diff(agg.row_offsets) >= 0) # monotone + assert len(agg.values) == int(agg.row_offsets[-1]) + # per-ordinal slice length == non-NaN count for that ordinal + np.testing.assert_array_equal( + np.diff(agg.row_offsets).astype(np.intp), agg.counts.astype(np.intp) + ) + + +def test_csr_value_store_invariants_across_mutations(): + """The CSR store stays consistent (offsets integer & monotone, values length + == last offset, window slice == masked selection) through build, + append_predictions, append_observations and trim_to_last.""" + from mlforecast.pooled import PooledState + + def slice_matches_rows(agg, time_index, y): + for k, o in enumerate(agg.unique_times): + got = np.sort(agg.values[agg.row_offsets[k] : agg.row_offsets[k + 1]]) + want = np.sort(y[(time_index == o) & ~np.isnan(y)]) + np.testing.assert_array_equal(got, want) + + n_series = 4 + ids = np.repeat([f"s{i}" for i in range(n_series)], 5) + times = np.tile(np.arange(5), n_series) + y = np.arange(n_series * 5, dtype=float) + df = pd.DataFrame({"unique_id": ids, "ds": times, "y": y}) + + def fresh_state(): + return PooledState.from_global( + df, + id_col="unique_id", + time_col="ds", + target_col="y", + ga_data_dtype=np.dtype("float64"), + n_series=n_series, + store_values=True, + ) + + # build + state = fresh_state() + assert state._store_values is True + _assert_csr_ok(state._ts_aggs[0]) + slice_matches_rows(state._ts_aggs[0], state.time_index, state.y) + + # append_predictions: one new ordinal whose slice is the step's non-NaN preds + state = fresh_state() + preds = np.array([100.0, np.nan, 102.0, 103.0]) + state.append_predictions(np.array([5]), preds, n_series) + agg = state._ts_aggs[0] + _assert_csr_ok(agg) + last = agg.values[agg.row_offsets[-2] : agg.row_offsets[-1]] + np.testing.assert_array_equal(np.sort(last), np.sort(preds[~np.isnan(preds)])) + slice_matches_rows(agg, state.time_index, state.y) + + # append_observations: rebuild keeps the store + state = fresh_state() + new = pd.DataFrame( + { + "unique_id": [f"s{i}" for i in range(n_series)], + "ds": [5] * n_series, + "y": [7.0, 8.0, 9.0, 10.0], + } + ) + state.append_observations( + new, + id_col="unique_id", + time_col="ds", + target_col="y", + ga_data_dtype=np.dtype("float64"), + ) + _assert_csr_ok(state._ts_aggs[0]) + slice_matches_rows(state._ts_aggs[0], state.time_index, state.y) + + # trim_to_last: suffix-slice keeps the store consistent (fit-time only, so a + # fresh state before any append -- bucket_df is still row-aligned) + state = fresh_state() + state.trim_to_last(3) + _assert_csr_ok(state._ts_aggs[0]) + slice_matches_rows(state._ts_aggs[0], state.time_index, state.y) + + +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +@pytest.mark.parametrize("lag", _LAGS) +def test_pooled_expanding_quantile_matches_numpy_oracle(engine, lag): + """Pooled ExpandingQuantile (global_) fit output equals a numpy + expanding-quantile oracle over the per-timestamp pooled series.""" + n_series, n_times = 4, 9 + rng = np.random.default_rng(5) + ids = np.repeat([f"s{i}" for i in range(n_series)], n_times) + times = np.tile(np.arange(n_times), n_series) + y = rng.standard_normal(n_series * n_times) + df = _make_df( + engine, + {"unique_id": ids.tolist(), "ds": times.tolist(), "y": y.tolist()}, + ) + p = 0.4 + tfm = ExpandingQuantile(p=p, global_=True) + col = tfm._get_name(lag) + ts = TimeSeries(freq=1, lag_transforms={lag: [tfm]}) + res = ts.fit_transform( + df, + id_col="unique_id", + time_col="ds", + target_col="y", + dropna=False, + static_features=[], + ) + # ordinals == ds (dense integer calendar); oracle read from res's own columns + rt = np.asarray(res["ds"], dtype=float) + ry = np.asarray(res["y"], dtype=float) + got = np.asarray(res[col], dtype=float) + exp = np.full(len(rt), np.nan) + for r in range(len(rt)): + upper = rt[r] - lag + if upper < 0: + continue + vals = ry[(rt <= upper) & ~np.isnan(ry)] + if len(vals) > 0: + exp[r] = float(np.quantile(vals, p)) + np.testing.assert_allclose(got, exp, atol=1e-12, equal_nan=True) + + # --------------------------------------------------------------------------- # time_agg: pre-aggregate rows sharing a timestamp within each bucket, then # apply the transform over that per-timestamp series. @@ -3631,10 +4292,7 @@ def test_slow_path_quantile_with_partition_by(engine, mode): from mlforecast.lag_transforms import ( # noqa: E402 Combine, - ExpandingQuantile, Offset, - RollingQuantile, - SeasonalRollingMean, ) from mlforecast.pooled import ( # noqa: E402 _build_ts_aggs, @@ -3798,6 +4456,28 @@ def test_time_agg_mean_differs_from_sum_and_rowpooled(engine): (lambda m: ExpandingMin(**m), "ExpandingMin"), (lambda m: ExpandingMax(**m), "ExpandingMax"), (lambda m: ExponentiallyWeightedMean(alpha=0.3, **m), "EWM"), + ( + lambda m: SeasonalRollingMean(season_length=3, window_size=2, **m), + "SeasonalRollingMean", + ), + ( + lambda m: SeasonalRollingStd(season_length=3, window_size=2, **m), + "SeasonalRollingStd", + ), + ( + lambda m: SeasonalRollingMin(season_length=3, window_size=2, **m), + "SeasonalRollingMin", + ), + ( + lambda m: SeasonalRollingMax(season_length=3, window_size=2, **m), + "SeasonalRollingMax", + ), + (lambda m: RollingQuantile(p=0.5, window_size=4, **m), "RollingQuantile"), + (lambda m: ExpandingQuantile(p=0.5, **m), "ExpandingQuantile"), + ( + lambda m: SeasonalRollingQuantile(p=0.5, season_length=3, window_size=2, **m), + "SeasonalRollingQuantile", + ), ] @@ -3930,8 +4610,10 @@ def test_ewm_time_agg_mean_is_noop(engine): @pytest.mark.parametrize("engine", ["pandas", "polars"]) -def test_time_agg_quantile_slow_path_literal(engine): - """RollingQuantile has no fast path: time_agg goes through row-collapse.""" +def test_time_agg_quantile_literal(engine): + """A time_agg RollingQuantile runs its aggregate fast path over the single + collapsed scalar per timestamp (the reaggregated value store, not a raw + O(rows) store); the hand-computed literals pin that it matches row-collapse.""" y_a = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] y_b = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0] # daily sums [11,22,33,44,55,66]; median over window=3 == middle value @@ -3947,7 +4629,9 @@ def test_time_agg_quantile_slow_path_literal(engine): @pytest.mark.parametrize("engine", ["pandas", "polars"]) -def test_time_agg_seasonal_slow_path_literal(engine): +def test_time_agg_seasonal_literal(engine): + """Hand-computed seasonal rolling mean of daily sums; the values are pinned + so the aggregate fast path must reproduce the row-collapse results exactly.""" y_a = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] y_b = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0] # daily sums [11,22,33,44,55,66]; season_length=2, window=2, lag=1 diff --git a/tests/test_pooled_keep_last_n_trim.py b/tests/test_pooled_keep_last_n_trim.py index d42f19b9..fc5f2af0 100644 --- a/tests/test_pooled_keep_last_n_trim.py +++ b/tests/test_pooled_keep_last_n_trim.py @@ -35,12 +35,16 @@ from mlforecast.lag_transforms import ( Combine, ExpandingMean, + ExpandingQuantile, ExponentiallyWeightedMean, Offset, RollingMax, RollingMean, RollingMin, + RollingQuantile, RollingStd, + SeasonalRollingMean, + SeasonalRollingQuantile, ) ID, TIME, TARGET = "unique_id", "ds", "y" @@ -76,6 +80,11 @@ def _finite_lag_transforms(): RollingStd(4, min_samples=2, global_=True), RollingMin(4, global_=True), RollingMax(4, global_=True), + SeasonalRollingMean(season_length=2, window_size=2, global_=True), + RollingQuantile(p=0.5, window_size=4, min_samples=1, global_=True), + SeasonalRollingQuantile( + p=0.5, season_length=2, window_size=2, min_samples=1, global_=True + ), ] } @@ -214,6 +223,24 @@ def test_g2_1_trimmed_predictions_match_untrimmed(keep_last_n): 1: [RollingMean(3, min_samples=1, groupby=["brand"], partition_by=["promo"])] }, "local+partition": {1: [RollingMean(3, min_samples=1, partition_by=["promo"])]}, + "global-seasonal": { + 1: [SeasonalRollingMean(season_length=2, window_size=2, global_=True)] + }, + "local+partition-seasonal": { + 1: [ + SeasonalRollingMean( + season_length=2, window_size=2, min_samples=1, partition_by=["promo"] + ) + ] + }, + "global-quantile": { + 1: [RollingQuantile(p=0.5, window_size=3, min_samples=1, global_=True)] + }, + "local+partition-quantile": { + 1: [ + RollingQuantile(p=0.5, window_size=3, min_samples=1, partition_by=["promo"]) + ] + }, } @@ -391,3 +418,134 @@ def test_g2_4_offset_and_combine_respect_inner_transform(): } state = _preprocess_states(df, 3, unbounded)[("global", (), ())] assert state.next_time_index_by_bucket[0] == T # not trimmed + + +# --------------------------------------------------------------------------- # +# Per-state retention: with an INFERRED keep_last_n, each finite pooled state +# keeps only its own window instead of the global max. Two partition columns +# with different windows (the on_display/promo example) exercise this. +# --------------------------------------------------------------------------- # +def _mixed_window_panel(T=30): + """Balanced panel with two dynamic partition columns that vary over time.""" + ids, ds, y, col_a, col_b = [], [], [], [], [] + for sid, base in {"a": 1.0, "b": 3.0, "c": 7.0, "d": 11.0}.items(): + for t in range(1, T + 1): + ids.append(sid) + ds.append(t) + y.append(base + 2.0 * t + ((t * 3) % 5)) + col_a.append(t % 2) # small-window partition key + col_b.append((t // 3) % 2) # large-window partition key + return pd.DataFrame({ID: ids, TIME: ds, TARGET: y, "col_a": col_a, "col_b": col_b}) + + +def _mixed_future_X(last_t, h): + rows = [] + for sid in ["a", "b", "c", "d"]: + for t in range(last_t + 1, last_t + 1 + h): + rows.append({ID: sid, TIME: t, "col_a": t % 2, "col_b": (t // 3) % 2}) + return pd.DataFrame(rows) + + +def _mixed_update_rows(t): + return pd.DataFrame( + { + ID: ["a", "b", "c", "d"], + TIME: [t] * 4, + TARGET: [50.0, 60.0, 70.0, 80.0], + "col_a": [t % 2] * 4, + "col_b": [(t // 3) % 2] * 4, + } + ) + + +def _fit_ts(df, keep_last_n, lag_transforms, static_features): + fcst = _build_fcst(lag_transforms) + fcst.preprocess( + df, + id_col=ID, + time_col=TIME, + target_col=TARGET, + keep_last_n=keep_last_n, + static_features=static_features, + dropna=False, + ) + return fcst.ts + + +_MIXED = { + 1: [ + RollingMean(3, min_samples=1, partition_by=["col_a"]), + RollingMean(8, min_samples=1, partition_by=["col_b"]), + ] +} + + +def test_per_state_inferred_trims_to_own_window(): + """Inferred keep_last_n: the small-window state keeps only its own window, + NOT the global max. Before the per-state fix this state was padded to the + global inferred keep_last_n (8).""" + T = 30 + df = _mixed_window_panel(T) + ts = _fit_ts(df, None, _MIXED, static_features=[]) # keep_last_n inferred + + # global inference is the widest window across all transforms (lag 1 + RM 8) + assert ts.keep_last_n == 8 + a_state = ts._pooled_states[("local", (), ("col_a",))] + b_state = ts._pooled_states[("local", (), ("col_b",))] + # each state keeps exactly its own window (= its transforms' update_samples) + assert set(a_state.next_time_index_by_bucket.values()) == {3} + assert set(b_state.next_time_index_by_bucket.values()) == {8} + # the small state keeps strictly less than the global keep_last_n + assert max(a_state.next_time_index_by_bucket.values()) < ts.keep_last_n + + +def test_per_state_inferred_predictions_and_update_match_full(): + """The per-state trim is prediction-neutral through predict + update().""" + T = 30 + h = 5 + df = _mixed_window_panel(T) + X1 = _mixed_future_X(T, h) + upd = _mixed_update_rows(T + 1) + X2 = _mixed_future_X(T + 1, h) + + def run(keep_last_n): + fcst = _build_fcst(_MIXED) + fcst.fit( + df, + id_col=ID, + time_col=TIME, + target_col=TARGET, + static_features=[], + keep_last_n=keep_last_n, + ) + p1 = _sorted_preds(fcst.predict(h=h, X_df=X1)) + fcst.update(upd) + p2 = _sorted_preds(fcst.predict(h=h, X_df=X2)) + return p1, p2 + + inf1, inf2 = run(None) # inferred -> per-state windows (3 and 8) + full1, full2 = run(_NO_TRIM) # full history retained everywhere + np.testing.assert_allclose(inf1, full1, rtol=1e-9, atol=1e-9) + np.testing.assert_allclose(inf2, full2, rtol=1e-9, atol=1e-9) + + +def test_boundary_unbounded_transform_disables_pooled_trim(): + """Explicit non-goal: when an unbounded transform makes keep_last_n + un-inferrable it stays None, so self.ga is not trimmed and finite pooled + states are NOT trimmed either (today's behavior, deliberately preserved).""" + T = 30 + df = _mixed_window_panel(T) + lag_transforms = { + 1: [ + RollingMean(3, min_samples=1, partition_by=["col_a"]), # finite + ExpandingQuantile(p=0.5, global_=True), # update_samples == -1 + ] + } + ts = _fit_ts(df, None, lag_transforms, static_features=[]) + + assert ts.keep_last_n is None # inference failed -> no resolved keep_last_n + # self.ga keeps full per-series history + assert int(np.diff(ts.ga.indptr).min()) == T + # the finite col_a state is NOT trimmed (no resolved keep_last_n) + a_state = ts._pooled_states[("local", (), ("col_a",))] + assert set(a_state.next_time_index_by_bucket.values()) == {T}