diff --git a/docs/api.md b/docs/api.md index a398a4cb..87ff4ed1 100644 --- a/docs/api.md +++ b/docs/api.md @@ -6,13 +6,19 @@ options: members: - pf2 - - export_factors - - load_factors - correct_conditions - order_components_by_energy - canonical_component_signs - match_components_across_ranks +## Factor Import/Export + +::: scrise.factor_io + options: + members: + - export_factors + - load_factors + ## Rank Selection ::: scrise.rank_selection @@ -24,6 +30,16 @@ ::: scrise.annotation_alignment +### Alignment Statistics + +::: scrise.alignment_stats + options: + members: + - compute_auroc_per_cell_type + - compute_tau + - compute_eta_squared + - compute_kruskal_epsilon_squared + ## Quantization & Compression ::: scrise.opq diff --git a/scrise/__init__.py b/scrise/__init__.py index a4231dfc..07a4b6b6 100644 --- a/scrise/__init__.py +++ b/scrise/__init__.py @@ -1,18 +1,17 @@ from parafac2.normalize import prepare_dataset from . import plotting +from .alignment_stats import compute_tau from .annotation_alignment import ( CellTypeAlignmentResults, ComponentAlignmentResult, cell_type_alignment, - compute_tau, score_cell_type_alignment, ) +from .factor_io import export_factors, load_factors from .factorization import ( canonical_component_signs, correct_conditions, - export_factors, - load_factors, match_components_across_ranks, order_components_by_energy, pf2, diff --git a/scrise/alignment_stats.py b/scrise/alignment_stats.py new file mode 100644 index 00000000..a3ba4d72 --- /dev/null +++ b/scrise/alignment_stats.py @@ -0,0 +1,283 @@ +from __future__ import annotations + +import numpy as np +import pandas as pd +import scipy.stats as sp + + +def compute_auroc_per_cell_type( + loadings: np.ndarray, + cell_type_codes: np.ndarray, + n_types: int, +) -> np.ndarray: + """Compute AUROC for each cell type vs all other cells. + + Parameters + ---------- + loadings : np.ndarray + 1D array of cell loadings of shape (n_cells,). + cell_type_codes : np.ndarray + 1D array of integer cell type assignments in [0, n_types - 1]. + n_types : int + Total number of unique cell types. + + Returns + ------- + np.ndarray + 1D array of AUROC values for each cell type of shape (n_types,). + """ + n_cells = loadings.size + if n_cells == 0 or n_types <= 1: + return np.full(n_types, 0.5, dtype=float) + + ranks = sp.rankdata(loadings, method="average") + counts = np.bincount(cell_type_codes, minlength=n_types).astype(float) + rank_sums = np.bincount(cell_type_codes, weights=ranks, minlength=n_types).astype( + float + ) + + aurocs = np.full(n_types, 0.5, dtype=float) + for k in range(n_types): + n1 = counts[k] + n0 = n_cells - n1 + if n1 > 0 and n0 > 0: + u1 = rank_sums[k] - n1 * (n1 + 1.0) / 2.0 + aurocs[k] = u1 / (n1 * n0) + + return aurocs + + +def compute_tau( + enrichment_scores: np.ndarray | pd.Series, + baseline: float = 0.0, +) -> float: + """Compute tissue specificity index tau (Yanai et al., 2005). + + tau = sum(1 - x_hat) / (n_types - 1), where x_hat = x / max(x). + + When computed over AUROC enrichment values: + - tau -> 1 indicates the component is specific/private to a single cell type. + - tau -> 0 indicates the component is evenly distributed across cell types. + + Parameters + ---------- + enrichment_scores : np.ndarray | pd.Series + 1D array of enrichment values (e.g. AUROC) across cell types. + baseline : float, optional (default: 0.0) + Baseline value subtracted before calculating tau. Values below baseline are clipped to 0. + + Returns + ------- + float + Tau index in [0.0, 1.0]. + """ + x = np.asarray(enrichment_scores, dtype=float) + n = x.size + if n <= 1: + return 0.0 + + if baseline > 0.0: + x = np.maximum(0.0, x - baseline) + + max_val = np.nanmax(x) + if max_val <= 0.0 or not np.isfinite(max_val): + return 0.0 + + x_hat = x / max_val + tau = np.sum(1.0 - x_hat) / (n - 1.0) + return float(np.clip(tau, 0.0, 1.0)) + + +def compute_eta_squared( + loadings: np.ndarray, + cell_type_codes: np.ndarray, + n_types: int, +) -> float: + """Compute omnibus eta-squared (loading ~ cell_type). + + Proportion of loading variance explained by cell-type identity. + + Parameters + ---------- + loadings : np.ndarray + 1D array of cell loadings (shape: n_cells,). + cell_type_codes : np.ndarray + 1D array of integer cell type assignments (shape: n_cells,). + n_types : int + Number of unique cell types. + + Returns + ------- + float + Eta-squared value in [0.0, 1.0]. + """ + y = np.asarray(loadings, dtype=float) + n_cells = y.size + if n_cells <= 1 or n_types <= 1: + return 0.0 + + y_mean = np.mean(y) + ss_total = np.sum((y - y_mean) ** 2) + if ss_total <= 0.0: + return 0.0 + + counts = np.bincount(cell_type_codes, minlength=n_types).astype(float) + valid = counts > 0 + sums = np.bincount(cell_type_codes, weights=y, minlength=n_types) + + means = np.zeros(n_types, dtype=float) + means[valid] = sums[valid] / counts[valid] + + ss_between = np.sum(counts[valid] * (means[valid] - y_mean) ** 2) + eta2 = ss_between / ss_total + return float(np.clip(eta2, 0.0, 1.0)) + + +def compute_kruskal_epsilon_squared( + loadings: np.ndarray, + cell_type_codes: np.ndarray, + n_types: int, +) -> float: + """Compute Kruskal-Wallis epsilon-squared effect size. + + Non-parametric measure of association between loading and cell type. + + Parameters + ---------- + loadings : np.ndarray + 1D array of cell loadings (shape: n_cells,). + cell_type_codes : np.ndarray + 1D array of integer cell type assignments (shape: n_cells,). + n_types : int + Number of unique cell types. + + Returns + ------- + float + Epsilon-squared value in [0.0, 1.0]. + """ + y = np.asarray(loadings, dtype=float) + n_cells = y.size + if n_cells <= 1 or n_types <= 1: + return 0.0 + + groups = [ + y[cell_type_codes == k] + for k in range(n_types) + if np.sum(cell_type_codes == k) > 0 + ] + if len(groups) <= 1: + return 0.0 + + try: + stat, _ = sp.kruskal(*groups) + if not np.isfinite(stat) or stat < 0: + return 0.0 + eps2 = stat / (n_cells - 1.0) + return float(np.clip(eps2, 0.0, 1.0)) + except (ValueError, ZeroDivisionError): + return 0.0 + + +def _validate_and_encode_cell_types( + cell_types: pd.Series | np.ndarray | list, +) -> tuple[np.ndarray, list[str]]: + """Encode cell types to integer codes and return category names.""" + if isinstance(cell_types, pd.Series) and isinstance( + cell_types.dtype, pd.CategoricalDtype + ): + categories = list(cell_types.cat.categories) + codes = cell_types.cat.codes.to_numpy() + # If there are unused categories or NaN, retain only observed categories + observed = np.unique(codes[codes >= 0]) + if len(observed) < len(categories): + # Remap to dense 0..K-1 + remap = {old: new for new, old in enumerate(observed)} + categories = [categories[old] for old in observed] + codes = np.array([remap.get(c, -1) for c in codes], dtype=int) + return codes, categories + + s = pd.Series(cell_types) + cat = pd.Categorical(s) + categories = [str(c) for c in cat.categories] + codes = cat.codes + return codes, categories + + +def _permutation_p_values( + y: np.ndarray, + codes: np.ndarray, + n_types: int, + n_cells: int, + aurocs: np.ndarray, + n_permutations: int, + rng: np.random.Generator, +) -> np.ndarray: + """Empirical one-sided p-values against a rank-permutation null.""" + ranks = sp.rankdata(y, method="average") + counts = np.bincount(codes, minlength=n_types).astype(float) + null_counts = np.zeros(n_types, dtype=int) + + for _ in range(n_permutations): + perm_ranks = rng.permutation(ranks) + perm_sums = np.bincount(codes, weights=perm_ranks, minlength=n_types).astype( + float + ) + for k in range(n_types): + n1 = counts[k] + n0 = n_cells - n1 + if n1 > 0 and n0 > 0: + u_null = perm_sums[k] - n1 * (n1 + 1.0) / 2.0 + auc_null = u_null / (n1 * n0) + if auc_null >= aurocs[k]: + null_counts[k] += 1 + + return (1.0 + null_counts) / (1.0 + n_permutations) + + +def _asymptotic_p_values(y: np.ndarray, codes: np.ndarray, n_types: int) -> np.ndarray: + """One-sided Mann-Whitney U p-values, used when ``n_permutations == 0``.""" + p_values = np.ones(n_types, dtype=float) + for k in range(n_types): + pos = y[codes == k] + neg = y[codes != k] + if pos.size > 0 and neg.size > 0: + res = sp.mannwhitneyu(pos, neg, alternative="greater") + p_values[k] = float(res.pvalue) + return p_values + + +def _component_p_values( + y: np.ndarray, + codes: np.ndarray, + n_types: int, + n_cells: int, + aurocs: np.ndarray, + n_permutations: int, + rng: np.random.Generator, +) -> np.ndarray: + """Per-cell-type p-values for one component's loadings. + + Shared by :func:`cell_type_alignment` and + :func:`score_cell_type_alignment`, which compute this identically. A + single cell type (or no cells) leaves every p-value at 1.0, since there + is nothing to be enriched against. + """ + if n_types <= 1 or n_cells == 0: + return np.ones(n_types, dtype=float) + if n_permutations > 0: + return _permutation_p_values( + y, codes, n_types, n_cells, aurocs, n_permutations, rng + ) + if n_permutations == 0: + return _asymptotic_p_values(y, codes, n_types) + return np.ones(n_types, dtype=float) + + +def _as_generator( + random_state: int | np.random.Generator | None, +) -> np.random.Generator: + """Accept a seed, an existing Generator, or None.""" + if isinstance(random_state, np.random.Generator): + return random_state + return np.random.default_rng(random_state) diff --git a/scrise/annotation_alignment.py b/scrise/annotation_alignment.py index 892af7f8..a54f4cee 100644 --- a/scrise/annotation_alignment.py +++ b/scrise/annotation_alignment.py @@ -17,6 +17,16 @@ import pandas as pd import scipy.stats as sp +from .alignment_stats import ( + _as_generator, + _component_p_values, + _validate_and_encode_cell_types, + compute_auroc_per_cell_type, + compute_eta_squared, + compute_kruskal_epsilon_squared, + compute_tau, +) + @dataclass class ComponentAlignmentResult: @@ -130,203 +140,62 @@ def summary(self) -> pd.DataFrame: return df -def compute_auroc_per_cell_type( - loadings: np.ndarray, - cell_type_codes: np.ndarray, - n_types: int, -) -> np.ndarray: - """Compute AUROC for each cell type vs all other cells. - - Parameters - ---------- - loadings : np.ndarray - 1D array of cell loadings of shape (n_cells,). - cell_type_codes : np.ndarray - 1D array of integer cell type assignments in [0, n_types - 1]. - n_types : int - Total number of unique cell types. - - Returns - ------- - np.ndarray - 1D array of AUROC values for each cell type of shape (n_types,). - """ - n_cells = loadings.size - if n_cells == 0 or n_types <= 1: - return np.full(n_types, 0.5, dtype=float) - - ranks = sp.rankdata(loadings, method="average") - counts = np.bincount(cell_type_codes, minlength=n_types).astype(float) - rank_sums = np.bincount(cell_type_codes, weights=ranks, minlength=n_types).astype( - float - ) - - aurocs = np.full(n_types, 0.5, dtype=float) - for k in range(n_types): - n1 = counts[k] - n0 = n_cells - n1 - if n1 > 0 and n0 > 0: - u1 = rank_sums[k] - n1 * (n1 + 1.0) / 2.0 - aurocs[k] = u1 / (n1 * n0) - - return aurocs - - -def compute_tau( - enrichment_scores: np.ndarray | pd.Series, - baseline: float = 0.0, -) -> float: - """Compute tissue specificity index tau (Yanai et al., 2005). - - tau = sum(1 - x_hat) / (n_types - 1), where x_hat = x / max(x). - - When computed over AUROC enrichment values: - - tau -> 1 indicates the component is specific/private to a single cell type. - - tau -> 0 indicates the component is evenly distributed across cell types. - - Parameters - ---------- - enrichment_scores : np.ndarray | pd.Series - 1D array of enrichment values (e.g. AUROC) across cell types. - baseline : float, optional (default: 0.0) - Baseline value subtracted before calculating tau. Values below baseline are clipped to 0. - - Returns - ------- - float - Tau index in [0.0, 1.0]. - """ - x = np.asarray(enrichment_scores, dtype=float) - n = x.size - if n <= 1: - return 0.0 - - if baseline > 0.0: - x = np.maximum(0.0, x - baseline) - - max_val = np.nanmax(x) - if max_val <= 0.0 or not np.isfinite(max_val): - return 0.0 - - x_hat = x / max_val - tau = np.sum(1.0 - x_hat) / (n - 1.0) - return float(np.clip(tau, 0.0, 1.0)) - - -def compute_eta_squared( - loadings: np.ndarray, - cell_type_codes: np.ndarray, - n_types: int, -) -> float: - """Compute omnibus eta-squared (loading ~ cell_type). - - Proportion of loading variance explained by cell-type identity. - - Parameters - ---------- - loadings : np.ndarray - 1D array of cell loadings (shape: n_cells,). - cell_type_codes : np.ndarray - 1D array of integer cell type assignments (shape: n_cells,). - n_types : int - Number of unique cell types. - - Returns - ------- - float - Eta-squared value in [0.0, 1.0]. - """ - y = np.asarray(loadings, dtype=float) - n_cells = y.size - if n_cells <= 1 or n_types <= 1: - return 0.0 - - y_mean = np.mean(y) - ss_total = np.sum((y - y_mean) ** 2) - if ss_total <= 0.0: - return 0.0 - - counts = np.bincount(cell_type_codes, minlength=n_types).astype(float) - valid = counts > 0 - sums = np.bincount(cell_type_codes, weights=y, minlength=n_types) - - means = np.zeros(n_types, dtype=float) - means[valid] = sums[valid] / counts[valid] - - ss_between = np.sum(counts[valid] * (means[valid] - y_mean) ** 2) - eta2 = ss_between / ss_total - return float(np.clip(eta2, 0.0, 1.0)) - - -def compute_kruskal_epsilon_squared( - loadings: np.ndarray, - cell_type_codes: np.ndarray, - n_types: int, -) -> float: - """Compute Kruskal-Wallis epsilon-squared effect size. - - Non-parametric measure of association between loading and cell type. +_CELL_TYPE_COLUMN_CANDIDATES = ( + "cell_type", + "CellType", + "cell_types", + "Cell_Type", + "celltype", +) + + +def _loadings_from_anndata(data: anndata.AnnData, projection_key: str) -> np.ndarray: + """Pull the cell-loading matrix out of ``obsm``, falling back to projections.""" + if projection_key in data.obsm: + return np.asarray(data.obsm[projection_key]) + if projection_key == "weighted_projections" and "projections" in data.obsm: + return np.asarray(data.obsm["projections"]) + raise KeyError(f"Could not find '{projection_key}' in data.obsm.") + + +def _cell_types_from_anndata( + data: anndata.AnnData, cell_types: pd.Series | np.ndarray | str | None +) -> pd.Series: + """Resolve cell-type labels for an AnnData: named column, guess, or literal.""" + if cell_types is None: + for candidate in _CELL_TYPE_COLUMN_CANDIDATES: + if candidate in data.obs: + return pd.Series(data.obs[candidate]) + raise KeyError( + "Cell-type column not specified and none of ['cell_type', 'CellType', 'cell_types'] found in data.obs." + ) + if isinstance(cell_types, str): + if cell_types not in data.obs: + raise KeyError(f"Column '{cell_types}' not found in data.obs.") + return pd.Series(data.obs[cell_types]) + return pd.Series(cell_types) - Parameters - ---------- - loadings : np.ndarray - 1D array of cell loadings (shape: n_cells,). - cell_type_codes : np.ndarray - 1D array of integer cell type assignments (shape: n_cells,). - n_types : int - Number of unique cell types. - Returns - ------- - float - Epsilon-squared value in [0.0, 1.0]. - """ - y = np.asarray(loadings, dtype=float) - n_cells = y.size - if n_cells <= 1 or n_types <= 1: - return 0.0 +def _resolve_loadings_and_cell_types( + data: anndata.AnnData | np.ndarray | pd.DataFrame, + cell_types: pd.Series | np.ndarray | str | None, + projection_key: str, +) -> tuple[np.ndarray, pd.Series]: + """Normalise the three accepted input shapes to (matrix, label series).""" + if isinstance(data, anndata.AnnData): + return ( + _loadings_from_anndata(data, projection_key), + _cell_types_from_anndata(data, cell_types), + ) - groups = [ - y[cell_type_codes == k] - for k in range(n_types) - if np.sum(cell_type_codes == k) > 0 - ] - if len(groups) <= 1: - return 0.0 - - try: - stat, _ = sp.kruskal(*groups) - if not np.isfinite(stat) or stat < 0: - return 0.0 - eps2 = stat / (n_cells - 1.0) - return float(np.clip(eps2, 0.0, 1.0)) - except (ValueError, ZeroDivisionError): - return 0.0 - - -def _validate_and_encode_cell_types( - cell_types: pd.Series | np.ndarray | list, -) -> tuple[np.ndarray, list[str]]: - """Encode cell types to integer codes and return category names.""" - if isinstance(cell_types, pd.Series) and isinstance( - cell_types.dtype, pd.CategoricalDtype - ): - categories = list(cell_types.cat.categories) - codes = cell_types.cat.codes.to_numpy() - # If there are unused categories or NaN, retain only observed categories - observed = np.unique(codes[codes >= 0]) - if len(observed) < len(categories): - # Remap to dense 0..K-1 - remap = {old: new for new, old in enumerate(observed)} - categories = [categories[old] for old in observed] - codes = np.array([remap.get(c, -1) for c in codes], dtype=int) - return codes, categories - - s = pd.Series(cell_types) - cat = pd.Categorical(s) - categories = [str(c) for c in cat.categories] - codes = cat.codes - return codes, categories + loadings_matrix = ( + data.to_numpy() if isinstance(data, pd.DataFrame) else np.asarray(data) + ) + if cell_types is None or isinstance(cell_types, str): + raise ValueError( + "cell_types must be provided when data is a DataFrame or array." + ) + return loadings_matrix, pd.Series(cell_types) def cell_type_alignment( @@ -381,41 +250,15 @@ def cell_type_alignment( # Compute observed AUROC per cell type aurocs = compute_auroc_per_cell_type(y, codes, n_types) - # Compute permutation p-values - p_values = np.ones(n_types, dtype=float) - if n_permutations > 0 and n_types > 1 and n_cells > 0: - rng = ( - random_state - if isinstance(random_state, np.random.Generator) - else np.random.default_rng(random_state) - ) - ranks = sp.rankdata(y, method="average") - counts = np.bincount(codes, minlength=n_types).astype(float) - null_counts = np.zeros(n_types, dtype=int) - - for _ in range(n_permutations): - perm_ranks = rng.permutation(ranks) - perm_sums = np.bincount( - codes, weights=perm_ranks, minlength=n_types - ).astype(float) - for k in range(n_types): - n1 = counts[k] - n0 = n_cells - n1 - if n1 > 0 and n0 > 0: - u_null = perm_sums[k] - n1 * (n1 + 1.0) / 2.0 - auc_null = u_null / (n1 * n0) - if auc_null >= aurocs[k]: - null_counts[k] += 1 - - p_values = (1.0 + null_counts) / (1.0 + n_permutations) - elif n_permutations == 0 and n_types > 1 and n_cells > 0: - # Asymptotic one-sided Mann-Whitney test - for k in range(n_types): - pos = y[codes == k] - neg = y[codes != k] - if pos.size > 0 and neg.size > 0: - res = sp.mannwhitneyu(pos, neg, alternative="greater") - p_values[k] = float(res.pvalue) + p_values = _component_p_values( + y, + codes, + n_types, + n_cells, + aurocs, + n_permutations, + _as_generator(random_state), + ) # BH FDR correction across cell types for this component if n_types > 1: @@ -491,50 +334,9 @@ def score_cell_type_alignment( CellTypeAlignmentResults Container with full results across all components. """ - # Extract loadings and cell_types - if isinstance(data, anndata.AnnData): - if projection_key in data.obsm: - loadings_matrix = np.asarray(data.obsm[projection_key]) - elif projection_key == "weighted_projections" and "projections" in data.obsm: - loadings_matrix = np.asarray(data.obsm["projections"]) - else: - raise KeyError(f"Could not find '{projection_key}' in data.obsm.") - - if cell_types is None: - for candidate in [ - "cell_type", - "CellType", - "cell_types", - "Cell_Type", - "celltype", - ]: - if candidate in data.obs: - cell_type_series = pd.Series(data.obs[candidate]) - break - else: - raise KeyError( - "Cell-type column not specified and none of ['cell_type', 'CellType', 'cell_types'] found in data.obs." - ) - elif isinstance(cell_types, str): - if cell_types not in data.obs: - raise KeyError(f"Column '{cell_types}' not found in data.obs.") - cell_type_series = pd.Series(data.obs[cell_types]) - else: - cell_type_series = pd.Series(cell_types) - elif isinstance(data, pd.DataFrame): - loadings_matrix = data.to_numpy() - if cell_types is None or isinstance(cell_types, str): - raise ValueError( - "cell_types must be provided when data is a DataFrame or array." - ) - cell_type_series = pd.Series(cell_types) - else: - loadings_matrix = np.asarray(data) - if cell_types is None or isinstance(cell_types, str): - raise ValueError( - "cell_types must be provided when data is a DataFrame or array." - ) - cell_type_series = pd.Series(cell_types) + loadings_matrix, cell_type_series = _resolve_loadings_and_cell_types( + data, cell_types, projection_key + ) if loadings_matrix.ndim == 1: loadings_matrix = loadings_matrix[:, np.newaxis] @@ -548,11 +350,7 @@ def score_cell_type_alignment( f"Length mismatch: data has {n_cells} cells but cell_types has {codes.size}." ) - rng = ( - random_state - if isinstance(random_state, np.random.Generator) - else np.random.default_rng(random_state) - ) + rng = _as_generator(random_state) component_labels = [i + 1 for i in range(n_comps)] aurocs_mat = np.zeros((n_comps, n_types), dtype=float) @@ -561,8 +359,6 @@ def score_cell_type_alignment( eta2_vec = np.zeros(n_comps, dtype=float) eps2_vec = np.zeros(n_comps, dtype=float) - counts = np.bincount(codes, minlength=n_types).astype(float) - for comp_idx in range(n_comps): y = loadings_matrix[:, comp_idx].astype(float) if signed: @@ -574,30 +370,9 @@ def score_cell_type_alignment( eta2_vec[comp_idx] = compute_eta_squared(y, codes, n_types) eps2_vec[comp_idx] = compute_kruskal_epsilon_squared(y, codes, n_types) - if n_permutations > 0 and n_types > 1 and n_cells > 0: - ranks = sp.rankdata(y, method="average") - null_counts = np.zeros(n_types, dtype=int) - for _ in range(n_permutations): - perm_ranks = rng.permutation(ranks) - perm_sums = np.bincount( - codes, weights=perm_ranks, minlength=n_types - ).astype(float) - for k in range(n_types): - n1 = counts[k] - n0 = n_cells - n1 - if n1 > 0 and n0 > 0: - u_null = perm_sums[k] - n1 * (n1 + 1.0) / 2.0 - auc_null = u_null / (n1 * n0) - if auc_null >= auc[k]: - null_counts[k] += 1 - p_vals_mat[comp_idx] = (1.0 + null_counts) / (1.0 + n_permutations) - elif n_permutations == 0 and n_types > 1 and n_cells > 0: - for k in range(n_types): - pos = y[codes == k] - neg = y[codes != k] - if pos.size > 0 and neg.size > 0: - res = sp.mannwhitneyu(pos, neg, alternative="greater") - p_vals_mat[comp_idx, k] = float(res.pvalue) + p_vals_mat[comp_idx] = _component_p_values( + y, codes, n_types, n_cells, auc, n_permutations, rng + ) # Joint BH FDR correction across all (component x cell_type) tests if aurocs_mat.size > 1: @@ -650,3 +425,11 @@ def score_cell_type_alignment( significant_cell_types=sig_dict, alpha=alpha, ) + + +__all__ = [ + "CellTypeAlignmentResults", + "ComponentAlignmentResult", + "cell_type_alignment", + "score_cell_type_alignment", +] diff --git a/scrise/factor_io.py b/scrise/factor_io.py new file mode 100644 index 00000000..2816af08 --- /dev/null +++ b/scrise/factor_io.py @@ -0,0 +1,309 @@ +"""Reading and writing RISE factors, without the raw expression matrix. + +Factors are exported to h5ad with the projection matrix OPQ-quantized and +the cell barcodes packed into a uint8 matrix, both of which have to be +undone on load. +""" + +import os +from collections.abc import Mapping, Sequence +from typing import Any, cast + +import anndata +import h5py +import hdf5plugin # noqa: F401 (registers the HDF5 filters these files use) +import numpy as np +import pandas as pd + +from .opq import OPQQuantizer, find_optimal_opq + + +def _floats_to_float32(mapping) -> dict: + """Downcast every floating ndarray in an AnnData attribute dict to float32.""" + return { + k: ( + v.astype(np.float32) + if isinstance(v, np.ndarray) and np.issubdtype(v.dtype, np.floating) + else v + ) + for k, v in mapping.items() + } + + +def _pack_obs_names(obs: pd.DataFrame) -> np.ndarray | None: + """Replace a string barcode index with a RangeIndex, returning the bytes.""" + orig_index = obs.index.to_numpy(dtype=str) + max_len = max((len(s) for s in orig_index), default=0) + if max_len == 0: + return None + s_arr = orig_index.astype(f"|S{max_len}") + byte_matrix = np.frombuffer(s_arr.tobytes(), dtype=np.uint8).reshape( + (len(orig_index), max_len) + ) + obs.index = pd.RangeIndex(len(obs)) + return byte_matrix + + +def _recompress_obs_names(filename: str, n_cells: int) -> None: + """Rewrite the packed barcode matrix chunked and gzipped, once written.""" + if n_cells <= 1000: + return + with h5py.File(filename, "r+") as f: + if "uns/_obs_names_bytes" in f: + d = f["uns/_obs_names_bytes"][()] + del f["uns/_obs_names_bytes"] + f.create_dataset( + "uns/_obs_names_bytes", + data=d, + chunks=(min(16384, len(d)), d.shape[1]), + compression="gzip", + compression_opts=6, + ) + + +def _restore_obs_names(adata: anndata.AnnData) -> None: + """Rebuild the string barcode index packed by :func:`_pack_obs_names`.""" + if "_obs_names_bytes" not in adata.uns: + return + byte_matrix = np.asarray(adata.uns["_obs_names_bytes"]) + max_len = byte_matrix.shape[1] + recon_barcodes = np.frombuffer(byte_matrix.tobytes(), dtype=f"|S{max_len}").astype( + str + ) + adata.obs.index = pd.Index(recon_barcodes) + del adata.uns["_obs_names_bytes"] + + +def _restore_projections(adata: anndata.AnnData) -> None: + """Decode OPQ-compressed projections and rebuild weighted_projections.""" + if "projections_opq_codes" in adata.obsm and "opq_rotation" in adata.uns: + quantizer = OPQQuantizer.from_saved( + R=adata.uns["opq_rotation"], + centroids_cat=adata.uns["opq_centroids"], + sub_dims=adata.uns["opq_subdims"], + ) + adata.obsm["projections"] = quantizer.decode( + np.asarray(adata.obsm["projections_opq_codes"]) + ) + elif "projections" in adata.obsm: + adata.obsm["projections"] = np.asarray( + adata.obsm["projections"], dtype=np.float32 + ) + + # Never stored on disk: recoverable as projections @ Pf2_B. + if "projections" in adata.obsm and "Pf2_B" in adata.uns: + adata.obsm["weighted_projections"] = ( + adata.obsm["projections"].astype(np.float32) + @ adata.uns["Pf2_B"].astype(np.float32) + ).astype(np.float32) + + if "embedding" in adata.obsm and "X_pf2_PaCMAP" not in adata.obsm: + adata.obsm["X_pf2_PaCMAP"] = adata.obsm["embedding"] + + +def _read_raw_dataset(raw_path: str) -> anndata.AnnData: + """Read raw expression data, preferring the IVCSR reader when it applies.""" + try: + import vsparse + + return vsparse.VCSCAnnData.read_h5ad(raw_path).to_anndata() + except (ImportError, AttributeError, KeyError, ValueError, OSError): + return anndata.read_h5ad(raw_path) + + +def _attach_raw_data(adata: anndata.AnnData, raw_path: str) -> None: + """Match the factors' cells and genes against raw data and attach ``X``.""" + raw = _read_raw_dataset(raw_path) + + if not isinstance(raw.obs, pd.DataFrame): + raise TypeError("raw.obs must be an in-memory pandas DataFrame.") + + # Match cells by index or cell_barcode column + if not np.all(adata.obs_names.isin(raw.obs_names)) and "cell_barcode" in raw.obs: + raw.obs.index = pd.Index(raw.obs["cell_barcode"].astype(str)) + + common_cells = adata.obs_names[adata.obs_names.isin(raw.obs_names)] + if len(common_cells) == 0: + raise ValueError( + "No matching cell barcodes found between factors and raw data." + ) + raw_sub = raw[adata.obs_names, :].copy() + + if not isinstance(raw_sub.var, pd.DataFrame): + raise TypeError("raw_sub.var must be an in-memory pandas DataFrame.") + + # Match genes + if ( + not np.all(adata.var_names.isin(raw_sub.var_names)) + and "gene_ids" in raw_sub.var + and "gene_ids" in adata.var + ): + raw_sub.var.index = pd.Index(raw_sub.var["gene_ids"].astype(str)) + + common_genes = adata.var_names[adata.var_names.isin(raw_sub.var_names)] + if len(common_genes) == 0: + raise ValueError("No matching gene names found between factors and raw data.") + raw_sub = raw_sub[:, adata.var_names].copy() + + from parafac2.normalize import prepare_dataset + + if "Condition" in adata.obs: + raw_prep = prepare_dataset(raw_sub, "Condition", geneThreshold=0.0) + adata.X = raw_prep.X + else: + adata.X = raw_sub.X + + +def _require_export_factors(X: anndata.AnnData) -> None: + """Reject an AnnData that has not been through a RISE fit.""" + if "Pf2_A" not in X.uns or "Pf2_B" not in X.uns or "Pf2_weights" not in X.uns: + raise KeyError( + "Input AnnData is missing required uns factors (Pf2_A, Pf2_B, Pf2_weights)." + ) + if "Pf2_C" not in X.varm: + raise KeyError("Input AnnData is missing required varm factor 'Pf2_C'.") + if "projections" not in X.obsm: + raise KeyError("Input AnnData is missing required obsm 'projections'.") + + +def _compress_projections( + X: anndata.AnnData, fidelity_threshold: float, random_state: int +) -> tuple[np.ndarray, dict]: + """OPQ-quantize the projections, returning the codes and the codebook.""" + projections = np.asarray(X.obsm["projections"], dtype=np.float32) + quantizer, codes, r2 = find_optimal_opq( + projections, + fidelity_threshold=fidelity_threshold, + random_state=random_state, + ) + + assert quantizer.R is not None + assert quantizer.centroids_cat is not None + assert quantizer.sub_dims is not None + return codes, { + "opq_rotation": quantizer.R.astype(np.float32), + "opq_centroids": quantizer.centroids_cat.astype(np.float32), + "opq_subdims": quantizer.sub_dims.astype(np.int32), + "opq_fidelity": float(r2), + } + + +def _embedding_for_export(X: anndata.AnnData) -> dict: + """The PaCMAP embedding under its on-disk name, if the fit produced one.""" + for key in ("X_pf2_PaCMAP", "embedding"): + if key in X.obsm: + return {"embedding": np.asarray(X.obsm[key], dtype=np.float32)} + return {} + + +def export_factors( + X: anndata.AnnData, + filename: str, + fidelity_threshold: float = 0.99, + random_state: int = 42, +) -> anndata.AnnData: + """Export RISE decomposition factors to an h5ad file without raw expression data. + + Compresses the projection matrix using Optimized Product Quantization (OPQ) + to meet or exceed the specified fidelity threshold (R^2 >= fidelity_threshold). + All factor matrices (Pf2_A, Pf2_B, Pf2_weights, Pf2_C) are stored in float32. + Weighted projections are never stored on disk because they can be reconstructed + deterministically as projections @ Pf2_B. + PaCMAP embeddings are optionally stored if present. + + Parameters + ---------- + X : anndata.AnnData + AnnData object containing RISE decomposition results. Must contain: + - X.uns["Pf2_A"], X.uns["Pf2_B"], X.uns["Pf2_weights"] + - X.varm["Pf2_C"] + - X.obsm["projections"] + filename : str + Output file path (.h5ad). + fidelity_threshold : float, optional (default: 0.99) + Target R^2 reconstruction accuracy threshold for projection compression. + random_state : int, optional (default: 42) + Random seed for reproducibility during OPQ codebook training. + + Returns + ------- + anndata.AnnData + The factor-only AnnData object written to disk. + """ + _require_export_factors(X) + + # Factor matrices in float32 + uns_dict = _floats_to_float32(X.uns) + uns_dict["Pf2_A"] = np.asarray(X.uns["Pf2_A"], dtype=np.float32) + uns_dict["Pf2_B"] = np.asarray(X.uns["Pf2_B"], dtype=np.float32) + uns_dict["Pf2_weights"] = np.asarray(X.uns["Pf2_weights"], dtype=np.float32) + + varm_dict = _floats_to_float32(X.varm) + varm_dict["Pf2_C"] = np.asarray(X.varm["Pf2_C"], dtype=np.float32) + + codes, opq_uns = _compress_projections(X, fidelity_threshold, random_state) + uns_dict.update(opq_uns) + + # obsm excludes weighted_projections and the uncompressed projections. + obsm_dict = {"projections_opq_codes": codes.astype(np.uint8)} + obsm_dict.update(_embedding_for_export(X)) + + obs_df, var_df = X.obs, X.var + if not isinstance(obs_df, pd.DataFrame) or not isinstance(var_df, pd.DataFrame): + raise TypeError( + "X.obs and X.var must be in-memory pandas DataFrames " + "(backed Dataset2D is not supported)." + ) + obs = obs_df.copy() + n_cells = len(obs) + packed_names = _pack_obs_names(obs) + if packed_names is not None: + uns_dict["_obs_names_bytes"] = packed_names + + factors_adata = anndata.AnnData( + obs=obs, + var=var_df.copy(), + uns=uns_dict, + varm=cast(Mapping[str, Sequence[Any]], varm_dict), + obsm=cast(Mapping[str, Sequence[Any]], obsm_dict), + ) + + out_dir = os.path.dirname(os.path.abspath(filename)) + if out_dir: + os.makedirs(out_dir, exist_ok=True) + factors_adata.write_h5ad(filename) + + if "_obs_names_bytes" in uns_dict: + _recompress_obs_names(filename, n_cells) + return factors_adata + + +def load_factors( + filename: str, + raw_path: str | None = None, +) -> anndata.AnnData: + """Load RISE decomposition factors from an h5ad file, decompressing OPQ projections + and optionally rebuilding the full dataset from raw expression data. + + Parameters + ---------- + filename : str + Path to factors .h5ad file. + raw_path : str, optional + Path to raw AnnData or IVCSR .h5ad / .h5 file. If provided, cells and genes + are matched based on cell barcodes and gene names to attach the expression matrix. + + Returns + ------- + anndata.AnnData + AnnData object with reconstructed projections, weighted_projections, + factors, and optionally the raw expression matrix X. + """ + adata = anndata.read_h5ad(filename) + _restore_obs_names(adata) + _restore_projections(adata) + + if raw_path is not None: + _attach_raw_data(adata, raw_path) + + return adata diff --git a/scrise/factorization.py b/scrise/factorization.py index b3f4dd08..4fbc54f9 100644 --- a/scrise/factorization.py +++ b/scrise/factorization.py @@ -1,9 +1,6 @@ -import os -from collections.abc import Mapping, Sequence from typing import Any, cast import anndata -import h5py import hdf5plugin # noqa: F401 import numpy as np import pandas as pd @@ -16,7 +13,6 @@ from tqdm import tqdm from ._pf2_utils import run_parafac2 -from .opq import OPQQuantizer, find_optimal_opq def correct_conditions(X: anndata.AnnData): @@ -405,256 +401,11 @@ def rise_pca_r2x( return r2x_rise, r2x_pca[np.array(ranks) - 1] -def export_factors( - X: anndata.AnnData, - filename: str, - fidelity_threshold: float = 0.99, - random_state: int = 42, -) -> anndata.AnnData: - """Export RISE decomposition factors to an h5ad file without raw expression data. - - Compresses the projection matrix using Optimized Product Quantization (OPQ) - to meet or exceed the specified fidelity threshold (R^2 >= fidelity_threshold). - All factor matrices (Pf2_A, Pf2_B, Pf2_weights, Pf2_C) are stored in float32. - Weighted projections are never stored on disk because they can be reconstructed - deterministically as projections @ Pf2_B. - PaCMAP embeddings are optionally stored if present. - - Parameters - ---------- - X : anndata.AnnData - AnnData object containing RISE decomposition results. Must contain: - - X.uns["Pf2_A"], X.uns["Pf2_B"], X.uns["Pf2_weights"] - - X.varm["Pf2_C"] - - X.obsm["projections"] - filename : str - Output file path (.h5ad). - fidelity_threshold : float, optional (default: 0.99) - Target R^2 reconstruction accuracy threshold for projection compression. - random_state : int, optional (default: 42) - Random seed for reproducibility during OPQ codebook training. - - Returns - ------- - anndata.AnnData - The factor-only AnnData object written to disk. - """ - if "Pf2_A" not in X.uns or "Pf2_B" not in X.uns or "Pf2_weights" not in X.uns: - raise KeyError( - "Input AnnData is missing required uns factors (Pf2_A, Pf2_B, Pf2_weights)." - ) - if "Pf2_C" not in X.varm: - raise KeyError("Input AnnData is missing required varm factor 'Pf2_C'.") - if "projections" not in X.obsm: - raise KeyError("Input AnnData is missing required obsm 'projections'.") - - # Factor matrices in float32 - uns_dict = { - k: ( - v.astype(np.float32) - if isinstance(v, np.ndarray) and np.issubdtype(v.dtype, np.floating) - else v - ) - for k, v in X.uns.items() - } - uns_dict["Pf2_A"] = np.asarray(X.uns["Pf2_A"], dtype=np.float32) - uns_dict["Pf2_B"] = np.asarray(X.uns["Pf2_B"], dtype=np.float32) - uns_dict["Pf2_weights"] = np.asarray(X.uns["Pf2_weights"], dtype=np.float32) - - varm_dict = { - k: ( - v.astype(np.float32) - if isinstance(v, np.ndarray) and np.issubdtype(v.dtype, np.floating) - else v - ) - for k, v in X.varm.items() - } - varm_dict["Pf2_C"] = np.asarray(X.varm["Pf2_C"], dtype=np.float32) - - # Compress projections using OPQ - projections = np.asarray(X.obsm["projections"], dtype=np.float32) - quantizer, codes, r2 = find_optimal_opq( - projections, - fidelity_threshold=fidelity_threshold, - random_state=random_state, - ) - - assert quantizer.R is not None - assert quantizer.centroids_cat is not None - assert quantizer.sub_dims is not None - uns_dict["opq_rotation"] = quantizer.R.astype(np.float32) - uns_dict["opq_centroids"] = quantizer.centroids_cat.astype(np.float32) - uns_dict["opq_subdims"] = quantizer.sub_dims.astype(np.int32) - uns_dict["opq_fidelity"] = float(r2) - - # Build obsm (excluding weighted_projections and uncompressed projections) - obsm_dict = { - "projections_opq_codes": codes.astype(np.uint8), - } - - # Optionally store PaCMAP embedding - if "X_pf2_PaCMAP" in X.obsm: - obsm_dict["embedding"] = np.asarray(X.obsm["X_pf2_PaCMAP"], dtype=np.float32) - elif "embedding" in X.obsm: - obsm_dict["embedding"] = np.asarray(X.obsm["embedding"], dtype=np.float32) - - obs_df, var_df = X.obs, X.var - if not isinstance(obs_df, pd.DataFrame) or not isinstance(var_df, pd.DataFrame): - raise TypeError( - "X.obs and X.var must be in-memory pandas DataFrames " - "(backed Dataset2D is not supported)." - ) - obs = obs_df.copy() - # Option B: Compress string barcode index into 2D uint8 ASCII character byte matrix - orig_index = obs.index.to_numpy(dtype=str) - max_len = max((len(s) for s in orig_index), default=0) - if max_len > 0: - s_arr = orig_index.astype(f"|S{max_len}") - byte_matrix = np.frombuffer(s_arr.tobytes(), dtype=np.uint8).reshape( - (len(orig_index), max_len) - ) - uns_dict["_obs_names_bytes"] = byte_matrix - obs.index = pd.RangeIndex(len(obs)) - - factors_adata = anndata.AnnData( - obs=obs, - var=var_df.copy(), - uns=uns_dict, - varm=cast(Mapping[str, Sequence[Any]], varm_dict), - obsm=cast(Mapping[str, Sequence[Any]], obsm_dict), - ) - - out_dir = os.path.dirname(os.path.abspath(filename)) - if out_dir: - os.makedirs(out_dir, exist_ok=True) - factors_adata.write_h5ad(filename) - - # Apply chunked gzip compression to _obs_names_bytes if present - if "_obs_names_bytes" in uns_dict and len(orig_index) > 1000: - with h5py.File(filename, "r+") as f: - if "uns/_obs_names_bytes" in f: - d = f["uns/_obs_names_bytes"][()] - del f["uns/_obs_names_bytes"] - f.create_dataset( - "uns/_obs_names_bytes", - data=d, - chunks=(min(16384, len(d)), d.shape[1]), - compression="gzip", - compression_opts=6, - ) - return factors_adata - - -def load_factors( - filename: str, - raw_path: str | None = None, -) -> anndata.AnnData: - """Load RISE decomposition factors from an h5ad file, decompressing OPQ projections - and optionally rebuilding the full dataset from raw expression data. - - Parameters - ---------- - filename : str - Path to factors .h5ad file. - raw_path : str, optional - Path to raw AnnData or IVCSR .h5ad / .h5 file. If provided, cells and genes - are matched based on cell barcodes and gene names to attach the expression matrix. - - Returns - ------- - anndata.AnnData - AnnData object with reconstructed projections, weighted_projections, - factors, and optionally the raw expression matrix X. - """ - adata = anndata.read_h5ad(filename) - - # Restore string barcode index if compressed with Option B - if "_obs_names_bytes" in adata.uns: - byte_matrix = np.asarray(adata.uns["_obs_names_bytes"]) - max_len = byte_matrix.shape[1] - recon_barcodes = np.frombuffer( - byte_matrix.tobytes(), dtype=f"|S{max_len}" - ).astype(str) - adata.obs.index = pd.Index(recon_barcodes) - del adata.uns["_obs_names_bytes"] - - # Decompress OPQ projections if present - if "projections_opq_codes" in adata.obsm and "opq_rotation" in adata.uns: - quantizer = OPQQuantizer.from_saved( - R=adata.uns["opq_rotation"], - centroids_cat=adata.uns["opq_centroids"], - sub_dims=adata.uns["opq_subdims"], - ) - adata.obsm["projections"] = quantizer.decode( - np.asarray(adata.obsm["projections_opq_codes"]) - ) - elif "projections" in adata.obsm: - adata.obsm["projections"] = np.asarray( - adata.obsm["projections"], dtype=np.float32 - ) - - # Reconstruct weighted_projections - if "projections" in adata.obsm and "Pf2_B" in adata.uns: - adata.obsm["weighted_projections"] = ( - adata.obsm["projections"].astype(np.float32) - @ adata.uns["Pf2_B"].astype(np.float32) - ).astype(np.float32) - - # Restore embedding alias if PaCMAP embedding was stored - if "embedding" in adata.obsm and "X_pf2_PaCMAP" not in adata.obsm: - adata.obsm["X_pf2_PaCMAP"] = adata.obsm["embedding"] - - # Optionally match and attach raw data - if raw_path is not None: - try: - import vsparse - - raw = vsparse.VCSCAnnData.read_h5ad(raw_path).to_anndata() - except (ImportError, AttributeError, KeyError, ValueError, OSError): - raw = anndata.read_h5ad(raw_path) - - if not isinstance(raw.obs, pd.DataFrame): - raise TypeError("raw.obs must be an in-memory pandas DataFrame.") - - # Match cells by index or cell_barcode column - if ( - not np.all(adata.obs_names.isin(raw.obs_names)) - and "cell_barcode" in raw.obs - ): - raw.obs.index = pd.Index(raw.obs["cell_barcode"].astype(str)) - - # Subset cells present in factors - common_cells = adata.obs_names[adata.obs_names.isin(raw.obs_names)] - if len(common_cells) == 0: - raise ValueError( - "No matching cell barcodes found between factors and raw data." - ) - raw_sub = raw[adata.obs_names, :].copy() - - if not isinstance(raw_sub.var, pd.DataFrame): - raise TypeError("raw_sub.var must be an in-memory pandas DataFrame.") - - # Match genes - if ( - not np.all(adata.var_names.isin(raw_sub.var_names)) - and "gene_ids" in raw_sub.var - and "gene_ids" in adata.var - ): - raw_sub.var.index = pd.Index(raw_sub.var["gene_ids"].astype(str)) - - common_genes = adata.var_names[adata.var_names.isin(raw_sub.var_names)] - if len(common_genes) == 0: - raise ValueError( - "No matching gene names found between factors and raw data." - ) - raw_sub = raw_sub[:, adata.var_names].copy() - - from parafac2.normalize import prepare_dataset - - if "Condition" in adata.obs: - raw_prep = prepare_dataset(raw_sub, "Condition", geneThreshold=0.0) - adata.X = raw_prep.X - else: - adata.X = raw_sub.X - - return adata +__all__ = [ + "canonical_component_signs", + "correct_conditions", + "match_components_across_ranks", + "order_components_by_energy", + "pf2", + "rise_pca_r2x", +] diff --git a/scrise/plotting/annotation_alignment.py b/scrise/plotting/annotation_alignment.py index 21014ec2..a4b98d28 100644 --- a/scrise/plotting/annotation_alignment.py +++ b/scrise/plotting/annotation_alignment.py @@ -19,6 +19,124 @@ cmap_enrichment = sns.diverging_palette(240, 10, as_cmap=True) +def _results_from_enrichment_frame( + enrichment: pd.DataFrame, alpha: float +) -> CellTypeAlignmentResults: + """Wrap a bare AUROC table as results, with no significance information. + + A raw frame carries no p-values, so q-values are all 1.0 and eta^2 is 0 -- + nothing will be starred. Tau is still computable from the AUROCs alone. + """ + q_values = pd.DataFrame(1.0, index=enrichment.index, columns=enrichment.columns) + tau = pd.Series( + [ + float(np.sum(1 - row / np.max(row)) / max(1, len(row) - 1)) + for _, row in enrichment.iterrows() + ], + index=enrichment.index, + ) + eta2 = pd.Series(0.0, index=enrichment.index) + return CellTypeAlignmentResults( + results=[], + enrichment=enrichment, + p_values=q_values, + q_values=q_values, + tau=tau, + eta_squared=eta2, + kruskal_epsilon_squared=eta2, + significant_cell_types={}, + alpha=alpha, + ) + + +def _coerce_alignment_results( + data: anndata.AnnData | CellTypeAlignmentResults | pd.DataFrame, + cell_type_col: str, + projection_key: str, + signed: bool, + n_permutations: int, + alpha: float, + random_state, +) -> CellTypeAlignmentResults: + """Accept already-scored results, an AnnData to score, or a raw AUROC table.""" + if isinstance(data, CellTypeAlignmentResults): + return data + if isinstance(data, anndata.AnnData): + return score_cell_type_alignment( + data=data, + cell_types=cell_type_col, + projection_key=projection_key, + signed=signed, + n_permutations=n_permutations, + alpha=alpha, + random_state=random_state, + ) + if isinstance(data, pd.DataFrame): + return _results_from_enrichment_frame(data, alpha) + raise TypeError(f"Unsupported data type: {type(data)}") + + +def _order_by_dominant_cell_type( + enrichment_df: pd.DataFrame, + q_values_df: pd.DataFrame, + tau_series: pd.Series, + eta2_series: pd.Series, +) -> tuple[pd.DataFrame, pd.DataFrame, pd.Series, pd.Series]: + """Group components by which cell type they peak on, strongest first.""" + vals = enrichment_df.to_numpy() + max_idx = np.argmax(vals, axis=1) + max_val = vals[np.arange(vals.shape[0]), max_idx] + order = np.lexsort((-max_val, max_idx)) + return ( + enrichment_df.iloc[order], + q_values_df.iloc[order], + tau_series.iloc[order], + eta2_series.iloc[order], + ) + + +def _significance_annotations( + enrichment_df: pd.DataFrame, q_values_df: pd.DataFrame, alpha: float +) -> pd.DataFrame: + """A frame of asterisks marking enriched, significant (component, cell type).""" + annot_mat = np.full(enrichment_df.shape, "", dtype=object) + for i, comp in enumerate(enrichment_df.index): + for j, ctype in enumerate(enrichment_df.columns): + if q_values_df.loc[comp, ctype] <= alpha and ( + enrichment_df.loc[comp, ctype] > 0.5 + ): + annot_mat[i, j] = "*" + return pd.DataFrame( + annot_mat, index=enrichment_df.index, columns=enrichment_df.columns + ) + + +def _resolve_alignment_axes( + ax: Axes | Sequence[Axes] | None, show_metrics: bool, enrichment_df: pd.DataFrame +) -> tuple[Axes, Axes | None]: + """Return (main, metrics) axes, creating a sized figure when none is given.""" + n_cols, n_rows = len(enrichment_df.columns), len(enrichment_df) + + if ax is None: + if show_metrics: + _, axes = plt.subplots( + 1, + 2, + figsize=(max(6, n_cols * 0.8 + 2), max(4, n_rows * 0.4 + 1)), + gridspec_kw={"width_ratios": [n_cols, 2], "wspace": 0.08}, + ) + return axes[0], axes[1] + _, ax_main = plt.subplots( + figsize=(max(5, n_cols * 0.8), max(4, n_rows * 0.4 + 1)) + ) + return ax_main, None + + if isinstance(ax, (list, tuple, np.ndarray)) and len(ax) >= 2: + ax_seq = cast(Sequence[Axes], ax) + return ax_seq[0], ax_seq[1] + return cast(Axes, ax), None + + def plot_cell_type_alignment( data: anndata.AnnData | CellTypeAlignmentResults | pd.DataFrame, ax: Axes | Sequence[Axes] | None = None, @@ -77,112 +195,34 @@ def plot_cell_type_alignment( Axes | tuple[Axes, ...] The plotted Matplotlib Axes. """ - if isinstance(data, CellTypeAlignmentResults): - results = data - elif isinstance(data, anndata.AnnData): - results = score_cell_type_alignment( - data=data, - cell_types=cell_type_col, - projection_key=projection_key, - signed=signed, - n_permutations=n_permutations, - alpha=alpha, - random_state=random_state, - ) - elif isinstance(data, pd.DataFrame): - # Raw enrichment dataframe - enrichment = data - q_values = pd.DataFrame(1.0, index=enrichment.index, columns=enrichment.columns) - tau = pd.Series( - [ - float(np.sum(1 - row / np.max(row)) / max(1, len(row) - 1)) - for _, row in enrichment.iterrows() - ], - index=enrichment.index, - ) - eta2 = pd.Series(0.0, index=enrichment.index) - results = CellTypeAlignmentResults( - results=[], - enrichment=enrichment, - p_values=q_values, - q_values=q_values, - tau=tau, - eta_squared=eta2, - kruskal_epsilon_squared=eta2, - significant_cell_types={}, - alpha=alpha, - ) - else: - raise TypeError(f"Unsupported data type: {type(data)}") + results = _coerce_alignment_results( + data, + cell_type_col, + projection_key, + signed, + n_permutations, + alpha, + random_state, + ) enrichment_df = results.enrichment.copy() q_values_df = results.q_values.copy() tau_series = results.tau.copy() eta2_series = results.eta_squared.copy() - # Reorder components by dominant cell type if requested if reorder and len(enrichment_df) > 1: - vals = enrichment_df.to_numpy() - max_idx = np.argmax(vals, axis=1) - max_val = vals[np.arange(vals.shape[0]), max_idx] - order = np.lexsort((-max_val, max_idx)) - enrichment_df = enrichment_df.iloc[order] - q_values_df = q_values_df.iloc[order] - tau_series = tau_series.iloc[order] - eta2_series = eta2_series.iloc[order] - - # Annotations for significance + enrichment_df, q_values_df, tau_series, eta2_series = ( + _order_by_dominant_cell_type( + enrichment_df, q_values_df, tau_series, eta2_series + ) + ) + annot_df = None if annotate_significance: - annot_mat = np.full(enrichment_df.shape, "", dtype=object) - for i, comp in enumerate(enrichment_df.index): - for j, ctype in enumerate(enrichment_df.columns): - q = q_values_df.loc[comp, ctype] - auc = enrichment_df.loc[comp, ctype] - if q <= results.alpha and auc > 0.5: - annot_mat[i, j] = "*" - annot_df = pd.DataFrame( - annot_mat, index=enrichment_df.index, columns=enrichment_df.columns - ) + annot_df = _significance_annotations(enrichment_df, q_values_df, results.alpha) - # Set up axes heatmap_cmap = cmap if cmap is not None else cmap_enrichment - - ax_main: Axes - ax_metrics: Axes | None - - if ax is None: - if show_metrics: - _, axes = plt.subplots( - 1, - 2, - figsize=( - max(6, len(enrichment_df.columns) * 0.8 + 2), - max(4, len(enrichment_df) * 0.4 + 1), - ), - gridspec_kw={ - "width_ratios": [len(enrichment_df.columns), 2], - "wspace": 0.08, - }, - ) - ax_main = axes[0] - ax_metrics = axes[1] - else: - _, ax_main = plt.subplots( - figsize=( - max(5, len(enrichment_df.columns) * 0.8), - max(4, len(enrichment_df) * 0.4 + 1), - ) - ) - ax_metrics = None - - elif isinstance(ax, (list, tuple, np.ndarray)) and len(ax) >= 2: - ax_seq = cast(Sequence[Axes], ax) - ax_main = ax_seq[0] - ax_metrics = ax_seq[1] - else: - ax_main = cast(Axes, ax) - ax_metrics = None + ax_main, ax_metrics = _resolve_alignment_axes(ax, show_metrics, enrichment_df) # Main AUROC heatmap sns.heatmap( diff --git a/scrise/plotting/factors.py b/scrise/plotting/factors.py index 4e359f69..7ec7c813 100644 --- a/scrise/plotting/factors.py +++ b/scrise/plotting/factors.py @@ -11,6 +11,63 @@ cmap = sns.diverging_palette(240, 10, as_cmap=True) +def _normalization_rows( + yt: pd.Series, + X: np.ndarray, + ThomsonNorm: bool, + control_pattern: str | None, + control_conditions: Sequence[str] | None, +) -> np.ndarray: + """The rows whose median and spread set the normalization. + + Explicit ``control_conditions`` win over a name pattern; ``ThomsonNorm`` + is shorthand for the 'CTRL' pattern. With none of them, every condition + contributes. + """ + if ThomsonNorm is True and control_pattern is None: + control_pattern = "CTRL" + + if control_conditions is not None: + return X[yt.isin(control_conditions)] + if control_pattern is not None: + return X[yt.str.contains(control_pattern)] + return X + + +def _draw_condition_group_labels( + ax: Axes, cond_group_labels: pd.Series, color_key +) -> None: + """Draw the colored row rail outside the heatmap, plus its legend.""" + ax.tick_params(axis="y", which="major", pad=20, length=0) + if color_key is None: + colors = sns.color_palette( + n_colors=pd.Series(cond_group_labels).nunique() + ).as_hex() + else: + colors = color_key + + lut = {} + legend_elements = [] + for index, group in enumerate(pd.unique(cond_group_labels)): + lut[group] = colors[index] + legend_elements.append(Patch(color=colors[index], label=group)) + + row_colors = pd.Series(cond_group_labels).map(lut) + for iii, color in enumerate(row_colors): + ax.add_patch( + plt.Rectangle( + xy=(-0.05, iii), + width=0.05, + height=1, + color=color, + lw=0, + transform=ax.get_yaxis_transform(), + clip_on=False, + ) + ) + ax.legend(handles=legend_elements, bbox_to_anchor=(0.18, 1.07)) + + def plot_condition_factors( data: anndata.AnnData, ax: Axes, @@ -61,17 +118,7 @@ def plot_condition_factors( if log_transform is True: X = np.log10(X) - 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 + XX = _normalization_rows(yt, X, ThomsonNorm, control_pattern, control_conditions) X -= np.median(XX, axis=0) X /= np.std(XX, axis=0) @@ -90,32 +137,7 @@ def plot_condition_factors( cond_group_labels = cond_group_labels.iloc[ind] X = X[ind] yt = yt.iloc[ind] - ax.tick_params(axis="y", which="major", pad=20, length=0) - if color_key is None: - colors = sns.color_palette( - n_colors=pd.Series(cond_group_labels).nunique() - ).as_hex() - else: - colors = color_key - lut = {} - legend_elements = [] - for index, group in enumerate(pd.unique(cond_group_labels)): - lut[group] = colors[index] - legend_elements.append(Patch(color=colors[index], label=group)) - row_colors = pd.Series(cond_group_labels).map(lut) - for iii, color in enumerate(row_colors): - ax.add_patch( - plt.Rectangle( - xy=(-0.05, iii), - width=0.05, - height=1, - color=color, - lw=0, - transform=ax.get_yaxis_transform(), - clip_on=False, - ) - ) - ax.legend(handles=legend_elements, bbox_to_anchor=(0.18, 1.07)) + _draw_condition_group_labels(ax, cond_group_labels, color_key) xticks = np.arange(1, X.shape[1] + 1) sns.heatmap( diff --git a/scrise/rank_selection.py b/scrise/rank_selection.py index bab4b4cb..4e3d93d9 100644 --- a/scrise/rank_selection.py +++ b/scrise/rank_selection.py @@ -162,6 +162,58 @@ def _bicv_trial( return 1.0 - ss_res / ss_tot +def _resolve_bicv_inputs( + X: anndata.AnnData | None, + adata: anndata.AnnData | None, + ranks: Sequence[int] | None, + n_repeats: int, + held_out_cell_frac: float, + held_out_gene_frac: float, + condition_key: str | None, +) -> tuple[anndata.AnnData, list[int]]: + """Validate `bicv`'s arguments and return the dataset and the rank list. + + Resolves the ``X``/``adata`` alias, fills in ``condition_unique_idxs`` from + ``condition_key`` when absent, brings a backed dataset into memory, and + rejects rank requests that cannot yield a well-posed trial at these + held-out fractions. + """ + 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 + + sorted_ranks = sorted({int(r) for r in ranks}) + max_rank = _max_feasible_rank(X, held_out_cell_frac, held_out_gene_frac) + if sorted_ranks[-1] > max_rank: + raise ValueError( + f"rank {sorted_ranks[-1]} exceeds the maximum feasible rank ({max_rank}) given " + f"held_out_cell_frac={held_out_cell_frac} and " + f"held_out_gene_frac={held_out_gene_frac}. Test lower ranks, or lower " + "the held-out fractions." + ) + return X, sorted_ranks + + def bicv( X: anndata.AnnData | None = None, ranks: Sequence[int] | None = None, @@ -232,39 +284,15 @@ def bicv( (one of "Fit R2X" or "BiCV R2X"), and "R2X". Ready to pass to :func:`scrise.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}) - max_rank = _max_feasible_rank(X, held_out_cell_frac, held_out_gene_frac) - if ranks[-1] > max_rank: - raise ValueError( - f"rank {ranks[-1]} exceeds the maximum feasible rank ({max_rank}) given " - f"held_out_cell_frac={held_out_cell_frac} and " - f"held_out_gene_frac={held_out_gene_frac}. Test lower ranks, or lower " - "the held-out fractions." - ) + X, ranks = _resolve_bicv_inputs( + X, + adata, + ranks, + n_repeats, + held_out_cell_frac, + held_out_gene_frac, + condition_key, + ) rng = np.random.default_rng(random_state) rows = [] diff --git a/scrise/tests/test_annotation_alignment.py b/scrise/tests/test_annotation_alignment.py index 3aa55c02..2e3f558c 100644 --- a/scrise/tests/test_annotation_alignment.py +++ b/scrise/tests/test_annotation_alignment.py @@ -18,7 +18,7 @@ compute_tau, score_cell_type_alignment, ) -from scrise.annotation_alignment import ( +from scrise.alignment_stats import ( compute_auroc_per_cell_type, compute_eta_squared, compute_kruskal_epsilon_squared, diff --git a/scrise/tests/test_contracts.py b/scrise/tests/test_contracts.py index 9f5acee9..6b2ac40d 100644 --- a/scrise/tests/test_contracts.py +++ b/scrise/tests/test_contracts.py @@ -17,9 +17,9 @@ import pandas as pd import pytest +from ..factor_io import export_factors from ..factorization import ( correct_conditions, - export_factors, order_components_by_energy, pf2, )