fix: honor skipna in the stateful lag transform updates

[José Morales (jmoralez)]

Fixes #142.

skipna was ignored by update() in the five lag transforms that keep their accumulator in Python (ExpandingMean, ExpandingStd, ExpandingMin, ExpandingMax, ExponentiallyWeightedMean). A NaN arriving through the incremental path poisoned the group permanently even with skipna=True, while transform() over the same series was fine, so the features built while predicting recursively silently disagreed with the training ones.

Changes

• The five accumulators skip a NaN when skipna=True.
• A group whose whole lagged history was NaN seeds from the first value it gets instead of staying NaN. The driver never ran the kernel for it and filled its stats row with NaN, which is an empty accumulator rather than a poisoned one. This half applies to skipna=False too, since a leading run of NaNs is supported there.
• Seeding an empty accumulator is not the same as undoing a NaN that went into one, and ExpandingMin, ExpandingMax and ExponentiallyWeightedMean resume from a single value per group, so the two states look the same to them: both hold NaN. Their stats_ now carries the values-seen column ExpandingMean already had, and _skip and the seeding read emptiness off it. np.fmin and np.fmax are gone with it, since _skip already keeps the NaN inputs skipna is about out of the comparison. stats_ for these three is (n_groups, 2) instead of (n_groups,); take() and stack() handle it as they do for the other accumulators. A transform pickled by an earlier version is brought to the new layout when it's loaded (_ExpandingBase.__setstate__), reading a NaN as an empty accumulator since the old layout couldn't tell, so a saved model keeps working without a refit.
• Those three also read each group's last output from the position before its end, which for an empty group belongs to another group, so update() reported that group's statistic where transform() gives NaN; the same read raised an IndexError from them and from ExpandingMean on an array with no elements. _last_of_each gives an empty group NaN.
• ExpandingQuantile.update() with skipna=False sorted an interior NaN into its buffer and returned a quantile of whatever came out; it now returns NaN like the transform.
• The skipna docstrings said NaN "propagates" with skipna=False. That's true for transform() but the rolling updates recompute from the last window and recover, so they now say what the scalers' do: only a leading run of NaNs is supported without skipna.
• test_correctness drew unseeded data and had its float32 std tolerance widened to hide it; it's seeded now, float32 std comparisons use an absolute tolerance, and float64 is held to a relative bound alone (the worst of these transforms is off by 7e-13 of the value).

skipna=False stays a trusted precondition

The review asked whether the updates could match the transform for a NaN outside the leading run as well, since three expanding transforms propagate it and three didn't. For the accumulating transforms that is now the case and it costs nothing: the NaN is already in the state they resume from, so keeping it is a matter of not seeding over it.

The ten windowed transforms are a different trade. Their updates recompute from the last window, so a NaN outside the leading run stops showing once it leaves the window, while the transform reports NaN to the end of the group (of the season, for the seasonal ones). Reproducing that means finding out whether the group's values hold a NaN at all, which is a pass over them on every update: measured over 1000 groups of 1000 values, _rolling_mean_update goes from 0.034 to 0.53 ms and _seasonal_rolling_mean_update from 0.072 to 0.32 ms, O(window) to O(history) per horizon step. That is the work skipna=False exists to skip, and it would be paid by every caller that honours the precondition to catch input that breaks it. So the kernels are unchanged and the docstrings keep saying an interior NaN is unspecified without skipna.

The detection on the fit side is kept off the fast path the same way: _has_value checks the last value of each group first, which settles all of them unless one ends in NaN, so a transform over data without NaNs pays O(n_groups) (the EWM transform goes from 1.795 to 1.871 ms over a million rows, the expanding min from 9.601 to 9.748).

Tests

Three tests over the same helper, which steps update() and compares it with transform() read at the same position, for both lags and dtypes:

