moved sparse matrix generation code to test file out of utils

Author: nathanneike

ot/utils.py

@@ -12,7 +12,6 @@
 
 import numpy as np
 from scipy.spatial.distance import cdist
-from scipy.sparse import coo_array
 import sys
 import warnings
 from inspect import signature
@@ -435,39 +434,6 @@ def dist(
             return cdist(x1, x2, metric=metric)
 
 
-def get_sparse_test_matrices(n1, n2, k=2, seed=42, nx=None):
-    if nx is None:
-        nx = NumpyBackend()
-
-    rng = np.random.RandomState(seed)
-    M_orig = rng.rand(n1, n2)
-
-    mask = np.zeros((n1, n2))
-    for i in range(n1):
-        j_list = rng.choice(n2, min(k, n2), replace=False)
-        for j in j_list:
-            mask[i, j] = 1
-    for j in range(n2):
-        i_list = rng.choice(n1, min(k, n1), replace=False)
-        for i in i_list:
-            mask[i, j] = 1
-
-    M_sparse_np = coo_array(M_orig * mask)
-    rows, cols, data = M_sparse_np.row, M_sparse_np.col, M_sparse_np.data
-
-    if nx.__name__ == "numpy":
-        M_sparse = M_sparse_np
-    else:
-        rows_b = nx.from_numpy(rows.astype(np.int64))
-        cols_b = nx.from_numpy(cols.astype(np.int64))
-        data_b = nx.from_numpy(data)
-        M_sparse = nx.coo_matrix(data_b, rows_b, cols_b, shape=(n1, n2))
-
-    M_dense = nx.from_numpy(M_orig + 1e8 * (1 - mask))
-
-    return M_sparse, M_dense
-
-
 def dist0(n, method="lin_square"):
     r"""Compute standard cost matrices of size (`n`, `n`) for OT problems
 

test/test_ot.py

@@ -914,6 +914,43 @@ def test_dual_variables():
     assert constraint_violation.max() < 1e-8
 
 
+def _get_sparse_test_matrices(n1, n2, k=2, seed=42, nx=None):
+    """Helper function to create sparse and dense test matrices."""
+    from scipy.sparse import coo_array
+    from ot.backend import NumpyBackend
+
+    if nx is None:
+        nx = NumpyBackend()
+
+    rng = np.random.RandomState(seed)
+    M_orig = rng.rand(n1, n2)
+
+    mask = np.zeros((n1, n2))
+    for i in range(n1):
+        j_list = rng.choice(n2, min(k, n2), replace=False)
+        for j in j_list:
+            mask[i, j] = 1
+    for j in range(n2):
+        i_list = rng.choice(n1, min(k, n1), replace=False)
+        for i in i_list:
+            mask[i, j] = 1
+
+    M_sparse_np = coo_array(M_orig * mask)
+    rows, cols, data = M_sparse_np.row, M_sparse_np.col, M_sparse_np.data
+
+    if nx.__name__ == "numpy":
+        M_sparse = M_sparse_np
+    else:
+        rows_b = nx.from_numpy(rows.astype(np.int64))
+        cols_b = nx.from_numpy(cols.astype(np.int64))
+        data_b = nx.from_numpy(data)
+        M_sparse = nx.coo_matrix(data_b, rows_b, cols_b, shape=(n1, n2))
+
+    M_dense = nx.from_numpy(M_orig + 1e8 * (1 - mask))
+
+    return M_sparse, M_dense
+
+
 def test_emd_sparse_vs_dense(nx):
     """Test that sparse and dense EMD solvers produce identical results.
 
@@ -929,7 +966,7 @@ def test_emd_sparse_vs_dense(nx):
     n2 = 100
     k = 2
 
-    M_sparse, M_dense = ot.utils.get_sparse_test_matrices(n1, n2, k=k, seed=42, nx=nx)
+    M_sparse, M_dense = _get_sparse_test_matrices(n1, n2, k=k, seed=42, nx=nx)
 
     a = ot.utils.unif(n1, type_as=M_dense)
     b = ot.utils.unif(n2, type_as=M_dense)
@@ -971,7 +1008,7 @@ def test_emd2_sparse_vs_dense(nx):
     n2 = 150
     k = 2
 
-    M_sparse, M_dense = ot.utils.get_sparse_test_matrices(n1, n2, k=k, seed=43, nx=nx)
+    M_sparse, M_dense = _get_sparse_test_matrices(n1, n2, k=k, seed=43, nx=nx)
 
     a = ot.utils.unif(n1, type_as=M_dense)
     b = ot.utils.unif(n2, type_as=M_dense)