Repository: helmholtz-analytics/heat

Issue:
Runtime Error in cdist
**Description**
When running cdist on distributed tensors a runtime error occcurs.

Traceback (most recent call last):
  File "demo_knn.py", line 131, in <module>
    print(verify_algorithm(X, Y, 5, 30, 5))
  File "demo_knn.py", line 127, in verify_algorithm
    result_y = classifier.predict(verification_x)
  File "/home/jakob/neural/heat/heat/classification/knn.py", line 44, in predict
    distances = ht.spatial.cdist(X, self.x)
  File "/home/jakob/neural/heat/heat/spatial/distance.py", line 128, in cdist
    return _dist(X, Y, _euclidian)
  File "/home/jakob/neural/heat/heat/spatial/distance.py", line 383, in _dist
    d._DNDarray__array[:, cols[0] : cols[1]] = d_ij
RuntimeError: The expanded size of the tensor (8) must match the existing size (11) at non-singleton dimension 1.  Target sizes: [30, 8].  Tensor sizes: [27, 11]

**To Reproduce**
On branch https://github.com/helmholtz-analytics/heat/tree/features/556-assign_label
run "mpirun -n 4 demo_knn.py " (in folder heat/examples/classification)

**Version Info**
On branch https://github.com/helmholtz-analytics/heat/tree/features/556-assign_label



Discussion:
Uh, interessting. I will have a look at it, thanks for raising this
Ok, had a look into it... It is arguable whether this is a bug. 
The reason this breaks is due to the fact that your ``create_fold`` function does not yield balanced DNDarray. However, for the caculation of the tiles, ``_dist`` assume a balanced array (the tile locations/sizes are calculated via ``comm.counts_displs_shape`` )
Running ``.balance()`` on the fold and verification arrays solves the problem (this needs to be done after slicing indices in knn.py : predict l.65 as well, by the way, there are some more problems to this example).

However, we should discuss whether the distance functions should allow for unbalanced arrays. This would require additional communication overhead though. @Markus-Goetz what do you say?
I think calling balance is the way to go here, because we also want to ensure that computations are more or less equally distributed.

