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: the mean and std leave their count, sum and Welford state untouched, the min and max compare with np.fmin/np.fmax, and the EWM forward-fills its mean.
• 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 the accumulators read as a value they had taken in. This applies to skipna=False too, since a leading run of NaNs is supported there.
• stats_ keeps its layout for all five, so a transform pickled by an earlier version loads as is.

Without skipna, a NaN that reaches the min, max or EWM through update() is dropped by the next value where transform() keeps it, and the mean and std keep it. That is left unspecified, as it already is for the rolling updates, which recompute from the last window and recover once the NaN leaves it. Making the three propagate needs a second value per group, which changes stats_ and what mlforecast writes into it, so it stays out of this fix.

Tests

test_update_matches_transform and test_update_matches_transform_with_skipna step update() and compare it with transform() read at the same position, for every transform, both lags and both dtypes: the first over the series whose NaNs are a leading run only, with both skipna settings; the second over the series that put NaNs elsewhere, including through update(), with skipna=True. test_update_consistency_covers_every_transform enumerates lag_transforms.__all__ so a new transform can't be added without landing in the comparison. test_correctness is seeded and its float32 std tolerance is absolute, since that error doesn't scale with the value.

🤖 Generated with Claude Code

https://claude.ai/code/session_01MkfPDBgpQgPAfken98CBHE

[codspeed]

Merging this PR will not alter performance

✅ 20 untouched benchmarks

---

_{Comparing fix/skipna-in-stateful-updates (6e0063f) with main (17f84c4)}

[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)]

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)]

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).