feat: navigable 3D embedding view

README.md

@@ -19,6 +19,9 @@ Embedding Atlas is a tool that provides interactive visualizations for large emb
 - 🧊 **Order-independent transparency:**
   Ensure clear, accurate rendering of overlapping points.
 
+- 🧭 **Navigable 3D embedding view:**
+  Explore embeddings in 3D with orbit, pan, and zoom when a Z coordinate is available.
+
 - 🔍 **Real-time search & nearest neighbors:**
   Find similar data to a given query or existing data point.
 

package.json

@@ -14,7 +14,7 @@
   ],
   "scripts": {
     "build": "./scripts/build.sh",
-    "test:js": "npm run test -w @embedding-atlas/density-clustering -w @embedding-atlas/umap-wasm -w @embedding-atlas/utils -w @embedding-atlas/viewer",
+    "test:js": "npm run test -w @embedding-atlas/density-clustering -w @embedding-atlas/umap-wasm -w @embedding-atlas/utils -w @embedding-atlas/component -w @embedding-atlas/viewer",
     "test:rust": "cargo test --workspace",
     "test:python": "cd packages/backend && uv run pytest",
     "test": "npm run test:js && npm run test:python && npm run test:rust",

packages/backend/README.md

@@ -47,6 +47,12 @@ embedding-atlas path_to_dataset.parquet --x projection_x --y projection_y
 
 You may use the [SentenceTransformers](https://sbert.net/) package to compute high-dimensional embeddings from text data, and then use the [UMAP](https://umap-learn.readthedocs.io/en/latest/index.html) package to compute 2D projections.
 
+For a navigable 3D view, either store a pre-computed Z column and pass it with `--z` (alongside `--x`/`--y`), or generate a 3D projection with `--umap-n-components 3`:
+
+```bash
+embedding-atlas path_to_dataset.parquet --x projection_x --y projection_y --z projection_z
+```
+
 ### Using Pre-computed Vectors
 
 If you already have pre-computed embedding vectors (but not the 2D projections), you can specify the column containing the vectors with `--vector`:

packages/backend/embedding_atlas/cli.py

@@ -222,6 +222,11 @@ def import_modules(names: list[str]):
     "y_column",
     help="Column containing pre-computed Y coordinates for the embedding view.",
 )
+@click.option(
+    "--z",
+    "z_column",
+    help="Column containing pre-computed Z coordinates for the embedding view (requires --x and --y). To GENERATE a 3D projection instead, use --umap-n-components 3.",
+)
 @click.option(
     "--neighbors",
     "neighbors_column",
@@ -265,6 +270,14 @@ def import_modules(names: list[str]):
 @click.option(
     "--umap-random-state", type=int, help="Random seed for reproducible UMAP results."
 )
+@click.option(
+    "--umap-n-components",
+    "umap_n_components",
+    type=click.IntRange(2, 3),
+    default=2,
+    show_default=True,
+    help="Number of UMAP output dimensions (2 or 3). 3 enables a 3D embedding view.",
+)
 @click.option(
     "--duckdb",
     type=str,
@@ -359,6 +372,7 @@ def main(
     max_concurrency: int | None,
     x_column: str | None,
     y_column: str | None,
+    z_column: str | None,
     neighbors_column: str | None,
     pagerank_column: str | None,
     query: str | None,
@@ -367,6 +381,7 @@ def main(
     umap_min_dist: float | None,
     umap_metric: str | None,
     umap_random_state: int | None,
+    umap_n_components: int,
     static: str | None,
     duckdb: str,
     host: str,
@@ -383,13 +398,48 @@ def main(
 ):
     apply_logging_config()
 
+    # --z names a PRE-COMPUTED coordinate column, so it is only meaningful together
+    # with pre-computed --x/--y. Without them a projection is computed, and passing a
+    # z column there would (a) auto-enable 3D regardless of --umap-n-components and
+    # (b) get overwritten by the generated projection_z. To GENERATE 3D, callers use
+    # --umap-n-components 3 (which writes a fresh projection_z column).
+    if z_column is not None and (x_column is None or y_column is None):
+        raise click.UsageError(
+            "--z names a pre-computed column and requires both --x and --y. "
+            "To generate a 3D projection, use --umap-n-components 3 instead."
+        )
+
+    # --umap-n-components 3 only writes a projection_z when a projection is actually
+    # computed (projection enabled AND no pre-computed --x/--y). Combined with
+    # pre-computed coordinates or --disable-projection it would be silently ignored and
+    # the view would stay 2D, so fail loudly instead.
+    if umap_n_components == 3 and not (
+        enable_projection and (x_column is None or y_column is None)
+    ):
+        raise click.UsageError(
+            "--umap-n-components 3 generates a 3D projection and cannot be combined with "
+            "pre-computed --x/--y or --disable-projection. To display pre-computed 3D "
+            "coordinates, pass --x, --y and --z instead."
+        )
+
     if with_modules is not None:
         import_modules(with_modules)
 
     df = load_datasets(inputs, splits=split, query=query, sample=sample)
 
     print(df)
 
+    # A pre-computed --z column must exist and be numeric: the frontend casts it to FLOAT
+    # and fits the 3D camera to its extent, so a missing or non-numeric column would fail
+    # the query or collapse the view.
+    if z_column is not None:
+        if z_column not in df.columns:
+            raise click.UsageError(
+                f"--z column {z_column!r} was not found in the data."
+            )
+        if not pd.api.types.is_numeric_dtype(df[z_column]):
+            raise click.UsageError(f"--z column {z_column!r} must be numeric.")
+
     if enable_projection and (x_column is None or y_column is None):
         # No x, y column selected, first see if text/image/vectors column is specified, if not, ask for it
         if text is None and image is None and audio is None and vector is None:
@@ -407,6 +457,11 @@ def main(
             umap_args["random_state"] = umap_random_state
         if umap_metric is not None:
             umap_args["metric"] = umap_metric
+        # Only set n_components for 3D. Leaving it absent for the default 2D case
+        # preserves the existing projection cache key ({}), so default runs do not
+        # miss the cache and recompute (possibly paid) embeddings.
+        if umap_n_components != 2:
+            umap_args["n_components"] = umap_n_components
         # Run embedding and projection
         if (
             text is not None
@@ -419,6 +474,8 @@ def main(
 
             x_column = find_column_name(df.columns, "projection_x")
             y_column = find_column_name(df.columns, "projection_y")
+            if umap_n_components == 3:
+                z_column = find_column_name(df.columns, "projection_z")
             if neighbors_column is None:
                 neighbors_column = find_column_name(df.columns, "__neighbors")
                 new_neighbors_column = neighbors_column
@@ -462,6 +519,7 @@ def main(
                 modality=modality,
                 x=x_column,
                 y=y_column,
+                z=z_column,
                 neighbors=new_neighbors_column,
                 embedder=embedder,
                 model=model,
@@ -506,6 +564,7 @@ def main(
         row_id=id_column,
         x=x_column,
         y=y_column,
+        z=z_column,
         neighbors=neighbors_column,
         importance=pagerank_column,
         text=text,

packages/backend/embedding_atlas/options.py

@@ -14,6 +14,10 @@ class EmbeddingAtlasOptions(TypedDict, total=False):
     y:
         The column name for Y axis in the embedding.
 
+    z:
+        The column name for Z axis in the embedding. When specified (together with
+        ``x`` and ``y``), the embedding view becomes 3D-capable.
+
     text:
         The column name for the textual data.
 
@@ -66,6 +70,7 @@ class EmbeddingAtlasOptions(TypedDict, total=False):
     row_id: str | None
     x: str | None
     y: str | None
+    z: str | None
     text: str | None
     image: str | None
     importance: str | None
@@ -116,7 +121,10 @@ def set_prop(key: str, value):
     set_prop("data.table", options.get("table"))
     set_prop("data.id", options.get("row_id"))
     if options.get("x") is not None and options.get("y") is not None:
-        set_prop("data.projection", {"x": options.get("x"), "y": options.get("y")})
+        projection = {"x": options.get("x"), "y": options.get("y")}
+        if options.get("z") is not None:
+            projection["z"] = options.get("z")
+        set_prop("data.projection", projection)
     set_prop("data.text", options.get("text"))
     set_prop("data.image", options.get("image"))
     set_prop("data.importance", options.get("importance"))

packages/backend/embedding_atlas/projection.py

@@ -24,6 +24,7 @@ def compute_projection(
     modality: str = "auto",
     x: str = "projection_x",
     y: str = "projection_y",
+    z: str | None = None,
     neighbors: str | None = "neighbors",
     embedder: str | Callable | None = None,
     model: str | None = None,
@@ -34,7 +35,7 @@ def compute_projection(
     cache_root: str | Path | None = None,
 ) -> IntoDataFrameT:
     """
-    Compute embeddings and generate 2D projections for a DataFrame column.
+    Compute embeddings and generate 2D or 3D projections for a DataFrame column.
 
     This is a unified entry point that auto-detects the modality of the input
     data (text, image, audio, or vector) and delegates to the appropriate
@@ -51,6 +52,10 @@ def compute_projection(
             'text', 'image', 'audio', 'vector', or 'auto' (auto-detect).
         x: str, column name where the UMAP X coordinates will be stored.
         y: str, column name where the UMAP Y coordinates will be stored.
+        z: str or None, column name where the UMAP Z coordinates will be
+            stored. Set to a column name to request a 3D projection. When set
+            and ``n_components`` is not specified in ``umap_args``, UMAP defaults
+            to 3 output components so the z column "just works".
         neighbors: str or None, column name where nearest neighbor indices
             will be stored. Set to None to skip.
         embedder: the embedding backend to use. Can be:
@@ -80,6 +85,7 @@ def compute_projection(
             modality=modality,
             x=x,
             y=y,
+            z=z,
             neighbors=neighbors,
             embedder=embedder,
             model=model,
@@ -99,6 +105,7 @@ async def async_compute_projection(
     modality: str = "auto",
     x: str = "projection_x",
     y: str = "projection_y",
+    z: str | None = None,
     neighbors: str | None = "neighbors",
     embedder: str | Callable | None = None,
     model: str | None = None,
@@ -121,7 +128,13 @@ async def async_compute_projection(
     nw_frame = nw.from_native(data_frame, eager_only=True)
     series = nw_frame[inputs]
     embedder_args = embedder_args or {}
-    umap_args = umap_args or {}
+    umap_args = dict(umap_args or {})
+
+    # If a z column is requested, ensure UMAP produces at least 3 components so that
+    # requesting a z column "just works". We check the effective value rather than
+    # key presence, because callers (e.g. the CLI) may pass an explicit default of 2.
+    if z is not None and umap_args.get("n_components", 2) < 3:
+        umap_args["n_components"] = 3
 
     # 1. Infer modality
     if modality == "auto":
@@ -217,6 +230,18 @@ async def run() -> Projection:
         nw.new_series(x, proj.projection[:, 0].tolist(), nw.Float64, backend=backend),
         nw.new_series(y, proj.projection[:, 1].tolist(), nw.Float64, backend=backend),
     ]
+    if z is not None:
+        if proj.projection.shape[1] < 3:
+            raise ValueError(
+                "A z column was requested, but the UMAP projection only has "
+                f"{proj.projection.shape[1]} component(s). Pass "
+                'umap_args={"n_components": 3} to produce a 3D projection.'
+            )
+        new_columns.append(
+            nw.new_series(
+                z, proj.projection[:, 2].tolist(), nw.Float64, backend=backend
+            )
+        )
     if neighbors is not None:
         new_columns.append(
             nw.new_series(

packages/backend/embedding_atlas/widget.py

@@ -57,6 +57,11 @@ def __init__(
             y:
                 The column name for Y axis in the embedding.
 
+            z:
+                The column name for Z axis in the embedding. When specified
+                (together with ``x`` and ``y``), the embedding view becomes
+                3D-capable.
+
             text:
                 The column name for the textual data.
 

packages/backend/tests/test_cli.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+
+"""Regression tests for embedding-atlas CLI option handling."""
+
+from click.testing import CliRunner
+
+from embedding_atlas.cli import main
+
+
+def test_z_requires_x_and_y():
+    """--z names a pre-computed column, so it must not be accepted on its own (it
+    would otherwise be passed to projection generation, auto-enabling 3D and being
+    overwritten by the generated projection_z)."""
+    runner = CliRunner()
+
+    # --z alone is rejected before any data is loaded.
+    result = runner.invoke(main, ["dummy.csv", "--z", "depth"])
+    assert result.exit_code != 0
+    assert "--z" in result.output
+    assert "umap-n-components" in result.output
+
+    # --z with only --x is still rejected.
+    result = runner.invoke(main, ["dummy.csv", "--x", "px", "--z", "depth"])
+    assert result.exit_code != 0
+    assert "requires both --x and --y" in result.output
+
+
+def test_z_with_x_and_y_passes_validation():
+    """--z together with --x and --y is a valid pre-computed 3D input; it must pass
+    the option guard (failing later only because the dummy dataset cannot load)."""
+    runner = CliRunner()
+    result = runner.invoke(main, ["dummy.csv", "--x", "px", "--y", "py", "--z", "pz"])
+    # It should NOT fail on the --z guard (which is a click.UsageError mentioning z).
+    assert "requires both --x and --y" not in result.output
+
+
+def test_umap_3d_rejected_with_precomputed_xy():
+    """--umap-n-components 3 generates a projection; combined with pre-computed --x/--y it
+    would be silently ignored (a 2D view), so it must be rejected up front."""
+    runner = CliRunner()
+    result = runner.invoke(
+        main, ["dummy.csv", "--x", "px", "--y", "py", "--umap-n-components", "3"]
+    )
+    assert result.exit_code != 0
+    assert "umap-n-components 3" in result.output
+
+
+def test_umap_3d_rejected_with_disable_projection():
+    """--umap-n-components 3 with --disable-projection computes no projection, so it must
+    be rejected rather than silently produce a 2D view."""
+    runner = CliRunner()
+    result = runner.invoke(
+        main, ["dummy.csv", "--disable-projection", "--umap-n-components", "3"]
+    )
+    assert result.exit_code != 0
+    assert "umap-n-components 3" in result.output
+
+
+def test_z_column_must_exist():
+    """A pre-computed --z naming a missing column is rejected with a clear error."""
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        with open("data.csv", "w") as f:
+            f.write("px,py\n1,2\n3,4\n")
+        result = runner.invoke(
+            main, ["data.csv", "--x", "px", "--y", "py", "--z", "missing"]
+        )
+    assert result.exit_code != 0
+    assert "was not found" in result.output
+
+
+def test_z_column_must_be_numeric():
+    """A pre-computed --z naming a non-numeric column is rejected with a clear error."""
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        with open("data.csv", "w") as f:
+            f.write("px,py,depth\n1,2,a\n3,4,b\n")
+        result = runner.invoke(
+            main, ["data.csv", "--x", "px", "--y", "py", "--z", "depth"]
+        )
+    assert result.exit_code != 0
+    assert "must be numeric" in result.output