Implement balanced K-Means clustering; add new clustering type option in UI and update event handler to support balanced K-Means functionality.

clustering.js

@@ -1,5 +1,45 @@
 import { UMAP } from "https://cdn.jsdelivr.net/npm/umap-js@1.4.0/+esm";
 
+function balancedKMeans(emb, k, beta = 1, maxIter = 100) {
+    const n = emb.length, d = emb[0].length;
+    let cent = kmeansPlusPlusInit(emb, k);
+    const lab = new Uint32Array(n);
+    const cnt = new Uint32Array(k);
+
+    for (let iter = 0; iter < maxIter; ++iter) {
+        let moved = false;
+
+        // ── assignment with size penalty ──
+        cnt.fill(0);
+        for (let i = 0; i < n; ++i) {
+            let best = 0, bestCost = Infinity;
+            for (let c = 0; c < k; ++c) {
+                let dist = 0;
+                for (let j = 0; j < d; ++j) {
+                    const diff = emb[i][j] - cent[c][j];
+                    dist += diff * diff;
+                }
+                const sizePenalty = beta * (2 * cnt[c] + 1);
+                const cost = dist + sizePenalty;
+                if (cost < bestCost) { bestCost = cost; best = c; }
+            }
+            if (lab[i] !== best) { lab[i] = best; moved = true; }
+            cnt[lab[i]]++;
+        }
+
+        // ── update centroids ──
+        cent = Array.from({ length: k }, () => new Array(d).fill(0));
+        for (let i = 0; i < n; ++i)
+            for (let j = 0; j < d; ++j) cent[lab[i]][j] += emb[i][j];
+        for (let c = 0; c < k; ++c)
+            if (cnt[c]) for (let j = 0; j < d; ++j) cent[c][j] /= cnt[c];
+
+        if (!moved) break;
+    }
+    return { labels: Array.from(lab), centroids: cent };
+}
+
+
 function kmeansPlusPlusInit(embeddings, k) {
     const n = embeddings.length;
     const dim = embeddings[0].length;
@@ -102,3 +142,5 @@ export function runUMAP(embeddings, nNeighbors = 15) {
     });
     return umap.fit(embeddings);
 }
+
+export { balancedKMeans };

index.html

@@ -58,6 +58,11 @@
     </script>
     <label for="kmeans-k" style="margin-left:10px;">Clusters:</label>
     <input id="kmeans-k" type="number" min="2" max="100" value="7" style="width:60px;">
+    <label for="kmeans-type" style="margin-left:10px;">Clustering Type:</label>
+    <select id="kmeans-type" style="width:180px;">
+      <option value="kmeans">K-Means (standard)</option>
+      <option value="balancedKMeans">Balanced K-Means</option>
+    </select>
     <button id="kmeans-btn">K-Means Clustering</button>
     <button id="heatmap-btn">Similarity Heatmap</button>
     <div id="progress-bar">

main.js

@@ -1,5 +1,5 @@
 import { getGroupEmbeddings, getLineEmbeddings } from './embedding.js';
-import { kmeans } from './clustering.js';
+import { kmeans, balancedKMeans } from './clustering.js';
 import { plotHeatmap, plotScatter, updateScatter } from './plotting.js';
 import { nameCluster } from './cluster_naming.js';
 import { prompt_cluster } from './prompt_cluster.js';
@@ -57,8 +57,14 @@ document.getElementById("kmeans-btn").onclick = async () => {
     if (n < 2) return;
     const requestedK = parseInt(document.getElementById("kmeans-k").value) || 3;
     const k = Math.max(2, Math.min(requestedK, n));
-    // K-Means clustering
-    const { labels } = kmeans(embeddings, k);
+    // Read clustering type
+    const clusteringType = document.getElementById("kmeans-type").value;
+    let labels;
+    if (clusteringType === "balancedKMeans") {
+        labels = balancedKMeans(embeddings, k).labels;
+    } else {
+        labels = kmeans(embeddings, k).labels;
+    }
     // UMAP projection
     const { UMAP } = await import('https://cdn.jsdelivr.net/npm/umap-js@1.4.0/+esm');
     const nNeighbors = Math.max(1, Math.min(lines.length - 1, 15));