fix(cross-fitting): keep missing indicators of a wrapped Preprocessor

docs/core_concepts/target_awareness.md

@@ -73,9 +73,9 @@ held-out fold, so every training row is transformed by a model that never saw it
 the training data, since there is no leakage risk there.
 
 A wrapped `Preprocessor` is not refit from scratch on each fold: every fold reuses its all-data
-fit and refits only the target-aware representations. Column types, category codes and the
-blocks that do not use `y` are shared, so the out-of-fold features have exactly the columns and
-encoding `transform` produces.
+fit and refits only the target-aware representations. Column types, category codes, imputers
+with their missing-value indicators, scalers and the blocks that do not use `y` are shared, so
+the out-of-fold features have exactly the columns and encoding `transform` produces.
 
 ```python
 import numpy as np

pretab/compose/inspection.py

@@ -7,8 +7,12 @@
 :func:`build_transformer_summary` renders that metadata as an aligned table.
 """
 
+import copy
+from typing import Any, cast
+
 import numpy as np
-from sklearn.pipeline import FeatureUnion
+from sklearn.base import clone
+from sklearn.pipeline import FeatureUnion, Pipeline
 
 from ..core.logging import get_logger
 from ..core.representation import FeatureLineage
@@ -23,6 +27,7 @@
     "build_transformer_summary",
     "clean_feature_names",
     "get_output_slices",
+    "refit_block_representation",
 ]
 
 
@@ -255,32 +260,40 @@ def build_transformer_summary(numerical_info, categorical_info, embedding_info):
 }
 
 
-def _resolve_block_representation(pipeline, columns):
-    """Return ``(family, component, uses_target, is_interaction)`` for a block.
+def _locate_block_representation(pipeline, columns):
+    """Return ``(step_index, (family, component, uses_target, is_interaction))`` for a block.
 
     The representation-bearing step is the last pipeline step exposing a
     ``get_representation_spec`` (a PreTab transformer), a known scikit-learn
     step name, or a registered method that consumes the target; helper steps
     such as imputers and float casts are skipped. ``columns`` names the spec's
-    inputs; with None the step names its own.
+    inputs; with None the step names its own. ``step_index`` is that step's
+    position in ``pipeline.steps`` (0 for a bare transformer), or None when the
+    block has no representation step.
     """
     steps = pipeline.steps if hasattr(pipeline, "steps") else [("_", pipeline)]
-    for step_name, transformer in reversed(steps):
+    for index in reversed(range(len(steps))):
+        step_name, transformer = steps[index]
         if hasattr(transformer, "get_representation_spec"):
             spec = transformer.get_representation_spec(input_features=None if columns is None else list(columns))
-            return spec.family, spec.component_kind, spec.uses_target, spec.is_interaction
+            return index, (spec.family, spec.component_kind, spec.uses_target, spec.is_interaction)
         if step_name in _STEP_FAMILY:
             family, component = _STEP_FAMILY[step_name]
-            return family, component, False, False
+            return index, (family, component, False, False)
         registered = TRANSFORMER_REGISTRY.get(step_name)
         if registered is not None and registered.target_usage != "forbidden":
             # A registered class without a RepresentationSpec (e.g. scikit-learn's
             # TargetEncoder): its declared supervision says whether it used y, so
             # cross-fitting refits it per fold instead of reusing the all-data fit.
             uses_target = registered.target_usage == "required" or bool(getattr(transformer, "target_aware", False))
             component = "category" if registered.is_categorical else "basis"
-            return step_name, component, uses_target, registered.is_multivariate
-    return "passthrough", "raw", False, False
+            return index, (step_name, component, uses_target, registered.is_multivariate)
+    return None, ("passthrough", "raw", False, False)
+
+
+def _resolve_block_representation(pipeline, columns):
+    """Return ``(family, component, uses_target, is_interaction)`` for a block."""
+    return _locate_block_representation(pipeline, columns)[1]
 
 
 def block_uses_target(transformer) -> bool:
@@ -297,6 +310,35 @@ def block_uses_target(transformer) -> bool:
     return bool(_resolve_block_representation(representation, None)[2])
 
 
+def refit_block_representation(transformer, X, y):
+    """Return a copy of a fitted per-column block whose representation is refit on ``(X, y)``.
+
+    Only the representation step and the steps after it are refit. The steps
+    before it (the imputer with its missing indicator, scalers) never use the
+    target, so they are kept from the all-data fit and transform ``X`` for the
+    refit: a ``SimpleImputer(add_indicator=True)`` refit on rows without missing
+    values would drop its indicator column and so change the block's width. For
+    a separate-state union the missing branch is kept as well. The fitted
+    ``transformer`` itself is left untouched.
+    """
+    separate_state = _separate_state_branches(transformer)
+    if separate_state is not None:
+        representation = refit_block_representation(separate_state[0], X, y)
+        union = copy.copy(transformer)
+        union.transformer_list = [
+            (name, representation if name == "representation" else branch)
+            for name, branch in transformer.transformer_list
+        ]
+        return union
+    index = _locate_block_representation(transformer, None)[0]
+    if not isinstance(transformer, Pipeline) or not index:
+        return cast(Any, clone(transformer)).fit(X, y)
+    refit = cast(Any, clone(transformer[index:])).fit(transformer[:index].transform(X), y)
+    block = copy.copy(transformer)
+    block.steps = [*transformer.steps[:index], *refit.steps]
+    return block
+
+
 def _passthrough_source(columns, offset, feature_names_in):
     """Resolve the source feature name for a passthrough / remainder column."""
     column = columns[offset] if offset < len(columns) else columns[-1]

pretab/core/supervised.py

@@ -99,11 +99,12 @@ class CrossFittedTransformer(RepresentationSpecMixin, TransformerMixin, BaseEsti
     ----------
     transformer : estimator
         A supervised (target-aware) PreTab transformer to cross-fit. For a
-        :class:`~pretab.Preprocessor`, only its target-aware blocks are refit per
-        fold: column types, category codes and the blocks that do not use ``y``
-        come from the all-data fit, so the out-of-fold features use exactly the
-        encoding of :meth:`transform`. Any other estimator is cloned and refit
-        from scratch on each fold.
+        :class:`~pretab.Preprocessor`, only its target-aware representations are
+        refit per fold: column types, category codes, imputers with their
+        missing-value indicators, scalers and the blocks that do not use ``y``
+        come from the all-data fit, so the out-of-fold features have exactly the
+        columns and encoding of :meth:`transform`. Any other estimator is cloned
+        and refit from scratch on each fold.
     n_folds : int, default=5
         Number of cross-fitting folds. Must be at least 2.
     task : {"regression", "classification"}, default="regression"
@@ -205,8 +206,9 @@ def fit_transform(self, X, y=None):
                     raise IncompatibleParamsError(
                         "Cross-fitting requires a fixed output width across folds; expected "
                         f"{width}, got {fold_out.shape[1]}. The wrapped transformer's width "
-                        "depends on the training rows (e.g. adaptive sizing); fix it, for "
-                        "example by disabling adaptive sizing."
+                        "depends on the rows it is fit on (e.g. adaptive sizing, or category / "
+                        "indicator columns learned from the data); configure it to produce the "
+                        "same columns on every fold."
                     )
                 out[test_idx] = fold_out
         finally:

pretab/preprocessor.py

@@ -4,7 +4,7 @@
 import os
 import time
 import warnings
-from typing import Any, cast
+from typing import cast
 
 import numpy as np
 from scipy import sparse as sp
@@ -23,6 +23,7 @@
     build_transformer_summary,
     clean_feature_names,
     get_output_slices,
+    refit_block_representation,
 )
 from .compose.output import (
     compute_output_report,
@@ -875,16 +876,19 @@ def estimate_memory(self, X) -> int:
         return int(n_rows * n_cols * self._output_itemsize())
 
     def _cross_fit_fold(self, X, y):
-        """Return a copy of this fit whose target-aware blocks are refit on ``(X, y)``.
+        """Return a copy of this fit whose target-aware representations are refit on ``(X, y)``.
 
         Used by :class:`~pretab.CrossFittedTransformer` for its out-of-fold
-        features. Only the blocks that consume the target are refit on the fold;
-        column types, category vocabularies and every other block are shared
-        with this all-data fit, so an out-of-fold row is encoded exactly like
+        features. Only the representations that consume the target are refit on
+        the fold; column types, category vocabularies, imputers with their
+        missing-value indicators, scalers and every other block are shared with
+        this all-data fit, so an out-of-fold row is encoded exactly like
         ``transform`` encodes it, except that no target-aware placement has seen
         that row's target. Refitting the whole preprocessor per fold would
         re-detect column types and relearn categories on fewer rows, shifting
-        integer codes and one-hot columns between the folds and ``transform``.
+        integer codes and one-hot columns between the folds and ``transform``;
+        refitting a block's missing indicator would drop its column on a fold
+        without missing values.
         """
         check_is_fitted(self)
         # Each block sees the fold rows exactly as the internal ColumnTransformer
@@ -894,7 +898,7 @@ def _cross_fit_fold(self, X, y):
         column_transformer.transformers_ = [
             (
                 name,
-                cast(Any, clone(transformer)).fit(_safe_indexing(X, columns, axis=1), y)
+                refit_block_representation(transformer, _safe_indexing(X, columns, axis=1), y)
                 if name != "remainder" and block_uses_target(transformer)
                 else transformer,
                 columns,

tests/core/test_cross_fitted.py

@@ -6,14 +6,14 @@
 """
 
 import warnings
-from typing import cast
+from typing import Any, cast
 
 import numpy as np
 import pytest
 from sklearn.model_selection import KFold
 
 from pretab import CrossFittedTransformer, LeakageWarning
-from pretab.exceptions import IncompatibleParamsError, InvalidParamError
+from pretab.exceptions import IncompatibleParamsError, InvalidParamError, PretabDataError
 from pretab.transformers import PLETransformer
 
 
@@ -146,10 +146,32 @@ def test_wrapped_preprocessor_out_of_fold_codes_match_transform(array_with_rare_
     np.testing.assert_array_equal(np.unique(out[:, unsupervised[1]]), [1.0, 2.0, 3.0, 4.0])
 
 
+def _refit_representation_per_fold(pipeline, column, y):
+    """Expected out-of-fold rows of a fitted block pipeline whose last step is its
+    target-aware representation: the helper steps before it keep their all-data fit."""
+    from sklearn.base import clone
+
+    helpers, representation = pipeline[:-1], pipeline[-1]
+    blocks = {}
+    for train_idx, test_idx in KFold(n_splits=5, shuffle=True, random_state=0).split(column):
+        fold = cast(Any, clone(representation)).fit(helpers.transform(column.iloc[train_idx]), y[train_idx])
+        blocks[tuple(test_idx)] = fold.transform(helpers.transform(column.iloc[test_idx]))
+    out = np.empty((len(column), next(iter(blocks.values())).shape[1]))
+    for test_idx, rows in blocks.items():
+        out[list(test_idx)] = rows
+    return out
+
+
+def _fitted_block(cross_fitted, name):
+    from pretab import Preprocessor
+
+    return cast(Preprocessor, cross_fitted.estimator_).column_transformer_.named_transformers_[name]
+
+
 def test_wrapped_preprocessor_target_aware_blocks_stay_out_of_fold(array_with_rare_category):
-    """Only the target-aware block is refit, on each fold's training rows."""
+    """Only the target-aware representation is refit, on each fold's training rows; the
+    helper steps before it (imputer, scaler) keep their all-data fit."""
     import pandas as pd
-    from sklearn.base import clone
 
     from pretab import Preprocessor
 
@@ -159,12 +181,10 @@ def test_wrapped_preprocessor_target_aware_blocks_stay_out_of_fold(array_with_ra
     supervised = _split_columns(cross_fitted, "num_feature_0_ple")
     assert len(supervised) == 4
 
-    wrapped = cast(Preprocessor, cross_fitted.estimator_)
-    block = wrapped.column_transformer_.named_transformers_["num_feature_0"]
-    column = pd.DataFrame({"feature_0": X[:, 0]})
-    for train_idx, test_idx in KFold(n_splits=5, shuffle=True, random_state=0).split(X):
-        fold_block = clone(block).fit(column.iloc[train_idx], y[train_idx])
-        np.testing.assert_allclose(out[np.ix_(test_idx, supervised)], fold_block.transform(column.iloc[test_idx]))
+    block = _fitted_block(cross_fitted, "num_feature_0")
+    assert [name for name, _ in block.steps] == ["imputer", "minmax", "ple"]
+    expected = _refit_representation_per_fold(block, pd.DataFrame({"feature_0": X[:, 0]}), y)
+    np.testing.assert_array_equal(out[:, supervised], expected)
     assert not np.allclose(out[:, supervised], np.asarray(cross_fitted.transform(X))[:, supervised])
 
 
@@ -260,13 +280,113 @@ def test_wrapped_preprocessor_reads_the_target_use_of_an_indicator_block():
     assert not np.allclose(out, full)  # refit on every fold
 
 
-def test_width_mismatch_message_does_not_blame_adaptive_sizing_only():
+def test_width_mismatch_names_the_general_cause():
     from sklearn.preprocessing import OneHotEncoder
 
     x = np.array(["a"] * 99 + ["b"]).reshape(-1, 1)  # the fold holding out "b" never sees it
     y = np.arange(100.0)
     cross_fitted = CrossFittedTransformer(
         OneHotEncoder(handle_unknown="ignore", sparse_output=False), n_folds=5, random_state=0
     )
-    with pytest.raises(IncompatibleParamsError, match="depends on the training rows"):
+    with pytest.raises(IncompatibleParamsError, match="same columns on every fold"):
         cross_fitted.fit_transform(x, y)
+
+
+# --- a wrapped Preprocessor shares its missing indicators across folds (issue #89) --
+
+
+@pytest.fixture
+def array_with_one_missing_value():
+    rng = np.random.default_rng(0)
+    X = rng.normal(size=(200, 2))
+    X[7, 0] = np.nan  # the fold holding out row 7 trains on no missing value
+    y = np.nan_to_num(X[:, 0]) + rng.normal(scale=0.1, size=200)
+    return X, y
+
+
+@pytest.mark.parametrize(
+    ("method", "missing"),
+    [
+        ("ple", {"add_missing_indicator": True}),
+        ("ple", {"missing_policy": "impute_with_indicator"}),
+        ("rbf", {"add_missing_indicator": True}),
+        ("relu", {"add_missing_indicator": True}),
+        ("sigmoid", {"add_missing_indicator": True}),
+        ("tanh", {"add_missing_indicator": True}),
+    ],
+)
+def test_wrapped_preprocessor_keeps_its_missing_indicator_columns(array_with_one_missing_value, method, missing):
+    """Each fold refit the block's imputer, whose missing indicator emits no column when
+    the fold's training rows contain no missing value, so the width check failed."""
+    from pretab import Preprocessor
+
+    X, y = array_with_one_missing_value
+
+    def cross_fit(**missing_config):
+        wrapped = Preprocessor(numerical_method=method, output_dim=4, random_state=0, **missing_config)
+        cross_fitted = CrossFittedTransformer(wrapped, n_folds=5, random_state=0)
+        return cross_fitted, cross_fitted.fit_transform(X, y)
+
+    with_indicator, out = cross_fit(**missing)
+    _, imputed_only = cross_fit(missing_policy="impute")
+    names = [str(name) for name in with_indicator.get_feature_names_out()]
+    indicator = [i for i, name in enumerate(names) if "missingindicator" in name]
+
+    assert out.shape == np.asarray(with_indicator.transform(X)).shape
+    assert indicator  # the indicator columns of the all-data fit are kept on every fold
+    # The value's out-of-fold representation is the one cross-fitting gives without an indicator.
+    np.testing.assert_array_equal(np.delete(out, indicator, axis=1), imputed_only)
+
+
+def test_wrapped_preprocessor_refits_only_the_representation_after_the_imputer(array_with_one_missing_value):
+    """The all-data imputer (with its indicator) and scaler encode every fold; only the
+    PLE is refit, on the imputed value and the indicator of the fold's training rows."""
+    import pandas as pd
+
+    from pretab import Preprocessor
+
+    X, y = array_with_one_missing_value
+    wrapped = Preprocessor(numerical_method="ple", output_dim=4, add_missing_indicator=True, random_state=0)
+    cross_fitted = CrossFittedTransformer(wrapped, n_folds=5, random_state=0)
+    out = cross_fitted.fit_transform(X, y)
+
+    block = _fitted_block(cross_fitted, "num_feature_0")
+    assert [name for name, _ in block.steps] == ["imputer", "minmax", "ple"]
+    columns = _split_columns(cross_fitted, "num_feature_0")
+    assert len(columns) == 8  # the imputer feeds the value and its indicator to the PLE
+    expected = _refit_representation_per_fold(block, pd.DataFrame({"feature_0": X[:, 0]}), y)
+    np.testing.assert_array_equal(out[:, columns], expected)
+
+
+@pytest.mark.xfail(
+    strict=True,
+    raises=PretabDataError,
+    reason="Until #62 keeps the imputer's indicator out of the representation, a spline basis "
+    "is also fitted on the indicator, which is constant on a fold without missing values.",
+)
+def test_wrapped_preprocessor_keeps_the_indicator_next_to_a_target_aware_spline(array_with_one_missing_value):
+    from pretab import Preprocessor
+
+    X, y = array_with_one_missing_value
+    wrapped = Preprocessor(numerical_method="bspline", add_missing_indicator=True, random_state=0)
+    cross_fitted = CrossFittedTransformer(wrapped, n_folds=5, random_state=0)
+    assert cross_fitted.fit_transform(X, y).shape == np.asarray(cross_fitted.transform(X)).shape
+
+
+def test_wrapped_preprocessor_keeps_the_separate_state_missing_branch(array_with_one_missing_value):
+    """A separate-state block refits only its representation branch: the imputer imputes
+    the held-out missing value with the all-data median, as transform does."""
+    import pandas as pd
+
+    from pretab import Preprocessor
+
+    X, y = array_with_one_missing_value
+    wrapped = Preprocessor(numerical_method="ple", output_dim=4, missing_policy="separate_state", random_state=0)
+    cross_fitted = CrossFittedTransformer(wrapped, n_folds=5, random_state=0)
+    out = cross_fitted.fit_transform(X, y)
+
+    representation = dict(_fitted_block(cross_fitted, "num_feature_0").transformer_list)["representation"]
+    expected = _refit_representation_per_fold(representation, pd.DataFrame({"feature_0": X[:, 0]}), y)
+    np.testing.assert_array_equal(out[:, _split_columns(cross_fitted, "num_feature_0__representation")], expected)
+    missing = _split_columns(cross_fitted, "num_feature_0__missing")
+    np.testing.assert_array_equal(out[:, missing].ravel(), np.isnan(X[:, 0]).astype(float))