• test_update_matches_transform: every transform, both skipna settings, over the series whose NaNs are a leading run only, which is what skipna=False supports.
• test_update_matches_transform_with_skipna: every transform, over the series that put NaNs elsewhere.
• test_accumulating_updates_keep_a_nan_without_skipna: the transforms whose updates accumulate over the whole group, over those same series with skipna=False. 24 of these fail without the values-seen column, which is the reseed the review found.
• test_update_gives_nan_to_an_empty_group: every transform, both skipna settings, with empty groups in every position and an array with no elements.
• test_unpickling_state_from_an_earlier_version: the five accumulators, both skipna settings and dtypes, pickled with the layouts main stores and compared with a fresh fit after loading, on the state and on the next update.

test_update_consistency_covers_every_transform enumerates lag_transforms.__all__ and asserts each transform is either windowed or accumulating, so a new one can't be added without landing in one of the buckets above. The C++ transform stays the single definition of what skipna means.

No kernel changes other than QuantileUpdate. The Python accumulators stay O(1) per update; moving them to the stateless recompute-from-data update kernels would make them O(n) like ExpandingQuantile.

[codspeed]

Merging this PR will not alter performance

✅ 20 untouched benchmarks

---

_{Comparing fix/skipna-in-stateful-updates (8482285) with main (2722d38)}

[Copilot]

🟡 Changes recommended

A critical NaN-state handling issue remains unresolved in lag_transforms.py.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Fixes stateful lag-transform updates so skipna is honored and tests are deterministic.

Changes:

• Updates Python accumulator NaN handling.
• Adds update/transform consistency coverage.
• Stabilizes float32 comparisons and documents fixes.

File summaries

File	Summary
tests/test_lag_transforms.py	Adds comprehensive deterministic regression tests.
python/coreforecast/lag_transforms.py	Updates incremental accumulator behavior; critical issue remains with interior NaN state handling when skipna=False (3 votes).
CHANGELOG.md	Documents the behavioral fixes.

Review details

Suppressed comments (1)

python/coreforecast/lag_transforms.py:566

• This has the same state ambiguity for EWM. Under skipna=False, a valid prefix followed by an interior NaN leaves stats_ NaN, and the next valid update is incorrectly treated as the first observation and recovers to x; the transform contract requires the NaN to keep propagating. Keep a separate per-group empty/seen flag and only apply this seed path to groups whose lagged history was entirely skipped.

        ewm = np.where(np.isnan(self.stats_), x, ewm)

• Files reviewed: 3/3 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Saul (nasaul)]

Fetched the branch, built it, ran the suite (989 passed; the 20 errors are a missing pytest-benchmark plugin, pre-existing), then fuzzed update() against transform() across all seven series in consistency_series × both skipna settings × all 16 transforms.

The fix is correct for what it targets. Under skipna=True every expanding/EWM accumulator now agrees with the transform — zero mismatches. The QuantileUpdate guard in expanding.h correctly mirrors the transform's "NaN from the first one on". The count-column trick in ExpandingMean/ExpandingStd is a genuinely nice way to distinguish "driver skipped this group" from "poisoned accumulator", and it holds up.

The skipna=False carve-out looks avoidable

_ExpandingComp.update and ExponentiallyWeightedMean.update reseed whenever stats_ is NaN. Neither has a count column, so that one NaN state has to mean both "nothing seen yet" (must seed) and "poisoned" (must propagate) — and seeding won. The PR then declares the losing case unspecified in the docstrings and leaves the three (False, ...) rows out of update_consistency_cases.

Two things make that worth fixing rather than documenting:

1. It's a regression. np.minimum(nan, x) and the plain EWM recurrence on main both matched the transform for an interior NaN. And nan_via_update means the NaN arrives through the update path — the exact train/predict skew this PR exists to close, just under the other skipna setting.
2. Three expanding transforms propagate and three don't. ExpandingMean, ExpandingStd and ExpandingQuantile all stay NaN correctly; only min/max/EWM recover. That asymmetry reads as an implementation artifact rather than a decision.

