text: """
  ADMM Lasso
  
  @Authors: Aleksandar Armacki and Lidija Fodor
  @Affiliation: Faculty of Sciences, University of Novi Sad, Serbia
  
  This work is supported by the I-BiDaaS project, funded by the European
  Commission under Grant Agreement No. 780787.
  """
  
  try:
      import cvxpy as cp
  except ImportError:
      import warnings
      warnings.warn('Cannot import cvxpy module. ADMM estimator will not work.')
  import numpy as np
  from pycompss.api.api import compss_wait_on
  from pycompss.api.parameter import Type, Depth, COLLECTION_IN, COLLECTION_OUT
  from pycompss.api.task import task
  from sklearn.base import BaseEstimator
  
  import dislib as ds
  from dislib.data.array import Array
  from dislib.utils.base import _paired_partition
  
  
  class ADMM(BaseEstimator):
      """ Alternating Direction Method of Multipliers (ADMM) solver. ADMM is
      renowned for being well suited to the distributed settings [1]_, for its
      guaranteed convergence and general robustness with respect
      to the parameters. Additionally, the algorithm has a generic form that
      can be easily adapted to a wide range of machine learning problems with
      only minor tweaks in the code.
  
      Parameters
      ----------
      loss_fn : func
          Loss function.
      k : float
          Soft thresholding value.
      rho : float, optional (default=1)
          The penalty parameter for constraint violation.
      max_iter : int, optional (default=100)
          Maximum number of iterations to perform.
      atol : float, optional (default=1e-4)
          The absolute tolerance used to calculate the early stop criterion.
      rtol : float, optional (default=1e-2)
          The relative tolerance used to calculate the early stop criterion.
      verbose : boolean, optional (default=False)
          Whether to print information about the optimization process.
  
      Attributes
      ----------
      z_ : ds-array shape=(1, n_features)
          Computed z.
      n_iter_ : int
          Number of iterations performed.
      converged_ : boolean
          Whether the optimization converged.
  
      References
      ----------
      .. [1] S. Boyd, N. Parikh, E. Chu, B. Peleato, and J. Eckstein (2011).
          Distributed Optimization and Statistical Learning via the Alternating
          Direction Method of Multipliers. In Foundations and Trends in Machine
          Learning, 3(1):1–122.
      """
  
      def __init__(self, loss_fn, k, rho=1, max_iter=100, rtol=1e-2, atol=1e-4,
                   verbose=False):
          self.rho = rho
          self.atol = atol
          self.rtol = rtol
          self.loss_fn = loss_fn
          self.k = k
          self.max_iter = max_iter
          self.verbose = verbose
  
      def fit(self, x, y):
          """
          Fits the model with training data.
  
          Parameters
          ----------
          x : ds-array, shape=(n_samples, n_features)
              Training samples.
          y : ds-array, shape=(n_samples, 1)
              Class labels of x.
  
          Returns
          -------
          self : ADMM
          """
          if not x._is_regular():
              x_reg = x.rechunk(x._reg_shape)
          else:
              x_reg = x
  
          self._init_model(x_reg)
  
          while not self.converged_ and self.n_iter_ < self.max_iter:
              self._step(x_reg, y)
              self.n_iter_ += 1
  
              if self.verbose:
                  print("Iteration ", self.n_iter_)
  
          z_blocks = [object() for _ in range(x_reg._n_blocks[1])]
          _split_z(self._z, x._reg_shape[1], z_blocks)
          self.z_ = Array([z_blocks], (1, x._reg_shape[1]), (1, x._reg_shape[1]),
                          (1, x.shape[1]), False)
  
          return self
  
      def _init_model(self, x):
          n_features = x.shape[1]
  
          self.converged_ = False
          self.n_iter_ = 0
          self._z = np.zeros(n_features)
          # u has one row per each row-block in x
          self._u = ds.zeros((x._n_blocks[0], n_features), (1, x._reg_shape[1]))
  
      def _step(self, x, y):
          # update w
          self._w_step(x, y)
  
          z_old = self._z
  
          # update z
          self._z_step()
  
          # update u
          self._u_step()
  
          # after norm in axis=1 and sum in axis=0, these should be ds-arrays
          # of a single element, so we keep the only block
          nxstack = (self._w.norm(axis=1) ** 2).sum().sqrt()
          nystack = (self._u.norm(axis=1) ** 2).sum().sqrt()
  
          # termination check
          n_samples, n_features = self._u.shape
          dualres = _compute_dual_res(n_samples, self.rho, self._z, z_old)
          prires = self._compute_primal_res(z_old)
          n_total = n_samples * n_features
  
          self.converged_ = _check_convergence(prires._blocks[0][0], dualres,
                                               n_samples, n_total,
                                               nxstack._blocks[0][0],
                                               nystack._blocks[0][0],
                                               self.atol, self.rtol, self._z)
          self.converged_ = compss_wait_on(self.converged_)
  
      def _compute_primal_res(self, z_old):
          blocks = []
  
          for w_hblock in self._w._iterator():
              out_blocks = [object() for _ in range(self._w._n_blocks[1])]
              _substract(w_hblock._blocks, z_old, out_blocks)
              blocks.append(out_blocks)
  
          prires = Array(blocks, self._w._reg_shape, self._w._reg_shape,
                         self._w.shape, self._w._sparse)
  
          # this should be a ds-array of a single element. We return only the
          # block
          return (prires.norm(axis=1) ** 2).sum().sqrt()
  
      def _u_step(self):
          u_blocks = []
  
          for u_hblock, w_hblock in zip(self._u._iterator(),
                                        self._w._iterator()):
              out_blocks = [object() for _ in range(self._u._n_blocks[1])]
              _update_u(self._z, u_hblock._blocks, w_hblock._blocks, out_blocks)
              u_blocks.append(out_blocks)
  
          r_shape = self._u._reg_shape
          shape = self._u.shape
          self._u = Array(u_blocks, r_shape, r_shape, shape, self._u._sparse)
  
      def _z_step(self):
          w_mean = self._w.mean(axis=0)
          u_mean = self._u.mean(axis=0)
          self._z = _soft_thresholding(w_mean._blocks, u_mean._blocks, self.k)
  
      def _w_step(self, x, y):
          w_blocks = []
  
          for xy_hblock, u_hblock in zip(_paired_partition(x, y),
                                         self._u._iterator()):
              x_hblock, y_hblock = xy_hblock
              w_hblock = [object() for _ in range(x._n_blocks[1])]
              x_blocks = x_hblock._blocks
              y_blocks = y_hblock._blocks
              u_blocks = u_hblock._blocks
  
              _update_w(x_blocks, y_blocks, self._z, u_blocks, self.rho,
                        self.loss_fn, w_hblock)
              w_blocks.append(w_hblock)
  
          r_shape = self._u._reg_shape
          self._w = Array(w_blocks, r_shape, r_shape, self._u.shape, x._sparse)
  
  
  @task(z_blocks={Type: COLLECTION_OUT, Depth: 1})
  def _split_z(z, block_size, z_blocks):
      for i in range(len(z_blocks)):
          z_blocks[i] = z[i * block_size: (i + 1) * block_size]
  
  
  @task(x_blocks={Type: COLLECTION_IN, Depth: 2},
        y_blocks={Type: COLLECTION_IN, Depth: 2},
        u_blocks={Type: COLLECTION_IN, Depth: 2},
        w_blocks={Type: COLLECTION_OUT, Depth: 1})
  def _update_w(x_blocks, y_blocks, z, u_blocks, rho, loss, w_blocks):
      x_np = Array._merge_blocks(x_blocks)
      y_np = np.squeeze(Array._merge_blocks(y_blocks))
      u_np = np.squeeze(Array._merge_blocks(u_blocks))
  
      w_new = cp.Variable(x_np.shape[1])
  
      problem = cp.Problem(cp.Minimize(_objective(loss, x_np, y_np, w_new, z,
                                                  u_np, rho)))
      problem.solve()
      status = problem.status
  
      if 'infeasible' in status or 'unbounded' in status:
          raise Exception("Cannot solve the problem. CVXPY status: %s" % status)
  
      w_np = w_new.value
      n_cols = x_blocks[0][0].shape[1]
  
      for i in range(len(w_blocks)):
          w_blocks[i] = w_np[i * n_cols:(i + 1) * n_cols].reshape(1, -1)
  
  
  def _objective(loss, x, y, w, z, u, rho):
      reg = cp.norm(w - z + u, p=2) ** 2
      return loss(x, y, w) + (rho / 2) * reg
  
  
  @task(w_blocks={Type: COLLECTION_IN, Depth: 2},
        u_blocks={Type: COLLECTION_IN, Depth: 2},
        returns=np.array)
  def _soft_thresholding(w_blocks, u_blocks, k):
      w_mean = np.squeeze(Array._merge_blocks(w_blocks))
      u_mean = np.squeeze(Array._merge_blocks(u_blocks))
      v = w_mean + u_mean
  
      z = np.zeros(v.shape)
      for i in range(z.shape[0]):
          if np.abs(v[i]) <= k:
              z[i] = 0
          else:
              if v[i] > k:
                  z[i] = v[i] - k
              else:
                  z[i] = v[i] + k
      return z
  
  
  @task(u_blocks={Type: COLLECTION_IN, Depth: 2},
        w_blocks={Type: COLLECTION_IN, Depth: 2},
        out_blocks={Type: COLLECTION_OUT, Depth: 1})
  def _update_u(z, u_blocks, w_blocks, out_blocks):
      u_np = np.squeeze(Array._merge_blocks(u_blocks))
      w_np = np.squeeze(Array._merge_blocks(w_blocks))
      u_new = u_np + w_np - z
      n_cols = u_blocks[0][0].shape[1]
  
      for i in range(len(out_blocks)):
          out_blocks[i] = u_new[i * n_cols: (i + 1) * n_cols].reshape(1, -1)
  
  
  @task(returns=1)
  def _compute_dual_res(n_samples, rho, z, z_old):
      return np.sqrt(n_samples) * rho * np.linalg.norm(z - z_old)
  
  
  @task(blocks={Type: COLLECTION_IN, Depth: 2},
        out_blocks={Type: COLLECTION_OUT, Depth: 1})
  def _substract(blocks, z, out_blocks):
      w_np = Array._merge_blocks(blocks) - z
      n_cols = blocks[0][0].shape[1]
  
      for i in range(len(out_blocks)):
          out_blocks[i] = w_np[i * n_cols: (i + 1) * n_cols].reshape(1, -1)
  
  
  @task(returns=bool)
  def _check_convergence(prires, dualres, n_samples, n_total, nxstack,
                         nystack, abstol, reltol, z):
      eps_pri = (np.sqrt(n_total)) * abstol + reltol * (
          max(nxstack, np.sqrt(n_samples) * np.linalg.norm(z)))
      eps_dual = np.sqrt(n_total) * abstol + reltol * nystack
  
      if prires <= eps_pri and dualres <= eps_dual:
          return True
  
      return False
  