fix(binning): learn bin edges at fit instead of per transform call

pretab/transformers/binning/binning.py

@@ -3,8 +3,9 @@
 import numpy as np
 import pandas as pd
 from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.utils.validation import check_is_fitted
 
-from ...core.exceptions import InsufficientSamplesError, InvalidParamError, PretabDataError
+from ...core.exceptions import InvalidParamError, PretabDataError
 from ...core.params import UNSET, AliasResolverMixin
 
 
@@ -29,6 +30,11 @@ class CustomBinTransformer(AliasResolverMixin, TransformerMixin, BaseEstimator):
     n_features_in_ : int
         The number of input features seen during ``fit`` (expected to be 1).
 
+    bin_edges_ : list of ndarray
+        Bin edges learned during ``fit``, one array per input column. Reused
+        verbatim by every ``transform`` call so a given value always receives the
+        same code.
+
     total_output_dim_ : int
         Total number of output columns (fitted). Always ``1`` because the output
         is a single ordinal column.
@@ -37,10 +43,13 @@ class CustomBinTransformer(AliasResolverMixin, TransformerMixin, BaseEstimator):
     -----
     This transformer operates on a single feature of shape ``(n_samples, 1)``. When
     ``output_dim`` is an integer, equal-width bin edges are computed from the data
-    range; when it is an array-like, the provided edges are used directly. The
-    output contains integer bin indices in a single column, so its width is ``1``
-    regardless of ``output_dim`` -- this is a documented exception to the
-    exact-width contract that the fixed-basis families follow.
+    range **at fit time**; when it is an array-like, the provided edges are used
+    directly. The output contains integer bin indices in a single column, so its
+    width is ``1`` regardless of ``output_dim`` -- this is a documented exception
+    to the exact-width contract that the fixed-basis families follow.
+
+    Values outside the fitted range are clipped into the first or last bin, so
+    ``transform`` never emits ``NaN`` into what is otherwise an integer column.
 
     The input must be numeric: binning is performed with :func:`pandas.cut`, so
     string / categorical data cannot be processed and raises a
@@ -80,12 +89,47 @@ def fit(self, X, y=None):
         self : object
             Fitted transformer.
         """
-        # Fit doesn't need to do anything as we are directly using provided bins
-        X = np.asarray(X)
+        X = self._as_numeric(np.asarray(X))
         self.n_features_in_ = X.shape[1] if X.ndim > 1 else 1
         self.total_output_dim_ = 1
+
+        bins_spec = self._resolve_param("output_dim", default=UNSET)
+        if bins_spec is UNSET:
+            raise InvalidParamError("CustomBinTransformer requires 'output_dim'.")
+
+        # Edges are learned here and reused verbatim by ``transform``. Deriving
+        # them inside ``transform`` instead made the encoding depend on the batch
+        # being transformed, so the same value got different codes depending on
+        # which other rows travelled with it.
+        self.bin_edges_ = [
+            self._edges_for(X[:, j], bins_spec) for j in range(X.shape[1])
+        ]
         return self
 
+    @staticmethod
+    def _as_numeric(X: np.ndarray) -> np.ndarray:
+        """Return ``X`` as a float array, raising a clear error on non-numeric input."""
+        if np.issubdtype(X.dtype, np.number):
+            return X
+        try:
+            return X.astype(np.float64)
+        except (ValueError, TypeError) as exc:
+            raise PretabDataError(
+                "CustomBinTransformer requires numeric input: it bins continuous "
+                "values with pandas.cut and cannot process string/categorical "
+                "data. Encode string columns with a categorical method (e.g. "
+                "'int' or 'one-hot') before binning."
+            ) from exc
+
+    @staticmethod
+    def _edges_for(column: np.ndarray, bins_spec) -> np.ndarray:
+        """Resolve the bin edges for one column from an int count or explicit edges."""
+        if isinstance(bins_spec, int):
+            _, edges = pd.cut(column, bins=bins_spec, retbins=True)
+        else:
+            edges = bins_spec
+        return np.sort(np.unique(np.asarray(edges)))
+
     def transform(self, X):
         """
         Transform the data using the specified binning strategy.
@@ -101,43 +145,26 @@ def transform(self, X):
             Binned data with integer bin indices.
         """
 
+        check_is_fitted(self, "bin_edges_")
+
         X = np.asarray(X)  # Ensures squeeze works and consistent input
         if X.ndim != 2 or X.shape[1] != 1:
             raise PretabDataError("Input must be a 2D array with shape (n_samples, 1).")
 
-        if X.shape[0] <= 2:
-            raise InsufficientSamplesError("Input must have more than 2 observations.")
-
-        if not np.issubdtype(X.dtype, np.number):
-            try:
-                X = X.astype(np.float64)
-            except (ValueError, TypeError) as exc:
-                raise PretabDataError(
-                    "CustomBinTransformer requires numeric input: it bins continuous "
-                    "values with pandas.cut and cannot process string/categorical "
-                    "data. Encode string columns with a categorical method (e.g. "
-                    "'int' or 'one-hot') before binning."
-                ) from exc
+        X = self._as_numeric(X)
 
-        bins_spec = self._resolve_param("output_dim", default=UNSET)
-        if bins_spec is UNSET:
-            raise InvalidParamError("CustomBinTransformer requires 'output_dim'.")
-
-        if isinstance(bins_spec, int):
-            # Calculate equal width bins based on the range of the data and number of bins
-            _, bins = pd.cut(X.squeeze(), bins=bins_spec, retbins=True)
-        else:
-            # Use predefined bins
-            bins = bins_spec
+        edges = self.bin_edges_[0]
+        # Clip into the fitted range so unseen extremes land in the first / last
+        # bin rather than becoming NaN in an otherwise integer column.
+        values = np.clip(X[:, 0], edges[0], edges[-1])
 
-        # Apply the bins to the data
         binned_data = pd.cut(  # type: ignore
-            X.squeeze(),
-            bins=np.sort(np.unique(bins)),  # type: ignore
+            values,
+            bins=edges,
             labels=False,
             include_lowest=True,
         )
-        return np.expand_dims(np.array(binned_data), 1)
+        return np.expand_dims(np.asarray(binned_data, dtype=int), 1)
 
     def get_feature_names_out(self, input_features=None):
         """Return the names of the transformed features.

