Typing isse solved

Author: vm-aifluence-jro

qolmat/imputations/rpca/rpca.py

@@ -57,11 +57,21 @@ def _prepare_data(self, X: NDArray) -> NDArray:
             else:
                 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
-        X = X.flatten()[: np.prod(shape)]
-        return X.reshape(shape)
+    def get_shape_original(self, M: NDArray, X: NDArray) -> NDArray:
+        """Shapes an output matrix from the RPCA algorithm into the original shape.
+
+        Parameters
+        ----------
+        M : NDArray
+            Matrix to reshape
+        X : NDArray
+            Matrix of the desired shape
+
+        Returns
+        -------
+        NDArray
+            Reshaped matrix
+        """
+        size = X.size
+        M_flat = M.flatten()[:size]
+        return M_flat.reshape(X.shape)

qolmat/imputations/rpca/rpca_noisy.py

@@ -366,7 +366,7 @@ def decompose_rpca_signal(
         elif self.norm == "L2":
             M, A, U, V, errors = self.decompose_rpca_L2(D_proj, Omega, lam, tau, rank)
 
-        M_final = self.get_shape_original(M, X.shape)
-        A_final = self.get_shape_original(A, X.shape)
+        M_final = self.get_shape_original(M, X)
+        A_final = self.get_shape_original(A, X)
 
         return M_final, A_final

qolmat/imputations/rpca/rpca_pcp.py

@@ -100,6 +100,6 @@ def decompose_rpca_signal(
         D = self._prepare_data(X)
         M, A = self.decompose_rpca(D)
 
-        M_final = self.get_shape_original(M, X.shape)
-        A_final = self.get_shape_original(A, X.shape)
+        M_final = self.get_shape_original(M, X)
+        A_final = self.get_shape_original(A, X)
         return M_final, A_final