test_batch_kmeans.py 文件源码

def test_step_M(self):
        points = np.random.randn(self.n_points,self.dim)
        KM = Kmeans(self.n_components)
        KM._initialize(points)
        assignements = KM._step_E(points)

        expected_means = KM.means.copy()
        for i in range(self.n_components):
            assignements_i = assignements[:,i:i+1]
            n_set = np.sum(assignements_i)
            idx_set,_ = np.where(assignements_i==1)
            sets = points[idx_set]
            if n_set > 0:
                expected_means[i] = np.asarray(np.sum(sets, axis=0)/n_set)


        KM._step_M(points,assignements)

        assert_almost_equal(expected_means,KM.means)