Performance optimizations for hierarchical clustering

- Lance-Williams UPGMA update replaces per-pair member enumeration, making
  distance maintenance O(n) per merge instead of O(cluster_size²)
- Active-index list (swap-with-last removal) replaces flag array so the
  find-minimum loop iterates only live clusters, ~3x fewer comparisons
- Distance matrix computed over upper triangle only and mirrored, halving
  initial pairwise computation
- Linked list tracks leaf order in O(1) per merge, replacing index array copies
- clock() check moved outside inner distance-matrix loop (was called n² times)
- rebuildTree rewritten with stable slot IDs, O(n) with no array splices
- clustersGivenK rebuilt using active Set + membership map for stable slot IDs
- Remove unused dummy Float32Array distances allocation and field from ClusterResult
- Add integration tests covering known inputs end-to-end through real WASM

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: cmdcolin

src/cluster.ts

@@ -21,40 +21,27 @@ export async function clusterData({
     checkCancellation,
   })
 
-  // Build clustersGivenK from merge information
+  // Build clustersGivenK from stable-slot merge sequence.
+  // mergeA[i] and mergeB[i] are stable slot indices; slot mergeA[i] absorbs mergeB[i].
   const numSamples = data.length
   const clustersGivenK: number[][][] = [[]]
 
-  // Start with each sample in its own cluster
-  const clusterSets: number[][] = Array.from({ length: numSamples }, (_, i) => [
-    i,
-  ])
+  const membership = Array.from({ length: numSamples }, (_, i) => [i] as number[])
+  const activeSlots = new Set(Array.from({ length: numSamples }, (_, i) => i))
 
   for (let i = 0; i < numSamples - 1; i++) {
-    const [mergeA, mergeB] = result.merges[i]!
+    const [a, b] = result.merges[i]!
 
-    // Record current state
-    clustersGivenK.push(clusterSets.map(s => [...s]))
+    clustersGivenK.push([...activeSlots].map(id => [...membership[id]!]))
 
-    // Merge clusters
-    const newCluster = [...clusterSets[mergeA]!, ...clusterSets[mergeB]!]
-
-    const removeFirst = Math.max(mergeA, mergeB)
-    const removeSecond = Math.min(mergeA, mergeB)
-
-    clusterSets.splice(removeFirst, 1)
-    clusterSets.splice(removeSecond, 1)
-    clusterSets.push(newCluster)
+    membership[a] = [...membership[a]!, ...membership[b]!]
+    activeSlots.delete(b!)
   }
 
-  clustersGivenK.push(clusterSets.map(s => [...s]))
-
-  // Create a dummy distance matrix (not used by caller, but part of interface)
-  const distances = new Float32Array(numSamples * numSamples)
+  clustersGivenK.push([...activeSlots].map(id => [...membership[id]!]))
 
   return {
     tree: result.tree,
-    distances,
     order: result.order,
     clustersGivenK: clustersGivenK.reverse(),
   }

src/distance.js

