Label clusters

Co-authored-by: Ioannis Panagiotas <ioannis.panagiotas@neotechnology.com>

Author: vnickolov

…org/neo4j/gds/hdbscan/StabilityStep.java …org/neo4j/gds/hdbscan/LabellingStep.javaalgo/src/main/java/org/neo4j/gds/hdbscan/StabilityStep.java renamed to algo/src/main/java/org/neo4j/gds/hdbscan/LabellingStep.java algo/src/main/java/org/neo4j/gds/hdbscan/StabilityStep.java renamed to algo/src/main/java/org/neo4j/gds/hdbscan/LabellingStep.java

@@ -21,14 +21,24 @@
 
 import com.carrotsearch.hppc.BitSet;
 import org.neo4j.gds.collections.ha.HugeDoubleArray;
+import org.neo4j.gds.collections.ha.HugeLongArray;
 
-class StabilityStep {
-    HugeDoubleArray computeStabilities(CondensedTree condensedTree, long nodeCount) {
+class LabellingStep {
+
+    private final CondensedTree condensedTree;
+    private final long nodeCount;
+
+    LabellingStep(CondensedTree condensedTree, long nodeCount) {
+        this.condensedTree = condensedTree;
+        this.nodeCount = nodeCount;
+    }
+
+    HugeDoubleArray computeStabilities() {
         var result = HugeDoubleArray.newArray(nodeCount - 1);
 
-        var condensedTreeRoot = condensedTree.root();
+        var condensedTreeRoot = this.condensedTree.root();
         // process the leaves of the tree
-        for (var p = 0; p < nodeCount; p++) {
+        for (var p = 0; p < this.nodeCount; p++) {
             var lambdaP = 1. / condensedTree.lambda(p);
             var birthPoint = condensedTree.fellOutOf(p);
             var lambdaBirth = birthPoint == condensedTreeRoot
@@ -51,7 +61,7 @@ HugeDoubleArray computeStabilities(CondensedTree condensedTree, long nodeCount)
         return result;
     }
 
-    BitSet selectedClusters(CondensedTree condensedTree, HugeDoubleArray stabilities, long nodeCount) {
+    BitSet selectedClusters(HugeDoubleArray stabilities) {
 
         var selectedClusters = new BitSet(nodeCount);
 
@@ -81,4 +91,34 @@ BitSet selectedClusters(CondensedTree condensedTree, HugeDoubleArray stabilities
 
         return selectedClusters;
     }
+
+    HugeLongArray computeLabels(BitSet selectedClusters) {
+        var labels = HugeLongArray.newArray(nodeCount);
+        labels.fill(-1L);
+        var nodeCountLabels = HugeLongArray.newArray(nodeCount);
+        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);
+            } else if (selectedClusters.get(adaptedIndex)) {
+                labels.set(adaptedIndex, adaptedIndex);
+            }
+        }
+
+        for (var n = 0; n < nodeCount; n++) {
+            nodeCountLabels.set(n, labels.get(condensedTree.fellOutOf(n) - nodeCount));
+        }
+
+        return nodeCountLabels;
+    }
+
+    HugeLongArray label() {
+        var stabilities = computeStabilities();
+        var selectedClusters = selectedClusters(stabilities);
+        return computeLabels(selectedClusters);
+    }
 }

…neo4j/gds/hdbscan/StabilityStepTest.java …org/neo4j/gds/hdbscan/LabellingTest.javaalgo/src/test/java/org/neo4j/gds/hdbscan/StabilityStepTest.java renamed to algo/src/test/java/org/neo4j/gds/hdbscan/LabellingTest.java algo/src/test/java/org/neo4j/gds/hdbscan/StabilityStepTest.java renamed to algo/src/test/java/org/neo4j/gds/hdbscan/LabellingTest.java

@@ -19,14 +19,15 @@
  */
 package org.neo4j.gds.hdbscan;
 
+import com.carrotsearch.hppc.BitSet;
 import org.assertj.core.data.Offset;
 import org.junit.jupiter.api.Test;
 import org.neo4j.gds.collections.ha.HugeDoubleArray;
 import org.neo4j.gds.collections.ha.HugeLongArray;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
-class StabilityStepTest {
+class LabellingTest {
 
     @Test
     void clusterStability() {
@@ -39,9 +40,9 @@ void clusterStability() {
         var maximumClusterId = 6;
 
         var condensedTree = new CondensedTree(root, parent, lambda, size, maximumClusterId, nodeCount);
-        var stabilityStep = new StabilityStep();
+        var stabilityStep = new LabellingStep(condensedTree, nodeCount);
 
-        var stabilities = stabilityStep.computeStabilities(condensedTree, nodeCount);
+        var stabilities = stabilityStep.computeStabilities();
 
 
         assertThat(stabilities.toArray()).containsExactly(
@@ -65,9 +66,9 @@ void clusterStabilityBiggerTest() {
 
         var condensedTree = new CondensedTree(root, parent, lambda, size, maximumClusterId, nodeCount);
 
-        var stabilityStep = new StabilityStep();
+        var stabilityStep = new LabellingStep(condensedTree, nodeCount);
 
-        var stabilities = stabilityStep.computeStabilities(condensedTree, nodeCount);
+        var stabilities = stabilityStep.computeStabilities();
 
         assertThat(stabilities.toArray()).containsExactly(
             new double[] {
@@ -101,9 +102,9 @@ void clusterSelectionOfChildClusters() {
         var stabilities = HugeDoubleArray.of(3., 4., 5.);
 
         var condensedTree = new CondensedTree(root, parent, lambda, size, maximumClusterId, nodeCount);
-        var stabilityStep = new StabilityStep();
+        var stabilityStep = new LabellingStep(condensedTree, nodeCount);
 
-        var selectedClusters = stabilityStep.selectedClusters(condensedTree, stabilities, nodeCount);
+        var selectedClusters = stabilityStep.selectedClusters(stabilities);
 
         assertThat(selectedClusters.get(0))
             .withFailMessage("Root should be unselected")
@@ -132,9 +133,9 @@ void clusterSelectionOfParentCluster() {
         var stabilities = HugeDoubleArray.of(10., 4., 5.);
 
         var condensedTree = new CondensedTree(root, parent, lambda, size, maximumClusterId, nodeCount);
-        var stabilityStep = new StabilityStep();
+        var stabilityStep = new LabellingStep(condensedTree, nodeCount);
 
-        var selectedClusters = stabilityStep.selectedClusters(condensedTree, stabilities, nodeCount);
+        var selectedClusters = stabilityStep.selectedClusters(stabilities);
 
         assertThat(selectedClusters.get(0))
             .withFailMessage("Root should be selected")
@@ -146,4 +147,59 @@ void clusterSelectionOfParentCluster() {
             .withFailMessage("Second child should be selected")
             .isTrue();
     }
+
+    @Test
+    void labelling() {
+        var parent = HugeLongArray.of(8, 8, 10, 10, 11, 11, 11, 0, 7, 7, 9, 9, 0);
+        var lambda = HugeDoubleArray.of(11.0, 11.0, 9.0, 9.0, 8.0, 7.0, 7.0, 0.0, 12.0, 12.0, 10.0, 10.0, 0.0);
+        var size = HugeLongArray.of(7, 2, 5, 2, 3, 0, 0);
+        var maximumClusterId = 11;
+        var nodeCount = 7;
+        var root = 7;
+
+        var condensedTree = new CondensedTree(root, parent, lambda, size, maximumClusterId, nodeCount);
+        var selectedClusters = new BitSet(5);
+        // selects cluster `8`
+        selectedClusters.set(1);
+        // selects cluster `11`
+        selectedClusters.set(4);
+
+        var stabilityStep = new LabellingStep(condensedTree, nodeCount);
+
+        var labels = stabilityStep.computeLabels(selectedClusters);
+
+        assertThat(labels.size()).isEqualTo(nodeCount);
+
+        assertThat(labels.get(0)).isEqualTo(1L);
+        assertThat(labels.get(1)).isEqualTo(1L);
+
+        assertThat(labels.get(2)).isEqualTo(-1L);
+        assertThat(labels.get(3)).isEqualTo(-1L);
+
+        assertThat(labels.get(4)).isEqualTo(4L);
+        assertThat(labels.get(5)).isEqualTo(4L);
+        assertThat(labels.get(6)).isEqualTo(4L);
+
+    }
+
+    @Test
+    void labellingWhenAllClustersAreSelected() {
+        var parent = HugeLongArray.of(8, 8, 10, 10, 11, 11, 11, 0, 7, 7, 9, 9, 0);
+        var lambda = HugeDoubleArray.of(11.0, 11.0, 9.0, 9.0, 8.0, 7.0, 7.0, 0.0, 12.0, 12.0, 10.0, 10.0, 0.0);
+        var size = HugeLongArray.of(7, 2, 5, 2, 3, 0, 0);
+        var maximumClusterId = 11;
+        var nodeCount = 7;
+        var root = 7;
+
+        var condensedTree = new CondensedTree(root, parent, lambda, size, maximumClusterId, nodeCount);
+        var selectedClusters = new BitSet(5);
+        selectedClusters.set(0, 5);
+
+        var stabilityStep = new LabellingStep(condensedTree, nodeCount);
+
+        var labels = stabilityStep.computeLabels(selectedClusters);
+
+        assertThat(labels.size()).isEqualTo(nodeCount);
+        assertThat(labels.toArray()).containsOnly(0L);
+    }
 }