EvalClusterBatchOp extends BatchOperator<EvalClusterBatchOp> implements EvalClusterParams<EvalClusterBatchOp>, EvaluationMetricsCollector<ClusterMetrics, EvalClusterBatchOp> { @Override public EvalClusterBatchOp linkFrom(BatchOperator<?>... inputs) { BatchOperator in = checkAndGetFirst(inputs); String labelColName = this.getLabelCol(); String predResultColName = this.getPredictionCol(); String vectorColName = this.getVectorCol(); DistanceType distanceType = getDistanceType(); FastDistance distance = distanceType.getFastDistance(); DataSet<Params> empty = MLEnvironmentFactory.get(getMLEnvironmentId()).getExecutionEnvironment().fromElements( new Params()); DataSet<Params> labelMetrics = empty, vectorMetrics; if (null != labelColName) { DataSet<Row> data = in.select(new String[] {labelColName, predResultColName}).getDataSet(); labelMetrics = calLocalPredResult(data) .reduce(new ReduceFunction<LongMatrix>() { @Override public LongMatrix reduce(LongMatrix value1, LongMatrix value2) { value1.plusEqual(value2); return value1; } }) .map(new MapFunction<LongMatrix, Params>() { @Override public Params map(LongMatrix value) { return ClusterEvaluationUtil.extractParamsFromConfusionMatrix(value); } }); } if (null != vectorColName) { DataSet<Row> data = in.select(new String[] {predResultColName, vectorColName}).getDataSet(); DataSet<BaseMetricsSummary> metricsSummary = data .groupBy(0) .reduceGroup(new CalcClusterMetricsSummary(distance)) .reduce(new EvaluationUtil.ReduceBaseMetrics()); DataSet<Tuple1<Double>> silhouetteCoefficient = data.map( new RichMapFunction<Row, Tuple1<Double>>() { @Override public Tuple1<Double> map(Row value) { List<BaseMetricsSummary> list = getRuntimeContext().getBroadcastVariable(METRICS_SUMMARY); return ClusterEvaluationUtil.calSilhouetteCoefficient(value, (ClusterMetricsSummary)list.get(0)); } }).withBroadcastSet(metricsSummary, METRICS_SUMMARY) .aggregate(Aggregations.SUM, 0); vectorMetrics = metricsSummary.map(new ClusterEvaluationUtil.SaveDataAsParams()).withBroadcastSet( silhouetteCoefficient, SILHOUETTE_COEFFICIENT); } else { vectorMetrics = in.select(predResultColName) .getDataSet() .reduceGroup(new BasicClusterParams()); } DataSet<Row> out = labelMetrics .union(vectorMetrics) .reduceGroup(new GroupReduceFunction<Params, Row>() { @Override public void reduce(Iterable<Params> values, Collector<Row> out) { Params params = new Params(); for (Params p : values) { params.merge(p); } out.collect(Row.of(params.toJson())); } }); this.setOutputTable(DataSetConversionUtil.toTable(getMLEnvironmentId(), out, new TableSchema(new String[] {EVAL_RESULT}, new TypeInformation[] {Types.STRING}) )); return this; } EvalClusterBatchOp();  EvalClusterBatchOp(Params params); @Override ClusterMetrics createMetrics(List<Row> rows); @Override EvalClusterBatchOp linkFrom(BatchOperator<?>... inputs); static final String SILHOUETTE_COEFFICIENT; }
@Test public void test() throws Exception { MemSourceBatchOp inOp = new MemSourceBatchOp(Arrays.asList(rows), new String[] {"label", "Y"}); KMeans train = new KMeans() .setVectorCol("Y") .setPredictionCol("pred") .setK(2); ClusterMetrics metrics = new EvalClusterBatchOp(new Params()) .setPredictionCol("pred") .setVectorCol("Y") .setLabelCol("label") .linkFrom(train.fit(inOp).transform(inOp)) .collectMetrics(); Assert.assertEquals(metrics.getVrc(), 12150.00, 0.01); Assert.assertEquals(metrics.getCp(), 0.115, 0.01); Assert.assertEquals(metrics.getCount().intValue(), 6); Assert.assertEquals(metrics.getDb(), 0.014, 0.01); Assert.assertEquals(metrics.getSp(), 15.58, 0.01); Assert.assertEquals(metrics.getK().intValue(), 2); Assert.assertEquals(metrics.getSsb(), 364.5, 0.01); Assert.assertEquals(metrics.getSsw(), 0.119, 0.01); Assert.assertEquals(metrics.getPurity(), 1.0, 0.01); Assert.assertEquals(metrics.getNmi(), 1.0, 0.01); Assert.assertEquals(metrics.getAri(), 1.0, 0.01); Assert.assertEquals(metrics.getRi(), 1.0, 0.01); Assert.assertEquals(metrics.getSilhouetteCoefficient(), 0.99,0.01); }
@Test public void testNoVector() throws Exception { MemSourceBatchOp inOp = new MemSourceBatchOp(Arrays.asList(rows), new String[] {"label", "Y"}); KMeans train = new KMeans() .setVectorCol("Y") .setPredictionCol("pred") .setK(2); ClusterMetrics metrics = new EvalClusterBatchOp() .setPredictionCol("pred") .linkFrom(train.fit(inOp).transform(inOp)) .collectMetrics(); Assert.assertEquals(metrics.getCount().intValue(), 6); Assert.assertArrayEquals(metrics.getClusterArray(), new String[]{"0", "1"}); }