Select clusters

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

Author: vnickolov

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

@@ -19,6 +19,7 @@
  */
 package org.neo4j.gds.hdbscan;
 
+import com.carrotsearch.hppc.BitSet;
 import org.neo4j.gds.collections.ha.HugeDoubleArray;
 
 class StabilityStep {
@@ -49,4 +50,35 @@ HugeDoubleArray computeStabilities(CondensedTree condensedTree, long nodeCount)
 
         return result;
     }
+
+    BitSet selectedClusters(CondensedTree condensedTree, HugeDoubleArray stabilities, long nodeCount) {
+
+        var selectedClusters = new BitSet(nodeCount);
+
+        var condensedTreeRoot = condensedTree.root();
+        var condensedTreeMaxClusterId = condensedTree.maximumClusterId();
+
+        var stabilitySums = HugeDoubleArray.newArray(nodeCount);
+        for (var p = condensedTreeMaxClusterId; p >= condensedTreeRoot; p--) {
+            var adaptedPIndex = p - nodeCount;
+            var stabilityP = stabilities.get(adaptedPIndex);
+            var childrenStabilitySum = stabilitySums.get(adaptedPIndex);
+            double stabilityToAdd;
+            if (childrenStabilitySum > stabilityP) {
+                stabilityToAdd = childrenStabilitySum;
+                selectedClusters.clear(adaptedPIndex);
+            } else {
+                stabilityToAdd = stabilityP;
+                selectedClusters.set(adaptedPIndex);
+                // Selected clusters below `p` are implicitly unselected - they will be ignored during- `labeling`
+            }
+            if (p == condensedTreeRoot) {
+                continue;
+            }
+            var parent = condensedTree.parent(p);
+            stabilitySums.addTo(parent - nodeCount, stabilityToAdd);
+        }
+
+        return selectedClusters;
+    }
 }

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

@@ -84,4 +84,66 @@ void clusterStabilityBiggerTest() {
                 0.0
             }, Offset.offset(1e-10));
     }
+
+    @Test
+    void clusterSelectionOfChildClusters() {
+        //    3
+        //   4  5
+
+        var nodeCount = 3;
+        var root = 3;
+
+        var parent = HugeLongArray.of(-1, -1, -1, 0, 3, 3);
+        var lambda = HugeDoubleArray.of(-1, -1, -1, -1, -1, -1);
+        var size = HugeLongArray.of(-1, -1, -1);
+        var maximumClusterId = 5;
+
+        var stabilities = HugeDoubleArray.of(3., 4., 5.);
+
+        var condensedTree = new CondensedTree(root, parent, lambda, size, maximumClusterId, nodeCount);
+        var stabilityStep = new StabilityStep();
+
+        var selectedClusters = stabilityStep.selectedClusters(condensedTree, stabilities, nodeCount);
+
+        assertThat(selectedClusters.get(0))
+            .withFailMessage("Root should be unselected")
+            .isFalse();
+        assertThat(selectedClusters.get(1))
+            .withFailMessage("First child should be selected cluster")
+            .isTrue();
+        assertThat(selectedClusters.get(2))
+            .withFailMessage("Second child should be selected cluster")
+            .isTrue();
+    }
+
+    @Test
+    void clusterSelectionOfParentCluster() {
+        //    3
+        //   4  5
+
+        var nodeCount = 3;
+        var root = 3;
+
+        var parent = HugeLongArray.of(-1, -1, -1, 0, 3, 3);
+        var lambda = HugeDoubleArray.of(-1, -1, -1, -1, -1, -1);
+        var size = HugeLongArray.of(-1, -1, -1);
+        var maximumClusterId = 5;
+
+        var stabilities = HugeDoubleArray.of(10., 4., 5.);
+
+        var condensedTree = new CondensedTree(root, parent, lambda, size, maximumClusterId, nodeCount);
+        var stabilityStep = new StabilityStep();
+
+        var selectedClusters = stabilityStep.selectedClusters(condensedTree, stabilities, nodeCount);
+
+        assertThat(selectedClusters.get(0))
+            .withFailMessage("Root should be selected")
+            .isTrue();
+        assertThat(selectedClusters.get(1))
+            .withFailMessage("First child should be selected")
+            .isTrue();
+        assertThat(selectedClusters.get(2))
+            .withFailMessage("Second child should be selected")
+            .isTrue();
+    }
 }