zmap-tools 0.2.0__tar.gz → 0.2.2__tar.gz

{zmap_tools-0.2.0 → zmap_tools-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: zmap-tools
-Version: 0.2.0
+Version: 0.2.2
 Summary: API for working with ZMAP (Zebrafish Multi-Atlas Project)
 Author-email: Daniel Wagner <Daniel.Wagner@ucsf.edu>
 License: BSD 3-Clause License
@@ -50,7 +50,7 @@ License-File: LICENSE
 Requires-Dist: anndata>=0.10
 Requires-Dist: scanpy>=1.9
 Requires-Dist: numpy>=1.24
-Requires-Dist: pandas>=2.0
+Requires-Dist: pandas<3.0,>=2.0
 Requires-Dist: matplotlib>=3.7
 Requires-Dist: requests>=2.31
 Requires-Dist: scipy>=1.10

{zmap_tools-0.2.0 → zmap_tools-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "zmap-tools"
-version = "0.2.0"
+version = "0.2.2"
 description = "API for working with ZMAP (Zebrafish Multi-Atlas Project)"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -30,7 +30,7 @@ dependencies = [
   "anndata>=0.10",
   "scanpy>=1.9",
   "numpy>=1.24",
-  "pandas>=2.0",
+  "pandas>=2.0,<3.0",
   "matplotlib>=3.7",
   "requests>=2.31",
   "scipy>=1.10",

{zmap_tools-0.2.0 → zmap_tools-0.2.2}/src/zmap/predict/knn_backend.py

@@ -18,10 +18,12 @@ _FAISS_INDEX_CACHE_ORDER: list[tuple[str, str, int, int]] = []
 
 
 def _as_float32_c(x: np.ndarray) -> np.ndarray:
+    """Cast array to float32 and ensure C-contiguous memory layout (required by FAISS)."""
     return np.ascontiguousarray(np.asarray(x, dtype=np.float32))
 
 
 def _l2_normalize(x: np.ndarray) -> np.ndarray:
+    """Row-wise L2-normalize an array to unit vectors (used for cosine similarity via inner product)."""
     x = _as_float32_c(x)
     norms = np.linalg.norm(x, axis=1, keepdims=True)
     norms = np.where(norms == 0, 1.0, norms)
@@ -29,6 +31,7 @@ def _l2_normalize(x: np.ndarray) -> np.ndarray:
 
 
 def _parse_cuda_device(device: str) -> int:
+    """Parse a device string like 'cuda:1' into an integer GPU index. Returns 0 on failure."""
     d = str(device).strip().lower()
     if ":" in d:
         try:
@@ -44,6 +47,17 @@ def _search_sklearn(
     n_neighbors: int,
     metric: str,
 ) -> tuple[np.ndarray, np.ndarray]:
+    """
+    kNN search using scikit-learn's NearestNeighbors.
+
+    Used as the fallback when FAISS is unavailable or fails. Slower than FAISS
+    on large references but supports any sklearn-compatible metric.
+
+    Returns
+    -------
+    neighbor_indices : np.ndarray, shape (n_query, n_neighbors)
+    distances : np.ndarray, shape (n_query, n_neighbors)
+    """
     nn = NearestNeighbors(n_neighbors=n_neighbors, metric=metric)
     nn.fit(ref)
     distances, neighbor_indices = nn.kneighbors(query, return_distance=True)
@@ -59,6 +73,44 @@ def _search_faiss(
     nprobe: int | None = None,
     cache_key: str | None = None,
 ) -> tuple[np.ndarray, np.ndarray, dict[str, Any]]:
+    """
+    kNN search using FAISS with optional GPU acceleration and index caching.
+
+    Builds an IVF-Flat index by default for fast approximate search on large references.
+    Falls back to an exact Flat index if IVF training fails. Cosine similarity is
+    implemented as inner product after L2-normalization; euclidean uses L2 distance.
+
+    Parameters
+    ----------
+    ref : np.ndarray, shape (n_ref, d)
+        Reference embedding matrix.
+    query : np.ndarray, shape (n_query, d)
+        Query embedding matrix.
+    n_neighbors : int
+        Number of nearest neighbors to return.
+    metric : {'euclidean', 'cosine'}
+        Distance metric.
+    device : str
+        Target device — 'cpu', 'auto', 'cuda', or 'cuda:<id>'.
+    nprobe : int, optional
+        Number of IVF cells to probe during search. Higher values increase recall
+        at the cost of speed. Defaults to DEFAULT_FAISS_NPROBE (16).
+    cache_key : str, optional
+        If provided, the built FAISS index is cached in-memory under this key and
+        reused on subsequent calls with matching (key, metric, n_ref, d). Up to
+        MAX_FAISS_INDEX_CACHE (16) indices are retained (LRU eviction).
+
+    Returns
+    -------
+    indices : np.ndarray, shape (n_query, n_neighbors), dtype int64
+        Indices into ref for each neighbor. Entries are -1 (set to NaN in distances)
+        when FAISS cannot fill all k neighbors (rare with Flat indices).
+    distances : np.ndarray, shape (n_query, n_neighbors), dtype float32
+        Euclidean distances or cosine-distance proxies (1 - cosine_similarity).
+    meta : dict
+        Diagnostic info: backend_used, device_used, index_type, nlist, nprobe,
+        cache_key_used, cache_hit.
+    """
     try:
         import faiss  # type: ignore
     except Exception as e:
@@ -233,21 +285,47 @@ def knn_search(
     cache_key: str | None = None,
 ) -> tuple[np.ndarray, np.ndarray, dict[str, Any]]:
     """
-    Unified kNN search for ZMAP API.
-
-    backend:
-      - auto: try faiss first, fallback to sklearn
-      - faiss: require faiss, fallback to sklearn with warning
-      - sklearn: force sklearn
-    device:
-      - auto | cpu | cuda | cuda:<id>
-      - used only by faiss backend
-    nprobe:
-      - IVF probe count for FAISS (None -> default)
-      - used only by faiss backend
-    cache_key:
-      - optional key to reuse FAISS index build across repeated calls
-      - used only by faiss backend
+    Unified kNN search supporting FAISS and scikit-learn backends.
+
+    Tries FAISS first (fast, approximate) and falls back to sklearn (exact) if
+    FAISS is unavailable or raises. Both backends return consistent output shapes
+    and distance semantics.
+
+    Parameters
+    ----------
+    ref : np.ndarray, shape (n_ref, d)
+        Reference embedding matrix.
+    query : np.ndarray, shape (n_query, d)
+        Query embedding matrix.
+    n_neighbors : int
+        Number of nearest neighbors to return per query point.
+    metric : {'cosine', 'euclidean'}, default 'cosine'
+        Distance metric. Cosine is computed as 1 - cosine_similarity.
+    backend : {'auto', 'faiss', 'sklearn'}, default 'auto'
+        Search backend. 'auto' tries FAISS first and falls back to sklearn.
+        'faiss' also falls back to sklearn with a warning if FAISS fails.
+        'sklearn' skips FAISS entirely.
+    device : {'auto', 'cpu', 'cuda', 'cuda:<id>'}, default 'auto'
+        Device for FAISS index. Ignored by the sklearn backend. 'auto' uses
+        GPU if FAISS GPU support is available, otherwise CPU.
+    nprobe : int, optional
+        Number of IVF partitions to probe during FAISS search. Higher values
+        improve recall at the cost of speed. Defaults to DEFAULT_FAISS_NPROBE.
+        Ignored by the sklearn backend and FAISS Flat indices.
+    cache_key : str, optional
+        If provided, the FAISS index built for `ref` is cached in-memory under
+        this key and reused on subsequent calls with the same (key, metric,
+        n_ref, d). Useful when the same reference is queried repeatedly.
+
+    Returns
+    -------
+    indices : np.ndarray, shape (n_query, n_neighbors), dtype int64
+        Row indices into `ref` for each neighbor, ordered by ascending distance.
+    distances : np.ndarray, shape (n_query, n_neighbors), dtype float32
+        Corresponding distances. NaN where FAISS returns invalid indices (-1).
+    meta : dict
+        Diagnostic info including backend_used, device_used, index_type, and
+        cache_hit. Useful for verifying which backend and device were actually used.
     """
     metric = str(metric).lower()
     backend_req = str(backend).lower()

{zmap_tools-0.2.0 → zmap_tools-0.2.2}/src/zmap/predict/predict.py

@@ -962,7 +962,7 @@ def predict_labels_kNN(
         )
 
     # compact digest to validate cache
-    _mask_digest = [int(n_ref_keep), int(np.sum(np.flatnonzero(ref_keep_mask.to_numpy()) % 1048573))]
+    _mask_digest = [str(int(n_ref_keep)), str(int(np.sum(np.flatnonzero(ref_keep_mask.to_numpy()) % 1048573)))]
 
     # filtered reference arrays
     X_ref_all = adata_ref.obsm[ref_basis]
@@ -991,7 +991,7 @@ def predict_labels_kNN(
             cache.get('query_basis') == query_basis and
             cache.get('metric') == metric and
             cache.get('n_neighbors') == n_neighbors and
-            cache.get('mask_digest') == _mask_digest and
+            list(cache.get('mask_digest', [])) == list(_mask_digest) and
             cache.get('tissue_mode', 'none') == expected_cache_mode and
             cache.get('knn_backend_requested', 'auto') == knn_backend and
             cache.get('knn_device_requested', 'auto') == knn_device and
@@ -1101,7 +1101,7 @@ def predict_labels_kNN(
 
     predicted_labels = sorted_classes[np.argmax(probabilities_sorted, axis=1)]
     predicted_labels = predicted_labels.astype(object)
-    predicted_labels[~has_votes] = pd.NA
+    predicted_labels[~has_votes] = np.nan
 
     # ---------- outputs ----------
     if omit_labels:
@@ -1224,7 +1224,7 @@ def predict_labels_kNN(
         adata_query.obs[col_reject] = ~accept
 
         # Mask rejected predictions
-        adata_query.obs.loc[~accept, col_main] = pd.NA
+        adata_query.obs.loc[~accept, col_main] = np.nan
 
         n_total = len(accept)
         n_accept = int(accept.sum())
@@ -1285,7 +1285,7 @@ def predict_labels_kNN(
         rare_labels = label_counts[label_counts < min_cells_per_label].index
         if len(rare_labels) > 0:
             adata_query.obs[col_rareflag] = adata_query.obs[col_main].isin(rare_labels)
-            adata_query.obs.loc[adata_query.obs[col_rareflag], col_main] = pd.NA
+            adata_query.obs.loc[adata_query.obs[col_rareflag], col_main] = np.nan
             adata_query.uns.setdefault('zmap_labels', {}).setdefault(space, {})
             adata_query.uns['zmap_labels'][space]['Rare Labels'] = list(rare_labels)
             _zlog(f"Filtered {len(rare_labels)} rare labels: {list(rare_labels[:10])}{'...' if len(rare_labels) > 10 else ''}")
@@ -1671,8 +1671,8 @@ def predict_labels_tissue_kNN(
         )
 
     mask_digest = [
-        int(n_ref_keep),
-        int(np.sum(np.flatnonzero(ref_keep_mask.to_numpy()) % 1048573)),
+        str(int(n_ref_keep)),
+        str(int(np.sum(np.flatnonzero(ref_keep_mask.to_numpy()) % 1048573))),
     ]
 
     p_thresh_use = p_thresh
@@ -2061,6 +2061,8 @@ def summarize_knn_run(adata_query, label_key):
         ("pct_assigned", cov.get("pct_assigned")),
     ], columns=["Key", "Value"])
 
+    df["Value"] = df["Value"].astype(str)
+
     return df
 
 
@@ -2227,6 +2229,17 @@ def aggregate_by_cluster(
 
     df = pd.DataFrame(records)
 
+    # Ensure object columns are h5ad-serializable in uns:
+    # - numeric columns demoted to object by pd.NA → coerce back to float (pd.NA → np.nan)
+    # - string columns with pd.NA → fillna("") and cast to str
+    _numeric_cols = {"top_fraction", "mean_prob", "margin", "second_fraction"}
+    for col in df.columns:
+        if df[col].dtype == object:
+            if col in _numeric_cols:
+                df[col] = pd.to_numeric(df[col], errors="coerce")
+            else:
+                df[col] = df[col].fillna("").astype(str)
+
     if save_csv:
         os.makedirs(output_dir, exist_ok=True)
         out_path = os.path.join(output_dir, f"{label_space}_cluster_summary.csv")
@@ -2305,6 +2318,11 @@ def build_cell_annotations_table(
     df.index.name = "cell_id"
     df = df.reset_index()
 
+    # Ensure object columns are pure str (np.nan in object dtype is not h5ad-serializable in uns)
+    for col in df.columns:
+        if df[col].dtype == object:
+            df[col] = df[col].fillna("").astype(str)
+
     if save_csv:
         os.makedirs(output_dir, exist_ok=True)
         out_path = os.path.join(output_dir, f"{label_space}_cell_annotations.csv")
@@ -2557,6 +2575,7 @@ def plot_embedding_with_ondata_labels(
     # ---- data/keys ----
     color_key: str = "ZMAP_Tissue_predicted",
     basis: str = "X_umap",
+    ref_basis: str | None = None,
     filter_na: bool = True,
 
     # ---- palette handling ----
@@ -2572,10 +2591,14 @@ def plot_embedding_with_ondata_labels(
     # ---- figure style ----
     figsize: tuple[float, float] = (6, 6),
     dpi: int = 200,
+    title_fontsize: float = 10,
 
-    # ---- point/legend style ----
+    # ---- point style ----
     ref_size: float = 2,
+    ref_alpha: float = 0.3,
     test_size: float = 2,
+    test_alpha: float = 1.0,
+    show_ref: bool = True,
     cmap: str = "jet",
     frameon: bool = False,
     sort_order: bool = True,
@@ -2583,6 +2606,9 @@ def plot_embedding_with_ondata_labels(
     legend_fontsize: float = 5,
     legend_fontweight: str = "normal",
 
+    # ---- label visibility ----
+    show_labels: bool = True,
+
     # ---- label styling ----
     recolor_labels_from_palette: bool = True,
     text_stroke_width: float = 1.0,
@@ -2613,9 +2639,10 @@ def plot_embedding_with_ondata_labels(
     Plot a query dataset overlaid on the reference embedding, with on-data labels
     and an optional vertical time distribution strip.
 
-    Renders two layers: (1) the full reference embedding as a faint grey
-    background for spatial context, and (2) the query cells colored by a
-    predicted label column. Labels are drawn directly on the embedding using
+    Renders up to two layers: (1) optionally, the full reference embedding
+    as a faint grey background for spatial context (controlled by
+    ``show_ref``), and (2) the query cells colored by a predicted label
+    column. Labels are drawn directly on the embedding using
     ``adjustText`` to minimize overlap. A vertical colorbar histogram of
     predicted developmental time (``ZMAP_time_id``) can optionally be added
     as a strip on the right side of the figure.
@@ -2630,7 +2657,12 @@ def plot_embedding_with_ondata_labels(
         Column in ``adata_test.obs`` containing the categorical labels to color
         and annotate. Typically a ``_predicted`` column from ``predict_labels_kNN``.
     basis : str, default ``"X_umap"``
-        ``obsm`` key used for the 2D embedding coordinates in both datasets.
+        ``obsm`` key used for the 2D embedding coordinates in the query dataset.
+    ref_basis : str or None, default ``None``
+        ``obsm`` key for the reference background embedding. When ``None``,
+        falls back to ``basis``. Useful when the query embedding lives in a
+        different key than the reference (e.g. ``basis="X_umap_zmap"`` for
+        the query while the reference uses ``"X_umap"``).
     filter_na : bool, default ``True``
         Drop query cells with ``NaN`` in ``color_key`` before plotting.
     palette : dict or None, default ``None``
@@ -2653,15 +2685,34 @@ def plot_embedding_with_ondata_labels(
         Figure size in inches ``(width, height)``.
     dpi : int, default ``200``
         Figure resolution.
-    ref_size, test_size : float, default ``2``
-        Scatter point sizes for reference background and query cells respectively.
+    ref_size : float, default ``2``
+        Scatter point size for reference background cells.
+    ref_alpha : float, default ``0.3``
+        Opacity of reference background points. Lower values push the
+        reference further into the background.
+    test_size : float, default ``2``
+        Scatter point size for query (projected) cells.
+    test_alpha : float, default ``1.0``
+        Opacity of query overlay points.
+    show_ref : bool, default ``True``
+        If ``True``, draw the reference embedding as a faint grey background
+        for spatial context. If ``False``, only the query cells are plotted
+        (useful when ``do_ingest=False``). When ``False``, ``adata_ref``
+        may be ``None``.
     cmap : str, default ``"jet"``
         Colormap used for the reference background scatter.
     legend_loc : str, default ``"on data"``
         Where to place the category legend. ``"on data"`` draws labels directly
         at centroid positions; other values follow matplotlib legend conventions.
+        Ignored when ``show_labels=False`` (forced to ``"none"``).
     legend_fontsize, legend_fontweight : float and str, default ``5`` and ``"normal"``
         Font size and weight for on-data legend labels.
+    show_labels : bool, default ``True``
+        If ``True``, draw on-data text labels at category centroids with
+        ``adjustText`` repositioning and optional arrow connectors. If
+        ``False``, suppress all text labels and arrows — only the colored
+        scatter is shown, which is useful for clean figures or when the
+        number of categories is too large for readable labels.
     replace_underscores : bool, default ``True``
         Replace underscores in label strings with line breaks for cleaner
         on-data annotation.
@@ -2669,17 +2720,23 @@ def plot_embedding_with_ondata_labels(
         ``(x_expand, y_expand)`` passed to ``adjustText`` for label placement.
     match_arrow_color_to_text : bool, default ``True``
         Color annotation arrows to match their corresponding text label.
+    ref_kwargs : dict or None, default ``None``
+        Extra keyword arguments forwarded to the reference ``sc.pl.embedding``
+        call. Explicit ``ref_alpha`` takes priority over ``alpha`` in this dict.
+    test_kwargs : dict or None, default ``None``
+        Extra keyword arguments forwarded to the query ``sc.pl.embedding``
+        call. Explicit ``test_alpha`` takes priority over ``alpha`` in this dict.
     show : bool, default ``False``
         Call ``plt.show()`` after rendering.
     save : bool, default ``True``
-        Save the figure as PNG and PDF to ``./zmap/predict/``.
+        Save the figure as PNG and PDF to ``output_dir``.
     return_ax : bool, default ``False``
         Return the main ``matplotlib.axes.Axes`` object.
 
     Returns
     -------
-    matplotlib.axes.Axes or None
-        The main axes when ``return_ax=True``, otherwise ``None``.
+    tuple or None
+        ``(fig, ax_umap, ax_strip)`` when ``return_ax=True``, otherwise ``None``.
     """
     # ---- prepare test AnnData (drop NAs on requested key, cast to categorical) ----
     if filter_na:
@@ -2698,7 +2755,7 @@ def plot_embedding_with_ondata_labels(
     # Only attempt sync if there's no _color_map dict already available
     # (annotate_with_zmap populates _color_map from reference colormaps)
     cmap_dict_key = f"{base_obs}_color_map"
-    if cmap_dict_key not in adata_test_plot.uns:
+    if cmap_dict_key not in adata_test_plot.uns and adata_ref is not None:
         try:
             sync_zmap_colors(adata_ref, obs_key=base_obs)
             sync_zmap_colors(adata_test_plot, obs_key=color_key, ref_adata=adata_ref, ref_obs_key=base_obs)
@@ -2745,6 +2802,13 @@ def plot_embedding_with_ondata_labels(
     test_kwargs = {} if test_kwargs is None else dict(test_kwargs)
     time_strip_kwargs = {} if time_strip_kwargs is None else dict(time_strip_kwargs)
 
+    # ---- Inject alpha into kwargs (explicit params take priority) ----
+    ref_kwargs.setdefault("alpha", ref_alpha)
+    test_kwargs.setdefault("alpha", test_alpha)
+
+    # ---- Resolve legend_loc: suppress labels when show_labels=False ----
+    legend_loc_use = legend_loc if show_labels else "none"
+
     # ---- Check if we actually have a time vector ----
     has_time = (
         show_time_strip and
@@ -2756,6 +2820,9 @@ def plot_embedding_with_ondata_labels(
     ax_umap = None
     ax_strip = None
 
+    # ---- Resolve ref_basis: defaults to basis if not provided ----
+    ref_basis_use = ref_basis if ref_basis is not None else basis
+
     # ---- figure + axes layout ----
     with plt.rc_context({'figure.figsize': figsize, 'figure.dpi': dpi}):
         if has_time:
@@ -2773,15 +2840,18 @@ def plot_embedding_with_ondata_labels(
             fig, ax_umap = plt.subplots()
             ax_strip = None
 
-        # ---- reference embedding in UMAP axis ----
-        ax = sc.pl.embedding(
-            adata_ref,
-            basis=basis,
-            show=False,
-            s=ref_size,
-            ax=ax_umap,
-            **ref_kwargs,
-        )
+        # ---- reference embedding in UMAP axis (optional) ----
+        if show_ref:
+            ax = sc.pl.embedding(
+                adata_ref,
+                basis=ref_basis_use,
+                show=False,
+                s=ref_size,
+                ax=ax_umap,
+                **ref_kwargs,
+            )
+        else:
+            ax = ax_umap
 
         # ---- query overlay in UMAP axis ----
         ax = sc.pl.embedding(
@@ -2793,7 +2863,7 @@ def plot_embedding_with_ondata_labels(
             basis=basis,
             sort_order=sort_order,
             size=test_size,
-            legend_loc=legend_loc,
+            legend_loc=legend_loc_use,
             legend_fontweight=legend_fontweight,
             ax=ax,
             show=False,
@@ -2801,42 +2871,44 @@ def plot_embedding_with_ondata_labels(
             palette=palette,
             **test_kwargs,
         )
+        ax.title.set_fontsize(title_fontsize)
 
-        # --- obtain existing on-data labels ---
-        texts = [t for t in ax.texts]
+        # --- on-data label styling (only when show_labels=True) ---
+        if show_labels:
+            texts = [t for t in ax.texts]
 
-        # --- recolor existing on-data text labels by palette ---
-        if recolor_labels_from_palette:
-            for t in texts:
-                label = t.get_text()
-                if label in palette:
-                    t.set_color(palette[label])
-
-        # --- add white border to text labels ---
-        for t in texts:
-            t.set_path_effects([pe.withStroke(linewidth=text_stroke_width, foreground='white')])
+            # --- recolor existing on-data text labels by palette ---
+            if recolor_labels_from_palette:
+                for t in texts:
+                    label = t.get_text()
+                    if label in palette:
+                        t.set_color(palette[label])
 
-        # --- replace underscores with line breaks ---
-        if replace_underscores:
+            # --- add white border to text labels ---
             for t in texts:
-                t.set_text(t.get_text().replace(linebreak_from, linebreak_to))
-
-        # --- adjust label positions with slim arrows ---
-        if len(texts) > 0:
-            adjust_text(
-                texts,
-                ax=ax,
-                expand=adjust_expand,
-                arrowprops=arrowprops,
-                min_arrow_len=min_arrow_len,
-            )
+                t.set_path_effects([pe.withStroke(linewidth=text_stroke_width, foreground='white')])
+
+            # --- replace underscores with line breaks ---
+            if replace_underscores:
+                for t in texts:
+                    t.set_text(t.get_text().replace(linebreak_from, linebreak_to))
+
+            # --- adjust label positions with slim arrows ---
+            if len(texts) > 0:
+                adjust_text(
+                    texts,
+                    ax=ax,
+                    expand=adjust_expand,
+                    arrowprops=arrowprops,
+                    min_arrow_len=min_arrow_len,
+                )
 
-        # --- match arrow color to label color ---
-        if match_arrow_color_to_text and len(texts) > 0:
-            arrows = [p for p in ax.patches if isinstance(p, FancyArrowPatch)]
-            for t, a in zip(texts, arrows[-len(texts):]):
-                a.set_color(t.get_color())
-                a.set_alpha(arrow_alpha)
+            # --- match arrow color to label color ---
+            if match_arrow_color_to_text and len(texts) > 0:
+                arrows = [p for p in ax.patches if isinstance(p, FancyArrowPatch)]
+                for t, a in zip(texts, arrows[-len(texts):]):
+                    a.set_color(t.get_color())
+                    a.set_alpha(arrow_alpha)
 
         # ---- time distribution strip on the right (rotated) ----
         if has_time and ax_strip is not None:
@@ -2920,7 +2992,6 @@ def plot_embedding_with_ondata_labels(
     return None
 
 
-
 # ================================================================
 #  6. Overlap matrix & plot
 # ================================================================
@@ -3078,20 +3149,26 @@ def map_query_labels(
     fig = None
     if show_plot or save_plots:
         plt.rcParams["axes.grid"] = False
-        fig, ax = plt.subplots(figsize=(figsize, figsize))
+        n_rows, n_cols = plot_table.shape
+        tick_fontsize = 8
+        cell_size = 0.25          # inches per cell — fits 8pt labels without crowding
+        fig_w = min(28, max(figsize, n_cols * cell_size))
+        fig_h = min(28, max(figsize, n_rows * cell_size))
+        cb_shrink = min(0.8, 3.0 / fig_h)   # keep colorbar ~3 inches tall
+        fig, ax = plt.subplots(figsize=(fig_w, fig_h))
 
         im = ax.imshow(plot_table.to_numpy(), cmap=cmap, vmin=vmin, vmax=vmax)
 
         ax.set_aspect("equal")
         ax.set_xticks(np.arange(plot_table.shape[1]))
         ax.set_yticks(np.arange(plot_table.shape[0]))
-        ax.set_xticklabels(plot_table.columns, rotation=90)
-        ax.set_yticklabels(plot_table.index)
+        ax.set_xticklabels(plot_table.columns, rotation=90, fontsize=tick_fontsize)
+        ax.set_yticklabels(plot_table.index, fontsize=tick_fontsize)
         ax.set_title(title)
         ax.set_xlabel(obs_A)
         ax.set_ylabel(obs_B)
 
-        cb = fig.colorbar(im, ax=ax, shrink=0.5)
+        cb = fig.colorbar(im, ax=ax, shrink=cb_shrink)
         cb.ax.set_ylabel(colorbar_label)
 
         if overlay_values:
@@ -3315,7 +3392,7 @@ def annotate_with_zmap(
         - ``.obs[f"{label_space}_predicted"]``                       — transferred cell labels.
         - ``.obs[f"{label_space}_prob"]``                            — label confidence (0–1).
         - ``.obs["ZMAP_time_id_predicted"]``                         — predicted time (hpf).
-        - ``.obsm["X_umap"]``                                        — UMAP coordinates (if ingested).
+        - ``.obsm["X_umap_zmap"]``                                   — projected UMAP coordinates (if ingested).
         - ``.uns['zmap_labels']['_last_space']``                     — most recent label_space.
         - ``.uns['zmap_labels'][label_space]['_run_config']``        — stored run parameters
           for zero-arg on-demand plot accessors.
@@ -3478,9 +3555,29 @@ def annotate_with_zmap(
         if do_ingest:
             if v >= 1:
                 _zlog("Ingesting query into reference UMAP...")
+
+            # Stash pre-existing X_umap so symphony doesn't clobber it
+            _had_existing_umap = "X_umap" in adata_query.obsm
+            if _had_existing_umap:
+                _existing_umap = adata_query.obsm["X_umap"].copy()
+
             with warnings.catch_warnings():
                 warnings.simplefilter("ignore", FutureWarning)
                 sp.tl.ingest(adata_query=adata_query, adata_ref=adata_ref)
+
+            # Move ingested coords to X_umap_zmap (non-destructive)
+            adata_query.obsm["X_umap_zmap"] = adata_query.obsm["X_umap"].copy()
+
+            # Restore original X_umap if one existed, otherwise remove
+            if _had_existing_umap:
+                adata_query.obsm["X_umap"] = _existing_umap
+                if v >= 1:
+                    _zlog("Restored pre-existing X_umap; ingested coords in X_umap_zmap.")
+            else:
+                del adata_query.obsm["X_umap"]
+                if v >= 1:
+                    _zlog("Ingested coords stored in X_umap_zmap.")
+
             if v >= 1:
                 _zlog("Ingestion complete.")
 
@@ -3554,15 +3651,26 @@ def annotate_with_zmap(
     adata_query.uns.setdefault("zmap_labels", {})
     adata_query.uns["zmap_labels"]["_last_space"] = space
     adata_query.uns["zmap_labels"].setdefault(space, {})
+    # Determine embedding basis and whether to show reference background
+    if do_ingest:
+        _plot_basis = "X_umap_zmap"
+        _show_ref = True
+    else:
+        # No ingest → query lives in its own UMAP space
+        _plot_basis = "X_umap" if "X_umap" in adata_query.obsm else None
+        _show_ref = False
+
     adata_query.uns["zmap_labels"][space]["_run_config"] = {
         "label_space": space,
         "ref_label_col": ref_label_col,
         "query_label_col": query_label_col,
         "time_col": time_col_actual,
         "output_dir": space_dir,
+        "basis": _plot_basis,
+        "show_ref": _show_ref,
     }
-    # Store reference UMAP for on-demand plot_embedding (avoids needing adata_ref)
-    if "X_umap" in adata_ref.obsm:
+    # Store reference UMAP for on-demand plot_embedding (only when ingested)
+    if do_ingest and "X_umap" in adata_ref.obsm:
         adata_query.uns["zmap_labels"][space]["_ref_umap"] = (
             np.asarray(adata_ref.obsm["X_umap"]).copy()
         )
@@ -3652,15 +3760,24 @@ def annotate_with_zmap(
     # ------------------------------------------------------------------
     # 9. UMAP overlay figure (v >= 2)
     # ------------------------------------------------------------------
-    if v >= 2:
+    if _plot_basis is None:
+        if v >= 1:
+            _zlog(
+                "Skipping UMAP overlay: no embedding found "
+                "(do_ingest=False and no pre-existing X_umap)."
+            )
+    elif v >= 2:
         try:
             if v >= 1:
                 _zlog("Plotting UMAP overlay with predicted labels...")
             plot_embedding_with_ondata_labels(
-                adata_ref,
+                adata_ref if _show_ref else None,
                 adata_query,
                 color_key=f"{space}_predicted",
                 time_key=time_col_actual,
+                basis=_plot_basis,
+                ref_basis="X_umap",
+                show_ref=_show_ref,
                 show=True,
                 save=save_outputs,
                 output_dir=space_dir,
@@ -3675,10 +3792,13 @@ def annotate_with_zmap(
         # Still save the figure even if not showing it
         try:
             plot_embedding_with_ondata_labels(
-                adata_ref,
+                adata_ref if _show_ref else None,
                 adata_query,
                 color_key=f"{space}_predicted",
                 time_key=time_col_actual,
+                basis=_plot_basis,
+                ref_basis="X_umap",
+                show_ref=_show_ref,
                 show=False,
                 save=True,
                 output_dir=space_dir,
@@ -3946,6 +4066,7 @@ def validate_markers(
             groupby=groupby,
             method=method,
             corr_method=corr_method,
+            use_raw=False,
         )
 
     # ---- Extract and filter DE results per group ----
@@ -4186,8 +4307,14 @@ def plot_embedding(
     label_space: str | None = None,
     *,
     save: bool = False,
+    show_labels: bool = True,
+    ref_size: float = 2,
+    ref_alpha: float = 0.3,
+    test_size: float = 2,
+    test_alpha: float = 1.0,
+    return_ax: bool = False,
     **kwargs,
-) -> None:
+) -> tuple | None:
     """
     Re-plot the UMAP overlay with on-data labels and time strip.
 
@@ -4202,37 +4329,75 @@ def plot_embedding(
         Label namespace. When ``None``, uses the most recent run.
     save : bool, default ``False``
         Save the figure to the run's output directory.
+    show_labels : bool, default ``True``
+        If ``True``, draw on-data text labels at category centroids.
+        If ``False``, suppress all text labels and arrows.
+    ref_size : float, default ``2``
+        Scatter point size for reference background cells.
+    ref_alpha : float, default ``0.3``
+        Opacity of reference background points.
+    test_size : float, default ``2``
+        Scatter point size for query (projected) cells.
+    test_alpha : float, default ``1.0``
+        Opacity of query overlay points.
+    return_ax : bool, default ``False``
+        Return ``(fig, ax_umap, ax_strip)`` instead of ``None``.
     **kwargs
         Extra keyword arguments forwarded to ``plot_embedding_with_ondata_labels``.
+
+    Returns
+    -------
+    tuple or None
+        ``(fig, ax_umap, ax_strip)`` when ``return_ax=True``, otherwise ``None``.
     """
     label_space, config = _resolve_config(adata_query, label_space)
     space_store = adata_query.uns.get("zmap_labels", {}).get(label_space, {})
 
-    ref_umap = space_store.get("_ref_umap")
-    if ref_umap is None:
+    # Resolve basis and show_ref from stored config
+    plot_basis = config.get("basis", "X_umap_zmap")
+    show_ref = config.get("show_ref", True)
+
+    if plot_basis is None:
         raise KeyError(
-            f"No stored reference UMAP found at "
-            f"adata.uns['zmap_labels']['{label_space}']['_ref_umap']. "
-            "Run annotate_with_zmap first."
+            f"No embedding basis stored for label_space='{label_space}'. "
+            "The original run had do_ingest=False and no pre-existing X_umap."
         )
 
-    # Build a minimal AnnData shell with just the reference UMAP
-    adata_ref_mini = ad.AnnData(
-        obs=pd.DataFrame(index=[f"ref_{i}" for i in range(ref_umap.shape[0])]),
-    )
-    adata_ref_mini.obsm["X_umap"] = ref_umap
+    # Build reference shell only when needed
+    adata_ref_mini = None
+    if show_ref:
+        ref_umap = space_store.get("_ref_umap")
+        if ref_umap is None:
+            raise KeyError(
+                f"No stored reference UMAP found at "
+                f"adata.uns['zmap_labels']['{label_space}']['_ref_umap']. "
+                "Run annotate_with_zmap first."
+            )
+        adata_ref_mini = ad.AnnData(
+            obs=pd.DataFrame(index=[f"ref_{i}" for i in range(ref_umap.shape[0])]),
+        )
+        adata_ref_mini.obsm["X_umap"] = ref_umap
 
     color_key = f"{label_space}_predicted"
     time_key = config.get("time_col", "ZMAP_time_id_predicted")
     space_dir = config.get("output_dir", os.path.join("zmap_predict", label_space))
 
-    plot_embedding_with_ondata_labels(
+    return plot_embedding_with_ondata_labels(
         adata_ref_mini,
         adata_query,
         color_key=color_key,
         time_key=time_key,
+        basis=plot_basis,
+        ref_basis="X_umap",
+        show_ref=show_ref,
         show=True,
         save=save,
+        show_labels=show_labels,
+        ref_size=ref_size,
+        ref_alpha=ref_alpha,
+        test_size=test_size,
+        test_alpha=test_alpha,
+        return_ax=return_ax,
         output_dir=space_dir,
         **kwargs,
     )

{zmap_tools-0.2.0 → zmap_tools-0.2.2}/src/zmap_tools.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: zmap-tools
-Version: 0.2.0
+Version: 0.2.2
 Summary: API for working with ZMAP (Zebrafish Multi-Atlas Project)
 Author-email: Daniel Wagner <Daniel.Wagner@ucsf.edu>
 License: BSD 3-Clause License
@@ -50,7 +50,7 @@ License-File: LICENSE
 Requires-Dist: anndata>=0.10
 Requires-Dist: scanpy>=1.9
 Requires-Dist: numpy>=1.24
-Requires-Dist: pandas>=2.0
+Requires-Dist: pandas<3.0,>=2.0
 Requires-Dist: matplotlib>=3.7
 Requires-Dist: requests>=2.31
 Requires-Dist: scipy>=1.10

{zmap_tools-0.2.0 → zmap_tools-0.2.2}/src/zmap_tools.egg-info/requires.txt

@@ -1,7 +1,7 @@
 anndata>=0.10
 scanpy>=1.9
 numpy>=1.24
-pandas>=2.0
+pandas<3.0,>=2.0
 matplotlib>=3.7
 requests>=2.31
 scipy>=1.10