Machine-Learning-Explained/Algorithms/kmeans/code/kmeans.py at master · TannerGilbert/Machine-Learning-Explained · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
from __future__ import annotations
from typing import Union
import numpy as np
import random


class KMeans:
    """KMeans
    Parameters:
    -----------
    n_clusters: int = 2
        Number of clusters
    tol: float = 0.0001
        Relative tolerance
    max_iter: int = 300
        Maximum number of iterations of the k-means algorithm for a single run
    """
    def __init__(self, n_clusters: int = 2, tol: float = 0.0001, max_iter: int = 300) -> None:
        self.n_clusters = n_clusters
        self.tol = tol
        self.max_iter = max_iter
        self.centroids = {}
        self.classifications = {}

    def fit(self, data: Union[list, np.ndarray]) -> KMeans:
        self.centroids = {}

        # Randomly initialize centroids
        for i in range(self.n_clusters):
            self.centroids[i] = data[random.randint(0, len(data)-1)]

        for _ in range(self.max_iter):
            self.classifications = {}

            for i in range(self.n_clusters):
                self.classifications[i] = []

            # put row into the right cluster
            for row in data:
                distances = [np.linalg.norm(
                    row-self.centroids[centroid]) for centroid in self.centroids]
                classification = distances.index(min(distances))
                self.classifications[classification].append(row)

            prev_centroids = dict(self.centroids)

            for classification in self.classifications:
                self.centroids[classification] = np.average(
                    self.classifications[classification], axis=0)

            # Check if we are finished
            optimized = True
            for c in self.centroids:
                original_centroid = prev_centroids[c]
                current_centroid = self.centroids[c]
                if np.sum((current_centroid - original_centroid) / original_centroid * 100.0) > self.tol:
                    optimized = False

            if optimized:
                break
        return self

    def predict(self, data: Union[list, np.ndarray]) -> list:
        classifications = []
        for row in data:
            distances = [np.linalg.norm(row-self.centroids[centroid])
                         for centroid in self.centroids]
            classification = distances.index(min(distances))
            classifications.append(classification)
        return classifications


if __name__ == '__main__':
    import matplotlib.pyplot as plt
    from matplotlib import style
    from sklearn.datasets import make_blobs
    style.use('ggplot')

    X, y = make_blobs(n_samples=30, centers=3, n_features=2)

    model = KMeans(n_clusters=3)
    model.fit(X)

    centroids = model.centroids

    colors = 10 * ['r', 'g', 'b', 'c', 'k', 'y']

    for classification, featureset in zip(model.predict(X), X):
        color = colors[classification]
        plt.scatter(featureset[0], featureset[1], marker="x",
                    color=color, s=150, linewidths=5, zorder=10)

    for c in centroids:
        plt.scatter(centroids[c][0], centroids[c][1],
                    color='k', marker="*", s=150, linewidths=5)

    plt.show()