ClusterEvaluationUtil implements AllWindowFunction<Row, Row, TimeWindow> { public static Params extractParamsFromConfusionMatrix(LongMatrix longMatrix) { long[][] matrix = longMatrix.getMatrix(); long[] actualLabel = longMatrix.getColSums(); long[] predictLabel = longMatrix.getRowSums(); long total = longMatrix.getTotal(); double entropyActual = 0.0; double entropyPredict = 0.0; double mutualInfor = 0.0; double purity = 0.0; long tp = 0L; long tpFpSum = 0L; long tpFnSum = 0L; for (long anActualLabel : actualLabel) { entropyActual += entropy(anActualLabel, total); tpFpSum += combination(anActualLabel); } entropyActual /= -Math.log(2); for (long aPredictLabel : predictLabel) { entropyPredict += entropy(aPredictLabel, total); tpFnSum += combination(aPredictLabel); } entropyPredict /= -Math.log(2); for (int i = 0; i < matrix.length; i++) { long max = 0; for (int j = 0; j < matrix[0].length; j++) { max = Math.max(max, matrix[i][j]); mutualInfor += (0 == matrix[i][j] ? 0.0 : 1.0 * matrix[i][j] / total * Math.log(1.0 * total * matrix[i][j] / predictLabel[i] / actualLabel[j])); tp += combination(matrix[i][j]); } purity += max; } purity /= total; mutualInfor /= Math.log(2); long fp = tpFpSum - tp; long fn = tpFnSum - tp; long totalCombination = combination(total); long tn = totalCombination - tp - fn - fp; double expectedIndex = 1.0 * tpFpSum * tpFnSum / totalCombination; double maxIndex = 1.0 * (tpFpSum + tpFnSum) / 2; double ri = 1.0 * (tp + tn) / (tp + tn + fp + fn); return new Params() .set(ClusterMetrics.NMI, 2.0 * mutualInfor / (entropyActual + entropyPredict)) .set(ClusterMetrics.PURITY, purity) .set(ClusterMetrics.RI, ri) .set(ClusterMetrics.ARI, (tp - expectedIndex) / (maxIndex - expectedIndex)); } ClusterEvaluationUtil(ContinuousDistance distance); static Params extractParamsFromConfusionMatrix(LongMatrix longMatrix); static Tuple1<Double> calSilhouetteCoefficient(Row row, ClusterMetricsSummary clusterMetricsSummary); static Params getBasicClusterStatistics(Iterable<Row> rows); static ClusterMetricsSummary getClusterStatistics(Iterable<Row> rows, ContinuousDistance distance); @Override void apply(TimeWindow timeWindow, Iterable<Row> rows, Collector<Row> collector); static int COUNT; static int MEAN; static int SUM_2; }
@Test public void matrixToParamsTest(){ long[][] matrix = new long[][]{{5, 1, 2}, {1, 4, 0}, {0, 1, 3}}; LongMatrix longMatrix = new LongMatrix(matrix); Params params = ClusterEvaluationUtil.extractParamsFromConfusionMatrix(longMatrix); Assert.assertEquals(params.get(ClusterMetrics.NMI), 0.364, 0.001); Assert.assertEquals(params.get(ClusterMetrics.PURITY), 0.705, 0.001); Assert.assertEquals(params.get(ClusterMetrics.RI), 0.68, 0.01); Assert.assertEquals(params.get(ClusterMetrics.ARI), 0.24, 0.01); }