The create folds function could be improved here. Instead of taking a global slice, leaving the last processes empty handed, we could simply do a local slice, and sticht the results back together into a global DNDarray. This way we also avoid heavy communication.
That is what @Inzlinger is doing actually (calling the ``ht.array(torch.tensor, is_split=0)`` However, this does not result in balanced ``DNDarrays`` (something I also discovered when doing the kmeans test enhancement).
So maybe that is the point to look at. 
I agree though, that we should leave ``cdist`` as is, since the official heat policy is "balanced in, balanced out" as far as I recall.
Sorry my bad. Did not look into it deep enough. In this case calling ``balance_()`` should be sufficient and actually not encompass to much communication.

Correct about the balanced in, balanced out part.

Resolving patch:
diff --git a/examples/classification/demo_knn.py b/examples/classification/demo_knn.py
--- a/examples/classification/demo_knn.py
+++ b/examples/classification/demo_knn.py
@@ -1,5 +1,6 @@
 import sys
 import os
+import random
 
 # Fix python path if run from terminal
 curdir = os.path.dirname(os.path.abspath(__file__))
@@ -19,20 +20,20 @@
     keys.append(1)
 for i in range(100, 150):
     keys.append(2)
-Y = ht.array(keys)
+Y = ht.array(keys, split=0)
 
 
 def calculate_accuracy(new_y, verification_y):
     """
-    Calculates the accuracy of classification/clustering-Algorithms.
-    Note this only works with integer/discrete classes. For Algorithms that give Approximations an error function is
+    Calculates the accuracy of classification/clustering-algorithms.
+    Note this only works with integer/discrete classes. For algorithms that give approximations an error function is
     required.
 
     Parameters
     ----------
-    new_y : ht.tensor of shape (n_samples,), required
-        The new labels that are generated
-    verification_y : ht.tensor of shape (n_samples,), required
+    new_y : ht.tensor of shape (n_samples, n_features), required
+        The new labels that where generated
+    verification_y : ht.tensor of shape (n_samples, n_features), required
         Known labels
 
     Returns
@@ -41,17 +42,16 @@ def calculate_accuracy(new_y, verification_y):
         the accuracy, number of properly labeled samples divided by amount of labels.
     """
 
-    if len(new_y) != len(verification_y):
+    if new_y.gshape != verification_y.gshape:
         raise ValueError(
-            "Expecting results of same length, got {}, {}".format(len(new_y), len(verification_y))
+            "Expecting results of same length, got {}, {}".format(
+                new_y.gshape, verification_y.gshape
+            )
         )
 
-    length = len(new_y)
-    count = 0
-    for index in range(length):
-        if new_y[index] == verification_y[index].item():
-            count += 1
-    return count / length
+    count = ht.sum(ht.where(new_y == verification_y, 1, 0))
+
+    return count / new_y.gshape[0]
 
 
 def create_fold(dataset_x, dataset_y, size, seed=None):
@@ -84,23 +84,26 @@ def create_fold(dataset_x, dataset_y, size, seed=None):
     assert size < len(dataset_x)
 
     data_length = len(dataset_x)
+
     if seed:
-        ht.random.seed(seed)
-    indices = ht.random.randint(low=0, high=data_length - 1, size=(size,))
-    indices = ht.unique(indices, sorted=True)
-    while len(indices) < size:
-        diff = size - len(indices)
-        additional = ht.random.randint(low=0, high=data_length, size=(diff,))
-        indices = ht.concatenate((indices, additional))
-        indices = ht.unique(indices, sorted=True)
-    indices = indices.tolist()
-    all_indices = [index for index in range(data_length)]
-    verification_indices = [index for index in all_indices if index not in indices]
-
-    fold_x = ht.array(dataset_x[indices], is_split=0)
-    fold_y = ht.array(dataset_y[indices], is_split=0)
+        random.seed(seed)
+    indices = [i for i in range(data_length)]
+    random.shuffle(indices)
+
+    data_indices = ht.array(indices[0:size], split=0)
+    verification_indices = ht.array(indices[size:], split=0)
+
+    fold_x = ht.array(dataset_x[data_indices], is_split=0)
+    fold_y = ht.array(dataset_y[data_indices], is_split=0)
     verification_y = ht.array(dataset_y[verification_indices], is_split=0)
     verification_x = ht.array(dataset_x[verification_indices], is_split=0)
+
+    # Balance arrays
+    fold_x.balance_()
+    fold_y.balance_()
+    verification_y.balance_()
+    verification_x.balance_()
+
     return fold_x, fold_y, verification_x, verification_y
 
 
@@ -135,8 +138,8 @@ def verify_algorithm(x, y, split_number, split_size, k, seed=None):
         fold_x, fold_y, verification_x, verification_y = create_fold(x, y, split_size, seed)
         classifier = KNN(fold_x, fold_y, k)
         result_y = classifier.predict(verification_x)
-        accuracies.append(calculate_accuracy(result_y, verification_y))
+        accuracies.append(calculate_accuracy(result_y, verification_y).item())
     return accuracies
 
 
-print(verify_algorithm(X, Y, 5, 30, 5))
+print(verify_algorithm(X, Y, 1, 30, 5, 1))
diff --git a/heat/classification/knn.py b/heat/classification/knn.py
--- a/heat/classification/knn.py
+++ b/heat/classification/knn.py
@@ -16,20 +16,38 @@ class KNN(ht.ClassificationMixin, ht.BaseEstimator):
         Array of shape (n_samples, sample_length,), required
         The training vectors
     y : ht.DNDarray
-        Array of shape (n_samples,), required
-        Labels for the training set
+            Labels for the data. Can be either of shape (n_samples,) as label array, or in one-hot-encoding with shape
+            (n_samples, n_classes).
     num_neighbours: int
         Number of neighbours to consider when choosing label
-
+    is_one_hot: bool
+        Whether y is a label array or in one-hot-encoding
     References
     --------
     [1] T. Cover and P. Hart, "Nearest neighbor pattern classification," in IEEE Transactions on Information Theory,
         vol. 13, no. 1, pp. 21-27, January 1967, doi: 10.1109/TIT.1967.1053964.
     """
 
-    def __init__(self, x, y, num_neighbours):
+    def __init__(self, x, y, num_neighbours, is_one_hot=False):
+
+        if x.shape[0] != y.shape[0]:
+            raise ValueError(
+                "Number of samples and labels needs to be the same, got {}, {}".format(
+                    x.shape[0], y.shape[0]
+                )
+            )
+
         self.x = x
-        self.y = y
+        if len(y.shape) == 1:
+            self.y = self.label_to_one_hot(y)
+        elif len(y.shape) == 2:
+            self.y = y
+        else:
+            raise ValueError(
+                "Expected labels of shape (n_samples,) or (n_samples, n_classes) but got {}".format(
+                    y.shape
+                )
+            )
         self.num_neighbours = num_neighbours
 
     def fit(self, X, Y):
@@ -39,7 +57,8 @@ def fit(self, X, Y):
         X : ht.DNDarray
             Data vectors used for prediction
         Y : ht.DNDarray
-            Labels for the data
+            Labels for the data. Can be either of shape (n_samples,) as label array, or in one-hot-encoding with shape
+            (n_samples, n_classes).
         """
 
         if X.shape[0] != Y.shape[0]:
@@ -50,7 +69,16 @@ def fit(self, X, Y):
             )
 
         self.x = X
