text: import os
  import unittest
  
  import heat as ht
  import numpy as np
  import torch
  
  from ...core.tests.test_suites.basic_test import TestCase
  
  
  class TestKMedians(TestCase):
      def create_spherical_dataset(
          self, num_samples_cluster, radius=1.0, offset=4.0, dtype=ht.float32, random_state=1
      ):
          """
          Creates k=4 sperical clusters in 3D space along the space-diagonal
  
          Parameters
          ----------
          num_samples_cluster: int
              Number of samples per cluster. Each process will create n // MPI_WORLD.size elements for each cluster
          radius: float
              Radius of the sphere
          offset: float
              Shift of the clusters along the axes. The 4 clusters will be positioned centered around c1=(offset, offset,offset),
              c2=(2*offset,2*offset,2*offset), c3=(-offset, -offset, -offset) and c4=(2*offset, -2*offset, -2*offset)
          dtype: ht.datatype
          random_state: int
              seed of the torch random number generator
          """
          # contains num_samples
  
          p = ht.MPI_WORLD.size
          # create k sperical clusters with each n elements per cluster. Each process creates k * n/p elements
          num_ele = num_samples_cluster // p
          ht.random.seed(random_state)
          # radius between 0 and 1
          r = ht.random.rand(num_ele, split=0) * radius
          # theta between 0 and pi
          theta = ht.random.rand(num_ele, split=0) * 3.1415
          # phi between 0 and 2pi
          phi = ht.random.rand(num_ele, split=0) * 2 * 3.1415
          # Cartesian coordinates
          x = r * ht.sin(theta) * ht.cos(phi)
          x.astype(dtype, copy=False)
          y = r * ht.sin(theta) * ht.sin(phi)
          y.astype(dtype, copy=False)
          z = r * ht.cos(theta)
          z.astype(dtype, copy=False)
  
          cluster1 = ht.stack((x + offset, y + offset, z + offset), axis=1)
          cluster2 = ht.stack((x + 2 * offset, y + 2 * offset, z + 2 * offset), axis=1)
          cluster3 = ht.stack((x - offset, y - offset, z - offset), axis=1)
          cluster4 = ht.stack((x - 2 * offset, y - 2 * offset, z - 2 * offset), axis=1)
  
          data = ht.concatenate((cluster1, cluster2, cluster3, cluster4), axis=0)
          # Note: enhance when shuffel is available
          return data
  
      def test_clusterer(self):
          kmedian = ht.cluster.KMedians()
          self.assertTrue(ht.is_estimator(kmedian))
          self.assertTrue(ht.is_clusterer(kmedian))
  
      def test_get_and_set_params(self):
          kmedian = ht.cluster.KMedians()
          params = kmedian.get_params()
  
          self.assertEqual(
              params,
              {"n_clusters": 8, "init": "random", "max_iter": 300, "tol": 1e-4, "random_state": None},
          )
  
          params["n_clusters"] = 10
          kmedian.set_params(**params)
          self.assertEqual(10, kmedian.n_clusters)
  
      def test_fit_iris_unsplit(self):
          split = 0
          # get some test data
          iris = ht.load("heat/datasets/iris.csv", sep=";", split=split)
  
          # fit the clusters
          k = 3
          kmedian = ht.cluster.KMedians(n_clusters=k)
          kmedian.fit(iris)
  
          # check whether the results are correct
          self.assertIsInstance(kmedian.cluster_centers_, ht.DNDarray)
          self.assertEqual(kmedian.cluster_centers_.shape, (k, iris.shape[1]))
          # same test with init=kmedians++
          kmedian = ht.cluster.KMedians(n_clusters=k, init="kmedians++")
          kmedian.fit(iris)
  
          # check whether the results are correct
          self.assertIsInstance(kmedian.cluster_centers_, ht.DNDarray)
          self.assertEqual(kmedian.cluster_centers_.shape, (k, iris.shape[1]))
  
      def test_exceptions(self):
          # get some test data
          iris_split = ht.load("heat/datasets/iris.csv", sep=";", split=1)
  
          # build a clusterer
          k = 3
          kmedian = ht.cluster.KMedians(n_clusters=k)
  
          with self.assertRaises(NotImplementedError):
              kmedian.fit(iris_split)
          with self.assertRaises(ValueError):
              kmedian.set_params(foo="bar")
          with self.assertRaises(ValueError):
              kmedian = ht.cluster.KMedians(n_clusters=k, init="random_number")
              kmedian.fit(iris_split)
  
      def test_spherical_clusters(self):
          seed = 1
          n = 20 * ht.MPI_WORLD.size
          data = self.create_spherical_dataset(
              num_samples_cluster=n, radius=1.0, offset=4.0, dtype=ht.float32, random_state=seed
          )
          kmedians = ht.cluster.KMedians(n_clusters=4, init="kmedians++")
          kmedians.fit(data)
          self.assertIsInstance(kmedians.cluster_centers_, ht.DNDarray)
          self.assertEqual(kmedians.cluster_centers_.shape, (4, 3))
  
          # More Samples
          n = 100 * ht.MPI_WORLD.size
          data = self.create_spherical_dataset(
              num_samples_cluster=n, radius=1.0, offset=4.0, dtype=ht.float32, random_state=seed
          )
          kmedians = ht.cluster.KMedians(n_clusters=4, init="kmedians++")
          kmedians.fit(data)
          self.assertIsInstance(kmedians.cluster_centers_, ht.DNDarray)
          self.assertEqual(kmedians.cluster_centers_.shape, (4, 3))
  
          # different datatype
          n = 20 * ht.MPI_WORLD.size
          data = self.create_spherical_dataset(
              num_samples_cluster=n, radius=1.0, offset=4.0, dtype=ht.float64, random_state=seed
          )
          kmedians = ht.cluster.KMedians(n_clusters=4, init="kmedians++")
          kmedians.fit(data)
          self.assertIsInstance(kmedians.cluster_centers_, ht.DNDarray)
          self.assertEqual(kmedians.cluster_centers_.shape, (4, 3))
  
          # on Ints (different radius, offset and datatype
          data = self.create_spherical_dataset(
              num_samples_cluster=n, radius=10.0, offset=40.0, dtype=ht.int32, random_state=seed
          )
          kmedians = ht.cluster.KMedians(n_clusters=4, init="kmedians++")
          kmedians.fit(data)
          self.assertIsInstance(kmedians.cluster_centers_, ht.DNDarray)
          self.assertEqual(kmedians.cluster_centers_.shape, (4, 3))
  