Details and numbers are in the two inline comments. The fix is small — track "has this group seen a value" explicitly instead of overloading NaN, which is the same discriminator ExpandingMean gets for free from its count column. I prototyped it: all expanding mismatches go to zero across all seven series and both skipna settings, and the full suite still passes.

Worth noting the docstring change stays accurate either way — the rolling transforms genuinely cannot propagate, since their updates recompute from the last window and recover once the NaN falls out of it. That carve-out is real; the expanding one does not appear to be.

Smaller notes

• If you take the fix, _skip()'s docstring ("the running comparisons and the EWM read emptiness off their NaN state instead") needs updating.
• test_update_consistency_covers_every_transform asserting against __all__ is a good guard against a future transform silently skipping the property.
• Seeding test_correctness is the right call. Minor: for float64 expanding_std, atol=1e-4 is a touch looser than the old rtol=1e-5 at these magnitudes — the float32 reasoning in the comment is sound, but float64 relaxed slightly along with it.

---

Generated by Claude Code

[Saul (nasaul)]

This reseed is unconditional, so with skipna=False (the default) a NaN that is not part of the leading run makes update() recover where transform() stays NaN. Confirmed on this branch with ExpandingMin, lag=1, float64:

series	update()	transform()
nan_via_update	[1.52, 1.52, nan, 10.321, 2.149, 2.149]	[1.52, 1.52, nan, nan, nan, nan]
long_nan_run	[1.52, nan, nan, nan, 2.149, 2.149]	[1.52, nan, nan, nan, nan, nan]

ExpandingMax diverges identically. On main, np.minimum(nan, x) matched the transform here, so this is a regression — and nan_via_update means the NaN arrives through the update path, which is the same train/predict feature skew this PR is closing under the other skipna setting.

The root cause is that the NaN state has to mean two things at once here — "nothing seen yet" (must seed) and "poisoned" (must propagate) — because _ExpandingComp has no count column to tell them apart, unlike ExpandingMean/ExpandingStd. Tracking that explicitly resolves it:

def _any_value(out: np.ndarray, indptr: np.ndarray) -> np.ndarray:
    """Whether each group's transform produced any non-NaN value."""
    c = np.append(0, np.cumsum(~np.isnan(out)))
    return c[indptr[1:]] - c[indptr[:-1]] > 0

Set self.seen_ = _any_value(out, ga.indptr) in transform, then:

self.stats_ = np.where(~self.seen_, x, comp_fn(self.stats_, x))
self.seen_ |= ~np.isnan(x)

and slice seen_ in take(). I prototyped exactly this (same change for the EWM below): every expanding transform then agrees with transform() across all seven series in consistency_series under both skipna settings — zero mismatches — and the suite still passes. The three (False, ...) cases currently excluded from update_consistency_cases can move in rather than being carved out.

---

Generated by Claude Code

[Saul (nasaul)]

Same unconditional reseed as _ExpandingComp.update above, with the same consequence under skipna=False. Confirmed on this branch, lag=1, float64:

series	update()	transform()
nan_via_update	[6.454, 7.4, nan, 10.321, 6.235, 7.263]	[6.454, 7.4, nan, nan, nan, nan]
long_nan_run	[6.454, nan, nan, nan, 2.149, 5.22]	[6.454, nan, nan, nan, nan, nan]

main propagated the NaN and matched. The seen_ fix from the other comment applies unchanged:

ewm = np.where(~self.seen_, x, ewm)
if self.skipna:
    ewm = np.where(np.isnan(x), self.stats_, ewm)
self.stats_ = ewm
self.seen_ |= ~np.isnan(x)

with self.seen_ = _any_value(out, ga.indptr) in transform and out.seen_ = self.seen_[idxs].copy() in take.

---

Generated by Claude Code

[Saul (nasaul)]

The seen-flag + accumulator state model is the right fix for #142, and keeping the kernel as the single definition of skipna is the correct call. Three things before merge, all in _keep_last / ExpandingMean.transform.

