Repository: benchopt/benchopt
Pull request #285 (2022-01-04 08:32:52)

ENH Add n_tasks parameter to make_correlated_data


Patch:
diff --git a/benchopt/datasets/simulated.py b/benchopt/datasets/simulated.py
index f8da89f..b81a8b2 100644
--- a/benchopt/datasets/simulated.py
+++ b/benchopt/datasets/simulated.py
@@ -5,9 +5,9 @@ from numpy.linalg import norm
 from ..utils.checkers import check_random_state
 
 
-def make_correlated_data(n_samples=100, n_features=50, rho=0.6, snr=3,
-                         w_true=None, density=0.2, random_state=None):
-
+def make_correlated_data(
+        n_samples=100, n_features=50, n_tasks=1, rho=0.6, snr=3,
+        w_true=None, density=0.2, random_state=None):
     r"""Generate a linear regression with decaying correlation for the design
     matrix :math:`\rho^{|i-j|}`.
 
@@ -32,19 +32,21 @@ def make_correlated_data(n_samples=100, n_features=50, rho=0.6, snr=3,
         Number of samples in the design matrix.
     n_features: int
         Number of features in the design matrix.
+    n_tasks: int
+        Number of tasks.
     rho: float
         Correlation :math:`\rho` between successive features. The cross
         correlation :math:`C_{i, j}` between feature i and feature j will be
         :math:`\rho^{|i-j|}`. This parameter should be selected in
         :math:`[0, 1[`.
-    snr : float
+    snr: float or np.inf
         Signal-to-noise ratio.
-    w_true: np.array, shape (n_features,) | None
+    w_true: np.array, shape (n_features,) or (n_features, n_tasks)| None
         True regression coefficients. If None, a sparse array with standard
         Gaussian non zero entries is simulated.
     density: float
         Proportion of non zero elements in w_true if the latter is simulated.
-    random_state : int | RandomState instance | None (default)
+    random_state: int | RandomState instance | None (default)
         Determines random number generation for data generation. Use an int to
         make the randomness deterministic.
 
@@ -52,10 +54,10 @@ def make_correlated_data(n_samples=100, n_features=50, rho=0.6, snr=3,
     -------
     X: ndarray, shape (n_samples, n_features)
         A design matrix with Toeplitz covariance.
-    y: ndarray, shape (n_samples,)
-        Observation vector.
-    w_true: ndarray, shape (n_features,)
-        True regression vector of the model.
+    y: ndarray, shape (n_samples,) or (n_samples, n_tasks)
+        Observation vector/matrix.
+    w_true: ndarray, shape (n_features,) or (n_features, n_tasks)
+        True regression vector/matrix of the model.
     """
     if not 0 <= rho < 1:
         raise ValueError("The correlation `rho` should be chosen in [0, 1[.")
@@ -83,15 +85,21 @@ def make_correlated_data(n_samples=100, n_features=50, rho=0.6, snr=3,
         X = rng.randn(n_samples, n_features)
 
     if w_true is None:
-        w_true = np.zeros(n_features)
+        w_true = np.zeros((n_features, n_tasks))
         support = rng.choice(n_features, nnz, replace=False)
-        w_true[support] = rng.randn(nnz)
+        w_true[support, :] = rng.randn(nnz, n_tasks)
+    else:
+        if w_true.ndim == 1:
+            w_true = w_true[:, None]
 
-    y = X @ w_true
-    noise = rng.randn(n_samples)
+    Y = X @ w_true
+    noise = rng.randn(n_samples, n_tasks)
     if snr not in [0, np.inf]:
-        y += noise / norm(noise) * norm(y) / snr
+        Y += noise / norm(noise) * norm(Y) / snr
     elif snr == 0:
-        y = noise
+        Y = noise
 
-    return X, y, w_true
+    if n_tasks == 1:
+        return X, Y.flatten(), w_true.flatten()
+    else:
+        return X, Y, w_true


Test patch:
diff --git a/benchopt/datasets/tests/test_dataset_simulated.py b/benchopt/datasets/tests/test_dataset_simulated.py
index f83cf3c..f117ea8 100644
--- a/benchopt/datasets/tests/test_dataset_simulated.py
+++ b/benchopt/datasets/tests/test_dataset_simulated.py
@@ -49,6 +49,20 @@ def test_correlated_w_true(density):
     assert all(w_new == w_true)
 
 
+@pytest.mark.parametrize('n_tasks', [1, 7])
+def test_correlated_n_tasks(n_tasks):
+    X, y, w_true = make_correlated_data(
+        n_tasks=n_tasks, random_state=42
+    )
+
+    if n_tasks == 1:
+        assert y.ndim == 1
+        assert w_true.ndim == 1
+    else:
+        assert y.ndim == 2
+        assert y.shape[1] == w_true.shape[1]
+
+
 @pytest.mark.parametrize('param_name, p_range', [
     ('rho', [0, 1]), ('density', [0, 1]), ('snr', [0, np.inf])
 ])