-        self.y = Y
+        if len(Y.shape) == 1:
+            self.y = self.label_to_one_hot(Y)
+        elif len(Y.shape) == 2:
+            self.y = Y
+        else:
+            raise ValueError(
+                "Expected labels of shape (n_samples,) or (n_samples, n_classes) but got {}".format(
+                    Y.shape
+                )
+            )
 
     def predict(self, X) -> ht.dndarray:
         """
@@ -63,9 +91,21 @@ def predict(self, X) -> ht.dndarray:
         distances = ht.spatial.cdist(X, self.x)
         _, indices = ht.topk(distances, self.num_neighbours, largest=False)
 
-        labels = self.y[indices._DNDarray__array]
-        uniques = ht.unique(labels, sorted=True)
-        uniques = ht.resplit(ht.expand_dims(uniques, axis=0), axis=0)
-        labels = ht.expand_dims(labels, axis=2)
+        labels = self.y[indices.flatten()]
+        labels.balance_()
+        labels = ht.reshape(labels, (indices.gshape + (self.y.gshape[1],)))
+        labels = ht.sum(labels, axis=1)
+        maximums = ht.argmax(labels, axis=1)
+
+        return maximums
+
+    @staticmethod
+    def label_to_one_hot(a):
+        max_label = ht.max(a)
+        a = a.expand_dims(1)
+
+        items = ht.arange(0, max_label.item() + 1)
+        one_hot = ht.stack([items for i in range(a.shape[0])], axis=0)
+        one_hot = ht.where(one_hot == a, 1, 0)
 
-        return ht.argmax(ht.sum(labels == uniques, axis=1), axis=1)
+        return one_hot


Test patch:
diff --git a/heat/classification/tests/__init__.py b/heat/classification/tests/__init__.py
new file mode 100644
diff --git a/heat/classification/tests/test_knn.py b/heat/classification/tests/test_knn.py
new file mode 100644
--- /dev/null
+++ b/heat/classification/tests/test_knn.py
@@ -0,0 +1,118 @@
+import os
+import unittest
+
+import heat as ht
+
+from heat.core.tests.test_suites.basic_test import TestCase
+
+from heat.classification.knn import KNN
+
+
+class TestKNN(TestCase):
+    def test_split_none(self):
+        X = ht.load_hdf5("heat/datasets/data/iris.h5", dataset="data")
+
+        # Generate keys for the iris.h5 dataset
+        keys = []
+        for i in range(50):
+            keys.append(0)
+        for i in range(50, 100):
+            keys.append(1)
+        for i in range(100, 150):
+            keys.append(2)
+        Y = ht.array(keys)
+
+        knn = KNN(X, Y, 5)
+
+        result = knn.predict(X)
+
+        self.assertTrue(ht.is_estimator(knn))
+        self.assertTrue(ht.is_classifier(knn))
+        self.assertIsInstance(result, ht.DNDarray)
+        self.assertEqual(result.shape, Y.shape)
+
+    def test_split_zero(self):
+        X = ht.load_hdf5("heat/datasets/data/iris.h5", dataset="data", split=0)
+
+        # Generate keys for the iris.h5 dataset
+        keys = []
+        for i in range(50):
+            keys.append(0)
+        for i in range(50, 100):
+            keys.append(1)
+        for i in range(100, 150):
+            keys.append(2)
+        Y = ht.array(keys, split=0)
+
+        knn = KNN(X, Y, 5)
+
+        result = knn.predict(X)
+
+        self.assertIsInstance(result, ht.DNDarray)
+        self.assertEqual(result.shape, Y.shape)
+
+    def test_exception(self,):
+        a = ht.zeros((3,))
+        b = ht.zeros((1,))
+        c = ht.zeros((3,))
+        d = ht.zeros((2, 2, 2))
+
+        with self.assertRaises(ValueError):
+            knn = KNN(a, b, 1)
+
+        with self.assertRaises(ValueError):
+            knn = KNN(a, b, 1)
+
+        knn = KNN(a, c, 1)
+        with self.assertRaises(ValueError):
+            knn.fit(a, b)
+
+        knn = KNN(a, c, 1)
+        with self.assertRaises(ValueError):
+            knn.fit(a, d)
+
+        with self.assertRaises(ValueError):
+            knn = KNN(a, d, 1)
+
+    def test_utility(self,):
+        a = ht.array([1, 2, 3, 4])
+        b = ht.array([[0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [0, 0, 0, 0, 1]])
+
+        one_hot = KNN.label_to_one_hot(a)
+
+        self.assertTrue((one_hot == b).all())
+
+    def test_fit_one_hot(self,):
+        X = ht.load_hdf5("heat/datasets/data/iris.h5", dataset="data")
+
+        # Keys as label array
+        keys = []
+        for i in range(50):
+            keys.append(0)
+        for i in range(50, 100):
+            keys.append(1)
+        for i in range(100, 150):
+            keys.append(2)
+        labels = ht.array(keys, split=0)
+
+        # Keys as one_hot
+        keys = []
+        for i in range(50):
+            keys.append([1, 0, 0])
+        for i in range(50, 100):
+            keys.append([0, 1, 0])
+        for i in range(100, 150):
+            keys.append([0, 0, 1])
+
+        Y = ht.array(keys)
+
+        knn = KNN(X, Y, 5)
+
+        knn.fit(X, Y)
+
+        result = knn.predict(X)
+
+        self.assertTrue(ht.is_estimator(knn))
+        self.assertTrue(ht.is_classifier(knn))
+        self.assertIsInstance(result, ht.DNDarray)
+        self.assertEqual(result.shape, labels.shape)