sampledisco 0.2.0__tar.gz → 0.3.0__tar.gz

{sampledisco-0.2.0/src/sampledisco.egg-info → sampledisco-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sampledisco
-Version: 0.2.0
+Version: 0.3.0
 Summary: Cross-omics, cross-condition sample embedding for single-cell data (RNA + ATAC)
 Author: Hongkai Ji
 Author-email: Harry Jiang <hjiang55@jh.edu>
@@ -61,15 +61,15 @@ For each modality (RNA, ATAC, or integrated multi-omics) the pipeline produces t
 
 | Key | Role | Source |
 |---|---|---|
-| **`Z_clust`** | sample-removed embedding — used for clustering and composition blocks | Harmony (single-omics) / Harmony post-pass on scGLUE (multi-omics) |
-| **`Z_cmd`**  | sample-preserved embedding — used for the counterfactual displacement (CMD) block | second Harmony pass (single-omics) / scGLUE primary output (multi-omics) |
+| **`Z_comp`** | sample-removed embedding — used for cell typing and the composition blocks | Harmony (single-omics) / Harmony post-pass on scGLUE (multi-omics) |
+| **`Z_rmd`**  | sample-preserved embedding — used for the reference-relative mean displacement (RMD) block | second Harmony pass (single-omics) / scGLUE primary output (multi-omics) |
 
 It then assembles **four blocks** per sample (or per sample × modality for multi-omics):
 
 1. **A1** — one-hot cell-type composition
 2. **A2** — soft k-means composition at K_med (≈120)
 3. **A3** — soft k-means composition at K_fine (≈300)
-4. **CMD** — leave-one-out cell-type-resolved displacement on `Z_cmd`
+4. **RMD** — leave-one-out cell-type-resolved displacement on `Z_rmd`
 
 The four blocks are inverse-variance weighted, Frobenius-stacked, PCA-reduced to 10 dimensions, and Harmony-corrected at sample level. The result is stored as `adata.uns['X_DR_sample']` and feeds every downstream module.
 
@@ -87,20 +87,20 @@ code/
 │   ├── atac_wrapper.py
 │   └── multiomics_wrapper.py
 ├── preparation/               # Preprocessing
-│   ├── rna_preprocess_{cpu,gpu}.py   # QC → HVG → PCA → dual Harmony → Z_clust + Z_cmd
-│   ├── atac_preprocess_{cpu,gpu}.py  # QC → TF-IDF → HVF → LSI → dual Harmony → Z_clust + Z_cmd
-│   ├── cell_type_{cpu,gpu}.py        # Leiden clustering on Z_clust (RNA or ATAC)
+│   ├── rna_preprocess_{cpu,gpu}.py   # QC → HVG → PCA → dual Harmony → Z_comp + Z_rmd
+│   ├── atac_preprocess_{cpu,gpu}.py  # QC → TF-IDF → HVF → LSI → dual Harmony → Z_comp + Z_rmd
+│   ├── cell_type_{cpu,gpu}.py        # Leiden clustering on Z_comp (RNA or ATAC)
 │   ├── ATAC_cell_type{,_gpu}.py      # ATAC-specific cell typing variants
 │   ├── multi_omics_glue.py           # scGLUE integration (cross-modality VAE + guidance graph)
-│   ├── multi_omics_batch_correction.py # Harmony post-pass on X_glue → Z_clust
+│   ├── multi_omics_batch_correction.py # Harmony post-pass on X_glue → Z_comp
 │   ├── multi_omics_merge.py          # post-GLUE merge + per-modality preprocess/slimming
 │   └── multi_omics_cell_type_{cpu,gpu}.py  # RNA-Leiden + k-NN label transfer to ATAC
 ├── sample_embedding/          # Core method
-│   ├── blocks.py              # composition, CMD, weighting, Frobenius stack, final PCA + Harmony
+│   ├── blocks.py              # composition, RMD, weighting, Frobenius stack, final PCA + Harmony
 │   ├── sample_embedding.py    # CPU pipeline
 │   └── sample_embedding_gpu.py # GPU pipeline (cuML + cupy)
 ├── parameter_selection/
-│   └── autotune.py            # Bayesian GP sweep over CMD α; adaptive proxy ensemble
+│   └── autotune.py            # Bayesian GP sweep over RMD α; adaptive proxy ensemble
 ├── sample_distance/           # Pairwise sample distances (DR / EMD / chi-square / JS)
 ├── sample_clustering/         # Hierarchical (HRA / HRC / NN / UPGMA / consensus), K-means, proportion test, RAISIN
 ├── sample_trajectory/         # CCA (supervised) and TSCAN (unsupervised) + GAM-based trajectory DGE
@@ -165,7 +165,7 @@ For multi-omics, the pipeline takes two separate h5ads (RNA + ATAC) and integrat
 ├── atac/   (parallel structure)
 ├── multiomics/
 │   ├── integration/glue/{rna-pp,atac-pp,guidance.graphml.gz}
-│   ├── preprocess/adata_sample.h5ad           # post-GLUE merged adata with Z_clust + Z_cmd
+│   ├── preprocess/adata_sample.h5ad           # post-GLUE merged adata with Z_comp + Z_rmd
 │   ├── sample_embedding/sample_embedding.csv
 │   └── (same downstream subdirs as single-omics)
 └── sys_log/main_process_status.json           # which stages completed

{sampledisco-0.2.0 → sampledisco-0.3.0}/README.md

@@ -14,15 +14,15 @@ For each modality (RNA, ATAC, or integrated multi-omics) the pipeline produces t
 
 | Key | Role | Source |
 |---|---|---|
-| **`Z_clust`** | sample-removed embedding — used for clustering and composition blocks | Harmony (single-omics) / Harmony post-pass on scGLUE (multi-omics) |
-| **`Z_cmd`**  | sample-preserved embedding — used for the counterfactual displacement (CMD) block | second Harmony pass (single-omics) / scGLUE primary output (multi-omics) |
+| **`Z_comp`** | sample-removed embedding — used for cell typing and the composition blocks | Harmony (single-omics) / Harmony post-pass on scGLUE (multi-omics) |
+| **`Z_rmd`**  | sample-preserved embedding — used for the reference-relative mean displacement (RMD) block | second Harmony pass (single-omics) / scGLUE primary output (multi-omics) |
 
 It then assembles **four blocks** per sample (or per sample × modality for multi-omics):
 
 1. **A1** — one-hot cell-type composition
 2. **A2** — soft k-means composition at K_med (≈120)
 3. **A3** — soft k-means composition at K_fine (≈300)
-4. **CMD** — leave-one-out cell-type-resolved displacement on `Z_cmd`
+4. **RMD** — leave-one-out cell-type-resolved displacement on `Z_rmd`
 
 The four blocks are inverse-variance weighted, Frobenius-stacked, PCA-reduced to 10 dimensions, and Harmony-corrected at sample level. The result is stored as `adata.uns['X_DR_sample']` and feeds every downstream module.
 
@@ -40,20 +40,20 @@ code/
 │   ├── atac_wrapper.py
 │   └── multiomics_wrapper.py
 ├── preparation/               # Preprocessing
-│   ├── rna_preprocess_{cpu,gpu}.py   # QC → HVG → PCA → dual Harmony → Z_clust + Z_cmd
-│   ├── atac_preprocess_{cpu,gpu}.py  # QC → TF-IDF → HVF → LSI → dual Harmony → Z_clust + Z_cmd
-│   ├── cell_type_{cpu,gpu}.py        # Leiden clustering on Z_clust (RNA or ATAC)
+│   ├── rna_preprocess_{cpu,gpu}.py   # QC → HVG → PCA → dual Harmony → Z_comp + Z_rmd
+│   ├── atac_preprocess_{cpu,gpu}.py  # QC → TF-IDF → HVF → LSI → dual Harmony → Z_comp + Z_rmd
+│   ├── cell_type_{cpu,gpu}.py        # Leiden clustering on Z_comp (RNA or ATAC)
 │   ├── ATAC_cell_type{,_gpu}.py      # ATAC-specific cell typing variants
 │   ├── multi_omics_glue.py           # scGLUE integration (cross-modality VAE + guidance graph)
-│   ├── multi_omics_batch_correction.py # Harmony post-pass on X_glue → Z_clust
+│   ├── multi_omics_batch_correction.py # Harmony post-pass on X_glue → Z_comp
 │   ├── multi_omics_merge.py          # post-GLUE merge + per-modality preprocess/slimming
 │   └── multi_omics_cell_type_{cpu,gpu}.py  # RNA-Leiden + k-NN label transfer to ATAC
 ├── sample_embedding/          # Core method
-│   ├── blocks.py              # composition, CMD, weighting, Frobenius stack, final PCA + Harmony
+│   ├── blocks.py              # composition, RMD, weighting, Frobenius stack, final PCA + Harmony
 │   ├── sample_embedding.py    # CPU pipeline
 │   └── sample_embedding_gpu.py # GPU pipeline (cuML + cupy)
 ├── parameter_selection/
-│   └── autotune.py            # Bayesian GP sweep over CMD α; adaptive proxy ensemble
+│   └── autotune.py            # Bayesian GP sweep over RMD α; adaptive proxy ensemble
 ├── sample_distance/           # Pairwise sample distances (DR / EMD / chi-square / JS)
 ├── sample_clustering/         # Hierarchical (HRA / HRC / NN / UPGMA / consensus), K-means, proportion test, RAISIN
 ├── sample_trajectory/         # CCA (supervised) and TSCAN (unsupervised) + GAM-based trajectory DGE
@@ -118,7 +118,7 @@ For multi-omics, the pipeline takes two separate h5ads (RNA + ATAC) and integrat
 ├── atac/   (parallel structure)
 ├── multiomics/
 │   ├── integration/glue/{rna-pp,atac-pp,guidance.graphml.gz}
-│   ├── preprocess/adata_sample.h5ad           # post-GLUE merged adata with Z_clust + Z_cmd
+│   ├── preprocess/adata_sample.h5ad           # post-GLUE merged adata with Z_comp + Z_rmd
 │   ├── sample_embedding/sample_embedding.csv
 │   └── (same downstream subdirs as single-omics)
 └── sys_log/main_process_status.json           # which stages completed

{sampledisco-0.2.0 → sampledisco-0.3.0}/pyproject.toml

@@ -9,7 +9,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "sampledisco"
-version = "0.2.0"
+version = "0.3.0"
 description = "Cross-omics, cross-condition sample embedding for single-cell data (RNA + ATAC)"
 readme = "README.md"
 requires-python = ">=3.10"

{sampledisco-0.2.0 → sampledisco-0.3.0}/src/sampledisco/config/config_demo.yaml

@@ -95,7 +95,7 @@ rna_autotune_enable: false
 rna_autotune_search: "bayesian"
 rna_autotune_scoring: "auto"
 rna_autotune_scope: "alpha_only"
-rna_autotune_alpha_bounds: [0.1, 10.0]
+rna_autotune_alpha_bounds: [0.1, 100.0]
 rna_autotune_grouping_col:
 
 # --- Trajectory Analysis Parameters ---
@@ -216,7 +216,7 @@ atac_autotune_enable: false
 atac_autotune_search: "bayesian"
 atac_autotune_scoring: "auto"
 atac_autotune_scope: "alpha_only"
-atac_autotune_alpha_bounds: [0.1, 10.0]
+atac_autotune_alpha_bounds: [0.1, 100.0]
 atac_autotune_grouping_col:
 
 # --- Trajectory Analysis Parameters ---
@@ -332,9 +332,9 @@ multiomics_treat_sample_as_batch: false  # GLUE preserves per-sample variance (V
 
 # --- Sample-removed cluster embedding (X_glue from scGLUE is the RMD/sample-preserved emb) ---
 # By default, one Harmony pass on X_glue removes per-sample variance to
-# produce X_glue_harmony (used by cell typing + A1/A2/A3). Alternative:
+# produce Z_comp (used by cell typing + A1/A2/A3). Alternative:
 # train scGLUE a second time with treat_sample_as_batch=True; that run's
-# embedding is merged in as X_glue_harmony and the Harmony pass auto-skips.
+# embedding is merged in as Z_comp and the Harmony pass auto-skips.
 multiomics_harmonize_xglue_max_iter: 30
 multiomics_run_glue_twice_for_sample_removal: false
 multiomics_save_prefix: "glue"
@@ -355,7 +355,7 @@ multiomics_metric: "cosine"
 multiomics_existing_cell_types: false
 multiomics_n_target_clusters: 8
 multiomics_cluster_resolution: 0.8
-# null → auto-resolve to X_glue_harmony (sample-removed) if present, else X_glue
+# null → auto-resolve to Z_comp (sample-removed); legacy key names accepted
 multiomics_use_rep_celltype:
 multiomics_markers:
 multiomics_generate_umap_celltype: true
@@ -384,7 +384,7 @@ multiomics_autotune_enable: false
 multiomics_autotune_search: "bayesian"
 multiomics_autotune_scoring: "auto"
 multiomics_autotune_scope: "alpha_only"
-multiomics_autotune_alpha_bounds: [0.1, 10.0]
+multiomics_autotune_alpha_bounds: [0.1, 100.0]
 multiomics_autotune_grouping_col:
 
 # --- Trajectory Analysis Parameters ---

{sampledisco-0.2.0 → sampledisco-0.3.0}/src/sampledisco/parameter_selection/autotune.py

@@ -20,6 +20,7 @@ from __future__ import annotations
 import math
 import os
 import time
+import warnings
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -45,9 +46,17 @@ from sampledisco.sample_embedding.blocks import (
     loo_rmd,
     soft_assign,
 )
+from sampledisco.utils.embedding_keys import resolve_comp_key, resolve_rmd_key
 
 
-DEFAULT_ALPHA_BOUNDS = (0.1, 10.0)
+# Only the ceiling moved (10.0 -> 100.0). Datasets whose signal lives in the RMD
+# block were being truncated by the old ceiling: a dense sweep puts the optimum
+# at alpha = 16 on the 1M-scBloodNL stimulation time-course, and several saved
+# runs (covid ATAC, covid 279, unpaired paper) had returned alpha pinned at 10.0.
+# The floor stays at 0.1: no dataset has ever optimised against it, and on a
+# composition-dominated cohort the objective is flat below ~1.5, so a lower floor
+# only lets alpha drift down until the RMD block is numerically absent.
+DEFAULT_ALPHA_BOUNDS = (0.1, 100.0)
 
 
 # ============================================================ #
@@ -57,7 +66,7 @@ def build_blocks(
     adata: AnnData,
     sample_col: str,
     celltype_col: str,
-    cluster_emb_key: str,
+    comp_emb_key: Optional[str] = None,
     rmd_emb_key: Optional[str] = None,
     modality_col: Optional[str] = None,
     batch_col: Optional[str] = None,
@@ -69,12 +78,12 @@ def build_blocks(
     verbose: bool = True,
 ) -> Dict:
     """Build composition + RMD blocks once. Returns a dict the inner loop reuses."""
-    rmd_key = rmd_emb_key if rmd_emb_key and rmd_emb_key in adata.obsm else (
-        "Z_rmd" if "Z_rmd" in adata.obsm else cluster_emb_key
-    )
+    comp_key = resolve_comp_key(adata, comp_emb_key, context="autotune")
+    rmd_key = resolve_rmd_key(adata, rmd_emb_key, comp_key=comp_key,
+                              context="autotune")
 
-    units, unit_cellids, unit_ids, unit_groups, unit_batches, all_cellids, Z_clust = \
-        assemble_units(adata, sample_col, cluster_emb_key,
+    units, unit_cellids, unit_ids, unit_groups, unit_batches, all_cellids, Z_comp = \
+        assemble_units(adata, sample_col, comp_key,
                        modality_col=modality_col, batch_col=batch_col)
     n_units = len(units)
     cellid_idx = {cid: i for i, cid in enumerate(all_cellids)}
@@ -85,26 +94,26 @@ def build_blocks(
 
     # A1: hard cell-type-label composition (one-hot, no clustering)
     L1 = {ct: i for i, ct in enumerate(unique_cts)}
-    soft1 = np.zeros((Z_clust.shape[0], K_c), dtype=np.float32)
+    soft1 = np.zeros((Z_comp.shape[0], K_c), dtype=np.float32)
     for i, ct in enumerate(cell_type):
         soft1[i, L1[ct]] = 1.0
     unit_cellids_list = [unit_cellids[uid] for uid in unit_ids]
     A1 = composition_per_unit(unit_cellids_list, soft1, cellid_idx)
 
-    K_med = min(medium_K, max(2, Z_clust.shape[0] // 200))
+    K_med = min(medium_K, max(2, Z_comp.shape[0] // 200))
     if verbose:
         print(f"[autotune.build_blocks] K-means K={K_med}...")
     km_med = MiniBatchKMeans(n_clusters=K_med, random_state=seed,
-                              batch_size=4096, n_init=5, max_iter=200).fit(Z_clust)
-    soft2 = soft_assign(Z_clust, km_med.cluster_centers_)
+                              batch_size=4096, n_init=5, max_iter=200).fit(Z_comp)
+    soft2 = soft_assign(Z_comp, km_med.cluster_centers_)
     A2 = composition_per_unit(unit_cellids_list, soft2, cellid_idx)
 
-    K_fine = min(fine_K, max(2, Z_clust.shape[0] // 100))
+    K_fine = min(fine_K, max(2, Z_comp.shape[0] // 100))
     if verbose:
         print(f"[autotune.build_blocks] K-means K={K_fine}...")
     km_fine = MiniBatchKMeans(n_clusters=K_fine, random_state=seed + 1,
-                                batch_size=4096, n_init=5, max_iter=200).fit(Z_clust)
-    soft3 = soft_assign(Z_clust, km_fine.cluster_centers_)
+                                batch_size=4096, n_init=5, max_iter=200).fit(Z_comp)
+    soft3 = soft_assign(Z_comp, km_fine.cluster_centers_)
     A3 = composition_per_unit(unit_cellids_list, soft3, cellid_idx)
 
     # RMD
@@ -152,7 +161,7 @@ def build_blocks(
         unit_ids=unit_ids, unit_groups=unit_groups, unit_batches=unit_batches,
         n_units=n_units, grouping=grouping_arr, batch=batch_arr,
         has_batch=has_batch, has_grouping=has_grouping,
-        cluster_emb_key=cluster_emb_key, rmd_emb_key=rmd_key,
+        comp_emb_key=comp_key, rmd_emb_key=rmd_key,
     )
 
 
@@ -441,8 +450,29 @@ def make_scorer(name: str, meta: Dict, lam: float = 0.5) -> Callable[[np.ndarray
 # ============================================================ #
 # Search strategies                                              #
 # ============================================================ #
-def search_grid(objective: Callable, alpha_grid: List[float]):
-    """Exhaustive grid search over ``alpha_grid``. Returns (best_alpha, best_score, trace)."""
+def _log10_bounds(bounds: Tuple[float, float]) -> Tuple[float, float]:
+    """Map an alpha interval to log10 space, guarding against a non-positive low end."""
+    lo, hi = float(bounds[0]), float(bounds[1])
+    if hi <= 0:
+        raise ValueError(f"alpha_bounds upper limit must be positive, got {hi}")
+    lo = max(lo, hi * 1e-6)
+    return math.log10(lo), math.log10(hi)
+
+
+def search_grid(objective: Callable, alpha_grid: Optional[List[float]] = None,
+                bounds: Optional[Tuple[float, float]] = None, n: int = 15):
+    """Exhaustive grid search. Returns (best_alpha, best_score, trace).
+
+    With ``alpha_grid`` the caller's points are used verbatim. Otherwise a
+    log-spaced grid of ``n`` points is built from ``bounds`` — alpha is a scale
+    parameter, so equal spacing in log10 gives each order of magnitude the same
+    number of probes.
+    """
+    if alpha_grid is None:
+        if bounds is None:
+            raise ValueError("search_grid needs alpha_grid or bounds")
+        t_lo, t_hi = _log10_bounds(bounds)
+        alpha_grid = [float(10.0 ** t) for t in np.linspace(t_lo, t_hi, n)]
     trace = []
     best = None
     for a in alpha_grid:
@@ -453,31 +483,34 @@ def search_grid(objective: Callable, alpha_grid: List[float]):
     return best[0], best[1], trace
 
 
-def search_golden(objective: Callable, bounds=(0.1, 10.0), max_iter=12):
-    """Golden-section search for a unimodal objective over ``bounds``. Returns (best_alpha, best_score, trace)."""
+def search_golden(objective: Callable, bounds=DEFAULT_ALPHA_BOUNDS, max_iter=12):
+    """Golden-section search over log10(alpha) for a unimodal objective.
+
+    Returns (best_alpha, best_score, trace); the trace is in alpha, not log10.
+    """
     phi = (1 + math.sqrt(5)) / 2
-    a, b = bounds
+    a, b = _log10_bounds(bounds)
     resphi = 2 - phi
     x1 = a + resphi * (b - a)
     x2 = b - resphi * (b - a)
-    f1 = objective(x1)
-    f2 = objective(x2)
-    trace = [(x1, f1), (x2, f2)]
+    f1 = objective(10.0 ** x1)
+    f2 = objective(10.0 ** x2)
+    trace = [(10.0 ** x1, f1), (10.0 ** x2, f2)]
     for _ in range(max_iter):
         if f1 > f2:
             b = x2
             x2 = x1
             f2 = f1
             x1 = a + resphi * (b - a)
-            f1 = objective(x1)
-            trace.append((x1, f1))
+            f1 = objective(10.0 ** x1)
+            trace.append((10.0 ** x1, f1))
         else:
             a = x1
             x1 = x2
             f1 = f2
             x2 = b - resphi * (b - a)
-            f2 = objective(x2)
-            trace.append((x2, f2))
+            f2 = objective(10.0 ** x2)
+            trace.append((10.0 ** x2, f2))
     best = max(trace, key=lambda x: x[1])
     return best[0], best[1], trace
 
@@ -493,45 +526,53 @@ def _gp_ei(gp, X_grid, y_best, xi=0.01):
     return ei
 
 
-def search_bayesian(objective: Callable, bounds=(0.1, 10.0),
+def search_bayesian(objective: Callable, bounds=DEFAULT_ALPHA_BOUNDS,
                      n_init=5, n_iter=10, seed=42):
-    """GP-EI Bayesian optimisation over a 1-D ``bounds`` interval.
-    Seeds with ``n_init`` equally-spaced evaluations, then acquires via Expected Improvement.
-    Returns (best_alpha, best_score, trace)."""
+    """GP-EI Bayesian optimisation over log10(alpha).
+
+    Seeds with ``n_init`` log-spaced evaluations, then acquires via Expected
+    Improvement on a log-spaced grid. Returns (best_alpha, best_score, trace),
+    with the trace reported in alpha rather than log10.
+
+    The search runs in log space because alpha is a scale parameter: seeding
+    linearly over, say, [0.01, 1000] would place the first five probes at
+    0.01/250/500/750/1000 and never examine the sub-unit region where the
+    composition-dominated datasets have their optimum.
+    """
     rng = np.random.default_rng(seed)
-    lo, hi = bounds
-    X_init = np.linspace(lo, hi, n_init).reshape(-1, 1)
+    t_lo, t_hi = _log10_bounds(bounds)
     trace = []
-    for x in X_init:
-        s = objective(float(x[0]))
-        trace.append((float(x[0]), s))
-    X = np.array([[a] for a, _ in trace])
+    for t in np.linspace(t_lo, t_hi, n_init):
+        a = float(10.0 ** t)
+        trace.append((a, objective(a)))
+    X = np.array([[math.log10(a)] for a, _ in trace])
     y = np.array([s for _, s in trace])
     kernel = (ConstantKernel(1.0, (1e-3, 1e3))
               * Matern(length_scale=1.0, length_scale_bounds=(1e-2, 1e2), nu=2.5)
               + WhiteKernel(noise_level=1e-3, noise_level_bounds=(1e-5, 1e-1)))
-    grid = np.linspace(lo, hi, 500).reshape(-1, 1)
+    grid = np.linspace(t_lo, t_hi, 500).reshape(-1, 1)
     for _ in range(n_iter):
         try:
             gp = GaussianProcessRegressor(kernel=kernel, normalize_y=True,
                                             n_restarts_optimizer=2,
                                             random_state=seed).fit(X, y)
             ei = _gp_ei(gp, grid, y.max())
-            x_next = float(grid[int(np.argmax(ei))][0])
+            t_next = float(grid[int(np.argmax(ei))][0])
         except Exception:
-            x_next = float(rng.uniform(lo, hi))
-        if any(abs(x_next - x) < 1e-3 for x, _ in trace):
-            x_next = float(rng.uniform(lo, hi))
-        s = objective(x_next)
-        trace.append((x_next, s))
-        X = np.vstack([X, [[x_next]]])
+            t_next = float(rng.uniform(t_lo, t_hi))
+        if any(abs(t_next - math.log10(a)) < 1e-3 for a, _ in trace):
+            t_next = float(rng.uniform(t_lo, t_hi))
+        a_next = float(10.0 ** t_next)
+        s = objective(a_next)
+        trace.append((a_next, s))
+        X = np.vstack([X, [[t_next]]])
         y = np.append(y, s)
     best = max(trace, key=lambda x: x[1])
     return best[0], best[1], trace
 
 
 SEARCH_FUNCS = {
-    "grid":            lambda obj, b: search_grid(obj, [0.1, 0.5, 1.0, 2.0, 4.0, 6.0, 8.0]),
+    "grid":            lambda obj, b: search_grid(obj, bounds=b, n=15),
     "golden_section":  lambda obj, b: search_golden(obj, bounds=b, max_iter=12),
     "bayesian":        lambda obj, b: search_bayesian(obj, bounds=b, n_init=5, n_iter=10),
 }
@@ -546,7 +587,7 @@ def run_autotune(
     *,
     sample_col: str = "sample",
     celltype_col: str = "cell_type",
-    cluster_emb_key: str = "Z_clust",
+    comp_emb_key: Optional[str] = None,
     rmd_emb_key: Optional[str] = None,
     modality_col: Optional[str] = None,
     batch_col: Optional[Union[str, List[str]]] = None,
@@ -564,6 +605,7 @@ def run_autotune(
     save: bool = True,
     verbose: bool = True,
     tune_on_modality: Optional[str] = None,
+    cluster_emb_key: Optional[str] = None,
 ) -> Dict:
     """Run autotune and return the best params + final sample-AnnData.
 
@@ -574,6 +616,12 @@ def run_autotune(
     how does the other modality fare under that α?". Set to ``None`` (default)
     to score on every unit (current behaviour).
     """
+    if cluster_emb_key is not None:
+        warnings.warn("cluster_emb_key= is deprecated and will be removed in 1.0; "
+                      "use comp_emb_key=.", FutureWarning, stacklevel=2)
+        if comp_emb_key is None:
+            comp_emb_key = cluster_emb_key
+
     t0 = time.time()
     primary_batch = batch_col[0] if isinstance(batch_col, (list, tuple)) and batch_col else batch_col
     if isinstance(primary_batch, list):
@@ -581,7 +629,7 @@ def run_autotune(
 
     blocks = build_blocks(
         adata, sample_col=sample_col, celltype_col=celltype_col,
-        cluster_emb_key=cluster_emb_key, rmd_emb_key=rmd_emb_key,
+        comp_emb_key=comp_emb_key, rmd_emb_key=rmd_emb_key,
         modality_col=modality_col, batch_col=primary_batch,
         grouping_col=grouping_col, medium_K=medium_K, fine_K=fine_K,
         rmd_dim=rmd_dim, seed=seed, verbose=verbose,
@@ -755,7 +803,7 @@ def _format_autotune_report(*, best_params, best_score, trace, weights,
     lines.append(f"  K_c   (cell types) : {int(blocks['K_c'])}")
     lines.append(f"  K_med (k-means)    : {int(blocks['K_med'])}")
     lines.append(f"  K_fine (k-means)   : {int(blocks['K_fine'])}")
-    lines.append(f"  cluster_emb_key    : {blocks.get('cluster_emb_key', '?')}")
+    lines.append(f"  comp_emb_key       : {blocks.get('comp_emb_key', '?')}")
     lines.append(f"  rmd_emb_key        : {blocks.get('rmd_emb_key', '?')}")
     lines.append("")
     lines.append("Active scoring proxies (gated by data availability)")
@@ -782,6 +830,20 @@ def _format_autotune_report(*, best_params, best_score, trace, weights,
                   + "  (A1, A2, A3, RMD)")
     lines.append(f"  total evaluations : {len(trace)}")
     lines.append(f"  wall time         : {elapsed_s:.2f} s")
+    # A flat objective makes the single winning α arbitrary within the plateau;
+    # report the plateau so the number is not read as more precise than it is.
+    near = [a for a, s in trace
+            if best_score == best_score and s >= best_score - 0.01 * abs(best_score)]
+    if len(near) > 1:
+        lines.append(f"  α within 1% of best: [{min(near):g}, {max(near):g}]  "
+                     f"({len(near)}/{len(trace)} evaluations) — the objective is "
+                     f"flat over this range")
+    alpha_best = best_params.get("rmd_weight")
+    if alpha_best is not None and any(
+            abs(alpha_best - b) <= 1e-9 * max(1.0, abs(b)) for b in alpha_bounds):
+        lines.append(f"  NOTE: α equals a search bound; the optimum may lie outside "
+                     f"[{alpha_bounds[0]:g}, {alpha_bounds[1]:g}] — widen "
+                     f"alpha_bounds and re-run.")
     lines.append("")
     lines.append(f"Search trace ({len(trace)} evals, top 25 by score)")
     lines.append("-" * 68)
@@ -821,7 +883,9 @@ def _finalize(adata, blocks, final_emb, weights, *,
         "K_c": int(blocks["K_c"]),
         "K_med": int(blocks["K_med"]),
         "K_fine": int(blocks["K_fine"]),
-        "cluster_emb_key": str(blocks.get("cluster_emb_key", "")),
+        "comp_emb_key": str(blocks.get("comp_emb_key", "")),
+        # Deprecated 0.2.0 spelling; removed in 1.0.
+        "cluster_emb_key": str(blocks.get("comp_emb_key", "")),
         "rmd_emb_key": str(blocks.get("rmd_emb_key", "")),
         "n_evals": len(trace),
         "autotuned": True,

{sampledisco-0.2.0 → sampledisco-0.3.0}/src/sampledisco/preparation/ATAC_cell_type.py

@@ -8,6 +8,7 @@ import scanpy as sc
 import time
 import matplotlib.pyplot as plt
 from sampledisco.utils.safe_save import safe_h5ad_write  # Importing the new safe save method
+from sampledisco.utils.embedding_keys import resolve_comp_key
 
 def clean_obs_for_saving(adata, verbose=True):
     """
@@ -214,7 +215,7 @@ def cell_types_atac(
     Save=False,
     output_dir=None,
     cluster_resolution=0.8,
-    use_rep='Z_clust',
+    use_rep=None,
     peaks=None, 
     method='average', 
     metric='euclidean', 
@@ -234,7 +235,7 @@ def cell_types_atac(
     Assigns cell types based on existing annotations or performs Leiden clustering if no annotation exists.
     Uses recursive strategy to adaptively find optimal clustering resolution when target clusters specified.
     
-    ATAC VERSION: Uses dimension reduction (Z_clust) for dendrogram construction and differential peaks.
+    ATAC VERSION: Uses dimension reduction (Z_comp) for dendrogram construction and differential peaks.
 
     Parameters:
     - adata: AnnData object
@@ -245,7 +246,7 @@ def cell_types_atac(
     - Save: Boolean, whether to save the output
     - output_dir: Directory to save the output if Save=True
     - cluster_resolution: Starting resolution for Leiden clustering
-    - use_rep: Representation to use for neighborhood graph (default: 'Z_clust')
+    - use_rep: Representation to use for neighborhood graph (default: 'Z_comp')
     - peaks: List of peak names for mapping numeric IDs to names
     - method, metric, distance_mode: Parameters for hierarchical clustering
     - num_DMs: Number of diffusion map components for neighborhood graph
@@ -264,6 +265,8 @@ def cell_types_atac(
     start_time = time.time() if verbose else None
     from sampledisco.utils.random_seed import set_global_seed
     set_global_seed(seed = 42)
+    use_rep = resolve_comp_key(adata, use_rep, fallbacks=('X_lsi',),
+                               context="cell_types_atac")
     
     if _recursion_depth > 10:
         raise RuntimeError(f"Maximum recursion depth exceeded. Could not achieve {n_target_clusters} clusters.")

{sampledisco-0.2.0 → sampledisco-0.3.0}/src/sampledisco/preparation/ATAC_cell_type_gpu.py

@@ -8,6 +8,7 @@ import rapids_singlecell as rsc
 import scanpy as sc
 from sklearn.neighbors import KNeighborsTransformer
 import time
+from sampledisco.utils.embedding_keys import resolve_comp_key
 
 def standardize_cell_type_column_gpu(adata, verbose=True):
     """Coerce obs['cell_type'] to string categorical, handling Leiden integer output."""
@@ -56,7 +57,7 @@ def cell_types_atac_gpu(
     Save=False,
     output_dir=None,
     cluster_resolution=0.8,
-    use_rep='Z_clust',
+    use_rep=None,
     peaks=None, 
     method='average', 
     metric='euclidean', 
@@ -75,6 +76,8 @@ def cell_types_atac_gpu(
     _recursion_depth is internal; do not set manually.
     """
     start_time = time.time() if verbose else None
+    use_rep = resolve_comp_key(adata, use_rep, fallbacks=('X_lsi',),
+                               context="cell_types_atac_gpu")
 
     if _recursion_depth > 10:
         raise RuntimeError(f"Maximum recursion depth exceeded. Could not achieve {n_target_clusters} clusters.")
@@ -284,7 +287,7 @@ def cell_type_dendrogram_atac_gpu(
     method='average',
     metric='euclidean',
     distance_mode='centroid',
-    use_rep='Z_clust',
+    use_rep=None,
     num_DMs=20,
     verbose=True
 ):
@@ -297,8 +300,8 @@ def cell_type_dendrogram_atac_gpu(
         print(f'=== Preparing data for dendrogram (using {use_rep}) ===')
     if groupby not in adata.obs.columns:
         raise ValueError(f"The groupby key '{groupby}' is not present in adata.obs.")
-    if use_rep not in adata.obsm:
-        raise ValueError(f"The representation '{use_rep}' is not present in adata.obsm.")
+    use_rep = resolve_comp_key(adata, use_rep, fallbacks=('X_lsi',),
+                               context="cell_type_dendrogram_atac_gpu")
 
     # Convert from GPU array if needed.
     if hasattr(adata.obsm[use_rep], 'get'):

{sampledisco-0.2.0 → sampledisco-0.3.0}/src/sampledisco/preparation/atac_preprocess_cpu.py

@@ -8,6 +8,7 @@ import scanpy as sc
 import muon as mu
 from muon import atac as ac
 from sampledisco.utils.harmony_compat import harmonize_embedding
+from sampledisco.utils.embedding_keys import COMP_KEY, LEGACY_COMP_KEY, write_comp_key
 from scipy.sparse import issparse
 
 from sampledisco.utils.safe_save import safe_h5ad_write
@@ -20,7 +21,7 @@ def _run_lsi_on_hvg(adata, n_comps, drop_first, hvg_col="highly_variable", out_k
 
     muon's `ac.tl.lsi` doesn't accept `use_highly_variable`, so we run on a
     temporary HVF-only copy. varm['LSI'] and uns['lsi'] are only written for the
-    canonical out_key='X_lsi' to avoid clobbering the Z_clust basis.
+    canonical out_key='X_lsi' to avoid clobbering the Z_comp basis.
     """
     hvg_mask = adata.var[hvg_col].values
     if hvg_mask.sum() < n_comps + 1:
@@ -62,7 +63,7 @@ def anndata_cluster(
       - `.X`                             TF-IDF + log1p normalized
       - `.var['highly_variable']`        HVF flag (no subsetting)
       - `.obsm['X_lsi']`                 LSI on HVF subset (drop_first applied if set)
-      - `.obsm['Z_clust']`         sample-removed Harmony
+      - `.obsm['Z_comp']`         sample-removed Harmony
       - `.obsm['Z_rmd']`  sample-preserved Harmony (RMD)
 
     Writes a single file: `<output_dir>/adata_preprocessed.h5ad`.
@@ -102,18 +103,18 @@ def anndata_cluster(
         print("=== [CPU] Harmony pass 1: WITH sample (sample-removed) ===")
         print("  batch keys:", ", ".join(cell_level_batch_key_for_harmony or []))
     if cell_level_batch_key_for_harmony:
-        adata.obsm["Z_clust"] = harmonize_embedding(
+        write_comp_key(adata, harmonize_embedding(
             adata.obsm["X_lsi"], adata.obs,
             batch_key=cell_level_batch_key_for_harmony,
             max_iter_harmony=num_harmony_iterations,
             use_gpu=False,
-        )
+        ))
     else:
-        adata.obsm["Z_clust"] = adata.obsm["X_lsi"].copy()
+        write_comp_key(adata, adata.obsm["X_lsi"].copy())
 
     # HVF2 for RMD basis: sample-naive (batch_key=None). ATAC has no raw-counts
     # layer post TF-IDF, so HVF2 runs on the same TF-IDF+log1p .X as HVF1.
-    # Placed after Z_clust Harmony to preserve RNG ordering.
+    # Placed after Z_comp Harmony to preserve RNG ordering.
     adata.var["highly_variable_clust"] = adata.var["highly_variable"].to_numpy().copy()
     if verbose:
         print("Running HVF2 (sample-naive) selection...")
@@ -146,7 +147,9 @@ def anndata_cluster(
             adata.obsm["X_lsi_rmd"], dtype=np.float32)
 
     if verbose:
-        print(f"  Z_clust   shape: {adata.obsm['Z_clust'].shape}")
+        print(f"  Z_comp    shape: {adata.obsm[COMP_KEY].shape}")
+        if LEGACY_COMP_KEY in adata.obsm:
+            print(f"  (legacy alias obsm['{LEGACY_COMP_KEY}'] also written; removed in 1.0)")
         print(f"  X_lsi_rmd shape: {adata.obsm['X_lsi_rmd'].shape}")
         print(f"  Z_rmd     shape: {adata.obsm['Z_rmd'].shape}")