Expand label result with some stats

Co-authored-by: Veselin Nikolov <[email protected]>

Author: IoannisPanagiotas

algo/src/main/java/org/neo4j/gds/hdbscan/HDBScan.java

@@ -22,15 +22,14 @@
 import org.neo4j.gds.Algorithm;
 import org.neo4j.gds.api.IdMap;
 import org.neo4j.gds.api.properties.nodes.NodePropertyValues;
-import org.neo4j.gds.collections.ha.HugeLongArray;
 import org.neo4j.gds.collections.ha.HugeObjectArray;
 import org.neo4j.gds.core.concurrency.Concurrency;
 import org.neo4j.gds.core.concurrency.ParallelUtil;
 import org.neo4j.gds.core.utils.paged.HugeSerialObjectMergeSort;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.termination.TerminationFlag;
 
-public class HDBScan extends Algorithm<HugeLongArray> {
+public class HDBScan extends Algorithm<Labels> {
 
     private final IdMap nodes;
     private final NodePropertyValues nodePropertyValues;
@@ -61,7 +60,7 @@ public HDBScan(
     }
 
     @Override
-    public HugeLongArray compute() {
+    public Labels compute() {
         progressTracker.beginSubTask();
         var kdTree = buildKDTree();
 

algo/src/main/java/org/neo4j/gds/hdbscan/LabellingStep.java

@@ -29,6 +29,7 @@ class LabellingStep {
     private final CondensedTree condensedTree;
     private final long nodeCount;
     private final ProgressTracker progressTracker;
+    private static final long NOISE = -1;
 
     LabellingStep(CondensedTree condensedTree, long nodeCount, ProgressTracker progressTracker) {
         this.condensedTree = condensedTree;
@@ -97,36 +98,42 @@ BitSet selectedClusters(HugeDoubleArray stabilities) {
         return selectedClusters;
     }
 
-    HugeLongArray computeLabels(BitSet selectedClusters) {
+    Labels computeLabels(BitSet selectedClusters) {
         progressTracker.beginSubTask();
+        var treeLabels = HugeLongArray.newArray(nodeCount);
         var labels = HugeLongArray.newArray(nodeCount);
-        labels.fill(-1L);
-        var nodeCountLabels = HugeLongArray.newArray(nodeCount);
+        treeLabels.fill(NOISE);
+        long clusters=0;
         var root = condensedTree.root();
         var maximumClusterId = condensedTree.maximumClusterId();
         for (var p = root; p <= maximumClusterId; p++) {
             var adaptedIndex = p - nodeCount;
             var parent = condensedTree.parent(p);
-            long parentLabel = p == root ? -1L : labels.get(parent - nodeCount);
-            if (parentLabel != -1L) {
-                labels.set(adaptedIndex, parentLabel);
+            long parentLabel = p == root ? NOISE : treeLabels.get(parent - nodeCount);
+            if (parentLabel != NOISE) {
+                treeLabels.set(adaptedIndex, parentLabel);
             } else if (selectedClusters.get(adaptedIndex)) {
-                labels.set(adaptedIndex, adaptedIndex);
+                clusters++;
+                treeLabels.set(adaptedIndex, adaptedIndex);
             }
             progressTracker.logProgress();
         }
 
+        long noisePoints=0;
         for (var n = 0; n < nodeCount; n++) {
-            nodeCountLabels.set(n, labels.get(condensedTree.fellOutOf(n) - nodeCount));
+            long label = treeLabels.get(condensedTree.fellOutOf(n) - nodeCount);
+            if (label == NOISE)  {
+                noisePoints++;
+            }
+            labels.set(n, label);
             progressTracker.logProgress();
         }
         progressTracker.endSubTask();
 
-
-        return nodeCountLabels;
+        return new Labels(labels,noisePoints,clusters);
     }
 
-    HugeLongArray labels() {
+    Labels labels() {
         progressTracker.beginSubTask();
         var stabilities = computeStabilities();
         var selectedClusters = selectedClusters(stabilities);

algo/src/main/java/org/neo4j/gds/hdbscan/Labels.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.neo4j.gds.hdbscan;
+
+import org.neo4j.gds.collections.ha.HugeLongArray;
+
+public record Labels(HugeLongArray labels, long numberOfNoisePoints, long numberOfClusters) {}

algo/src/test/java/org/neo4j/gds/hdbscan/HDBScanE2ETest.java

@@ -73,7 +73,9 @@ void hdbscan() {
             TerminationFlag.RUNNING_TRUE
         );
 
-        var labelsWithOffset = hdbScan.compute();
+        var result =hdbScan.compute();
+
+        var labelsWithOffset = result.labels();
 
         var labels = new long[10];
         for (char letter='a'; letter<='j';++letter){
@@ -83,6 +85,9 @@ void hdbscan() {
 
         var expectedLabels = new long[] {2, 2, 1, 2, 2, 2, 1, 1, 1, 1};
 
+        assertThat(result.numberOfClusters()).isEqualTo(2L);
+        assertThat(result.numberOfNoisePoints()).isEqualTo(0);
+
         assertThat(labels).containsExactly(expectedLabels);
     }
 

algo/src/test/java/org/neo4j/gds/hdbscan/LabellingTest.java

@@ -176,8 +176,9 @@ void labelling() {
 
         var stabilityStep = new LabellingStep(condensedTree, nodeCount,ProgressTracker.NULL_TRACKER);
 
-        var labels = stabilityStep.computeLabels(selectedClusters);
+        var labelsResult = stabilityStep.computeLabels(selectedClusters);
 
+        var labels=labelsResult.labels();
         assertThat(labels.size()).isEqualTo(nodeCount);
 
         assertThat(labels.get(0)).isEqualTo(1L);
@@ -190,6 +191,10 @@ void labelling() {
         assertThat(labels.get(5)).isEqualTo(4L);
         assertThat(labels.get(6)).isEqualTo(4L);
 
+        assertThat(labelsResult.numberOfNoisePoints()).isEqualTo(2L);
+        assertThat(labelsResult.numberOfClusters()).isEqualTo(2L);
+
+
     }
 
     @Test
@@ -207,10 +212,15 @@ void labellingWhenAllClustersAreSelected() {
 
         var stabilityStep = new LabellingStep(condensedTree, nodeCount,ProgressTracker.NULL_TRACKER);
 
-        var labels = stabilityStep.computeLabels(selectedClusters);
+        var labelsResult = stabilityStep.computeLabels(selectedClusters);
 
+        var labels= labelsResult.labels();
         assertThat(labels.size()).isEqualTo(nodeCount);
         assertThat(labels.toArray()).containsOnly(0L);
+
+        assertThat(labelsResult.numberOfClusters()).isEqualTo(1L);
+        assertThat(labelsResult.numberOfNoisePoints()).isEqualTo(0L);
+
     }
 
     @Test