Refactor balanced K-Means implementation; update beta parameter default value, handle empty embeddings, and ensure proper centroid updates. Enhance UI to include beta input and adjust event handler to read beta value for clustering.

clustering.js

@@ -1,16 +1,20 @@
 import { UMAP } from "https://cdn.jsdelivr.net/npm/umap-js@1.4.0/+esm";
 
-function balancedKMeans(emb, k, beta = 1, maxIter = 100) {
+function balancedKMeans(emb, k, beta = 0.01, maxIter = 100) {
+    if (!emb.length) return { labels: [], centroids: [] };
+
     const n = emb.length, d = emb[0].length;
+    k = Math.max(2, Math.min(k, n));                 // guard k ≤ n
+
     let cent = kmeansPlusPlusInit(emb, k);
-    const lab = new Uint32Array(n);
+    const lab = new Uint32Array(n).fill(k);          // start “unassigned”
     const cnt = new Uint32Array(k);
 
     for (let iter = 0; iter < maxIter; ++iter) {
         let moved = false;
+        cnt.fill(0);
 
         // ── 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) {
@@ -24,18 +28,23 @@ function balancedKMeans(emb, k, beta = 1, maxIter = 100) {
                 if (cost < bestCost) { bestCost = cost; best = c; }
             }
             if (lab[i] !== best) { lab[i] = best; moved = true; }
-            cnt[lab[i]]++;
+            cnt[best]++;
         }
 
         // ── 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 (cnt[c]) {
+                const inv = 1 / cnt[c];
+                for (let j = 0; j < d; ++j) cent[c][j] *= inv;
+            }
 
-        if (!moved) break;
+        if (!moved) break;                            // converged
     }
+
     return { labels: Array.from(lab), centroids: cent };
 }
 

index.html

@@ -53,17 +53,19 @@
     <h1>Embedding Similarity Heatmap</h1>
     <textarea id="input"></textarea>
     <script type="module">
-      import { sentences } from './sentences.js';
-      document.getElementById("input").value = sentences.join("\n");
+        import { sentences } from './sentences.js';
+        document.getElementById("input").value = sentences.join("\n");
     </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-beta" style="margin-left:10px;">Beta (balance):</label>
+    <input id="kmeans-beta" type="number" min="0" max="10" step="0.0001" value="0.01" style="width:80px;">
     <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>
+        <option value="balancedKMeans">Balanced K-Means</option>
+        <option value="kmeans">K-Means (standard)</option>
     </select>
-    <button id="kmeans-btn">K-Means Clustering</button>
+    <button id="kmeans-btn">Clustering</button>
     <button id="heatmap-btn">Similarity Heatmap</button>
     <div id="progress-bar">
         <div id="progress-bar-inner"></div>

main.js

@@ -57,11 +57,12 @@ 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));
-    // Read clustering type
+    // Read clustering type and beta
     const clusteringType = document.getElementById("kmeans-type").value;
+    const beta = parseFloat(document.getElementById("kmeans-beta").value) || 0.01;
     let labels;
     if (clusteringType === "balancedKMeans") {
-        labels = balancedKMeans(embeddings, k).labels;
+        labels = balancedKMeans(embeddings, k, beta).labels;
     } else {
         labels = kmeans(embeddings, k).labels;
     }