Performance enhancement

optimask/_optimask.py

@@ -141,7 +141,7 @@ def _preprocess(x):
         """
         m, n = x.shape
         iy, ix = np.empty(m * n, dtype=np.uint32), np.empty(m * n, dtype=np.uint32)
-        cols_index_mapper = -np.ones(n, dtype=np.int16)
+        cols_index_mapper = -np.ones(n, dtype=np.int32)
         rows_with_nan = np.empty(m, dtype=np.uint32)
         n_rows_with_nan = 0
         n_cols_with_nan = 0
@@ -172,11 +172,17 @@ def _preprocess(x):
         ].astype(np.uint32)
         return iy, ix, rows_with_nan, cols_with_nan
 
-    def _trial(self, rng, m_nan, n_nan, iy, ix, m, n):
-        p_rows = rng.permutation(m_nan).astype(np.uint32)
-        p_cols = rng.permutation(n_nan).astype(np.uint32)
-        iy_trial = self.apply_permutation(p_rows, iy, inplace=False)
-        ix_trial = self.apply_permutation(p_cols, ix, inplace=False)
+    def _trial(self, k, rng, m_nan, n_nan, iy, ix, m, n):
+        if k:
+            p_rows = rng.permutation(m_nan).astype(np.uint32)
+            p_cols = rng.permutation(n_nan).astype(np.uint32)
+            iy_trial = self.apply_permutation(p_rows, iy, inplace=False)
+            ix_trial = self.apply_permutation(p_cols, ix, inplace=False)
+        else:
+            p_rows = np.arange(m_nan, dtype=np.uint32)
+            p_cols = np.arange(n_nan, dtype=np.uint32)
+            iy_trial = iy.copy()
+            ix_trial = ix.copy()
 
         hy = self.groupby_max(iy_trial, ix_trial, m_nan)
         step = 0
@@ -271,7 +277,7 @@ def _solve(self, x):
             rng = np.random.default_rng(seed=self.random_state)
             area_max = -1
             for k in range(self.n_tries):
-                area, i0, j0, p_rows, p_cols = self._trial(rng, m_nan, n_nan, iy, ix, m, n)
+                area, i0, j0, p_rows, p_cols = self._trial(k, rng, m_nan, n_nan, iy, ix, m, n)
                 self._verbose(f"\tTrial {k + 1} : submatrix of size {m - j0}x{n - i0} ({area} elements) found.")
                 if area > area_max:
                     area_max = area

tests/test_optimask.py

@@ -115,14 +115,10 @@ def test_seed():
     X = generate_random(m=10_000, n=500, ratio=0.025)
     om1 = OptiMask(random_state=99)
     om2 = OptiMask(random_state=99)
-    om3 = OptiMask(random_state=999)
     rows1, cols1 = om1.solve(X, check_result=True)
     rows2, cols2 = om2.solve(X, check_result=True)
-    rows3, cols3 = om3.solve(X, check_result=True)
     assert np.allclose(rows1, rows2)
     assert np.allclose(cols1, cols2)
-    assert (len(rows1) != len(rows3)) or (not np.allclose(rows1, rows3))
-    assert (len(cols1) != len(cols3)) or (not np.allclose(cols1, cols3))
 
 
 def test_speed(opti_mask_instance):
@@ -132,3 +128,11 @@ def test_speed(opti_mask_instance):
         start = perf_counter()
         opti_mask_instance.solve(X=x)
         print(f"{1e3 * (perf_counter() - start):.2f}ms")
+
+
+def test_large_arrays(opti_mask_instance):
+    x = generate_random(m=100_000, n=1_000, ratio=0.02)
+    opti_mask_instance.solve(x)
+
+    x = generate_random(m=1_000, n=100_000, ratio=0.02)
+    opti_mask_instance.solve(x)