TST add test_multi_task_lasso_vs_skglm (#31957)

lorentzenchr · web-flow · commit 3883ba73ac4a · 2025-08-19T13:50:45.000+05:00
diff --git a/sklearn/linear_model/_cd_fast.pyx b/sklearn/linear_model/_cd_fast.pyx
@@ -786,6 +786,12 @@ def enet_coordinate_descent_multi_task(
 
         0.5 * norm(Y - X W.T, 2)^2 + l1_reg ||W.T||_21 + 0.5 * l2_reg norm(W.T, 2)^2
 
+    The algorithm follows
+    Noah Simon, Jerome Friedman, Trevor Hastie. 2013.
+    A Blockwise Descent Algorithm for Group-penalized Multiresponse and Multinomial
+    Regression
+    https://doi.org/10.48550/arXiv.1311.6529
+
     Returns
     -------
     W : ndarray of shape (n_tasks, n_features)
diff --git a/sklearn/linear_model/tests/test_coordinate_descent.py b/sklearn/linear_model/tests/test_coordinate_descent.py
@@ -510,6 +510,46 @@ def test_uniform_targets():
             assert_array_equal(model.alphas_, [np.finfo(float).resolution] * 3)
 
 
+@pytest.mark.filterwarnings("error::sklearn.exceptions.ConvergenceWarning")
+def test_multi_task_lasso_vs_skglm():
+    """Test that MultiTaskLasso gives same results as the one from skglm.
+
+    To reproduce numbers, just use
+    from skglm import MultiTaskLasso
+    """
+    # Numbers are with skglm version 0.5.
+    n_samples, n_features, n_tasks = 5, 4, 3
+    X = np.vander(np.arange(n_samples), n_features)
+    Y = np.arange(n_samples * n_tasks).reshape(n_samples, n_tasks)
+
+    def obj(W, X, y, alpha):
+        intercept = W[:, -1]
+        W = W[:, :-1]
+        l21_norm = np.sqrt(np.sum(W**2, axis=0)).sum()
+        return (
+            np.linalg.norm(Y - X @ W.T - intercept, ord="fro") ** 2 / (2 * n_samples)
+            + alpha * l21_norm
+        )
+
+    alpha = 0.1
+    # TODO: The high number of iterations are required for convergence and show room
+    # for improvement of the CD algorithm.
+    m = MultiTaskLasso(alpha=alpha, tol=1e-10, max_iter=5000).fit(X, Y)
+    assert_allclose(
+        obj(np.c_[m.coef_, m.intercept_], X, Y, alpha=alpha),
+        0.4965993692547902,
+        rtol=1e-10,
+    )
+    assert_allclose(
+        m.intercept_, [0.219942959407, 1.219942959407, 2.219942959407], rtol=1e-7
+    )
+    assert_allclose(
+        m.coef_,
+        np.tile([-0.032075014794, 0.25430904614, 2.44785152982, 0], (n_tasks, 1)),
+        rtol=1e-6,
+    )
+
+
 def test_multi_task_lasso_and_enet():
     X, y, X_test, y_test = build_dataset()
     Y = np.c_[y, y]
