LocalKmeansFunc implements Serializable { static FastDistanceMatrixData kmeans(int k, long[] sampleWeights, FastDistanceVectorData[] samples, FastDistance distance, int vectorSize) { Random random = new Random(0); List<FastDistanceVectorData> initCentroidsList = sampleInitialCentroids(k, sampleWeights, samples, distance, random); FastDistanceMatrixData initCentroids = KMeansUtil.buildCentroidsMatrix(initCentroidsList, distance, vectorSize); boolean converge = false; int iteration = 0; DenseMatrix sumMatrix = new DenseMatrix(vectorSize + 1, k); DenseMatrix distanceMatrix = new DenseMatrix(k, 1); double[] sumMatrixData = sumMatrix.getData(); double[] initCentroidsData = initCentroids.getVectors().getData(); int[] indices = new int[samples.length]; while (!converge && iteration < LOCAL_MAX_ITER) { iteration++; converge = true; for (int i = 0; i < samples.length; i++) { int clusterId = KMeansUtil.updateSumMatrix(samples[i], sampleWeights[i], initCentroids, vectorSize, sumMatrixData, k, distance, distanceMatrix); if (clusterId != indices[i]) { indices[i] = clusterId; converge = false; } } Arrays.fill(initCentroidsData, 0.0); for (int i = 0; i < k; i++) { int initCentroidsStartIndex = i * vectorSize; int sumMatrixStartIndex = initCentroidsStartIndex + i; double weight = sumMatrixData[sumMatrixStartIndex + vectorSize]; if (weight > 0) { BLAS.axpy(vectorSize, 1.0 / weight, sumMatrixData, sumMatrixStartIndex, initCentroidsData, initCentroidsStartIndex); } else { int index = random.nextInt(samples.length); MatVecOp.appendVectorToMatrix(initCentroids.getVectors(), false, i, samples[index].getVector()); } distance.updateLabel(initCentroids); } } if (iteration != LOCAL_MAX_ITER) { LOG.info("Local kmeans converge with {} steps.", iteration); } else { LOG.info("Local kmeans reach max iteration number!"); } return initCentroids; }    }
@Test public void kmeansTest() { EuclideanDistance distance = new EuclideanDistance(); long[] sampleWeights = new long[10]; FastDistanceVectorData[] samples = new FastDistanceVectorData[sampleWeights.length]; for (int i = 0; i < 10; i++) { sampleWeights[i] = i; samples[i] = distance.prepareVectorData(Tuple2.of(DenseVector.ones(2).scale(i), null)); } FastDistanceMatrixData initCentroid = kmeans(2, sampleWeights, samples, distance, 2); DenseMatrix initCentroidData = initCentroid.getVectors(); Assert.assertEquals(initCentroidData.numCols(), 2); Assert.assertEquals(new DenseVector(initCentroidData.getColumn(0)).normL2(), 10.842, 0.001); Assert.assertEquals(new DenseVector(initCentroidData.getColumn(1)).normL2(), 5.185, 0.001); }
@Test public void kmeansSparseTest() { int len = 10; int k = 2; int size = 20; EuclideanDistance distance = new EuclideanDistance(); long[] sampleWeights = new long[len]; FastDistanceVectorData[] samples = new FastDistanceVectorData[len]; for (int i = 0; i < 10; i++) { sampleWeights[i] = i; samples[i] = distance.prepareVectorData(Tuple2.of(new SparseVector(size, new int[]{i, i + 1}, new double[]{i, i}), null)); } FastDistanceMatrixData initCentroid = kmeans(k, sampleWeights, samples, distance, size); DenseMatrix initCentroidData = initCentroid.getVectors(); Assert.assertEquals(initCentroidData.numCols(), k); Assert.assertEquals(new DenseVector(initCentroidData.getColumn(0)).normL2(), 8.615, 0.001); Assert.assertEquals(new DenseVector(initCentroidData.getColumn(1)).normL2(), 4.128, 0.001); }