tests/test_custombin_transformer.py

@@ -53,11 +53,17 @@ def test_custom_bin_transformer_invalid_input():
 
 
 def test_custom_bin_transformer_raises_on_invalid_shape():
-    transformer = CustomBinTransformer(output_dim=3)
-    X = np.array([[0.1]])  # This will become scalar after squeeze()
+    transformer = CustomBinTransformer(output_dim=3).fit(np.linspace(0, 1, 10).reshape(-1, 1))
+
+    with pytest.raises(ValueError, match=r"2D array with shape \(n_samples, 1\)"):
+        transformer.transform(np.zeros((5, 2)))
+
+
+def test_custom_bin_transformer_raises_before_fit():
+    from sklearn.exceptions import NotFittedError
 
-    with pytest.raises(ValueError, match="Input must have more than 2 observations."):
-        transformer.transform(X)
+    with pytest.raises(NotFittedError):
+        CustomBinTransformer(output_dim=3).transform(np.linspace(0, 1, 10).reshape(-1, 1))
 
 
 def test_custom_bin_transformer_invalid_bins_type():
@@ -80,3 +86,74 @@ def test_custom_bin_transformer_feature_names_out_raises():
 
 def test_custom_bin_transformer_is_sklearn_compatible():
     assert isinstance(CustomBinTransformer(output_dim=3), (BaseEstimator, TransformerMixin))
+
+
+# --------------------------------------------------------------------------- #
+# Bin edges must be learned at fit and reused verbatim.
+#
+# ``fit`` recorded only ``n_features_in_``; ``transform`` re-derived equal-width
+# edges from whatever batch it was handed, so the same value received different
+# codes depending on which other rows travelled with it.
+# --------------------------------------------------------------------------- #
+def test_edges_are_learned_at_fit():
+    transformer = CustomBinTransformer(output_dim=4).fit(np.linspace(0, 10, 100).reshape(-1, 1))
+
+    assert hasattr(transformer, "bin_edges_")
+    assert len(transformer.bin_edges_[0]) == 5  # 4 bins -> 5 edges
+    assert transformer.bin_edges_[0][0] <= 0.0
+    assert transformer.bin_edges_[0][-1] >= 10.0
+
+
+def test_codes_do_not_depend_on_the_transform_batch():
+    transformer = CustomBinTransformer(output_dim=4).fit(np.linspace(0, 10, 100).reshape(-1, 1))
+    rows = np.array([[0.0], [2.5], [5.0], [7.5], [10.0]])
+
+    alone = transformer.transform(rows)
+    with_outlier = transformer.transform(np.vstack([rows, [[100.0]]]))
+
+    np.testing.assert_array_equal(alone.ravel(), with_outlier[:-1].ravel())
+
+
+def test_codes_reflect_the_fitted_range_not_the_batch():
+    transformer = CustomBinTransformer(output_dim=4).fit(np.linspace(0, 10, 100).reshape(-1, 1))
+
+    codes = transformer.transform(np.array([[0.0], [2.5], [5.0], [7.5], [10.0]])).ravel()
+
+    # Evenly spaced points across the fitted [0, 10] range must span the bins.
+    np.testing.assert_array_equal(codes, [0, 0, 1, 2, 3])
+
+
+def test_single_row_transform_is_supported():
+    transformer = CustomBinTransformer(output_dim=4).fit(np.linspace(0, 10, 100).reshape(-1, 1))
+
+    out = transformer.transform(np.array([[5.0]]))
+
+    assert out.shape == (1, 1)
+    assert out[0, 0] == transformer.transform(np.array([[0.0], [5.0], [10.0]]))[1, 0]
+
+
+def test_out_of_range_values_clip_into_the_end_bins():
+    transformer = CustomBinTransformer(output_dim=4).fit(np.linspace(0, 10, 100).reshape(-1, 1))
+
+    codes = transformer.transform(np.array([[-50.0], [5.0], [50.0]])).ravel()
+
+    assert codes[0] == 0
+    assert codes[2] == 3
+    assert not np.isnan(codes).any()
+
+
+def test_explicit_edges_are_stored_and_reused():
+    edges = [0.0, 0.5, 1.0]
+    transformer = CustomBinTransformer(output_dim=edges).fit(np.linspace(0, 1, 20).reshape(-1, 1))
+
+    np.testing.assert_allclose(transformer.bin_edges_[0], edges)
+    np.testing.assert_array_equal(
+        transformer.transform(np.array([[0.1], [0.9]])).ravel(), [0, 1]
+    )
+
+
+def test_string_input_raises_at_fit():
+    from pretab.core.exceptions import PretabDataError
+
+    with pytest.raises(PretabDataError, match="requires numeric input"):
+        CustomBinTransformer(output_dim=3).fit(np.array([["a"], ["b"], ["c"]], dtype=object))