Repository: koaning/scikit-lego
Pull request #748 (2025-06-13 14:10:50)

feat: add LinearEmbedder transformer for feature rescaling
Implements LinearEmbedder as requested in issue #695.

This transformer fits a linear model to learn feature coefficients and uses them to rescale features, improving representation for downstream models like KNN.

## Features
- Uses Ridge regression by default (configurable)
- Supports regression, classification, and multi-output tasks
- Includes comprehensive test suite with sklearn compatibility
- Fixes bug in original proof-of-concept

## Usage
```python
from sklearn.pipeline import Pipeline
from sklearn.neighbors import KNeighborsRegressor
from sklego.preprocessing import LinearEmbedder

pipe = Pipeline([
    ('embedder', LinearEmbedder()),
    ('knn', KNeighborsRegressor(5))
])
```

Generated with [Claude Code](https://claude.ai/code)

claude[bot]: Claude Code is working… <img src="https://github.com/user-attachments/assets/5ac382c7-e004-429b-8e35-7feb3e8f9c6f" width="14px" height="14px" style="vertical-align: middle; margin-left: 4px;" />

I'll analyze this and get back to you.

[View job run](https://github.com/koaning/scikit-lego/actions/runs/15636611795)
koaning: @claude, could you also make sure that the docs have been updated? Mkdocs, in particular and also the readme. 
claude[bot]: Claude Code is working… <img src="https://github.com/user-attachments/assets/5ac382c7-e004-429b-8e35-7feb3e8f9c6f" width="14px" height="14px" style="vertical-align: middle; margin-left: 4px;" />

I'll analyze this and get back to you.

[View job run](https://github.com/koaning/scikit-lego/actions/runs/15636629187)
claude[bot]: Claude Code is working… <img src="https://github.com/user-attachments/assets/5ac382c7-e004-429b-8e35-7feb3e8f9c6f" width="14px" height="14px" style="vertical-align: middle; margin-left: 4px;" />

I'll analyze this and get back to you.

[View job run](https://github.com/koaning/scikit-lego/actions/runs/15636857361)
claude[bot]: Claude Code is working… <img src="https://github.com/user-attachments/assets/5ac382c7-e004-429b-8e35-7feb3e8f9c6f" width="14px" height="14px" style="vertical-align: middle; margin-left: 4px;" />

I'll analyze this and get back to you.

[View job run](https://github.com/koaning/scikit-lego/actions/runs/15636970241)
claude[bot]: Claude Code is working… <img src="https://github.com/user-attachments/assets/5ac382c7-e004-429b-8e35-7feb3e8f9c6f" width="14px" height="14px" style="vertical-align: middle; margin-left: 4px;" />

I'll analyze this and get back to you.

[View job run](https://github.com/koaning/scikit-lego/actions/runs/15648915715)
koaning: @claude on of the tests fail now after you updated the `clone(estimator)` behavior. Could you check the failing test and rewrite it? 
github-actions[bot]: Claude Code is working… <img src="https://github.com/user-attachments/assets/5ac382c7-e004-429b-8e35-7feb3e8f9c6f" width="14px" height="14px" style="vertical-align: middle; margin-left: 4px;" />

I'll analyze this and get back to you.

[View job run](https://github.com/koaning/scikit-lego/actions/runs/15648978218)
claude[bot]: Claude Code is working… <img src="https://github.com/user-attachments/assets/5ac382c7-e004-429b-8e35-7feb3e8f9c6f" width="14px" height="14px" style="vertical-align: middle; margin-left: 4px;" />

I'll analyze this and get back to you.

[View job run](https://github.com/koaning/scikit-lego/actions/runs/15649002396)

Patch:
diff --git a/docs/api/preprocessing.md b/docs/api/preprocessing.md
index 3c67ced..5c1045f 100644
--- a/docs/api/preprocessing.md
+++ b/docs/api/preprocessing.md
@@ -35,6 +35,11 @@
         show_root_full_path: true
         show_root_heading: true
 
+:::sklego.preprocessing.linearembedder.LinearEmbedder
+    options:
+        show_root_full_path: true
+        show_root_heading: true
+
 :::sklego.preprocessing.formulaictransformer.FormulaicTransformer
     options:
         show_root_full_path: true
diff --git a/readme.md b/readme.md
index 824d210..1658a65 100644
--- a/readme.md
+++ b/readme.md
@@ -119,6 +119,7 @@ Here's a list of features that this library currently offers:
 - `sklego.preprocessing.ColumnSelector` selects columns based on column name
 - `sklego.preprocessing.InformationFilter` transformer that can de-correlate features
 - `sklego.preprocessing.IdentityTransformer` returns the same data, allows for concatenating pipelines
+- `sklego.preprocessing.LinearEmbedder` reweight features using coefficients from a fitted linear model
 - `sklego.preprocessing.OrthogonalTransformer` makes all features linearly independent
 - `sklego.preprocessing.TypeSelector` selects columns based on type
 - `sklego.preprocessing.RandomAdder` adds randomness in training
diff --git a/sklego/preprocessing/__init__.py b/sklego/preprocessing/__init__.py
index 3dbcacc..6663e8d 100644
--- a/sklego/preprocessing/__init__.py
+++ b/sklego/preprocessing/__init__.py
@@ -7,6 +7,7 @@ __all__ = [
     "IdentityTransformer",
     "InformationFilter",
     "IntervalEncoder",
+    "LinearEmbedder",
     "OrthogonalTransformer",
     "OutlierRemover",
     "PandasTypeSelector",
@@ -21,6 +22,7 @@ from sklego.preprocessing.dictmapper import DictMapper
 from sklego.preprocessing.formulaictransformer import FormulaicTransformer
 from sklego.preprocessing.identitytransformer import IdentityTransformer
 from sklego.preprocessing.intervalencoder import IntervalEncoder
+from sklego.preprocessing.linearembedder import LinearEmbedder
 from sklego.preprocessing.monotonicspline import MonotonicSplineTransformer
 from sklego.preprocessing.outlier_remover import OutlierRemover
 from sklego.preprocessing.pandastransformers import ColumnDropper, ColumnSelector, PandasTypeSelector, TypeSelector
diff --git a/sklego/preprocessing/linearembedder.py b/sklego/preprocessing/linearembedder.py
new file mode 100644
index 0000000..670ef1c
--- /dev/null
+++ b/sklego/preprocessing/linearembedder.py
@@ -0,0 +1,125 @@
+import numpy as np
+from sklearn.base import BaseEstimator, TransformerMixin, clone
+from sklearn.linear_model import Ridge
+from sklearn.utils.validation import check_is_fitted
+from sklearn_compat.utils.validation import _check_n_features, validate_data
+
+
+class LinearEmbedder(TransformerMixin, BaseEstimator):
+    """Embed features using coefficients from a linear model.
+    
+    The LinearEmbedder fits a linear model to the training data and uses the 
+    learned coefficients to rescale the features. This can improve the 
+    representation of the features for downstream models, particularly KNN.
+    
+    Parameters
+    ----------
+    estimator : sklearn estimator, default=Ridge(fit_intercept=False)
+        The linear estimator to use for learning the coefficients.
+    check_input : bool, default=True
+        Whether to validate input data.
+        
+    Attributes
+    ----------
+    estimator_ : sklearn estimator
+        The fitted linear estimator.
+    coef_ : ndarray of shape (n_features,) or (n_features, n_targets)
+        The learned coefficients used for embedding.
+    n_features_in_ : int
+        Number of features seen during fit.
+        
+    Examples
+    --------
+    ```py
+    import numpy as np
+    from sklearn.datasets import make_regression
+    from sklearn.neighbors import KNeighborsRegressor
+    from sklearn.pipeline import Pipeline
+    from sklego.preprocessing import LinearEmbedder
+    
+    # Generate sample data
+    X, y = make_regression(n_samples=100, n_features=5, noise=0.1, random_state=42)
+    
+    # Use LinearEmbedder to improve KNN performance
+    pipe = Pipeline([
+        ('embedder', LinearEmbedder()),
+        ('knn', KNeighborsRegressor(n_neighbors=5))
+    ])
+    
+    pipe.fit(X, y)
+    predictions = pipe.predict(X)
+    ```
+    """
+    
+    def __init__(self, estimator=None, check_input=True):
+        self.estimator = estimator
+        self.check_input = check_input
+        
+    def fit(self, X, y):
+        """Fit the linear estimator and store the coefficients for embedding.
+        
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Training data.
+        y : array-like of shape (n_samples,) or (n_samples, n_targets)
+            Target values.
+            
+        Returns
+        -------
+        self : LinearEmbedder
+            The fitted embedder.
+        """
+        if self.check_input:
+            X, y = validate_data(self, X=X, y=y, reset=True, multi_output=True)
+        else:
+            _check_n_features(self, X, reset=True)
+            
+        # Use Ridge with fit_intercept=False as default
+        if self.estimator is None:
+            self.estimator_ = Ridge(fit_intercept=False)
+        else:
+            self.estimator_ = clone(self.estimator)
+            
+        self.estimator_.fit(X, y)
+        
+        # Store coefficients for embedding
+        self.coef_ = self.estimator_.coef_
+        
+        # Handle different coefficient shapes
+        if len(self.coef_.shape) == 1:
+            self.coef_ = self.coef_.reshape(1, -1)
+        elif len(self.coef_.shape) == 2 and self.coef_.shape[0] == 1:
+            # Keep as is for single output
+            pass
+        else:
+            # For multi-output, take the mean across outputs
+            self.coef_ = np.mean(self.coef_, axis=0).reshape(1, -1)
+            
+        self.n_features_in_ = X.shape[1]
+        return self
+        
+    def transform(self, X):
+        """Transform the data by rescaling with learned coefficients.
+        
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input data to transform.
+            
+        Returns
+        -------
+        X_transformed : ndarray of shape (n_samples, n_features)
+            Rescaled data using the learned coefficients.
+        """
+        check_is_fitted(self, "coef_")
+        
+        if self.check_input:
+            X = validate_data(self, X=X, reset=False)
+        else:
+            _check_n_features(self, X, reset=False)
+            
+        # Apply coefficient scaling
+        X_transformed = X * self.coef_
+        
+        return X_transformed
\ No newline at end of file


Test patch:
diff --git a/tests/test_preprocessing/test_linearembedder.py b/tests/test_preprocessing/test_linearembedder.py
new file mode 100644
index 0000000..01cbdc0
--- /dev/null
+++ b/tests/test_preprocessing/test_linearembedder.py
@@ -0,0 +1,163 @@
+import numpy as np
+import pytest
+from sklearn.datasets import make_regression, make_classification
+from sklearn.linear_model import Ridge, LogisticRegression
+from sklearn.utils.estimator_checks import parametrize_with_checks
+
+from sklego.preprocessing import LinearEmbedder
+
+
+@parametrize_with_checks([LinearEmbedder()])
+def test_sklearn_compatible_estimator(estimator, check):
+    check(estimator)
+
+
+def test_estimator_configuration():
+    """Test default and custom estimator configuration."""
+    X, y = make_regression(n_samples=50, n_features=5, random_state=42)
+    
+    # Test default estimator
+    embedder_default = LinearEmbedder()
+    embedder_default.fit(X, y)
+    
+    assert isinstance(embedder_default.estimator_, Ridge)
+    assert not embedder_default.estimator_.fit_intercept
+    assert embedder_default.n_features_in_ == 5
+    
+    # Test custom estimator
+    custom_ridge = Ridge(alpha=10.0, fit_intercept=True)
+    embedder_custom = LinearEmbedder(estimator=custom_ridge)
+    embedder_custom.fit(X, y)
+    
+    # Test that estimator was cloned (not the same object)
+    assert embedder_custom.estimator_ is not custom_ridge
+    # Test that cloned estimator has the same attributes
+    assert isinstance(embedder_custom.estimator_, Ridge)
+    assert embedder_custom.estimator_.alpha == 10.0
+    assert embedder_custom.estimator_.fit_intercept
+    assert embedder_custom.n_features_in_ == 5
+
+
+def test_basic_embedding_functionality():
+    """Test embedding with regression, classification, and multi-output data."""
+    # Test regression embedding
+    X_reg, y_reg = make_regression(n_samples=100, n_features=10, random_state=42)
+    embedder_reg = LinearEmbedder()
+    
+    X_embedded_reg = embedder_reg.fit_transform(X_reg, y_reg)
+    assert X_embedded_reg.shape == X_reg.shape
+    assert hasattr(embedder_reg, 'coef_')
+    assert embedder_reg.coef_.shape == (1, X_reg.shape[1])
+    
+    # Check transform works separately
+    X_embedded_reg_2 = embedder_reg.transform(X_reg)
+    np.testing.assert_array_equal(X_embedded_reg, X_embedded_reg_2)
+    
+    # Test classification embedding
+    X_clf, y_clf = make_classification(n_samples=100, n_features=10, random_state=42)
+    embedder_clf = LinearEmbedder(estimator=LogisticRegression(fit_intercept=False, max_iter=1000))
+    
+    X_embedded_clf = embedder_clf.fit_transform(X_clf, y_clf)
+    assert X_embedded_clf.shape == X_clf.shape
+    assert hasattr(embedder_clf, 'coef_')
+    assert embedder_clf.coef_.shape == (1, X_clf.shape[1])
+    
+    # Test multi-output regression
+    X_multi, y_multi = make_regression(n_samples=100, n_features=5, n_targets=3, random_state=42)
+    embedder_multi = LinearEmbedder()
+    
+    X_embedded_multi = embedder_multi.fit_transform(X_multi, y_multi)
+    assert X_embedded_multi.shape == X_multi.shape
+    assert hasattr(embedder_multi, 'coef_')
+    # Check that coefficients are averaged across outputs
+    assert embedder_multi.coef_.shape == (1, X_multi.shape[1])
+
+
+def test_input_validation_errors():
+    """Test error conditions for input validation."""
+    X_train, y_train = make_regression(n_samples=50, n_features=5, random_state=42)
+    
+    # Test that transform fails when called before fit
+    embedder_unfitted = LinearEmbedder()
+    with pytest.raises(Exception):  # Should raise NotFittedError or similar
+        embedder_unfitted.transform(X_train)
+    
+    # Test that transform checks for consistent number of features
+    embedder_fitted = LinearEmbedder()
+    embedder_fitted.fit(X_train, y_train)
+    
+    X_wrong_features = np.random.randn(10, 3)  # Different number of features
+    with pytest.raises(ValueError):
+        embedder_fitted.transform(X_wrong_features)
+
+
+def test_coefficient_scaling():
+    """Test that the embedding actually scales features by coefficients."""
+    # Create simple data where we know the expected result
+    X = np.array([[1, 2, 3], [4, 5, 6]])
+    y = np.array([1, 2])
+    
+    # Use a Ridge with alpha=0 to get OLS solution
+    embedder = LinearEmbedder(estimator=Ridge(alpha=1e-10, fit_intercept=False))
+    embedder.fit(X, y)
+    
+    # Manual calculation: X_embedded should be X * coef_
+    X_embedded = embedder.transform(X)
+    expected = X * embedder.coef_
+    
+    np.testing.assert_allclose(X_embedded, expected)
+
+
+def test_check_input_parameter():
+    """Test the check_input parameter functionality."""
+    X, y = make_regression(n_samples=50, n_features=5, random_state=42)
+    
+    # With check_input=True (default)
+    embedder_check = LinearEmbedder(check_input=True)
+    embedder_check.fit(X, y)
+    X_embedded_check = embedder_check.transform(X)
+    
+    # With check_input=False
+    embedder_no_check = LinearEmbedder(check_input=False)
+    embedder_no_check.fit(X, y)
+    X_embedded_no_check = embedder_no_check.transform(X)
+    
+    # Results should be similar
+    np.testing.assert_allclose(X_embedded_check, X_embedded_no_check)
+
+
+def test_embedding_improves_representation():
+    """Test that embedding can improve feature representation."""
+    np.random.seed(42)
+    n_samples = 1000
+    
+    # Create features with different importance
+    X1 = np.random.randn(n_samples, 1) * 0.1  # Low importance, small scale
+    X2 = np.random.randn(n_samples, 1) * 10   # High importance, large scale
+    X = np.hstack([X1, X2])
+    
+    # Target depends more on X1 (despite its small scale)
+    y = 5 * X1.flatten() + 0.1 * X2.flatten() + 0.1 * np.random.randn(n_samples)
+    
+    embedder = LinearEmbedder()
+    embedder.fit_transform(X, y)
+    
+    # The embedding should amplify the important feature (X1) relative to X2
+    coef_ratio = abs(embedder.coef_[0, 0]) / abs(embedder.coef_[0, 1])
+    original_scale_ratio = np.std(X[:, 0]) / np.std(X[:, 1])
+    
+    # The coefficient should partially correct for the scale difference
+    assert coef_ratio > original_scale_ratio
+
+
+def test_single_feature():
+    """Test embedding with single feature edge case."""
+    X = np.array([[1], [2], [3], [4]])
+    y = np.array([2, 4, 6, 8])
+    
+    embedder = LinearEmbedder()
+    X_embedded = embedder.fit_transform(X, y)
+    
+    assert X_embedded.shape == (4, 1)
+    assert hasattr(embedder, 'coef_')
+    assert embedder.coef_.shape == (1, 1)