wehub-resource-sync
159 lines
5.2 KiB
Python
159 lines
5.2 KiB
Python
# coding:utf-8
|
|
|
|
import random
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import seaborn as sns
|
|
|
|
from mla.base import BaseEstimator
|
|
from mla.metrics.distance import euclidean_distance
|
|
|
|
random.seed(1111)
|
|
|
|
|
|
class KMeans(BaseEstimator):
|
|
"""Partition a dataset into K clusters.
|
|
|
|
Finds clusters by repeatedly assigning each data point to the cluster with
|
|
the nearest centroid and iterating until the assignments converge (meaning
|
|
they don't change during an iteration) or the maximum number of iterations
|
|
is reached.
|
|
|
|
Parameters
|
|
----------
|
|
|
|
K : int
|
|
The number of clusters into which the dataset is partitioned.
|
|
max_iters: int
|
|
The maximum iterations of assigning points to the nearest cluster.
|
|
Short-circuited by the assignments converging on their own.
|
|
init: str, default 'random'
|
|
The name of the method used to initialize the first clustering.
|
|
|
|
'random' - Randomly select values from the dataset as the K centroids.
|
|
'++' - Select a random first centroid from the dataset, then select
|
|
K - 1 more centroids by choosing values from the dataset with a
|
|
probability distribution proportional to the squared distance
|
|
from each point's closest existing cluster. Attempts to create
|
|
larger distances between initial clusters to improve convergence
|
|
rates and avoid degenerate cases.
|
|
"""
|
|
|
|
y_required = False
|
|
|
|
def __init__(self, K=5, max_iters=100, init="random"):
|
|
self.K = K
|
|
self.max_iters = max_iters
|
|
self.clusters = [[] for _ in range(self.K)]
|
|
self.centroids = []
|
|
self.init = init
|
|
|
|
def _initialize_centroids(self, init):
|
|
"""Set the initial centroids."""
|
|
|
|
if init == "random":
|
|
self.centroids = [
|
|
self.X[x] for x in random.sample(range(self.n_samples), self.K)
|
|
]
|
|
elif init == "++":
|
|
self.centroids = [random.choice(self.X)]
|
|
while len(self.centroids) < self.K:
|
|
self.centroids.append(self._choose_next_center())
|
|
else:
|
|
raise ValueError("Unknown type of init parameter")
|
|
|
|
def _predict(self, X=None):
|
|
"""Perform clustering on the dataset."""
|
|
self._initialize_centroids(self.init)
|
|
centroids = self.centroids
|
|
|
|
# Optimize clusters
|
|
for _ in range(self.max_iters):
|
|
self._assign(centroids)
|
|
centroids_old = centroids
|
|
centroids = [self._get_centroid(cluster) for cluster in self.clusters]
|
|
|
|
if self._is_converged(centroids_old, centroids):
|
|
break
|
|
|
|
self.centroids = centroids
|
|
|
|
return self._get_predictions()
|
|
|
|
def _get_predictions(self):
|
|
predictions = np.empty(self.n_samples)
|
|
|
|
for i, cluster in enumerate(self.clusters):
|
|
for index in cluster:
|
|
predictions[index] = i
|
|
return predictions
|
|
|
|
def _assign(self, centroids):
|
|
for row in range(self.n_samples):
|
|
for i, cluster in enumerate(self.clusters):
|
|
if row in cluster:
|
|
self.clusters[i].remove(row)
|
|
break
|
|
|
|
closest = self._closest(row, centroids)
|
|
self.clusters[closest].append(row)
|
|
|
|
def _closest(self, fpoint, centroids):
|
|
"""Find the closest centroid for a point."""
|
|
closest_index = None
|
|
closest_distance = None
|
|
for i, point in enumerate(centroids):
|
|
dist = euclidean_distance(self.X[fpoint], point)
|
|
if closest_index is None or dist < closest_distance:
|
|
closest_index = i
|
|
closest_distance = dist
|
|
return closest_index
|
|
|
|
def _get_centroid(self, cluster):
|
|
"""Get values by indices and take the mean."""
|
|
return [np.mean(np.take(self.X[:, i], cluster)) for i in range(self.n_features)]
|
|
|
|
def _dist_from_centers(self):
|
|
"""Calculate distance from centers."""
|
|
return np.array(
|
|
[min([euclidean_distance(x, c) for c in self.centroids]) for x in self.X]
|
|
)
|
|
|
|
def _choose_next_center(self):
|
|
distances = self._dist_from_centers()
|
|
squared_distances = distances**2
|
|
probs = squared_distances / squared_distances.sum()
|
|
ind = np.random.choice(self.X.shape[0], 1, p=probs)[0]
|
|
return self.X[ind]
|
|
|
|
def _is_converged(self, centroids_old, centroids):
|
|
"""Check if the distance between old and new centroids is zero."""
|
|
distance = 0
|
|
for i in range(self.K):
|
|
distance += euclidean_distance(centroids_old[i], centroids[i])
|
|
return distance == 0
|
|
|
|
def plot(self, ax=None, holdon=False):
|
|
sns.set(style="white")
|
|
palette = sns.color_palette("hls", self.K + 1)
|
|
data = self.X
|
|
|
|
if ax is None:
|
|
_, ax = plt.subplots()
|
|
|
|
for i, index in enumerate(self.clusters):
|
|
point = np.array(data[index]).T
|
|
ax.scatter(
|
|
*point,
|
|
c=[
|
|
palette[i],
|
|
],
|
|
)
|
|
|
|
for point in self.centroids:
|
|
ax.scatter(*point, marker="x", linewidths=10)
|
|
|
|
if not holdon:
|
|
plt.show()
|