+
+
+ {(['tsne', 'umap', 'pca'] as const).map((algo) => (
+
+ ))}
+
+
+ {/* t-SNE parameters */}
+ {reductionAlgorithm === 'tsne' && (
+
+
+
+ setTsnePerplexity(Number(e.target.value))}
+ className="w-full"
+ />
+
+
+
+
+
+
+ )}
+
+ {/* UMAP parameters */}
+ {reductionAlgorithm === 'umap' && (
+
+
+
+ setUmapNNeighbors(Number(e.target.value))}
+ className="w-full"
+ />
+
+
+
+ setUmapMinDist(Number(e.target.value))}
+ className="w-full"
+ />
+
+
+
+
+
+
+ )}
+
+ {/* PCA has no extra parameters */}
+ {reductionAlgorithm === 'pca' && (
+
+ PCA has no additional parameters.
+
+ )}
+
diff --git a/frontend/src/stores/plotStore.ts b/frontend/src/stores/plotStore.ts
index 3b6d9da..157cc9a 100644
--- a/frontend/src/stores/plotStore.ts
+++ b/frontend/src/stores/plotStore.ts
@@ -1,5 +1,5 @@
import { create } from 'zustand'
-import type { PlotResponse, SearchResult } from '../types'
+import type { PlotResponse, ReductionAlgorithm, SearchResult } from '../types'
interface PlotState {
plotData: PlotResponse | null
@@ -11,6 +11,12 @@ interface PlotState {
highlightedIds: Set
isSearching: boolean
queryPoint: { x: number; y: number; z: number } | null
+ reductionAlgorithm: ReductionAlgorithm
+ tsnePerplexity: number
+ tsneLearningRate: string
+ umapNNeighbors: number
+ umapMinDist: number
+ umapMetric: string
// actions
setPlotData: (data: PlotResponse | null) => void
toggleCluster: (index: number) => void
@@ -23,6 +29,12 @@ interface PlotState {
setIsSearching: (searching: boolean) => void
clearSearch: () => void
setQueryPoint: (point: { x: number; y: number; z: number } | null) => void
+ setReductionAlgorithm: (algorithm: ReductionAlgorithm) => void
+ setTsnePerplexity: (perplexity: number) => void
+ setTsneLearningRate: (rate: string) => void
+ setUmapNNeighbors: (n: number) => void
+ setUmapMinDist: (dist: number) => void
+ setUmapMetric: (metric: string) => void
}
export const CLUSTER_COLORS = [
@@ -42,6 +54,12 @@ export const usePlotStore = create((set) => ({
highlightedIds: new Set(),
isSearching: false,
queryPoint: null,
+ reductionAlgorithm: 'tsne',
+ tsnePerplexity: 30,
+ tsneLearningRate: 'auto',
+ umapNNeighbors: 15,
+ umapMinDist: 0.1,
+ umapMetric: 'cosine',
setPlotData: (data) => set({ plotData: data }),
@@ -80,4 +98,10 @@ export const usePlotStore = create((set) => ({
isSearching: false,
queryPoint: null,
}),
+ setReductionAlgorithm: (algorithm) => set({ reductionAlgorithm: algorithm }),
+ setTsnePerplexity: (perplexity) => set({ tsnePerplexity: perplexity }),
+ setTsneLearningRate: (rate) => set({ tsneLearningRate: rate }),
+ setUmapNNeighbors: (n) => set({ umapNNeighbors: n }),
+ setUmapMinDist: (dist) => set({ umapMinDist: dist }),
+ setUmapMetric: (metric) => set({ umapMetric: metric }),
}))
diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts
index b569381..a1839ab 100644
--- a/frontend/src/types/index.ts
+++ b/frontend/src/types/index.ts
@@ -58,14 +58,22 @@ export interface IndexStatusResponse {
}
// Plot
+export type ReductionAlgorithm = 'tsne' | 'umap' | 'pca'
+
export interface PlotRequest {
- chromadb_collection_name: string;
- num_clusters?: number;
- text_display_fields?: string[];
- image_field?: string;
- gpt_generate_cluster_name?: boolean;
- gpt_default_model?: string;
- gpt_default_temperature?: number;
+ chromadb_collection_name: string
+ num_clusters?: number
+ text_display_fields?: string[]
+ image_field?: string
+ gpt_generate_cluster_name?: boolean
+ gpt_default_model?: string
+ gpt_default_temperature?: number
+ reduction_algorithm?: ReductionAlgorithm
+ tsne_perplexity?: number
+ tsne_learning_rate?: string
+ umap_n_neighbors?: number
+ umap_min_dist?: number
+ umap_metric?: string
}
export interface PlotPoint {
diff --git a/pyproject.toml b/pyproject.toml
index e489a6c..61f6c42 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -35,6 +35,9 @@ dev = [
"pre-commit>=4,<5",
"httpx>=0.28,<1",
]
+umap = [
+ "umap-learn>=0.5",
+]
[tool.ruff]
target-version = "py313"
@@ -82,6 +85,7 @@ module = [
"PIL.*",
"openai.*",
"aiohttp.*",
+ "umap.*",
]
ignore_missing_imports = true
diff --git a/tests/test_scatter_plot.py b/tests/test_scatter_plot.py
index e3d322f..110ed83 100644
--- a/tests/test_scatter_plot.py
+++ b/tests/test_scatter_plot.py
@@ -230,16 +230,13 @@ def test_prepare_data(self, monkeypatch: pytest.MonkeyPatch) -> None:
"embedding_cluster.scatter_plot.load_chromadb_collection",
return_value=fake_collection,
),
- patch("embedding_cluster.scatter_plot.TSNE") as mock_tsne_cls,
+ patch(
+ "embedding_cluster.scatter_plot.reduce_dimensions",
+ return_value=np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]),
+ ),
patch("embedding_cluster.scatter_plot.KMeans") as mock_kmeans_cls,
patch("embedding_cluster.scatter_plot.StandardScaler") as mock_scaler_cls,
):
- mock_tsne = MagicMock()
- mock_tsne.fit_transform.return_value = np.array(
- [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]
- )
- mock_tsne_cls.return_value = mock_tsne
-
mock_kmeans = MagicMock()
mock_kmeans.fit_predict.return_value = [0, 0, 1, 1]
mock_kmeans_cls.return_value = mock_kmeans
@@ -251,7 +248,6 @@ def test_prepare_data(self, monkeypatch: pytest.MonkeyPatch) -> None:
mock_scaler_cls.return_value = mock_scaler
fig = prepare_data(settings)
-
assert isinstance(fig, go.Figure)
assert len(sp.cluster_images) == 2
assert len(sp.cluster_item_names) == 2
@@ -420,3 +416,100 @@ def test_large_dataset_uses_sample(self) -> None:
assert result["k_values"] == list(range(2, 6))
assert 2 <= result["suggested_k"] <= 5
+
+
+class TestReduceDimensions:
+ def test_tsne_output_shape(self) -> None:
+ from embedding_cluster.scatter_plot import reduce_dimensions
+
+ rng = np.random.default_rng(42)
+ embeddings = rng.random((50, 10))
+
+ result = reduce_dimensions(embeddings, algorithm="tsne", n_components=3)
+
+ assert result.shape == (50, 3)
+
+ def test_pca_output_shape(self) -> None:
+ from embedding_cluster.scatter_plot import reduce_dimensions
+
+ rng = np.random.default_rng(42)
+ embeddings = rng.random((30, 10))
+
+ result = reduce_dimensions(embeddings, algorithm="pca", n_components=3)
+
+ assert result.shape == (30, 3)
+
+ def test_pca_deterministic(self) -> None:
+ from embedding_cluster.scatter_plot import reduce_dimensions
+
+ rng = np.random.default_rng(42)
+ embeddings = rng.random((30, 10))
+
+ result1 = reduce_dimensions(embeddings, algorithm="pca", n_components=3)
+ result2 = reduce_dimensions(embeddings, algorithm="pca", n_components=3)
+
+ np.testing.assert_array_equal(result1, result2)
+
+ def test_tsne_custom_perplexity(self) -> None:
+ from embedding_cluster.scatter_plot import reduce_dimensions
+
+ rng = np.random.default_rng(42)
+ embeddings = rng.random((30, 10))
+
+ result = reduce_dimensions(
+ embeddings,
+ algorithm="tsne",
+ n_components=3,
+ perplexity=10.0,
+ )
+
+ assert result.shape == (30, 3)
+
+ def test_umap_output_shape(self) -> None:
+ from embedding_cluster.scatter_plot import reduce_dimensions
+
+ rng = np.random.default_rng(42)
+ embeddings = rng.random((30, 10))
+
+ try:
+ result = reduce_dimensions(
+ embeddings,
+ algorithm="umap",
+ n_components=3,
+ n_neighbors=5,
+ min_dist=0.1,
+ )
+ assert result.shape == (30, 3)
+ except ImportError:
+ pytest.skip("umap-learn not installed")
+
+ def test_umap_import_guard(self) -> None:
+ from embedding_cluster.scatter_plot import reduce_dimensions
+
+ rng = np.random.default_rng(42)
+ embeddings = rng.random((30, 10))
+
+ with (
+ patch.dict("sys.modules", {"umap": None}),
+ pytest.raises(ImportError, match="umap-learn is not installed"),
+ ):
+ reduce_dimensions(embeddings, algorithm="umap", n_components=3)
+
+ def test_invalid_algorithm_raises(self) -> None:
+ from embedding_cluster.scatter_plot import reduce_dimensions
+
+ rng = np.random.default_rng(42)
+ embeddings = rng.random((30, 10))
+
+ with pytest.raises(ValueError, match="Unknown reduction algorithm"):
+ reduce_dimensions(embeddings, algorithm="invalid", n_components=3)
+
+ def test_pca_two_components(self) -> None:
+ from embedding_cluster.scatter_plot import reduce_dimensions
+
+ rng = np.random.default_rng(42)
+ embeddings = rng.random((30, 10))
+
+ result = reduce_dimensions(embeddings, algorithm="pca", n_components=2)
+
+ assert result.shape == (30, 2)
diff --git a/tests/test_settings.py b/tests/test_settings.py
index a65063a..8d7dfed 100644
--- a/tests/test_settings.py
+++ b/tests/test_settings.py
@@ -47,6 +47,21 @@ def test_default_model_names(self) -> None:
assert s.image_model_name == "openai/clip-vit-base-patch32"
assert s.text_model_name == "BAAI/bge-small-en-v1.5"
+ def test_default_reduction_algorithm(self) -> None:
+ s = Settings()
+ assert s.reduction_algorithm == "tsne"
+
+ def test_default_tsne_params(self) -> None:
+ s = Settings()
+ assert s.tsne_perplexity == pytest.approx(30.0)
+ assert s.tsne_learning_rate == "auto"
+
+ def test_default_umap_params(self) -> None:
+ s = Settings()
+ assert s.umap_n_neighbors == 15
+ assert s.umap_min_dist == pytest.approx(0.1)
+ assert s.umap_metric == "cosine"
+
class TestSettingsEnvVars:
def test_running_mode_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
@@ -87,3 +102,22 @@ def test_local_csv_filename_from_env(self, monkeypatch: pytest.MonkeyPatch) -> N
monkeypatch.setenv("LOCAL_CSV_FILENAME", "/tmp/data.csv")
s = Settings()
assert s.local_csv_filename == "/tmp/data.csv"
+
+ def test_reduction_algorithm_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ monkeypatch.setenv("REDUCTION_ALGORITHM", "pca")
+ s = Settings()
+ assert s.reduction_algorithm == "pca"
+
+ def test_tsne_perplexity_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ monkeypatch.setenv("TSNE_PERPLEXITY", "50.0")
+ s = Settings()
+ assert s.tsne_perplexity == pytest.approx(50.0)
+
+ def test_umap_params_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ monkeypatch.setenv("UMAP_N_NEIGHBORS", "30")
+ monkeypatch.setenv("UMAP_MIN_DIST", "0.5")
+ monkeypatch.setenv("UMAP_METRIC", "euclidean")
+ s = Settings()
+ assert s.umap_n_neighbors == 30
+ assert s.umap_min_dist == pytest.approx(0.5)
+ assert s.umap_metric == "euclidean"
diff --git a/uv.lock b/uv.lock
index 94c2987..db32f7a 100644
--- a/uv.lock
+++ b/uv.lock
@@ -391,6 +391,9 @@ dev = [
{ name = "pytest-cov", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
{ name = "ruff", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
]
+umap = [
+ { name = "umap-learn", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+]
[package.metadata]
requires-dist = [
@@ -416,10 +419,11 @@ requires-dist = [
{ name = "sentence-transformers", specifier = ">=3.4,<4" },
{ name = "torch", specifier = ">=2.6,<3" },
{ name = "transformers", specifier = ">=4.48,<5" },
+ { name = "umap-learn", marker = "extra == 'umap'", specifier = ">=0.5" },
{ name = "uvicorn", specifier = ">=0.34,<1" },
{ name = "websockets", specifier = ">=14,<15" },
]
-provides-extras = ["dev"]
+provides-extras = ["dev", "umap"]
[[package]]
name = "fastapi"
@@ -769,6 +773,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/2b/39/191d3d28abc26c9099b19852e6c99f7f6d400b82fa5a4e80291bd3803e19/librt-0.8.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:cc3656283d11540ab0ea01978378e73e10002145117055e03722417aeab30994", size = 263001, upload-time = "2026-02-17T16:12:43.627Z" },
]
+[[package]]
+name = "llvmlite"
+version = "0.46.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/74/cd/08ae687ba099c7e3d21fe2ea536500563ef1943c5105bf6ab4ee3829f68e/llvmlite-0.46.0.tar.gz", hash = "sha256:227c9fd6d09dce2783c18b754b7cd9d9b3b3515210c46acc2d3c5badd9870ceb", size = 193456, upload-time = "2025-12-08T18:15:36.295Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/e6/ff/3eba7eb0aed4b6fca37125387cd417e8c458e750621fce56d2c541f67fa8/llvmlite-0.46.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:30b60892d034bc560e0ec6654737aaa74e5ca327bd8114d82136aa071d611172", size = 37232767, upload-time = "2025-12-08T18:15:13.22Z" },
+ { url = "https://files.pythonhosted.org/packages/0e/54/737755c0a91558364b9200702c3c9c15d70ed63f9b98a2c32f1c2aa1f3ba/llvmlite-0.46.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6cc19b051753368a9c9f31dc041299059ee91aceec81bd57b0e385e5d5bf1a54", size = 56275176, upload-time = "2025-12-08T18:15:16.339Z" },
+ { url = "https://files.pythonhosted.org/packages/95/ae/af0ffb724814cc2ea64445acad05f71cff5f799bb7efb22e47ee99340dbc/llvmlite-0.46.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:d252edfb9f4ac1fcf20652258e3f102b26b03eef738dc8a6ffdab7d7d341d547", size = 37232768, upload-time = "2025-12-08T18:15:25.055Z" },
+ { url = "https://files.pythonhosted.org/packages/c9/19/5018e5352019be753b7b07f7759cdabb69ca5779fea2494be8839270df4c/llvmlite-0.46.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:379fdd1c59badeff8982cb47e4694a6143bec3bb49aa10a466e095410522064d", size = 56275173, upload-time = "2025-12-08T18:15:28.109Z" },
+]
+
[[package]]
name = "markdown-it-py"
version = "4.0.0"
@@ -949,6 +965,22 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl", hash = "sha256:5bb13e3eed2923615535339b3c620e76779af4cb4c6a90deccc9e36b274d3827", size = 23438, upload-time = "2025-12-20T14:08:52.782Z" },
]
+[[package]]
+name = "numba"
+version = "0.64.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "llvmlite", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+ { name = "numpy", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/23/c9/a0fb41787d01d621046138da30f6c2100d80857bf34b3390dd68040f27a3/numba-0.64.0.tar.gz", hash = "sha256:95e7300af648baa3308127b1955b52ce6d11889d16e8cfe637b4f85d2fca52b1", size = 2765679, upload-time = "2026-02-18T18:41:20.974Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/52/80/2734de90f9300a6e2503b35ee50d9599926b90cbb7ac54f9e40074cd07f1/numba-0.64.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:3bab2c872194dcd985f1153b70782ec0fbbe348fffef340264eacd3a76d59fd6", size = 2683392, upload-time = "2026-02-18T18:41:06.563Z" },
+ { url = "https://files.pythonhosted.org/packages/42/e8/14b5853ebefd5b37723ef365c5318a30ce0702d39057eaa8d7d76392859d/numba-0.64.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:703a246c60832cad231d2e73c1182f25bf3cc8b699759ec8fe58a2dbc689a70c", size = 3812245, upload-time = "2026-02-18T18:41:07.963Z" },
+ { url = "https://files.pythonhosted.org/packages/3d/8a/77d26afe0988c592dd97cb8d4e80bfb3dfc7dbdacfca7d74a7c5c81dd8c2/numba-0.64.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:f565d55eaeff382cbc86c63c8c610347453af3d1e7afb2b6569aac1c9b5c93ce", size = 2683590, upload-time = "2026-02-18T18:41:12.897Z" },
+ { url = "https://files.pythonhosted.org/packages/8e/4b/600b8b7cdbc7f9cebee9ea3d13bb70052a79baf28944024ffcb59f0712e3/numba-0.64.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9b55169b18892c783f85e9ad9e6f5297a6d12967e4414e6b71361086025ff0bb", size = 3781163, upload-time = "2026-02-18T18:41:15.377Z" },
+]
+
[[package]]
name = "numpy"
version = "2.4.2"
@@ -1538,6 +1570,22 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217, upload-time = "2025-06-21T13:39:07.939Z" },
]
+[[package]]
+name = "pynndescent"
+version = "0.6.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "joblib", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+ { name = "llvmlite", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+ { name = "numba", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+ { name = "scikit-learn", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+ { name = "scipy", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/4a/fb/7f58c397fb31666756457ee2ac4c0289ef2daad57f4ae4be8dec12f80b03/pynndescent-0.6.0.tar.gz", hash = "sha256:7ffde0fb5b400741e055a9f7d377e3702e02250616834231f6c209e39aac24f5", size = 2992987, upload-time = "2026-01-08T21:29:58.943Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/b2/e6/94145d714402fd5ade00b5661f2d0ab981219e07f7db9bfa16786cdb9c04/pynndescent-0.6.0-py3-none-any.whl", hash = "sha256:dc8c74844e4c7f5cbd1e0cd6909da86fdc789e6ff4997336e344779c3d5538ef", size = 73511, upload-time = "2026-01-08T21:29:57.306Z" },
+]
+
[[package]]
name = "pypika"
version = "0.51.1"
@@ -2040,6 +2088,23 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" },
]
+[[package]]
+name = "umap-learn"
+version = "0.5.11"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "numba", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+ { name = "numpy", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+ { name = "pynndescent", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+ { name = "scikit-learn", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+ { name = "scipy", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+ { name = "tqdm", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'darwin'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/94/9a/a1e4a257a9aa979dac4f6d5781dac929cbb0949959e2003ed82657d10b0f/umap_learn-0.5.11.tar.gz", hash = "sha256:31566ffd495fbf05d7ab3efcba703861c0f5e6fc6998a838d0e2becdd00e54f5", size = 96409, upload-time = "2026-01-12T20:44:47.553Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/43/d2/fcf7192dd1cd8c090b6cfd53fa223c4fb2887a17c47e06bc356d44f40dfb/umap_learn-0.5.11-py3-none-any.whl", hash = "sha256:cb17adbde9d544ba79481b3ab4d81ac222e940f3d9219307bea6044f869af3cc", size = 90890, upload-time = "2026-01-12T20:44:46.511Z" },
+]
+
[[package]]
name = "urllib3"
version = "2.6.3"