text: from random import randint
  
  from sklearn.base import BaseEstimator, ClusterMixin
  
  import geomstats.backend as gs
  from geomstats.learning._template import TransformerMixin
  
  
  class RiemannianKMeans(TransformerMixin, ClusterMixin, BaseEstimator):
  
      def __init__(self, riemannian_metric, n_clusters=8, init='random',
                   tol=1e-2, mean_method='default', verbose=0):
          """ K-Means algorithm using Riemannian manifolds
  
          Parameters
          ----------
          n_clusters : Number of clusters (k value of the k-means)
  
          riemannian_metric : The geomstats riemmanian metric associate to
                              the space used
  
          init : How to init centroids at the beginning of the algorithm.
                 'random' : will select random uniformally train point as
                           initial centroids
  
          tol : convergence factor. If the difference of mean distance
               between two step is lower than tol
  
          verbose : if verbose > 0, information will be print during learning
  
          Returns
          -------
          self : object
              Returns the instance itself.
          """
          self.n_clusters = n_clusters
          self.init = init
          self.riemannian_metric = riemannian_metric
          self.tol = tol
          self.verbose = verbose
          self.mean_method = mean_method
  
      def fit(self, X, max_iter=100):
          """Predict for each data point the closest center in terms of
              riemannian_metric distance
  
          Parameters
          ----------
          X : array-like, shape=[n_samples, n_features]
              Training data, where n_samples is the number of samples
              and n_features is the number of features.
  
          max_iter : Maximum number of iterations
  
          Returns
          -------
          self : object
              Return centroids array
          """
  
          n_samples = X.shape[0]
          belongs = gs.zeros(n_samples)
          self.centroids = [gs.expand_dims(X[randint(0, n_samples - 1)], 0)
                            for i in range(self.n_clusters)]
          self.centroids = gs.concatenate(self.centroids)
          index = 0
          while index < max_iter:
              index += 1
  
              dists = [gs.to_ndarray(
                       self.riemannian_metric.dist(self.centroids[i], X), 2, 1)
                       for i in range(self.n_clusters)]
              dists = gs.hstack(dists)
              belongs = gs.argmin(dists, 1)
              old_centroids = gs.copy(self.centroids)
              for i in range(self.n_clusters):
                  fold = gs.squeeze(X[belongs == i])
  
                  if len(fold) > 0:
  
                      self.centroids[i] = self.riemannian_metric.mean(
                          fold,
                          mean_method=self.mean_method,
                          n_max_iterations=150)
  
                  else:
                      self.centroids[i] = X[randint(0, n_samples - 1)]
  
              centroids_distances = self.riemannian_metric.dist(old_centroids,
                                                                self.centroids)
  
              if gs.mean(centroids_distances) < self.tol:
                  if self.verbose > 0:
                      print("Convergence Reached after ", index, " iterations")
  
                  return gs.copy(self.centroids)
  
          if index == max_iter:
              print('K-means maximum number of iterations {} reached.'
                    'The mean may be inaccurate'.format(max_iter))
  
          return gs.copy(self.centroids)
  
      def predict(self, X):
  
          """Predict for each data point the closest center in terms of
              riemannian_metric distance
  
          Parameters
          ----------
          X : array-like, shape=[n_samples, n_features]
              data, where n_samples is the number of samples
              and n_features is the number of features.
  
          Returns
          -------
          self : object
              Return array containing for each point the cluster associated
          """
          dists = gs.hstack([self.riemannian_metric.dist(self.centroids[i], X)
                             for i in range(self.n_clusters)])
          belongs = gs.argmin(dists, -1)
          return belongs
  