feat: Add condition_key support, root API exports, and plotting improvements

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",
 ]

scrise/factorization.py

@@ -214,16 +214,18 @@ 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,
     max_iter: int = 100,
     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,

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]
 

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)

scrise/rank_selection.py

@@ -157,15 +157,17 @@ 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,
     random_state: int | None = None,
     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,13 +213,28 @@ 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."
         )
     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})

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__

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)

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)

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)