Verified as correct: the _has_value fast path and the reduceat slow path (including interleaved and trailing empty groups); the empty-vs-poisoned accumulator distinction against FirstNotNaN in the driver, including the skipna=True, valid_count == 0 case that would have poisoned ExpandingMean via 0 * NaN (unreachable for the same reason); ExpandingStd.update's Welford rewrite against both kernel paths including the n > 1 gate; fmin/fmax -> minimum/maximum under both skipna modes; take()/stack() 2-D handling; the C++ QuantileUpdate guard against the sticky has_nan_; view/copy safety and dtype preservation in the rewritten update()s. The accumulating-vs-windowed NaN asymmetry is deliberate and documented.

Test run: the new suite against this branch's Python with the main extension build gives 433 passed / 8 failed, and all 8 are ExpandingQuantile under test_accumulating_updates_keep_a_nan_without_skipna - i.e. exactly the cases the uncompiled QuantileUpdate change fixes. Correctness/quantiles/stack/take: 56 passed under the tightened rtol=1e-7, atol=0.

[Saul (nasaul)]

out[indptr[1:] - 1] reads one slot before the group's end for an empty group (indptr[i] == indptr[i+1], which is allowed - the C++ side only requires indptr non-decreasing, and Reduce documents "an empty remainder gets a NaN row"), so it picks up the previous group's last output.

_has_value was carefully written to exclude empty groups, but last wasn't: the seen flag ends up 0 while column 1 holds a real number. On the next update(), x is NaN (_index_from_end on an empty group), _skip returns True, and np.where(skip, prev, ...) hands back the neighbouring group's value.

Reproduced on this branch with indptr=[0, 4, 4], data [1, 2, 3, 9], lag 1:

min  update = [1.0, 1.0]     # 2nd group should be NaN
max  update = [9.0, 3.0]     # 2nd group should be NaN
ewm  update = [5.625, 2.25]  # 2nd group should be NaN
mean update = [3.75, nan]    # correct
std  update = [..., nan]     # correct

ExpandingMean/ExpandingStd NaN out empty groups via np.isnan(n); min/max/EWM now don't. Suggested fix:

Suggested change

	last = out[indptr[1:] - 1]
	last = np.where(_has_value(out, indptr), out[indptr[1:] - 1], np.nan)

(then reuse that for seen rather than calling _has_value twice.)

Secondly, the same expression raises IndexError: index -1 is out of bounds for axis 0 with size 0 on a zero-length GroupedArray (indptr=[0, 0]). _has_value returns cleanly for that input, so it's only the last indexing that's unguarded here - same one-line fix covers it.

This is arguably pre-existing in spirit (np.fmin(garbage, nan) returned the same garbage on main), but this PR is the one that introduces explicit empty-group handling in _has_value, so it'd be good to finish the job here.

[Saul (nasaul)]

Same zero-length-GroupedArray edge case as in _keep_last: out[ga.indptr[1:] - 1] raises IndexError: index -1 is out of bounds for axis 0 with size 0 when indptr=[0, 0]. The np.isnan(n) masking below handles empty groups correctly, but the indexing happens first.

[Saul (nasaul)]

Two docs/compat points on the new stats_ layout:

1. Breaking for persisted state. stats_ goes from shape (n_groups,) to (n_groups, 2) for ExpandingMin, ExpandingMax and ExponentiallyWeightedMean. Anything unpickling a transform fitted with an older version (mlforecast persists fitted ts objects) will hit IndexError: too many indices inside _seen()/update(), and stack() will hstack the 1-D arrays instead of vstack-ing them. Worth an explicit "breaking for persisted state" note in the CHANGELOG.

2. CHANGELOG has the column order backwards. It says the flag is "a second column", but the flag is column 0 here - the statistic is what moved to column 1. Anyone migrating code that reads stats_ directly will get it the wrong way round.

[Saul (nasaul)]

Approving. Both remaining items are docs-only.

Verified the earlier review threads are all closed, with the original repros:

