kmeans.py

# coding:utf-8

importrandom

importmatplotlib.pyplotasplt
importnumpyasnp
importseabornassns

frommla.baseimportBaseEstimator
frommla.metrics.distanceimporteuclidean_distance

random.seed(1111)


classKMeans(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_inrange(self.K)]
self.centroids= []
self.init=init

def_initialize_centroids(self, init):
"""Set the initial centroids."""

ifinit=="random":
self.centroids= [self.X[x] forxinrandom.sample(range(self.n_samples), self.K)]
elifinit=="++":
self.centroids= [random.choice(self.X)]
whilelen(self.centroids) <self.K:
self.centroids.append(self._choose_next_center())
else:
raiseValueError("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_inrange(self.max_iters):
self._assign(centroids)
centroids_old=centroids
centroids= [self._get_centroid(cluster) forclusterinself.clusters]

ifself._is_converged(centroids_old, centroids):
break

self.centroids=centroids

returnself._get_predictions()

def_get_predictions(self):
predictions=np.empty(self.n_samples)

fori, clusterinenumerate(self.clusters):
forindexincluster:
predictions[index] =i
returnpredictions

def_assign(self, centroids):

forrowinrange(self.n_samples):
fori, clusterinenumerate(self.clusters):
ifrowincluster:
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
fori, pointinenumerate(centroids):
dist=euclidean_distance(self.X[fpoint], point)
ifclosest_indexisNoneordist<closest_distance:
closest_index=i
closest_distance=dist
returnclosest_index

def_get_centroid(self, cluster):
"""Get values by indices and take the mean."""
return [np.mean(np.take(self.X[:, i], cluster)) foriinrange(self.n_features)]

def_dist_from_centers(self):
"""Calculate distance from centers."""
returnnp.array([min([euclidean_distance(x, c) forcinself.centroids]) forxinself.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]
returnself.X[ind]

def_is_converged(self, centroids_old, centroids):
"""Check if the distance between old and new centroids is zero."""
distance=0
foriinrange(self.K):
distance+=euclidean_distance(centroids_old[i], centroids[i])
returndistance==0

defplot(self, ax=None, holdon=False):
sns.set(style="white")
palette=sns.color_palette("hls", self.K+1)
data=self.X

ifaxisNone:
_, ax=plt.subplots()

fori, indexinenumerate(self.clusters):
point=np.array(data[index]).T
ax.scatter(*point, c=[palette[i], ])

forpointinself.centroids:
ax.scatter(*point, marker="x", linewidths=10)

ifnotholdon:
plt.show()