FEAT Add custom imputation strategy to SimpleImputer (scikit-learn#28053)

mark-thm · jnothman · web-flow · commit e2b3785b6a1f · 2024-01-04T16:04:30.000+01:00
Co-authored-by: Joel Nothman &lt;joeln@canva.com&gt;
diff --git a/doc/whats_new/v1.5.rst b/doc/whats_new/v1.5.rst
@@ -25,6 +25,12 @@ Changelog
     :pr:`123456` by :user:`Joe Bloggs <joeongithub>`.
     where 123455 is the *pull request* number, not the issue number.
 
+:mod:`sklearn.impute`
+.....................
+- |Enhancement| :class:`impute.SimpleImputer` now supports custom strategies
+  by passing a function in place of a strategy name.
+  :pr:`28053` by :user:`Mark Elliot <mark-thm>`.
+
 Code and Documentation Contributors
 -----------------------------------
 
diff --git a/sklearn/impute/_base.py b/sklearn/impute/_base.py
@@ -6,6 +6,7 @@
 import warnings
 from collections import Counter
 from functools import partial
+from typing import Callable
 
 import numpy as np
 import numpy.ma as ma
@@ -163,7 +164,7 @@ class SimpleImputer(_BaseImputer):
         nullable integer dtypes with missing values, `missing_values`
         can be set to either `np.nan` or `pd.NA`.
 
-    strategy : str, default='mean'
+    strategy : str or Callable, default='mean'
         The imputation strategy.
 
         - If "mean", then replace missing values using the mean along
@@ -175,10 +176,16 @@ class SimpleImputer(_BaseImputer):
           If there is more than one such value, only the smallest is returned.
         - If "constant", then replace missing values with fill_value. Can be
           used with strings or numeric data.
+        - If an instance of Callable, then replace missing values using the
+          scalar statistic returned by running the callable over a dense 1d
+          array containing non-missing values of each column.
 
         .. versionadded:: 0.20
            strategy="constant" for fixed value imputation.
 
+        .. versionadded:: 1.5
+           strategy=callable for custom value imputation.
+
     fill_value : str or numerical value, default=None
         When strategy == "constant", `fill_value` is used to replace all
         occurrences of missing_values. For string or object data types,
@@ -270,7 +277,10 @@ class SimpleImputer(_BaseImputer):
 
     _parameter_constraints: dict = {
         **_BaseImputer._parameter_constraints,
-        "strategy": [StrOptions({"mean", "median", "most_frequent", "constant"})],
+        "strategy": [
+            StrOptions({"mean", "median", "most_frequent", "constant"}),
+            callable,
+        ],
         "fill_value": "no_validation",  # any object is valid
         "copy": ["boolean"],
     }
@@ -456,6 +466,9 @@ def _sparse_fit(self, X, strategy, missing_values, fill_value):
                     elif strategy == "most_frequent":
                         statistics[i] = _most_frequent(column, 0, n_zeros)
 
+                    elif isinstance(strategy, Callable):
+                        statistics[i] = self.strategy(column)
+
         super()._fit_indicator(missing_mask)
 
         return statistics
@@ -518,6 +531,13 @@ def _dense_fit(self, X, strategy, missing_values, fill_value):
             # fill_value in each column
             return np.full(X.shape[1], fill_value, dtype=X.dtype)
 
+        # Custom
+        elif isinstance(strategy, Callable):
+            statistics = np.empty(masked_X.shape[1])
+            for i in range(masked_X.shape[1]):
+                statistics[i] = self.strategy(masked_X[:, i].compressed())
+            return statistics
+
     def transform(self, X):
         """Impute all missing values in `X`.
 
diff --git a/sklearn/impute/tests/test_impute.py b/sklearn/impute/tests/test_impute.py
@@ -1710,3 +1710,37 @@ def test_simple_imputer_keep_empty_features(strategy, array_type, keep_empty_fea
             assert_array_equal(constant_feature, 0)
         else:
             assert X_imputed.shape == (X.shape[0], X.shape[1] - 1)
+
+
+@pytest.mark.parametrize("csc_container", CSC_CONTAINERS)
+def test_imputation_custom(csc_container):
+    X = np.array(
+        [
+            [1.1, 1.1, 1.1],
+            [3.9, 1.2, np.nan],
+            [np.nan, 1.3, np.nan],
+            [0.1, 1.4, 1.4],
+            [4.9, 1.5, 1.5],
+            [np.nan, 1.6, 1.6],
+        ]
+    )
+
+    X_true = np.array(
+        [
+            [1.1, 1.1, 1.1],
+            [3.9, 1.2, 1.1],
+            [0.1, 1.3, 1.1],
+            [0.1, 1.4, 1.4],
+            [4.9, 1.5, 1.5],
+            [0.1, 1.6, 1.6],
+        ]
+    )
+
+    imputer = SimpleImputer(missing_values=np.nan, strategy=np.min)
+    X_trans = imputer.fit_transform(X)
+    assert_array_equal(X_trans, X_true)
+
+    # Sparse matrix
+    imputer = SimpleImputer(missing_values=np.nan, strategy=np.min)
+    X_trans = imputer.fit_transform(csc_container(X))
+    assert_array_equal(X_trans.toarray(), X_true)
