long EM/RPCA operations wrapped with tqdm

Author: Julien Roussel

pyproject.toml

@@ -45,6 +45,7 @@ statsmodels = ">= 0.14.0"
 typed-ast = { version = "*", optional = true }
 category-encoders = "^2.6.3"
 dcor = ">= 0.6"
+tqdm = "^0.0.1"
 
 [tool.poetry.group.torch.dependencies]
 torch = "< 2.5"

qolmat/imputations/em_sampler.py

@@ -11,6 +11,7 @@
 from scipy import optimize as spo
 from sklearn import utils as sku
 from sklearn.base import BaseEstimator, TransformerMixin
+from tqdm import tqdm
 
 from qolmat.utils import utils
 from qolmat.utils.utils import RandomSetting
@@ -433,7 +434,9 @@ def fit_X(self, X: NDArray) -> None:
 
         X = self._maximize_likelihood(X_imp, mask_na)
 
-        for iter_em in range(self.max_iter_em):
+        for iter_em in tqdm(
+            range(self.max_iter_em), desc="EM parameters estimation"
+        ):
             X = self._sample_ou(X, mask_na)
 
             self.combine_parameters()

qolmat/imputations/rpca/rpca_noisy.py

@@ -11,6 +11,7 @@
 from scipy.sparse import dok_matrix, identity
 from scipy.sparse.linalg import spsolve
 from sklearn import utils as sku
+from tqdm import tqdm
 
 from qolmat.imputations.rpca import rpca_utils
 from qolmat.imputations.rpca.rpca import RPCA
@@ -311,7 +312,9 @@ def minimise_loss(
         Ir = np.eye(rank)
         In = identity(n_rows)
 
-        for _ in range(max_iterations):
+        for _ in tqdm(
+            range(max_iterations), desc="Noisy RPCA loss minimization"
+        ):
             M_temp = M.copy()
             A_temp = A.copy()
             L_temp = L.copy()

qolmat/imputations/rpca/rpca_pcp.py

@@ -8,6 +8,7 @@
 import numpy as np
 from numpy.typing import NDArray
 from sklearn import utils as sku
+from tqdm import tqdm
 
 from qolmat.imputations.rpca import rpca_utils
 from qolmat.imputations.rpca.rpca import RPCA
@@ -125,7 +126,9 @@ def decompose(self, D: NDArray, Omega: NDArray) -> Tuple[NDArray, NDArray]:
         errors: NDArray = np.full((self.max_iterations,), fill_value=np.nan)
 
         M: NDArray = D - A
-        for iteration in range(self.max_iterations):
+        for iteration in tqdm(
+            range(self.max_iterations), desc="RPCA PCP decomposition"
+        ):
             M = rpca_utils.svd_thresholding(D - A + Y / mu, 1 / mu)
             A = rpca_utils.soft_thresholding(D - M + Y / mu, lam / mu)
             A[~Omega] = (D - M)[~Omega]

qolmat/imputations/softimpute.py

@@ -10,6 +10,7 @@
 from numpy.typing import NDArray
 from sklearn import utils as sku
 from sklearn.base import BaseEstimator, TransformerMixin
+from tqdm import tqdm
 
 from qolmat.imputations.rpca import rpca_utils
 from qolmat.utils import utils
@@ -146,7 +147,9 @@ def decompose(self, X: NDArray, Omega: NDArray) -> Tuple[NDArray, NDArray]:
         B = V * D
         M = A @ B.T
         cost_start = SoftImpute.cost_function(X, M, A, Omega, tau)
-        for iter_ in range(self.max_iterations):
+        for iter_ in tqdm(
+            range(self.max_iterations), desc="Soft Impute decomposition"
+        ):
             U_old = U
             V_old = V
             D_old = D