code: # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  from dataclasses import dataclass
  from typing import Any
  
  import numpy as np
  from sklearn.utils import check_array
  from sklearn.utils.extmath import randomized_svd
  
  from .common import masked_mae
  from .solver import Solver
  
  F32PREC = np.finfo(np.float32).eps
  
  
  class SoftImpute(Solver):
      """
      Implementation of the SoftImpute algorithm from:
      "Spectral Regularization Algorithms for Learning Large Incomplete Matrices"
      by Mazumder, Hastie, and Tibshirani.
      """
  
      def __init__(
              self,
              shrinkage_value=None,
              convergence_threshold=0.001,
              max_iters=100,
              max_rank=None,
              n_power_iterations=1,
              init_fill_method="zero",
              X_init=None,
              min_value=None,
              max_value=None,
              normalizer=None,
              verbose=True):
          """
          Parameters
          ----------
          shrinkage_value : float
              Value by which we shrink singular values on each iteration. If
              omitted then the default value will be the maximum singular
              value of the initialized matrix (zeros for missing values) divided
              by 50.
  
          convergence_threshold : float
              Minimum ration difference between iterations (as a fraction of
              the Frobenius norm of the current solution) before stopping.
  
          max_iters : int
              Maximum number of SVD iterations
  
          max_rank : int, optional
              Perform a truncated SVD on each iteration with this value as its
              rank.
  
          n_power_iterations : int
              Number of power iterations to perform with randomized SVD
  
          init_fill_method : str
              How to initialize missing values of data matrix, default is
              to fill them with zeros.
  
          min_value : float
              Smallest allowable value in the solution
  
          max_value : float
              Largest allowable value in the solution
  
          normalizer : object
              Any object (such as BiScaler) with fit() and transform() methods
  
          verbose : bool
              Print debugging info
          """
          Solver.__init__(
              self,
              fill_method=init_fill_method,
              min_value=min_value,
              max_value=max_value,
              normalizer=normalizer,
              X_init=X_init)
          self.shrinkage_value = shrinkage_value
          self.convergence_threshold = convergence_threshold
          self.max_iters = max_iters
          self.max_rank = max_rank
          self.n_power_iterations = n_power_iterations
          self.verbose = verbose
  
      def _converged(self, X_old, X_new, missing_mask):
          # check for convergence
          old_missing_values = X_old[missing_mask]
          new_missing_values = X_new[missing_mask]
          difference = old_missing_values - new_missing_values
          ssd = np.sum(difference ** 2)
          old_norm = np.sqrt((old_missing_values ** 2).sum())
          # edge cases
          if old_norm == 0 or (old_norm < F32PREC and np.sqrt(ssd) > F32PREC):
              return False
          else:
              return (np.sqrt(ssd) / old_norm) < self.convergence_threshold
  
      def _svd_step(self, X, shrinkage_value, max_rank=None):
          """
          Returns reconstructed X from low-rank thresholded SVD and
          the rank achieved.
          """
          if max_rank:
              # if we have a max rank then perform the faster randomized SVD
              (U, s, V) = randomized_svd(
                  X,
                  max_rank,
                  n_iter=self.n_power_iterations)
          else:
              # perform a full rank SVD using ARPACK
              (U, s, V) = np.linalg.svd(
                  X,
                  full_matrices=False,
                  compute_uv=True)
          s_thresh = np.maximum(s - shrinkage_value, 0)
          rank = (s_thresh > 0).sum()
          s_thresh = s_thresh[:rank]
          U_thresh = U[:, :rank]
          V_thresh = V[:rank, :]
          S_thresh = np.diag(s_thresh)
          X_reconstruction = np.dot(U_thresh, np.dot(S_thresh, V_thresh))
          return X_reconstruction, rank
  
      def _max_singular_value(self, X_filled):
          # quick decomposition of X_filled into rank-1 SVD
          _, s, _ = randomized_svd(
              X_filled,
              1,
              n_iter=5)
          return s[0]
  
      def max_singular_value(self, X):
          X_original, missing_mask = self.prepare_input_data(X)
          X = X_original.copy()
          X_filled = self.fill(X, missing_mask, inplace=True)
          return self._max_singular_value(X_filled), missing_mask
  
      def solve(self, X, missing_mask):
          X = check_array(X, force_all_finite=False)
  
          X_init = X.copy()
  
          X_filled = X
          observed_mask = ~missing_mask
          max_singular_value = self._max_singular_value(X_filled)
          if self.verbose:
              print("[SoftImpute] Max Singular Value of X_init = %f" % (
                  max_singular_value))
  
          if self.shrinkage_value:
              shrinkage_value = self.shrinkage_value
          else:
              # totally hackish heuristic: keep only components
              # with at least 1/50th the max singular value
              shrinkage_value = max_singular_value / 50.0
  
          for i in range(self.max_iters):
              X_reconstruction, rank = self._svd_step(
                  X_filled,
                  shrinkage_value,
                  max_rank=self.max_rank)
              X_reconstruction = self.clip(X_reconstruction)
  
              # print error on observed data
              if self.verbose:
                  mae = masked_mae(
                      X_true=X_init,
                      X_pred=X_reconstruction,
                      mask=observed_mask)
                  print(
                      "[SoftImpute] Iter %d: observed MAE=%0.6f rank=%d" % (
                          i + 1,
                          mae,
                          rank))
  
              converged = self._converged(
                  X_old=X_filled,
                  X_new=X_reconstruction,
                  missing_mask=missing_mask)
              X_filled[missing_mask] = X_reconstruction[missing_mask]
              if converged:
                  break
          if self.verbose:
              print("[SoftImpute] Stopped after iteration %d for lambda=%f" % (
                  i + 1,
                  shrinkage_value))
  
          return X_filled
  
  
  @dataclass
  class SoftImputeWarmStarts(object):
  
      max_rank: int = None
      min_value: float = None
      max_value: float = None
      normalizer: Any = None
      n_power_iterations: int = 1
      start_shrinkage_value: float = None
      verbose: bool = False
      init_fill_method: str = "zero"
      convergence_threshold: float = 0.001
      max_iters: int = 100
      shrinkage_values_number: int = 2
  
      def fit_transform(self, X):
          solver = SoftImpute(convergence_threshold=self.convergence_threshold,
                              max_iters=self.max_iters,
                              max_rank=self.max_rank,
                              n_power_iterations=self.n_power_iterations,
                              init_fill_method=self.init_fill_method,
                              min_value=self.min_value,
                              max_value=self.max_value,
                              normalizer=self.normalizer
                              )
          X = check_array(X, force_all_finite=False)
          max_shrinkage_value, missing_mask = solver.max_singular_value(X)
          observed_mask = ~missing_mask
          if not (self.start_shrinkage_value and self.start_shrinkage_value < max_shrinkage_value):
              self.start_shrinkage_value = max_shrinkage_value
          shrinkage_values = np.exp(np.linspace(start=np.log(self.start_shrinkage_value), stop=np.log(1), num=self.shrinkage_values_number))
          solver.shrinkage_value = self.start_shrinkage_value
          X_init = solver.fit_transform(X).copy()
          result = [(self.start_shrinkage_value, X_init)]
          for sv in shrinkage_values[1:]:
              X_start = X_init.copy()
              if self.verbose:
                  print("shrinkage value %s" % sv)
              solver = SoftImpute(convergence_threshold=self.convergence_threshold,
                                  max_iters=self.max_iters,
                                  max_rank=self.max_rank,
                                  n_power_iterations=self.n_power_iterations,
                                  min_value=self.min_value,
                                  max_value=self.max_value,
                                  normalizer=self.normalizer,
                                  X_init=X_start,
                                  shrinkage_value=sv,
                                  )
              X_out = solver.solve(X_start, missing_mask)
              # TODO project
              X_out[observed_mask] = X[observed_mask]
              result.append([sv, X_out])
              X_init = np.copy(X_out)
          return result
  
apis: sklearn.utils.extmath.randomized_svd sklearn.utils.check_array