• indptr=[0, 4, 4], data [1, 2, 3, 9], lag 1 now gives min [1, nan], max [9, nan], ewm [3.516, nan] — the empty group no longer picks up its neighbour's statistic.
• indptr=[0, 0] returns [nan] from all six expanding transforms instead of raising IndexError.
• nan_via_update and long_nan_run with skipna=False now match transform() for ExpandingMin, ExpandingMax and ExponentiallyWeightedMean — the reseed regression is gone.
• The CHANGELOG column order is right (flag in column 0, statistic in column 1), and __setstate__ goes further than the "note it's breaking" I asked for.

Two things to fix before merge:

1. The skipna=False caveat is over-broad — python/coreforecast/lag_transforms.py

"only a leading run of NaNs is supported; an interior one leaves the result unspecified" is on all sixteen transforms. It's only true of update() on the ten windowed ones. For the six expanding transforms an interior NaN is now fully specified in both directions — which is what this PR just achieved and what test_accumulating_updates_keep_a_nan_without_skipna asserts:

x = [5, 2, nan, 3, 9, 1], lag=1, skipna=False, transform():
  Mean  [nan 5.  3.5   nan nan nan]     Min   [nan 5.  2.  nan nan nan]
  Std   [nan nan 2.121 nan nan nan]     Max   [nan 5.  5.  nan nan nan]
  Quant [nan 5.  3.5   nan nan nan]     EWM   [nan 5.  3.5 nan nan nan]

It's also over-broad for transform() on the windowed ones, where propagation to the end of the group is specified too. As written, someone reads "unspecified" on ExpandingMean and switches to skipna=True, silently getting a different statistic. Suggest scoping the sentence to update() on _RollingBase / _SeasonalRollingBase and dropping it from the six expanding classes.

2. The stats_ shape change belongs under Breaking changes — CHANGELOG.md

ExpandingMin, ExpandingMax and ExponentiallyWeightedMean going from (n_groups,) to (n_groups, 2), with the statistic moving to column 1, is only under ### Bug fixes. stats_ is a public attribute, and the __setstate__ migration is forward-only — a pickle written by this version can't be loaded by an older coreforecast. It belongs in ### Breaking changes next to the int64 indptr entry.

While there, worth stating the migration's one lossy case: an old skipna=False state genuinely poisoned by an interior NaN is indistinguishable from an empty one in the old layout, so it loads as empty and reseeds. Unavoidable, but it means a loaded model changes answers on those groups.

Not blocking, your call: skipna=False is the default, so the path the docstrings now declare unsupported is the default path, and the windowed update() / transform() divergence is silent — the same shape of train/predict skew as #142. The PR's cost analysis rejects an O(history) scan per update, which is right, but there's a cheaper option it doesn't weigh: a sticky per-group (per-phase for seasonal) poisoned flag, derived from the transform output at fit time where the scan is already paid, then OR'd with isnan(x) on each update. That's O(n_groups) per update. The cost is new Python state on ten classes that currently hold none, plus take / stack / __setstate__ for it. Documenting the precondition is a defensible answer; just noting the option exists.

---

Verification:

• pytest tests/ on 8482285: 1135 passed. C++ doctest runner: 27 cases, 3768 assertions, all pass.
• Independent randomized oracle (not this PR's tests): 40 trials x 11 transforms x lag {1,2} x skipna {T,F} x {f32,f64}, with leading runs, interior NaNs, NaNs arriving through update(), and empty groups in every position — 21,760 update-vs-transform comparisons, 0 mismatches.
• take / stack round-trip and pickle / deepcopy identity checked for all five accumulators, float32 dtype preserved through __setstate__.
• The new test_correctness tolerances aren't brittle: float64 still passes at rtol=1e-11 against the 1e-7 set here, float32 std at atol=1e-4 against 1e-3. Deterministic across repeated runs.
• Fit-side cost claim holds: over 1M rows, EWM transform 1.66 ms and ExpandingMin 6.21 ms, unchanged whether _has_value takes the fast path or the reduceat path.
• Merges cleanly onto main (2722d38).