唯物是真 @Scaled_Wurm

プログラミング(主にPython2.7)とか機械学習とか

numpyを使ってK-meansを書いてみた

知り合いが書いたコードがまったくnumpyの機能を使ってなかったのでついカッとなって書いた。
今では反省している。


numpyを使うときはできるだけforループとインデックスによるアクセスを使わずに、行列・ベクトル単位で演算するのが基本的なやり方(Rとかと一緒です)。
ただ、下の新しいクラスタ割り当ての計算の式とか、慣れないとわりと読みづらい気がしないでもない……。

new_cluster = np.array([np.argmin(np.sum((d - centroids) ** 2, axis = 1)) for d in data])

適当に書いていたら、途中でどのデータ点も割り当てられないクラスタができてしまい、クラスタの重心がNaNになってしまうバグが出てしまいました。
そういう場合には重心の値を更新しないようにしたら一応最後まで回るようになりました。

ソースコード

import numpy as np
from optparse import OptionParser

data = np.array([[0,1],[3,8],[3,1],[5,10],[6,1],[11,2],[2,11]])

def scaling(data):
    mean = np.mean(data, axis = 0)
    std = np.std(data, axis = 0)
    return (data - mean) / std

def kmeans(data, K = 3, maxIter = 1000):
#    data = scaling(data)
    
    #random initialization
    cluster = np.hstack((np.arange(K), np.random.randint(0, K, data.shape[0] - K)))
    np.random.shuffle(cluster)
    
    centroids = np.empty((K, data.shape[1]))
    
    #iteration until convergence
    for i in xrange(maxIter):
        for k in xrange(K):
            if np.sum(cluster == k) > 0:
                centroids[k] = np.mean(data[cluster == k], axis = 0)
        
        new_cluster = np.array([np.argmin(np.sum((d - centroids) ** 2, axis = 1)) for d in data])
        
        if np.all(cluster == new_cluster):
            break
        else:
            cluster = new_cluster
    
    return [(centroids[k], data[cluster == k]) for k in xrange(K)]

if __name__ == '__main__':
    usage = 'usage: %prog [options]'
    parser = OptionParser(usage)
    parser.add_option('-k', dest = 'K', help = 'the number of the clusters', type = 'int', default = 2)
    (options, args) = parser.parse_args()

    for i, c in enumerate(kmeans(data, options.K)):
        print 'Cluster {0}:'.format(i)
        print '    Centroid:', c[0]
        for d in c[1]:
            print '    ', d