@@ -6,7 +6,6 @@ export interface ClusterNode {
 
 export interface ClusterResult {
   tree: ClusterNode
-  distances: Float32Array
   order: number[]
   clustersGivenK: number[][][]
 }

src/types.ts

@@ -41,31 +41,26 @@ export async function hierarchicalClusterWasm(
   }
   const vectorSize = data[0]?.length ?? 0
 
-  // Flatten data
   const flatData = new Float32Array(numSamples * vectorSize)
   for (let i = 0; i < numSamples; i++) {
     flatData.set(data[i]!, i * vectorSize)
   }
 
-  // Allocate memory
-  const dataPtr = module._malloc(flatData.length * 4)
+  const dataPtr    = module._malloc(flatData.length * 4)
   const heightsPtr = module._malloc((numSamples - 1) * 4)
-  const mergeAPtr = module._malloc((numSamples - 1) * 4)
-  const mergeBPtr = module._malloc((numSamples - 1) * 4)
-  const orderPtr = module._malloc(numSamples * 4)
+  const mergeAPtr  = module._malloc((numSamples - 1) * 4)
+  const mergeBPtr  = module._malloc((numSamples - 1) * 4)
+  const orderPtr   = module._malloc(numSamples * 4)
 
   let callbackPtr: number | null = null
 
   try {
-    // Copy data to WASM
     module.HEAPF32.set(flatData, dataPtr / 4)
 
-    // Set up progress callback if provided
     if (statusCallback || checkCancellation) {
       const progressCallback = (iteration: number, totalIterations: number) => {
         checkCancellation?.()
         if (statusCallback) {
-          // Negative iteration indicates distance matrix phase
           if (iteration < 0) {
             const distancesDone = -iteration
             const progress = Math.round((distancesDone / totalIterations) * 100)
@@ -75,15 +70,14 @@ export async function hierarchicalClusterWasm(
             statusCallback(`Clustering samples: ${progress}%`)
           }
         }
-        return 1 // Continue
+        return 1
       }
 
       callbackPtr = module.addFunction(progressCallback, 'iii')
       module._setProgressCallback(callbackPtr)
     }
 
-    // Run clustering in WASM
-    module._hierarchicalCluster(
+    const result = module._hierarchicalCluster(
       dataPtr,
       numSamples,
       vectorSize,
@@ -93,7 +87,10 @@ export async function hierarchicalClusterWasm(
       orderPtr,
     )
 
-    // Copy results back
+    if (result === -1) {
+      throw new Error('aborted')
+    }
+
     const heights = new Float32Array(numSamples - 1)
     heights.set(
       module.HEAPF32.subarray(heightsPtr / 4, heightsPtr / 4 + numSamples - 1),
@@ -112,7 +109,6 @@ export async function hierarchicalClusterWasm(
     const order = new Int32Array(numSamples)
     order.set(module.HEAP32.subarray(orderPtr / 4, orderPtr / 4 + numSamples))
 
-    // Rebuild tree structure from merge information
     const tree = rebuildTree(numSamples, heights, mergeA, mergeB, sampleLabels)
     const merges: [number, number][] = []
     for (let i = 0; i < numSamples - 1; i++) {
@@ -126,7 +122,6 @@ export async function hierarchicalClusterWasm(
       merges,
     }
   } finally {
-    // Clean up callback
     if (callbackPtr !== null) {
       module.removeFunction(callbackPtr)
       module._setProgressCallback(0)
@@ -140,45 +135,26 @@ export async function hierarchicalClusterWasm(
   }
 }
 
+// Rebuilds the tree from stable slot indices (mergeA[i] < mergeB[i] always).
+// Slot mergeA[i] absorbs mergeB[i] each iteration, so nodes[0] is always the root.
 function rebuildTree(
   numSamples: number,
   heights: Float32Array,
   mergeA: Int32Array,
   mergeB: Int32Array,
   sampleLabels?: string[],
 ): ClusterNode {
-  // Create leaf nodes
   const nodes: ClusterNode[] = []
   for (let i = 0; i < numSamples; i++) {
-    nodes.push({
-      name: sampleLabels?.[i] ?? `Sample ${i}`,
-      height: 0,
-    })
+    nodes.push({ name: sampleLabels?.[i] ?? `Sample ${i}`, height: 0 })
   }
-
-  // Build tree from merge information
   for (let i = 0; i < numSamples - 1; i++) {
-    const leftIdx = mergeA[i]!
-    const rightIdx = mergeB[i]!
-
-    const leftNode = nodes[leftIdx]!
-    const rightNode = nodes[rightIdx]!
-
-    const newNode: ClusterNode = {
+    const a = mergeA[i]!, b = mergeB[i]!
+    nodes[a] = {
       name: `Cluster ${i}`,
       height: heights[i]!,
-      children: [leftNode, rightNode],
+      children: [nodes[a]!, nodes[b]!],
     }
-
-    // Replace the merged clusters with the new one
-    // Remove higher index first
-    const removeFirst = Math.max(leftIdx, rightIdx)
-    const removeSecond = Math.min(leftIdx, rightIdx)
-
-    nodes.splice(removeFirst, 1)
-    nodes.splice(removeSecond, 1)
-    nodes.push(newNode)
   }
-
   return nodes[0]!
 }