From 217d13a42f1b47d97ba208640284eecfa59af0cf Mon Sep 17 00:00:00 2001 From: Sricharan Reddy Varra Date: Tue, 23 Jun 2026 14:34:52 -0700 Subject: [PATCH 1/2] fix(viscy-utils): assign unique uuid obs_names in embedding stores EmbeddingWriter wrote stores with duplicate obs_names: write_on_epoch_end concatenated per-batch index frames without ignore_index, so the positional index restarted each batch (0,1,..,0,1,..) and those labels became obs_names. Consumers worked around this with scattered obs_names_make_unique() calls, which only dedupe a single store and break once N stores are ad.concat'd. obs_names is never referenced directly (cell identity lives in the obs columns), so make it an opaque, globally-unique handle: a random uuid4 per observation, assigned at the single creation site. Stores stay unique within and across any number of concatenated stores. Also pass ignore_index=True at the per-batch concat to drop the duplicate intermediate index. Refs #475 Co-Authored-By: Claude Opus 4.8 (1M context) --- .../viscy_utils/callbacks/embedding_writer.py | 8 +++- .../tests/test_embedding_writer.py | 45 +++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) create mode 100644 packages/viscy-utils/tests/test_embedding_writer.py diff --git a/packages/viscy-utils/src/viscy_utils/callbacks/embedding_writer.py b/packages/viscy-utils/src/viscy_utils/callbacks/embedding_writer.py index 373507b8f..8bd7056e6 100644 --- a/packages/viscy-utils/src/viscy_utils/callbacks/embedding_writer.py +++ b/packages/viscy-utils/src/viscy_utils/callbacks/embedding_writer.py @@ -1,6 +1,7 @@ """Callback for writing embeddings to zarr store.""" import logging +import uuid from pathlib import Path from typing import Any, Dict, Literal, Optional, Sequence @@ -166,6 +167,11 @@ def write_embedding_dataset( elif hasattr(s, "cat") and isinstance(s.cat.categories.dtype, pd.StringDtype): ultrack_indices[col] = s.cat.rename_categories(s.cat.categories.astype(object)) + # obs_names are an opaque, unique handle — never referenced directly (cell identity + # lives in the obs columns). Random UUIDs stay unique within a store and across any + # number of concatenated stores, so consumers never need obs_names_make_unique(). + ultrack_indices.index = [str(uuid.uuid4()) for _ in range(len(ultrack_indices))] + if embedding_key == "projections": if projections is None: raise ValueError("embedding_key='projections' requires projections to be provided.") @@ -273,7 +279,7 @@ def write_on_epoch_end( """Write predictions and dimensionality reductions to a zarr store.""" features = _move_and_stack_embeddings(predictions, "features") projections = _move_and_stack_embeddings(predictions, "projections") - ultrack_indices = pd.concat([pd.DataFrame(p["index"]) for p in predictions]) + ultrack_indices = pd.concat([pd.DataFrame(p["index"]) for p in predictions], ignore_index=True) write_embedding_dataset( output_path=self.output_path, diff --git a/packages/viscy-utils/tests/test_embedding_writer.py b/packages/viscy-utils/tests/test_embedding_writer.py new file mode 100644 index 000000000..9939476e4 --- /dev/null +++ b/packages/viscy-utils/tests/test_embedding_writer.py @@ -0,0 +1,45 @@ +import anndata as ad +import numpy as np +import pandas as pd + +from viscy_utils.callbacks.embedding_writer import write_embedding_dataset + + +def _make_index_df(n: int, fov: str) -> pd.DataFrame: + # Mimic EmbeddingWriter.write_on_epoch_end: per-batch frames concatenated, so the + # positional index restarts each batch and the result has duplicate labels. + half = n // 2 + batches = [ + pd.DataFrame({"fov_name": [fov] * half, "track_id": range(half), "t": range(half)}), + pd.DataFrame( + {"fov_name": [fov] * (n - half), "track_id": range(n - half), "t": range(n - half)} + ), + ] + df = pd.concat(batches) + assert df.index.has_duplicates # precondition: the input that used to leak into obs_names + return df + + +def test_obs_names_unique_within_store(tmp_path): + n, d = 6, 4 + features = np.random.default_rng(0).standard_normal((n, d)).astype(np.float32) + out = tmp_path / "store.zarr" + + write_embedding_dataset(output_path=out, features=features, index_df=_make_index_df(n, "A/1")) + + adata = ad.read_zarr(out) + assert adata.obs_names.is_unique + + +def test_obs_names_unique_across_concatenated_stores(tmp_path): + rng = np.random.default_rng(0) + n, d = 6, 4 + paths = [] + for i, fov in enumerate(["A/1", "B/2"]): + out = tmp_path / f"store_{i}.zarr" + features = rng.standard_normal((n, d)).astype(np.float32) + write_embedding_dataset(output_path=out, features=features, index_df=_make_index_df(n, fov)) + paths.append(out) + + combined = ad.concat([ad.read_zarr(p) for p in paths]) + assert combined.obs_names.is_unique From 2b7c57028d74924d7f8703d9d204e57eb1d2a0a0 Mon Sep 17 00:00:00 2001 From: Sricharan Reddy Varra Date: Tue, 23 Jun 2026 14:36:41 -0700 Subject: [PATCH 2/2] refactor(dynaclr): drop obs_names_make_unique() workarounds EmbeddingWriter now assigns unique uuid obs_names at the source, so the per-consumer obs_names_make_unique() calls are dead no-ops. Remove all nine across the mmd, linear-classifier, and pseudotime code paths. Refs #475 Co-Authored-By: Claude Opus 4.8 (1M context) --- .../pseudotime/0-select_candidates/manual_candidates.py | 1 - .../pseudotime/0-select_candidates/select_candidates.py | 3 --- .../scripts/pseudotime/2-align_cells/align_embedding.py | 1 - .../dynaclr/scripts/pseudotime/2-align_cells/align_lc.py | 1 - .../pseudotime/3-organelle-remodeling/readout_common.py | 1 - .../src/dynaclr/evaluation/linear_classifiers/orchestrated.py | 1 - applications/dynaclr/src/dynaclr/evaluation/mmd/compute_mmd.py | 1 - 7 files changed, 9 deletions(-) diff --git a/applications/dynaclr/scripts/pseudotime/0-select_candidates/manual_candidates.py b/applications/dynaclr/scripts/pseudotime/0-select_candidates/manual_candidates.py index d837e0e2f..f5e35622e 100644 --- a/applications/dynaclr/scripts/pseudotime/0-select_candidates/manual_candidates.py +++ b/applications/dynaclr/scripts/pseudotime/0-select_candidates/manual_candidates.py @@ -117,7 +117,6 @@ def _validate_against_anndata(dataset_id: str, fov_name: str, tracks: list[dict] ds_meta = DATASETS[dataset_id] emb_path = ds_meta["embedding_zarr"] adata = ad.read_zarr(emb_path) - adata.obs_names_make_unique() obs = adata.obs fov_obs = obs[obs["fov_name"].astype(str) == fov_name] diff --git a/applications/dynaclr/scripts/pseudotime/0-select_candidates/select_candidates.py b/applications/dynaclr/scripts/pseudotime/0-select_candidates/select_candidates.py index 11cc21ba6..46990b560 100644 --- a/applications/dynaclr/scripts/pseudotime/0-select_candidates/select_candidates.py +++ b/applications/dynaclr/scripts/pseudotime/0-select_candidates/select_candidates.py @@ -195,7 +195,6 @@ def _select_productive_from_zarr( return pd.DataFrame(columns=OUTPUT_COLUMNS) adata = ad.read_zarr(matches[0]) - adata.obs_names_make_unique() if pred_column not in adata.obs.columns: _logger.warning(f"[{dataset_id}] {pred_column} not in {matches[0].name}; productive empty") return pd.DataFrame(columns=OUTPUT_COLUMNS) @@ -296,7 +295,6 @@ def _select_mock_from_zarr( return pd.DataFrame(columns=OUTPUT_COLUMNS) adata = ad.read_zarr(matches[0]) - adata.obs_names_make_unique() obs = adata.obs.copy() obs = obs[obs["fov_name"].astype(str).str.contains(fov_pattern, regex=False)] if obs.empty: @@ -413,7 +411,6 @@ def _load_lc_predictions( ) adata = ad.read_zarr(matches[0]) - adata.obs_names_make_unique() if pred_column not in adata.obs.columns: _logger.warning(f"{pred_column} not in {matches[0]} .obs; LC fallback") return pd.DataFrame() diff --git a/applications/dynaclr/scripts/pseudotime/2-align_cells/align_embedding.py b/applications/dynaclr/scripts/pseudotime/2-align_cells/align_embedding.py index 8618e0c71..e6f616dd3 100644 --- a/applications/dynaclr/scripts/pseudotime/2-align_cells/align_embedding.py +++ b/applications/dynaclr/scripts/pseudotime/2-align_cells/align_embedding.py @@ -240,7 +240,6 @@ def _load_query_embeddings( zarr_path = find_embedding_zarr(ds_cfg["pred_dir"], prefix + embedding_pattern) adata = ad.read_zarr(zarr_path) - adata.obs_names_make_unique() # FOV restriction from the dataset config (e.g. "C/2") — keeps us # out of control wells unless the user explicitly wants them. diff --git a/applications/dynaclr/scripts/pseudotime/2-align_cells/align_lc.py b/applications/dynaclr/scripts/pseudotime/2-align_cells/align_lc.py index 6a14e096b..edfa8246e 100644 --- a/applications/dynaclr/scripts/pseudotime/2-align_cells/align_lc.py +++ b/applications/dynaclr/scripts/pseudotime/2-align_cells/align_lc.py @@ -101,7 +101,6 @@ def _load_lc_predictions( f"{[m.name for m in matches]}; using first" ) adata = ad.read_zarr(matches[0]) - adata.obs_names_make_unique() if pred_column not in adata.obs.columns: _logger.warning(f"{pred_column} not in {matches[0]} .obs") return pd.DataFrame(columns=["fov_name", "track_id", "t", pred_column]) diff --git a/applications/dynaclr/scripts/pseudotime/3-organelle-remodeling/readout_common.py b/applications/dynaclr/scripts/pseudotime/3-organelle-remodeling/readout_common.py index 4930292c5..cbbb49a1b 100644 --- a/applications/dynaclr/scripts/pseudotime/3-organelle-remodeling/readout_common.py +++ b/applications/dynaclr/scripts/pseudotime/3-organelle-remodeling/readout_common.py @@ -89,7 +89,6 @@ def load_organelle_embeddings( _logger.warning(f"[{ds_id}] no embedding zarr matched {prefix + embedding_pattern}: {exc}") continue adata = ad.read_zarr(zarr_path) - adata.obs_names_make_unique() out[ds_id] = adata _logger.info(f"[{ds_id}] loaded {Path(zarr_path).name} ({adata.n_obs} cells)") return out diff --git a/applications/dynaclr/src/dynaclr/evaluation/linear_classifiers/orchestrated.py b/applications/dynaclr/src/dynaclr/evaluation/linear_classifiers/orchestrated.py index 70a0bd7b8..f4423eb57 100644 --- a/applications/dynaclr/src/dynaclr/evaluation/linear_classifiers/orchestrated.py +++ b/applications/dynaclr/src/dynaclr/evaluation/linear_classifiers/orchestrated.py @@ -73,7 +73,6 @@ def run_linear_classifiers( raise FileNotFoundError(f"No .zarr files found in {embeddings_path}") parts = [ad.read_zarr(p) for p in zarr_paths] adata = ad.concat(parts, join="outer") - adata.obs_names_make_unique() click.echo(f" Loaded {len(zarr_paths)} per-experiment zarrs") else: adata = ad.read_zarr(embeddings_path) diff --git a/applications/dynaclr/src/dynaclr/evaluation/mmd/compute_mmd.py b/applications/dynaclr/src/dynaclr/evaluation/mmd/compute_mmd.py index c08fdc40b..8a751f0fe 100644 --- a/applications/dynaclr/src/dynaclr/evaluation/mmd/compute_mmd.py +++ b/applications/dynaclr/src/dynaclr/evaluation/mmd/compute_mmd.py @@ -581,7 +581,6 @@ def run_mmd_pooled(config: MMDPooledConfig) -> pd.DataFrame: adatas = [ad.read_zarr(p) for p in config.input_paths] combined = ad.concat(adatas, join="outer", label="source_experiment") - combined.obs_names_make_unique() if config.obs_filter: mask = pd.Series([True] * len(combined), index=combined.obs.index)