From 0622e54f9c53d6ef0c4b25efbf1d053323c3068d Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Fri, 4 Sep 2026 11:18:53 -0700 Subject: [PATCH] feat: Add condition_key support, root API exports, and plotting improvements --- scrise/__init__.py | 10 +++++ scrise/factorization.py | 23 ++++++++++-- scrise/plotting/factors.py | 22 +++++++++-- scrise/plotting/stability.py | 8 ++-- scrise/rank_selection.py | 21 ++++++++++- scrise/tests/test_parafac2.py | 49 ++++++++++++++++++++++++- scrise/tests/test_plotting_factors.py | 29 +++++++++++++++ scrise/tests/test_plotting_stability.py | 17 ++++++++- scrise/tests/test_rank_selection.py | 32 ++++++++++++++++ 9 files changed, 196 insertions(+), 15 deletions(-) diff --git a/scrise/__init__.py b/scrise/__init__.py index 3eb18211..a4231dfc 100644 --- a/scrise/__init__.py +++ b/scrise/__init__.py @@ -1,3 +1,6 @@ +from parafac2.normalize import prepare_dataset + +from . import plotting from .annotation_alignment import ( CellTypeAlignmentResults, ComponentAlignmentResult, @@ -16,11 +19,16 @@ rise_pca_r2x, ) from .opq import OPQQuantizer, find_optimal_opq +from .rank_selection import bicv + +__version__ = "1.2.0" __all__ = [ "CellTypeAlignmentResults", "ComponentAlignmentResult", "OPQQuantizer", + "__version__", + "bicv", "canonical_component_signs", "cell_type_alignment", "compute_tau", @@ -31,6 +39,8 @@ "match_components_across_ranks", "order_components_by_energy", "pf2", + "plotting", + "prepare_dataset", "rise_pca_r2x", "score_cell_type_alignment", ] diff --git a/scrise/factorization.py b/scrise/factorization.py index d6983f78..6f7d2fb1 100644 --- a/scrise/factorization.py +++ b/scrise/factorization.py @@ -214,8 +214,8 @@ def match_components_across_ranks( def pf2( - X: anndata.AnnData, - rank: int, + X: anndata.AnnData | None = None, + rank: int | None = None, random_state=1, doEmbedding: bool = True, tolerance=1e-6, @@ -223,7 +223,9 @@ def pf2( normalize_slices: bool = False, backend: str | None = None, compress: int | tuple[int, int | None] | str | bool | None = None, -): + condition_key: str | None = None, + adata: anndata.AnnData | None = None, +) -> anndata.AnnData: """Perform PARAFAC2 tensor decomposition on single-cell RNA-seq data. This is the main function for running RISE analysis. It decomposes the @@ -284,6 +286,21 @@ def pf2( components -- typically the ones added when the rank is increased -- land at the high end of the ordering. """ + if X is None and adata is not None: + X = adata + if X is None: + raise ValueError("Either X or adata must be provided.") + if rank is None: + raise ValueError("rank must be provided.") + + if "condition_unique_idxs" not in X.obs: + if condition_key is not None and condition_key in X.obs: + X.obs["condition_unique_idxs"] = pd.Categorical(X.obs[condition_key]).codes + else: + raise KeyError( + "X.obs must contain 'condition_unique_idxs', or provide 'condition_key' pointing to a valid column in X.obs." + ) + pf_out, _ = parafac2_nd( X, rank=rank, diff --git a/scrise/plotting/factors.py b/scrise/plotting/factors.py index 4a61104b..4e359f69 100644 --- a/scrise/plotting/factors.py +++ b/scrise/plotting/factors.py @@ -1,3 +1,5 @@ +from collections.abc import Sequence + import anndata import numpy as np import pandas as pd @@ -15,9 +17,11 @@ def plot_condition_factors( cond: str = "Condition", log_transform: bool = True, cond_group_labels: pd.Series | None = None, - ThomsonNorm=False, + ThomsonNorm: bool = False, color_key=None, - group_cond=False, + group_cond: bool = False, + control_pattern: str | None = None, + control_conditions: Sequence[str] | None = None, ): """Plot condition factors as a heatmap showing how conditions contribute to components. @@ -57,8 +61,14 @@ def plot_condition_factors( if log_transform is True: X = np.log10(X) - if ThomsonNorm is True: - controls = yt.str.contains("CTRL") + if ThomsonNorm is True and control_pattern is None: + control_pattern = "CTRL" + + if control_conditions is not None: + controls = yt.isin(control_conditions) + XX = X[controls] + elif control_pattern is not None: + controls = yt.str.contains(control_pattern) XX = X[controls] else: XX = X @@ -184,6 +194,10 @@ def plot_gene_factors(data: anndata.AnnData, ax: Axes, weight=0.08, trim=True): if trim is True: max_weight = np.max(np.abs(X), axis=1) kept_idxs = max_weight > weight + if not np.any(kept_idxs): + raise ValueError( + f"No genes exceeded the weight threshold {weight}. Lower the threshold or set trim=False." + ) X = X[kept_idxs] yt = yt[kept_idxs] diff --git a/scrise/plotting/stability.py b/scrise/plotting/stability.py index d4c8e0ef..2a33c3a0 100644 --- a/scrise/plotting/stability.py +++ b/scrise/plotting/stability.py @@ -5,7 +5,6 @@ import anndata import numpy as np import pandas as pd -import scanpy as sc import seaborn as sns from matplotlib.axes import Axes from tensorly.cp_tensor import CPTensor @@ -80,9 +79,10 @@ def plot_fms_percent_drop( for j in range(runs): scores = [1.0] for i in percentList[1:]: - sampled_data: anndata.AnnData = sc.pp.subsample( - X, fraction=1 - (i / 100), random_state=j, copy=True - ) # type: ignore + n_cells = int(X.n_obs * (1 - (i / 100))) + rng = np.random.default_rng(j) + sampled_idx = rng.choice(X.n_obs, size=n_cells, replace=False) + sampled_data = X[sampled_idx].copy() sampledX = pf2(sampled_data, rank, random_state=j + 2, doEmbedding=False) fmsScore = calculateFMS(dataX, sampledX) diff --git a/scrise/rank_selection.py b/scrise/rank_selection.py index af1a07e3..1bc36ba7 100644 --- a/scrise/rank_selection.py +++ b/scrise/rank_selection.py @@ -157,8 +157,8 @@ def _bicv_trial( def bicv( - X: anndata.AnnData, - ranks: Sequence[int], + X: anndata.AnnData | None = None, + ranks: Sequence[int] | None = None, n_repeats: int = 3, held_out_cell_frac: float = 0.2, held_out_gene_frac: float = 0.2, @@ -166,6 +166,8 @@ def bicv( tolerance: float = 1e-6, max_iter: int = 200, compress: int | tuple[int, int | None] | str | bool | None = "auto", + condition_key: str | None = None, + adata: anndata.AnnData | None = None, ) -> pd.DataFrame: """Evaluate rank via bi-cross-validation (BiCV) and in-sample fit R2X. @@ -211,6 +213,13 @@ def bicv( (one of "Fit R2X" or "BiCV R2X"), and "R2X". Ready to pass to :func:`RISE.plotting.plot_bicv_r2x`. """ + if X is None and adata is not None: + X = adata + if X is None: + raise ValueError("Either X or adata must be provided.") + if ranks is None: + raise ValueError("ranks must be provided.") + if not (0 < held_out_cell_frac < 1) or not (0 < held_out_gene_frac < 1): raise ValueError( "held_out_cell_frac and held_out_gene_frac must both be between 0 and 1." @@ -218,6 +227,14 @@ def bicv( if n_repeats < 1: raise ValueError("n_repeats must be at least 1.") + if "condition_unique_idxs" not in X.obs: + if condition_key is not None and condition_key in X.obs: + X.obs["condition_unique_idxs"] = pd.Categorical(X.obs[condition_key]).codes + else: + raise KeyError( + "X.obs must contain 'condition_unique_idxs', or provide 'condition_key' pointing to a valid column in X.obs." + ) + X = X.to_memory() if hasattr(X, "to_memory") else X ranks = sorted({int(r) for r in ranks}) diff --git a/scrise/tests/test_parafac2.py b/scrise/tests/test_parafac2.py index 2272e3e6..20070b67 100644 --- a/scrise/tests/test_parafac2.py +++ b/scrise/tests/test_parafac2.py @@ -37,4 +37,51 @@ def test_factor_thomson_mlx_backend(): res = pf2(X, 5, doEmbedding=False, tolerance=1e-6, backend="mlx") - assert np.all(np.isfinite(res.varm["Pf2_C"])) + assert np.all(np.isfinite(np.asarray(res.varm["Pf2_C"]))) + + +def test_pf2_condition_key(): + import anndata + + from .conftest import make_synthetic_pf2_data + + orig = make_synthetic_pf2_data(n_cond=3, n_genes=15, rank=2, seed=42) + obs = pd.DataFrame({"custom_condition": orig.obs["Condition"].to_numpy()}) + X = anndata.AnnData(X=orig.X, obs=obs, var=pd.DataFrame(index=orig.var_names)) + + with pytest.raises(KeyError, match="condition_unique_idxs"): + pf2(X, 2, doEmbedding=False, max_iter=5, compress=None) + + res = pf2( + X, + 2, + condition_key="custom_condition", + doEmbedding=False, + max_iter=5, + compress=None, + ) + assert "condition_unique_idxs" in res.obs + assert "Pf2_A" in res.uns + + +def test_pf2_adata_alias(): + from .conftest import make_synthetic_pf2_data + + X = make_synthetic_pf2_data(n_cond=3, n_genes=15, rank=2, seed=42) + res = pf2(adata=X, rank=2, doEmbedding=False, max_iter=5, compress=None) + assert "Pf2_A" in res.uns + + +def test_root_api_exports(): + import scrise + from scrise import __version__, bicv, plotting, prepare_dataset + + assert callable(bicv) + assert callable(prepare_dataset) + assert hasattr(plotting, "plot_condition_factors") + assert hasattr(plotting, "plot_gene_factors") + assert __version__ == "1.2.0" + assert "bicv" in scrise.__all__ + assert "prepare_dataset" in scrise.__all__ + assert "plotting" in scrise.__all__ + assert "__version__" in scrise.__all__ diff --git a/scrise/tests/test_plotting_factors.py b/scrise/tests/test_plotting_factors.py index 3d20d4f5..7f48ed89 100644 --- a/scrise/tests/test_plotting_factors.py +++ b/scrise/tests/test_plotting_factors.py @@ -112,3 +112,32 @@ def test_plot_bicv_r2x_smoke_and_axis_labels(): assert ax.get_xlabel() == "Rank" assert ax.get_ylabel() == "R2X" assert ax.get_legend() is not None + + +def test_plot_condition_factors_control_pattern_and_conditions(): + adata = make_mock_factored_adata(n_conditions=6, rank=3) + adata.uns["Pf2_A"] = np.abs(adata.uns["Pf2_A"]) + 0.1 + # Assign conditions with CTRL and TRT + cond_codes = adata.obs["condition_unique_idxs"].cat.codes + cond_names = [f"CTRL_{c % 2}" if c < 2 else f"TRT_{c}" for c in cond_codes] + adata.obs["Condition"] = pd.Categorical(cond_names) + + fig1, ax1 = plt.subplots() + plot_condition_factors(adata, ax1, control_pattern="CTRL") + assert len(ax1.get_yticklabels()) > 0 + + fig2, ax2 = plt.subplots() + plot_condition_factors(adata, ax2, control_conditions=["CTRL_0", "CTRL_1"]) + assert len(ax2.get_yticklabels()) > 0 + + fig3, ax3 = plt.subplots() + plot_condition_factors(adata, ax3, ThomsonNorm=True) + assert len(ax3.get_yticklabels()) > 0 + + +def test_plot_gene_factors_no_genes_pass_weight_raises(): + adata = make_mock_factored_adata(n_genes=10, rank=2) + adata.varm["Pf2_C"] = np.full_like(adata.varm["Pf2_C"], 0.01) + fig, ax = plt.subplots() + with pytest.raises(ValueError, match="No genes exceeded the weight threshold"): + plot_gene_factors(adata, ax, weight=0.5, trim=True) diff --git a/scrise/tests/test_plotting_stability.py b/scrise/tests/test_plotting_stability.py index 4a3c14f3..ed376cb4 100644 --- a/scrise/tests/test_plotting_stability.py +++ b/scrise/tests/test_plotting_stability.py @@ -11,7 +11,12 @@ import matplotlib.pyplot as plt import numpy as np -from ..plotting.stability import calculateFMS, plot_fms_diff_ranks, resample +from ..plotting.stability import ( + calculateFMS, + plot_fms_diff_ranks, + plot_fms_percent_drop, + resample, +) from .conftest import make_synthetic_pf2_data @@ -97,3 +102,13 @@ def test_plot_fms_diff_ranks_smoke(): assert ax.get_xlabel() == "Component" assert ax.get_ylabel() == "FMS" assert ax.get_ylim() == (0.0, 1.0) + + +def test_plot_fms_percent_drop_smoke(): + X = make_synthetic_pf2_data( + n_cond=3, n_genes=10, rank=2, seed=0, cells_per_cond=(20, 30) + ) + fig, ax = plt.subplots() + plot_fms_percent_drop(X, ax, percentList=np.array([0, 10]), runs=1, rank=2) + assert ax.get_ylabel() == "FMS" + assert ax.get_ylim() == (0.0, 1.0) diff --git a/scrise/tests/test_rank_selection.py b/scrise/tests/test_rank_selection.py index 75089f52..9682ba47 100644 --- a/scrise/tests/test_rank_selection.py +++ b/scrise/tests/test_rank_selection.py @@ -86,3 +86,35 @@ def test_bicv_warns_when_best_rank_is_at_boundary(): with pytest.warns(UserWarning, match="edge of the tested ranks"): bicv(X, [8], n_repeats=1, random_state=0, max_iter=50) + + +def test_bicv_condition_key(): + orig = _make_test_data() + obs = pd.DataFrame( + {"custom_condition": orig.obs["condition_unique_idxs"].astype(str).to_numpy()} + ) + X = anndata.AnnData( + X=orig.X, + obs=obs, + var=pd.DataFrame({"means": np.zeros(orig.n_vars)}, index=orig.var_names), + ) + + with pytest.raises(KeyError, match="condition_unique_idxs"): + bicv(X, [2], n_repeats=1, random_state=0, max_iter=10) + + results = bicv( + X, + [2], + condition_key="custom_condition", + n_repeats=1, + random_state=0, + max_iter=10, + ) + assert isinstance(results, pd.DataFrame) + assert "condition_unique_idxs" in X.obs + + +def test_bicv_adata_alias(): + X = _make_test_data() + results = bicv(adata=X, ranks=[2], n_repeats=1, random_state=0, max_iter=10) + assert isinstance(results, pd.DataFrame)