Enhance heatmap event handling by incorporating search group functionality, allowing users to specify a search group for improved similarity calculations and visualization in the heatmap.

heatmap_event.js

@@ -58,17 +58,18 @@ export async function handleHeatmapEvent() {
     progressBarInner.style.width = "0%";
 
     const text = document.getElementById("input").value;
+    // Get search group from dedicated input (do not use ##search in main input)
+    const searchGroupText = document.getElementById("search-group-input")?.value.trim();
     // Get search sort mode from dropdown (either 'line' or 'group')
     const searchSortMode = document.getElementById("search-sort-mode")?.value || "group";
     const search_by_max_search_line = searchSortMode === "line";
     const search_by_max_search_group = searchSortMode === "group";
 
-    // Find the index of the search cluster (case-insensitive, looks for 'search' in header)
+    // Parse cluster names from main input (ignore any ##search)
     const clusterNames = text.split(/\n/)
         .map(x => x.trim())
         .filter(x => x && x.startsWith('##'))
         .map(x => x.replace(/^##\s*/, ''));
-    let searchIdx = clusterNames.findIndex(name => name.toLowerCase().includes('search'));
 
     const groups = text.split(/\n{3,}/);
     // Get group embeddings (removes ## lines internally)
@@ -91,53 +92,72 @@ export async function handleHeatmapEvent() {
         sim.push(row);
     }
     progressBarInner.style.width = "60%";
-    // Reorder clusters if search cluster is found, otherwise keep original order
     // Always use all group indices in order
     let order = Array.from({ length: n }, (_, i) => i);
-    if (searchIdx !== -1) {
-        // Compute similarity to search cluster and sort others by similarity
-        const simToSearch = sim[searchIdx].map((v, i) => ({ idx: i, sim: v }));
-        // Exclude searchIdx from sorting
-        const others = simToSearch.filter(x => x.idx !== searchIdx);
-        others.sort((a, b) => b.sim - a.sim); // descending similarity
-        order = [searchIdx, ...others.map(x => x.idx)];
+
+    // Only use search group if provided in search-group-input
+    let useSearchGroup = !!searchGroupText;
+    let searchIdx = -1;
+    let searchLines = [];
+    let searchEmbeds = [];
+    let refEmbed = null;
+    if (useSearchGroup) {
+        searchLines = searchGroupText.split(/\n/).map(l => l.trim()).filter(l => l);
+        if (searchLines.length > 0) {
+            searchEmbeds = await getLineEmbeddings(searchLines, task);
+            // For group similarity, use the mean embedding of the search group
+            refEmbed = searchEmbeds[0].map((_, i) => searchEmbeds.reduce((sum, e) => sum + e[i], 0) / searchEmbeds.length);
+            // Compute similarity to each group
+            const simToSearch = groupEmbeddings.map((emb, i) => ({ idx: i, sim: cosine(refEmbed, emb) }));
+            simToSearch.sort((a, b) => b.sim - a.sim);
+            order = [/* search group is not in groupEmbeddings, so just prepend -1 for heatmap */ -1, ...simToSearch.map(x => x.idx)];
+        }
     }
     // Reorder sim matrix and clusterNames for heatmap visualization
-    const simOrdered = order.map(i => order.map(j => sim[i][j]));
-    const xLabels = order.map(i => (clusterNames && clusterNames[i]) ? clusterNames[i] : `Group ${i + 1}`);
+    let simOrdered, xLabels;
+    if (useSearchGroup && searchLines.length > 0) {
+        // Insert search group as first row/col in heatmap, with similarity 1 to itself and to other groups
+        simOrdered = [
+            [1, ...order.slice(1).map(idx => idx === -1 ? 1 : cosine(refEmbed, groupEmbeddings[idx]))],
+            ...order.slice(1).map(i => [cosine(refEmbed, groupEmbeddings[i]), ...order.slice(1).map(j => sim[i][j])])
+        ];
+        xLabels = ["Search Group", ...order.slice(1).map(i => (clusterNames && clusterNames[i]) ? clusterNames[i] : `Group ${i + 1}`)];
+    } else {
+        simOrdered = order.map(i => order.map(j => sim[i][j]));
+        xLabels = order.map(i => (clusterNames && clusterNames[i]) ? clusterNames[i] : `Group ${i + 1}`);
+    }
     plotHeatmap(simOrdered, xLabels, xLabels);
 
-    // If search cluster is found and sorting by line, reorder lines in each group by similarity to search lines
-    if (searchIdx !== -1 && search_by_max_search_line) {
-        const searchLines = getCleanGroups(groups)[searchIdx];
-        const searchEmbeds = await getLineEmbeddings(searchLines, task);
+    // If search group is provided and sorting by line, reorder lines in each group by similarity to search lines
+    if (useSearchGroup && search_by_max_search_line && searchEmbeds.length > 0) {
         const cleanGroups = getCleanGroups(groups);
         const { allLines, allEmbeds } = await getAllLinesAndEmbeds(cleanGroups, task);
         const idxByGroup = getIdxByGroup(cleanGroups);
         const score = e => Math.max(...searchEmbeds.map(se => cosine(se, e)));
-        const sorted = order.map(g =>
+        // Skip -1 (search group) in order for main input reordering
+        const sorted = (order[0] === -1 ? order.slice(1) : order).map(g =>
             idxByGroup[g]
                 .map(i => ({ t: allLines[i], s: score(allEmbeds[i]) }))
                 .sort((a, b) => b.s - a.s)
                 .map(o => o.t)
         );
-        const finalText = buildFinalText(order, sorted, clusterNames, n);
+        const finalText = buildFinalText(order[0] === -1 ? order.slice(1) : order, sorted, clusterNames, n);
         document.getElementById("input").value = finalText;
     }
 
-    // If search cluster is found and sorting by group, reorder lines in each group by similarity to the search group embedding
-    if (searchIdx !== -1 && search_by_max_search_group) {
-        const refEmbed = groupEmbeddings[searchIdx];
+    // If search group is provided and sorting by group, reorder lines in each group by similarity to the search group embedding
+    if (useSearchGroup && search_by_max_search_group && refEmbed) {
         const cleanGroups = getCleanGroups(groups);
         const { allLines, allEmbeds } = await getAllLinesAndEmbeds(cleanGroups, task);
         const idxByGroup = getIdxByGroup(cleanGroups);
-        const sortedLines = order.map(gIdx =>
+        // Skip -1 (search group) in order for main input reordering
+        const sortedLines = (order[0] === -1 ? order.slice(1) : order).map(gIdx =>
             idxByGroup[gIdx]
                 .map(i => ({ t: allLines[i], s: cosine(refEmbed, allEmbeds[i]) }))
                 .sort((a, b) => b.s - a.s)
                 .map(o => o.t)
         );
-        const finalText = buildFinalText(order, sortedLines, clusterNames, n);
+        const finalText = buildFinalText(order[0] === -1 ? order.slice(1) : order, sortedLines, clusterNames, n);
         document.getElementById("input").value = finalText;
     }
     progressBarInner.style.width = "100%";