Include the total distance in the EMST result

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

Author: vnickolov

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

@@ -31,7 +31,7 @@
 import org.neo4j.gds.core.utils.paged.dss.HugeAtomicDisjointSetStruct;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
-public class DualTreeMSTAlgorithm extends Algorithm<HugeObjectArray<Edge>> {
+public class DualTreeMSTAlgorithm extends Algorithm<DualTreeMSTResult> {
 
     private final NodePropertyValues nodePropertyValues;
     private final KdTree kdTree;
@@ -44,6 +44,7 @@ public class DualTreeMSTAlgorithm extends Algorithm<HugeObjectArray<Edge>> {
 
     private final HugeObjectArray<Edge> edges;
     private long edgeCount = 0;
+    private double totalEdgeSum = 0d;
 
     public DualTreeMSTAlgorithm(
         NodePropertyValues nodePropertyValues,
@@ -69,15 +70,15 @@ public DualTreeMSTAlgorithm(
 
 
     @Override
-    public HugeObjectArray<Edge> compute() {
+    public DualTreeMSTResult compute() {
 
         var kdRoot = kdTree.root();
         var rootId = kdRoot.id();
         while (!kdNodeSingleComponent.get(rootId)) {
             kdNodeBound.fill(Double.MAX_VALUE);
             performIteration();
         }
-        return edges;
+        return new DualTreeMSTResult(edges, totalEdgeSum);
     }
 
     void baseCase(long p0, long p1, MutableDouble maxBound) {
@@ -252,11 +253,13 @@ void mergeComponents() {
                 continue;
             }
 
+            var distance = Math.sqrt(closestDistanceTracker.componentClosestDistance(componentId));
             this.edges.set(
                 edgeCount,
-                new Edge(u, v, Math.sqrt(closestDistanceTracker.componentClosestDistance(componentId)))
+                new Edge(u, v, distance)
             );
             this.edgeCount++;
+            this.totalEdgeSum += distance;
 
             unionFind.union(uComponent, vComponent);
         }

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

@@ -0,0 +1,25 @@
+/*
+ * 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.HugeObjectArray;
+
+public record DualTreeMSTResult(HugeObjectArray<Edge> edges, double totalDistance) {
+}

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

@@ -79,7 +79,7 @@ HugeObjectArray<Edge> dualTreeMSTPhase(KdTree kdTree, HugeDoubleArray coreValues
                 kdTree,coreValues,
                 nodes.nodeCount()
             );
-            return dualTreeMst.compute();
+            return dualTreeMst.compute().edges();
     }
     KdTree buildKDTree(){
         var builder =new KdTreeBuilder(nodes,nodePropertyValues,concurrency.value(),leafSize);

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

@@ -19,6 +19,7 @@
  */
 package org.neo4j.gds.hdbscan;
 
+import org.assertj.core.data.Offset;
 import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
 import org.neo4j.gds.collections.ha.HugeDoubleArray;
@@ -62,7 +63,7 @@ void shouldReturnEuclideanMSTWithZeroCoreValues() {
             var cores = HugeDoubleArray.newArray(graph.nodeCount());
 
             var dualTree = new DualTreeMSTAlgorithm(nodePropertyValues, kdTree, cores, graph.nodeCount());
-            var edges = dualTree.compute();
+            var result = dualTree.compute();
 
             var expected = List.of(
                 new Edge(graph.toMappedNodeId("g"), graph.toMappedNodeId("h"), 2.23606797749979),
@@ -75,9 +76,15 @@ void shouldReturnEuclideanMSTWithZeroCoreValues() {
                 new Edge(graph.toMappedNodeId("b"), graph.toMappedNodeId("d"), 3.1622776601683795)
             );
 
-            assertThat(edges.toArray())
+            assertThat(result.edges().toArray())
                 .usingElementComparator(new UndirectedEdgeComparator())
                 .containsExactlyInAnyOrderElementsOf(expected);
+
+            assertThat(result.totalDistance())
+                .isEqualTo(expected.stream()
+                    .mapToDouble(Edge::distance)
+                    .sum()
+                );
         }
 
     }
@@ -107,7 +114,7 @@ void shouldReturnEuclideanMSTWithZeroCoreValues() {
             var cores = HugeDoubleArray.newArray(graph.nodeCount());
 
             var dualTree = new DualTreeMSTAlgorithm(nodePropertyValues, kdTree, cores, graph.nodeCount());
-            var edges = dualTree.compute();
+            var result = dualTree.compute();
 
             var expected = List.of(
                 new Edge(graph.toMappedNodeId("a"), graph.toMappedNodeId("b"), 3.1622776601683795),
@@ -117,9 +124,17 @@ void shouldReturnEuclideanMSTWithZeroCoreValues() {
                 new Edge(graph.toMappedNodeId("e"), graph.toMappedNodeId("f"), 1.4142135623730951)
             );
 
-            assertThat(edges.toArray())
+            assertThat(result.edges().toArray())
                 .usingElementComparator(new UndirectedEdgeComparator())
                 .containsExactlyInAnyOrderElementsOf(expected);
+
+            assertThat(result.totalDistance())
+                .isCloseTo(
+                    expected.stream()
+                        .mapToDouble(Edge::distance)
+                        .sum(),
+                    Offset.offset(1e-12)
+                );
         }
 
     }
@@ -150,7 +165,7 @@ void shouldReturnEuclideanMSTWithZeroCoreValues() {
             var cores = HugeDoubleArray.newArray(graph.nodeCount());
 
             var dualTree = new DualTreeMSTAlgorithm(nodePropertyValues, kdTree, cores, graph.nodeCount());
-            var edges = dualTree.compute();
+            var result = dualTree.compute();
 
             var expected = List.of(
                 new Edge(graph.toMappedNodeId("a"), graph.toMappedNodeId("d"), 0.5),
@@ -160,9 +175,17 @@ void shouldReturnEuclideanMSTWithZeroCoreValues() {
                 new Edge(graph.toMappedNodeId("c"), graph.toMappedNodeId("e"), 997.0045135304052)
             );
 
-            assertThat(edges.toArray())
+            assertThat(result.edges().toArray())
                 .usingElementComparator(new UndirectedEdgeComparator())
                 .containsExactlyInAnyOrderElementsOf(expected);
+
+            assertThat(result.totalDistance())
+                .isCloseTo(
+                    expected.stream()
+                        .mapToDouble(Edge::distance)
+                        .sum(),
+                    Offset.offset(1e-12)
+                );
         }
 
     }