feat: dimension reduction algorithm selection (UMAP / PCA / t-SNE)

README.md

@@ -20,8 +20,8 @@ A full-stack tool for generating, indexing, and visualizing embedding
 clusters from CSV data. Feed it a CSV with image URLs or text fields,
 and it generates vector embeddings using CLIP (images) and
 SentenceTransformer (text) models, stores them in ChromaDB, clusters
-via k-means, and renders an interactive 3D scatter plot using t-SNE
-dimensionality reduction.
+via k-means, and renders an interactive 3D scatter plot with
+selectable dimensionality reduction (t-SNE, UMAP, or PCA).
 
 The project supports three running modes: CLI-based batch indexing,
 CLI-based Dash visualization, and a full web application with a
@@ -39,7 +39,8 @@ FastAPI backend and React frontend.
   curve) and silhouette score analysis across k=2..30, with an
   interactive chart to review the trade-off before applying.
 * Visualize clustering results in an interactive 3D scatter plot
-  using t-SNE dimensionality reduction.
+  with selectable dimensionality reduction: t-SNE, UMAP, or PCA,
+  each with configurable algorithm-specific parameters.
 * Web UI with FastAPI backend and React frontend for browser-based
   indexing and visualization.
 * Three switchable 3D render modes: colored particles, image sprites,
@@ -71,7 +72,7 @@ FastAPI backend and React frontend.
                        |
             StandardScaler + KMeans
                        |
-                t-SNE (3D projection)
+          t-SNE / UMAP / PCA (3D projection)
                        |
               3D Scatter Plot
            (Dash CLI / React Web UI)
