testing n_rows not divising X size

Author: vm-aifluence-jro

qolmat/imputations/rpca/rpca.py

@@ -58,8 +58,10 @@ def _prepare_data(self, X: NDArray) -> NDArray:
                 raise ValueError("`n_rows` should not be specified when imputing 2D data.")
 
     def get_shape_original(self, X: NDArray, shape: Tuple[int]) -> NDArray:
-        if len(shape) == 1 or shape[0] == 1:
-            n_values = sum(shape)
-            return X.reshape(1, -1)[:, :n_values]
-        else:
-            return X
+        # if len(shape) == 1 or shape[0] == 1:
+        #     n_values = sum(shape)
+        #     return X.reshape(1, -1)[:, :n_values]
+        # else:
+        #     return X
+        X = X.flatten()[: np.prod(shape)]
+        return X.reshape(shape)

tests/imputations/rpca/test_rpca.py

@@ -58,6 +58,7 @@ def test_rpca__prepare_data_1D_succeed(X: NDArray):
     [
         (2, X_flat, X_exp_nrows_2_pcp_decompose_rpca_signal),
         (3, X_flat, X_exp_nrows_3_pcp_decompose_rpca_signal),
+        (4, X_flat, X_exp_nrows_3_pcp_decompose_rpca_signal),
         (None, X_incomplete, X_exp_nrows_1_pcp_decompose_rpca_signal),
     ],
 )