code: """Template unit tests for scikit-learn estimators."""
  
  import pytest
  from sklearn.datasets import load_iris
  
  import geomstats.backend as gs
  import geomstats.tests
  from geomstats.learning._template import (
      TemplateClassifier,
      TemplateEstimator,
      TemplateTransformer,
  )
  
  ESTIMATORS = (TemplateClassifier, TemplateEstimator, TemplateTransformer)
  
  
  class TestEstimators(geomstats.tests.TestCase):
      _multiprocess_can_split_ = True
  
      def setup_method(self):
          self.data = load_iris(return_X_y=True)
  
      @geomstats.tests.np_and_autograd_only
      def test_template_estimator(self):
          est = TemplateEstimator()
          self.assertEqual(est.demo_param, "demo_param")
  
          X, y = self.data
  
          est.fit(X, y)
          self.assertTrue(hasattr(est, "is_fitted_"))
  
          y_pred = est.predict(X)
          self.assertAllClose(y_pred, gs.ones(gs.shape(X)[0]))
  
      @geomstats.tests.np_and_autograd_only
      def test_template_transformer_error(self):
          X, _ = self.data
          n_samples = gs.shape(X)[0]
          trans = TemplateTransformer()
          trans.fit(X)
          X_diff_size = gs.ones((n_samples, gs.shape(X)[1] + 1))
          with pytest.raises(ValueError):
              trans.transform(X_diff_size)
  
      def test_template_transformer(self):
          X, _ = self.data
          trans = TemplateTransformer()
          self.assertTrue(trans.demo_param == "demo")
  
          trans.fit(X)
          self.assertTrue(trans.n_features_ == X.shape[1])
  
          X_trans = trans.transform(X)
          self.assertAllClose(X_trans, gs.sqrt(X))
  
          X_trans = trans.fit_transform(X)
          self.assertAllClose(X_trans, gs.sqrt(X))
  
      @geomstats.tests.np_autograd_and_tf_only
      def test_template_classifier(self):
          X, y = self.data
          clf = TemplateClassifier()
          self.assertTrue(clf.demo_param == "demo")
  
          clf.fit(X, y)
          self.assertTrue(hasattr(clf, "classes_"))
          self.assertTrue(hasattr(clf, "X_"))
          self.assertTrue(hasattr(clf, "y_"))
  
          y_pred = clf.predict(X)
          self.assertTrue(y_pred.shape == (X.shape[0],))
   """Learning embedding of graph using Poincare Ball Model."""
  
  import matplotlib.patches as mpatches
  import matplotlib.pyplot as plt
  
  import geomstats.backend as gs
  import geomstats.visualization as visualization
  from geomstats.datasets.prepare_graph_data import HyperbolicEmbedding
  from geomstats.datasets.utils import load_karate_graph
  from geomstats.learning.kmeans import RiemannianKMeans
  from geomstats.learning.kmedoids import RiemannianKMedoids
  
  
  def main():
      """Learning Poincaré graph embedding.
  
      Learns Poincaré Ball embedding by using Riemannian
      gradient descent algorithm. Then K-means is applied
      to learn labels of each data sample.
      """
      gs.random.seed(1234)
  
      karate_graph = load_karate_graph()
      hyperbolic_embedding = HyperbolicEmbedding(max_epochs=3)
      embeddings = hyperbolic_embedding.embed(karate_graph)
  
      colors = {1: "b", 2: "r"}
      group_1 = mpatches.Patch(color=colors[1], label="Group 1")
      group_2 = mpatches.Patch(color=colors[2], label="Group 2")
  
      circle = visualization.PoincareDisk(point_type="ball")
  
      _, ax = plt.subplots(figsize=(8, 8))
      ax.axes.xaxis.set_visible(False)
      ax.axes.yaxis.set_visible(False)
      circle.set_ax(ax)
      circle.draw(ax=ax)
      for i_embedding, embedding in enumerate(embeddings):
          x_coords = embedding[0]
          y_coords = embedding[1]
          pt_id = i_embedding
          plt.scatter(x_coords, y_coords, c=colors[karate_graph.labels[pt_id][0]], s=150)
          ax.annotate(pt_id, (x_coords, y_coords))
  
      plt.tick_params(which="both")
      plt.title("Poincare Ball Embedding of the Karate Club Network")
      plt.legend(handles=[group_1, group_2])
      plt.show()
  
      n_clusters = 2
  
      kmeans = RiemannianKMeans(
          metric=hyperbolic_embedding.manifold.metric,
          n_clusters=n_clusters,
          init="random",
      )
  
      centroids = kmeans.fit(X=embeddings)
      labels = kmeans.predict(X=embeddings)
  
      colors = ["g", "c", "m"]
      circle = visualization.PoincareDisk(point_type="ball")
      _, ax2 = plt.subplots(figsize=(8, 8))
      circle.set_ax(ax2)
      circle.draw(ax=ax2)
      ax2.axes.xaxis.set_visible(False)
      ax2.axes.yaxis.set_visible(False)
      group_1_predicted = mpatches.Patch(color=colors[0], label="Predicted Group 1")
      group_2_predicted = mpatches.Patch(color=colors[1], label="Predicted Group 2")
      group_centroids = mpatches.Patch(color=colors[2], label="Cluster centroids")
  
      for _ in range(n_clusters):
          for i_embedding, embedding in enumerate(embeddings):
              x_coords = embedding[0]
              y_coords = embedding[1]
              pt_id = i_embedding
              if labels[i_embedding] == 0:
                  color = colors[0]
              else:
                  color = colors[1]
              plt.scatter(x_coords, y_coords, c=color, s=150)
              ax2.annotate(pt_id, (x_coords, y_coords))
  
      for _, centroid in enumerate(centroids):
          x_coords = centroid[0]
          y_coords = centroid[1]
          plt.scatter(
              x_coords,
              y_coords,
              c=colors[2],
              marker="*",
              s=150,
          )
  
      plt.title("K-means applied to Karate club embedding")
      plt.legend(handles=[group_1_predicted, group_2_predicted, group_centroids])
      plt.show()
  
      kmedoid = RiemannianKMedoids(
          metric=hyperbolic_embedding.manifold.metric,
          n_clusters=n_clusters,
          init="random",
          n_jobs=2,
      )
  
      centroids = kmedoid.fit(data=embeddings, max_iter=100)
      labels = kmedoid.predict(data=embeddings)
  
      colors = ["g", "c", "m"]
      circle = visualization.PoincareDisk(point_type="ball")
      _, ax2 = plt.subplots(figsize=(8, 8))
      circle.set_ax(ax2)
      circle.draw(ax=ax2)
      ax2.axes.xaxis.set_visible(False)
      ax2.axes.yaxis.set_visible(False)
      group_1_predicted = mpatches.Patch(color=colors[0], label="Predicted Group 1")
      group_2_predicted = mpatches.Patch(color=colors[1], label="Predicted Group 2")
      group_centroids = mpatches.Patch(color=colors[2], label="Cluster centroids")
  
      for _ in range(n_clusters):
          for i_embedding, embedding in enumerate(embeddings):
              x_coords = embedding[0]
              y_coords = embedding[1]
              pt_id = i_embedding
              if labels[i_embedding] == 0:
                  color = colors[0]
              else:
                  color = colors[1]
              plt.scatter(x_coords, y_coords, c=color, s=150)
              ax2.annotate(pt_id, (x_coords, y_coords))
  
      for _, centroid in enumerate(centroids):
          x_coords = centroid[0]
          y_coords = centroid[1]
          plt.scatter(
              x_coords,
              y_coords,
              c=colors[2],
              marker="*",
              s=150,
          )
  
      plt.title("K-Medoids applied to Karate club embedding")
      plt.legend(handles=[group_1_predicted, group_2_predicted, group_centroids])
      plt.show()
  
  
  if __name__ == "__main__":
      main()
  
apis: sklearn.datasets.load_iris
apis: matplotlib.pyplot.legend matplotlib.pyplot.tick_params matplotlib.pyplot.subplots matplotlib.pyplot.title matplotlib.pyplot.show matplotlib.patches.Patch matplotlib.pyplot.scatter