@@ -102,6 +103,9 @@ pre-commit hooks, commitizen
 
 > Models are downloaded from [HuggingFace](https://huggingface.co)
 > on first run. Ensure network access to `huggingface.co`.
+> UMAP support is optional. To install it:
+> `uv sync --extra umap` or `uv pip install umap-learn`.
+> t-SNE and PCA work without extra dependencies.
 
 ## Usage
 
@@ -172,6 +176,12 @@ collection name is the prefix combined with the embedded field name.
 | `GPT_GENERATE_CLUSTER_NAME` | Use GPT to name clusters (needs `OPENAI_API_KEY`) | No | `False` |
 | `GPT_DEFAULT_MODEL` | GPT model for cluster naming | No | `gpt-3.5-turbo` |
 | `GPT_DEFAULT_TEMPERATURE` | GPT temperature for cluster naming | No | `0.51` |
+| `REDUCTION_ALGORITHM` | Dimensionality reduction (`tsne`, `umap`, `pca`) | No | `tsne` |
+| `TSNE_PERPLEXITY` | t-SNE perplexity (5--50) | No | `30.0` |
+| `TSNE_LEARNING_RATE` | t-SNE learning rate (`auto` or numeric) | No | `auto` |
+| `UMAP_N_NEIGHBORS` | UMAP neighbor count (2--100) | No | `15` |
+| `UMAP_MIN_DIST` | UMAP minimum distance (0--1) | No | `0.1` |
+| `UMAP_METRIC` | UMAP distance metric | No | `cosine` |
 
 ```bash
 RUNNING_MODE=PLOT \

embedding_cluster/scatter_plot.py

@@ -10,6 +10,7 @@
 from dash import Dash, Input, Output, callback, dcc, html, no_update
 from openai import OpenAI
 from sklearn.cluster import KMeans
+from sklearn.decomposition import PCA
 from sklearn.manifold import TSNE
 from sklearn.metrics import silhouette_score
 from sklearn.preprocessing import StandardScaler
@@ -24,6 +25,52 @@
 logger = logging.getLogger(__name__)
 
 
+def reduce_dimensions(
+    embeddings: np.ndarray,
+    algorithm: str = "tsne",
+    n_components: int = 3,
+    random_state: int = 171,
+    **kwargs: Any,
+) -> np.ndarray:
+    """Reduce embedding dimensions using the specified algorithm."""
+    if algorithm == "tsne":
+        perplexity = kwargs.get("perplexity", 30.0)
+        learning_rate = kwargs.get("learning_rate", "auto")
+        reducer = TSNE(
+            n_components=n_components,
+            perplexity=perplexity,
+            learning_rate=learning_rate,
+            random_state=random_state,
+            verbose=1,
+            max_iter=1000,
+        )
+    elif algorithm == "umap":
+        try:
+            import umap
+        except ImportError as exc:
+            msg = (
+                "umap-learn is not installed. Install it with: uv pip install umap-learn"
+            )
+            raise ImportError(msg) from exc
+        n_neighbors = kwargs.get("n_neighbors", 15)
+        min_dist = kwargs.get("min_dist", 0.1)
+        metric = kwargs.get("metric", "cosine")
+        reducer = umap.UMAP(
+            n_components=n_components,
+            n_neighbors=n_neighbors,
+            min_dist=min_dist,
+            metric=metric,
+            random_state=random_state,
+        )
+    elif algorithm == "pca":
+        reducer = PCA(n_components=n_components)
+    else:
+        msg = f"Unknown reduction algorithm: '{algorithm}'. Supported: tsne, umap, pca"
+        raise ValueError(msg)
+    result: np.ndarray = reducer.fit_transform(embeddings)
+    return result
+
+
 def gpt_get_cluster_name(info: str, settings: Settings) -> str:
     openai_client = OpenAI()
     messages: list[dict[str, str]] = [
@@ -196,7 +243,7 @@ def generate_cluster_props(
 
 
 def compute_plot_data(settings: Settings) -> dict[str, Any]:
-    """Compute t-SNE + k-means and return raw data (no Plotly/Dash)."""
+    """Compute dimensionality reduction + k-means and return raw data."""
     random_state = 171
     n_iter = 1000
     collection_content_text_display: list[str] = []
@@ -211,15 +258,23 @@ def compute_plot_data(settings: Settings) -> dict[str, Any]:
     collection_content_vectors = collection_content["embeddings"]
     np_embeddings_arr = np.array(collection_content_vectors)
 
-    logger.info("Calculating t-SNE ...")
-    tsne = TSNE(
-        verbose=1,
-        learning_rate="auto",
-        max_iter=n_iter,
-        perplexity=30,
+    algorithm = settings.reduction_algorithm
+    logger.info("Calculating %s ...", algorithm.upper())
+    reduction_kwargs: dict[str, Any] = {}
+    if algorithm == "tsne":
+        reduction_kwargs["perplexity"] = settings.tsne_perplexity
+        reduction_kwargs["learning_rate"] = settings.tsne_learning_rate
+    elif algorithm == "umap":
+        reduction_kwargs["n_neighbors"] = settings.umap_n_neighbors
+        reduction_kwargs["min_dist"] = settings.umap_min_dist
+        reduction_kwargs["metric"] = settings.umap_metric
+    reduced = reduce_dimensions(
+        np_embeddings_arr,
+        algorithm=algorithm,
         n_components=3,
         random_state=random_state,
-    ).fit_transform(np_embeddings_arr)
+        **reduction_kwargs,
+    )
 
     common_params: dict[str, Any] = {
         "n_init": "auto",
@@ -262,9 +317,9 @@ def compute_plot_data(settings: Settings) -> dict[str, Any]:
             )
             points.append(
                 {
-                    "x": float(tsne[idx, 0]),
-                    "y": float(tsne[idx, 1]),
-                    "z": float(tsne[idx, 2]),
+                    "x": float(reduced[idx, 0]),
+                    "y": float(reduced[idx, 1]),
+                    "z": float(reduced[idx, 2]),
                     "cluster": cluster_i,
                     "metadata": metadata,
                     "id": point_id,

embedding_cluster/server/models.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-from pydantic import BaseModel
+from typing import Literal
+
+from pydantic import BaseModel, field_validator
 
 
 class CollectionInfo(BaseModel):
@@ -71,6 +73,24 @@ class PlotRequest(BaseModel):
     gpt_generate_cluster_name: bool = False
     gpt_default_model: str = "gpt-3.5-turbo"
     gpt_default_temperature: float = 0.51
+    reduction_algorithm: Literal["tsne", "umap", "pca"] = "tsne"
+    tsne_perplexity: float = 30.0
+    tsne_learning_rate: str = "auto"
+    umap_n_neighbors: int = 15
+    umap_min_dist: float = 0.1
+    umap_metric: str = "cosine"
+
+    @field_validator("reduction_algorithm")
+    @classmethod
+    def validate_algorithm(cls, v: str) -> str:
+        allowed = {"tsne", "umap", "pca"}
+        if v not in allowed:
+            msg = (
+                f"Invalid reduction algorithm: '{v}'. "
+                f"Must be one of: {', '.join(sorted(allowed))}"
+            )
+            raise ValueError(msg)
+        return v
 
 
 class SuggestClustersRequest(BaseModel):

embedding_cluster/server/routes/plot.py

@@ -36,6 +36,12 @@ async def _run_compute(task_state: TaskState, request: PlotRequest) -> None:
             gpt_generate_cluster_name=request.gpt_generate_cluster_name,
             gpt_default_model=request.gpt_default_model,
             gpt_default_temperature=request.gpt_default_temperature,
+            reduction_algorithm=request.reduction_algorithm,
+            tsne_perplexity=request.tsne_perplexity,
+            tsne_learning_rate=request.tsne_learning_rate,
+            umap_n_neighbors=request.umap_n_neighbors,
+            umap_min_dist=request.umap_min_dist,
+            umap_metric=request.umap_metric,
         )
         task_state.status = TaskStatus.RUNNING
         result = await asyncio.to_thread(compute_plot_data, settings)

embedding_cluster/settings.py

@@ -66,6 +66,16 @@ class Settings(BaseSettings):
         default=None, description="field name for the image to present on plot"
     )
 
+    reduction_algorithm: str = Field(
+        default="tsne",
+        description="Dimensionality reduction algorithm: tsne, umap, or pca",
+    )
+    tsne_perplexity: float = Field(default=30.0, description="t-SNE perplexity parameter")
+    tsne_learning_rate: str = Field(default="auto", description="t-SNE learning rate")
+    umap_n_neighbors: int = Field(default=15, description="UMAP number of neighbors")
+    umap_min_dist: float = Field(default=0.1, description="UMAP minimum distance")
+    umap_metric: str = Field(default="cosine", description="UMAP distance metric")
+
     gpt_generate_cluster_name: bool = Field(
         default=False, description="Generate cluster names using GPT"
     )

frontend/src/components/plot/PlotControls.tsx

@@ -3,7 +3,7 @@ import { useState, useEffect } from 'react'
 import { useSearchParams } from 'react-router-dom'
 import { fetchCollections, fetchCollection } from '../../api/collections'
 import { usePlotStore } from '../../stores/plotStore'
-import type { PlotRequest } from '../../types'
+import type { PlotRequest, ReductionAlgorithm } from '../../types'
 import ClusterSuggestion from './ClusterSuggestion'
 
 interface PlotControlsProps {
@@ -21,7 +21,15 @@ export default function PlotControls({ onCompute, isComputing }: PlotControlsPro
   const [gptModel, setGptModel] = useState('gpt-3.5-turbo')
   const [gptTemperature, setGptTemperature] = useState(0.51)
 
-  const { renderMode, setRenderMode, pointSize, setPointSize } = usePlotStore()
+  const {
+    renderMode, setRenderMode, pointSize, setPointSize,
+    reductionAlgorithm, setReductionAlgorithm,
+    tsnePerplexity, setTsnePerplexity,
+    tsneLearningRate, setTsneLearningRate,
+    umapNNeighbors, setUmapNNeighbors,
+    umapMinDist, setUmapMinDist,
+    umapMetric, setUmapMetric,
+  } = usePlotStore()
 
   // 1. Fetch collection list
   const { data: collections } = useQuery({
@@ -61,6 +69,16 @@ export default function PlotControls({ onCompute, isComputing }: PlotControlsPro
       gpt_generate_cluster_name: gptEnabled,
       gpt_default_model: gptEnabled ? gptModel : undefined,
       gpt_default_temperature: gptEnabled ? gptTemperature : undefined,
+      reduction_algorithm: reductionAlgorithm,
+      ...(reductionAlgorithm === 'tsne' && {
+        tsne_perplexity: tsnePerplexity,
+        tsne_learning_rate: tsneLearningRate,
+      }),
+      ...(reductionAlgorithm === 'umap' && {
+        umap_n_neighbors: umapNNeighbors,
+        umap_min_dist: umapMinDist,
+        umap_metric: umapMetric,
+      }),
     }
     onCompute(request)
   }
@@ -115,6 +133,115 @@ export default function PlotControls({ onCompute, isComputing }: PlotControlsPro
             />
           </div>
 
+          {/* Reduction Algorithm */}
+          <div className="space-y-2">
+            <label className="block text-sm font-medium text-gray-700">
+              Reduction Algorithm
+            </label>
+            <div className="flex space-x-2">
+              {(['tsne', 'umap', 'pca'] as const).map((algo) => (
+                <label key={algo} className="flex items-center text-xs cursor-pointer">
+                  <input
+                    type="radio"
+                    name="reductionAlgorithm"
+                    value={algo}
+                    checked={reductionAlgorithm === algo}
+                    onChange={() => setReductionAlgorithm(algo as ReductionAlgorithm)}
+                    className="mr-1"
+                  />
+                  {algo.toUpperCase()}
+                </label>
+              ))}
+            </div>
+
+            {/* t-SNE parameters */}
+            {reductionAlgorithm === 'tsne' && (
+              <div className="pl-4 space-y-2 border-l-2 border-blue-200 mt-2">
+                <div>
+                  <label className="block text-xs text-gray-500">
+                    Perplexity: {tsnePerplexity}
+                  </label>
+                  <input
+                    type="range"
+                    min="5"
+                    max="50"
+                    value={tsnePerplexity}
+                    onChange={(e) => setTsnePerplexity(Number(e.target.value))}
+                    className="w-full"
+                  />
+                </div>
+                <div>
+                  <label className="block text-xs text-gray-500">Learning Rate</label>
+                  <select
+                    className="w-full border border-gray-300 rounded px-2 py-1 text-sm"
+                    value={tsneLearningRate}
+                    onChange={(e) => setTsneLearningRate(e.target.value)}
+                  >
+                    <option value="auto">auto</option>
+                    <option value="50">50</option>
+                    <option value="100">100</option>
+                    <option value="200">200</option>
+                    <option value="500">500</option>
+                    <option value="1000">1000</option>
+                  </select>
+                </div>
+              </div>
+            )}
+
+            {/* UMAP parameters */}
+            {reductionAlgorithm === 'umap' && (
+              <div className="pl-4 space-y-2 border-l-2 border-green-200 mt-2">
+                <div>
+                  <label className="block text-xs text-gray-500">
+                    Neighbors: {umapNNeighbors}
+                  </label>
+                  <input
+                    type="range"
+                    min="2"
+                    max="100"
+                    value={umapNNeighbors}
+                    onChange={(e) => setUmapNNeighbors(Number(e.target.value))}
+                    className="w-full"
+                  />
+                </div>
+                <div>
+                  <label className="block text-xs text-gray-500">
+                    Min Distance: {umapMinDist}
+                  </label>
+                  <input
+                    type="range"
+                    min="0"
+                    max="1"
+                    step="0.05"
+                    value={umapMinDist}
+                    onChange={(e) => setUmapMinDist(Number(e.target.value))}
+                    className="w-full"
+                  />
+                </div>
+                <div>
+                  <label className="block text-xs text-gray-500">Metric</label>
+                  <select
+                    className="w-full border border-gray-300 rounded px-2 py-1 text-sm"
+                    value={umapMetric}
+                    onChange={(e) => setUmapMetric(e.target.value)}
+                  >
+                    <option value="cosine">cosine</option>
+                    <option value="euclidean">euclidean</option>
+                    <option value="manhattan">manhattan</option>
+                    <option value="correlation">correlation</option>
+                  </select>
+                </div>
+              </div>
+            )}
+
+            {/* PCA has no extra parameters */}
+            {reductionAlgorithm === 'pca' && (
+              <p className="text-xs text-gray-400 mt-1">
+                PCA has no additional parameters.
+              </p>
+            )}
+          </div>
+
           {/* Text Display Fields */}
           <div className="space-y-2">
             <label className="block text-sm font-medium text-gray-700">Display Fields</label>