ENH add n_jobs to mutual_info_regression and mutual_info_classif (#28085)

netomenoci · glemaitre · lesteve · web-flow · commit 06e566eb86cf · 2024-01-18T12:48:28.000+01:00
Co-authored-by: Guillaume Lemaitre &lt;g.lemaitre58@gmail.com&gt;
Co-authored-by: Loïc Estève &lt;loic.esteve@ymail.com&gt;
diff --git a/doc/whats_new/v1.5.rst b/doc/whats_new/v1.5.rst
@@ -54,6 +54,14 @@ Changelog
   by passing a function in place of a strategy name.
   :pr:`28053` by :user:`Mark Elliot <mark-thm>`.
 
+:mod:`sklearn.feature_selection`
+................................
+
+- |Enhancement| :func:`feature_selection.mutual_info_regression` and
+  :func:`feature_selection.mutual_info_classif` now support `n_jobs` parameter.
+  :pr:`28085` by :user:`Neto Menoci <netomenoci>` and
+  :user:`Florin Andrei <FlorinAndrei>`.
+
 :mod:`sklearn.metrics`
 ......................
 
diff --git a/sklearn/feature_selection/_mutual_info.py b/sklearn/feature_selection/_mutual_info.py
@@ -13,6 +13,7 @@
 from ..utils import check_random_state
 from ..utils._param_validation import Interval, StrOptions, validate_params
 from ..utils.multiclass import check_classification_targets
+from ..utils.parallel import Parallel, delayed
 from ..utils.validation import check_array, check_X_y
 
 
@@ -201,11 +202,13 @@ def _iterate_columns(X, columns=None):
 def _estimate_mi(
     X,
     y,
+    *,
     discrete_features="auto",
     discrete_target=False,
     n_neighbors=3,
     copy=True,
     random_state=None,
+    n_jobs=None,
 ):
     """Estimate mutual information between the features and the target.
 
@@ -242,6 +245,16 @@ def _estimate_mi(
         Pass an int for reproducible results across multiple function calls.
         See :term:`Glossary <random_state>`.
 
+    n_jobs : int, default=None
+        The number of jobs to use for computing the mutual information.
+        The parallelization is done on the columns of `X`.
+        ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
+        ``-1`` means using all processors. See :term:`Glossary <n_jobs>`
+        for more details.
+
+        .. versionadded:: 1.5
+
+
     Returns
     -------
     mi : ndarray, shape (n_features,)
@@ -301,10 +314,10 @@ def _estimate_mi(
             * rng.standard_normal(size=n_samples)
         )
 
-    mi = [
-        _compute_mi(x, y, discrete_feature, discrete_target, n_neighbors)
+    mi = Parallel(n_jobs=n_jobs)(
+        delayed(_compute_mi)(x, y, discrete_feature, discrete_target, n_neighbors)
         for x, discrete_feature in zip(_iterate_columns(X), discrete_mask)
-    ]
+    )
 
     return np.array(mi)
 
@@ -317,11 +330,19 @@ def _estimate_mi(
         "n_neighbors": [Interval(Integral, 1, None, closed="left")],
         "copy": ["boolean"],
         "random_state": ["random_state"],
+        "n_jobs": [Integral, None],
     },
     prefer_skip_nested_validation=True,
 )
 def mutual_info_regression(
-    X, y, *, discrete_features="auto", n_neighbors=3, copy=True, random_state=None
+    X,
+    y,
+    *,
+    discrete_features="auto",
+    n_neighbors=3,
+    copy=True,
+    random_state=None,
+    n_jobs=None,
 ):
     """Estimate mutual information for a continuous target variable.
 
@@ -367,6 +388,16 @@ def mutual_info_regression(
         Pass an int for reproducible results across multiple function calls.
         See :term:`Glossary <random_state>`.
 
+    n_jobs : int, default=None
+        The number of jobs to use for computing the mutual information.
+        The parallelization is done on the columns of `X`.
+
+        ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
+        ``-1`` means using all processors. See :term:`Glossary <n_jobs>`
+        for more details.
+
+        .. versionadded:: 1.5
+
     Returns
     -------
     mi : ndarray, shape (n_features,)
@@ -407,7 +438,16 @@ def mutual_info_regression(
     >>> mutual_info_regression(X, y)
     array([0.1..., 2.6...  , 0.0...])
     """
-    return _estimate_mi(X, y, discrete_features, False, n_neighbors, copy, random_state)
+    return _estimate_mi(
+        X,
+        y,
+        discrete_features=discrete_features,
+        discrete_target=False,
+        n_neighbors=n_neighbors,
+        copy=copy,
+        random_state=random_state,
+        n_jobs=n_jobs,
+    )
 
 
 @validate_params(
@@ -418,11 +458,19 @@ def mutual_info_regression(
         "n_neighbors": [Interval(Integral, 1, None, closed="left")],
         "copy": ["boolean"],
         "random_state": ["random_state"],
+        "n_jobs": [Integral, None],
     },
     prefer_skip_nested_validation=True,
 )
 def mutual_info_classif(
-    X, y, *, discrete_features="auto", n_neighbors=3, copy=True, random_state=None
+    X,
+    y,
+    *,
+    discrete_features="auto",
+    n_neighbors=3,
+    copy=True,
+    random_state=None,
+    n_jobs=None,
 ):
     """Estimate mutual information for a discrete target variable.
 
@@ -468,6 +516,15 @@ def mutual_info_classif(
         Pass an int for reproducible results across multiple function calls.
         See :term:`Glossary <random_state>`.
 
+    n_jobs : int, default=None
+        The number of jobs to use for computing the mutual information.
+        The parallelization is done on the columns of `X`.
+        ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
+        ``-1`` means using all processors. See :term:`Glossary <n_jobs>`
+        for more details.
+
+        .. versionadded:: 1.5
+
     Returns
     -------
     mi : ndarray, shape (n_features,)
@@ -511,4 +568,13 @@ def mutual_info_classif(
            0.     , 0.     , 0.     , 0.      , 0.        ])
     """
     check_classification_targets(y)
-    return _estimate_mi(X, y, discrete_features, True, n_neighbors, copy, random_state)
+    return _estimate_mi(
+        X,
+        y,
+        discrete_features=discrete_features,
+        discrete_target=True,
+        n_neighbors=n_neighbors,
+        copy=copy,
+        random_state=random_state,
+        n_jobs=n_jobs,
+    )
diff --git a/sklearn/feature_selection/tests/test_mutual_info.py b/sklearn/feature_selection/tests/test_mutual_info.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+from sklearn.datasets import make_classification, make_regression
 from sklearn.feature_selection import mutual_info_classif, mutual_info_regression
 from sklearn.feature_selection._mutual_info import _compute_mi
 from sklearn.utils import check_random_state
@@ -252,3 +253,18 @@ def test_mutual_info_regression_X_int_dtype(global_random_seed):
     expected = mutual_info_regression(X_float, y, random_state=global_random_seed)
     result = mutual_info_regression(X, y, random_state=global_random_seed)
     assert_allclose(result, expected)
+
+
+@pytest.mark.parametrize(
+    "mutual_info_func, data_generator",
+    [
+        (mutual_info_regression, make_regression),
+        (mutual_info_classif, make_classification),
+    ],
+)
+def test_mutual_info_n_jobs(global_random_seed, mutual_info_func, data_generator):
+    """Check that results are consistent with different `n_jobs`."""
+    X, y = data_generator(random_state=global_random_seed)
+    single_job = mutual_info_func(X, y, random_state=global_random_seed, n_jobs=1)
+    multi_job = mutual_info_func(X, y, random_state=global_random_seed, n_jobs=2)
+    assert_allclose(single_job, multi_job)
