def test_kmeans_constraint_weights(self): mat = numpy.array([[0, 0], [0.2, 0.2], [-0.1, -0.1], [1, 1]]) km = ConstraintKMeans(n_clusters=2, verbose=10, kmeans0=False, random_state=1, strategy='weights') buf = BufferedPrint() km.fit(mat, fLOG=buf.fprint) km = ConstraintKMeans(n_clusters=2, verbose=5, kmeans0=False, random_state=1, strategy='weights') km.fit(mat, fLOG=buf.fprint) self.assertEqual(km.cluster_centers_.shape, (2, 2)) self.assertLesser(km.inertia_, 4.55) self.assertEqual(km.cluster_centers_, numpy.array([[0.6, 0.6], [-0.05, -0.05]])) self.assertEqual(km.labels_, numpy.array([1, 0, 1, 0])) pred = km.predict(mat) self.assertEqual(pred, numpy.array([1, 1, 1, 0])) dist = km.transform(mat) self.assertEqual(dist.shape, (4, 2)) score = km.score(mat) self.assertEqual(score.shape, (4, )) self.assertIn("CKMeans", str(buf))
def test_kmeans_constraint_pickle(self): df = pandas.DataFrame(dict(y=[0, 1, 0, 1, 0, 1, 0, 1], X1=[0.5, 0.6, 0.52, 0.62, 0.5, 0.6, 0.51, 0.61], X2=[0.5, 0.6, 0.7, 0.5, 1.5, 1.6, 1.7, 1.8])) X = df.drop('y', axis=1) y = df['y'] model = ConstraintKMeans(n_clusters=2, strategy='distance') model.fit(X, y) pred = model.transform(X) st = BytesIO() pickle.dump(model, st) st = BytesIO(st.getvalue()) rec = pickle.load(st) pred2 = rec.transform(X) self.assertEqualArray(pred, pred2)