"""Dimensionality reduction and embedding functions for GEDI.
Provides SVD, PCA, and UMAP embeddings based on GEDI's factorized
decomposition, preserving the biological interpretability of the model.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Literal
import numpy as np
from threadpoolctl import threadpool_limits
if TYPE_CHECKING:
from anndata import AnnData
from .._logging import debug, info
from .._settings import settings
def _resolve_model_params(gedi_data: dict[str, Any], modality: str | None) -> dict[str, Any]:
"""Return the right model-params dict for single or multi-modal use."""
model = gedi_data.get("model", {})
if modality is None:
return model
if modality not in model:
available = [k for k in model if k != "joint"] + (["joint"] if "joint" in model else [])
raise ValueError(f"modality '{modality}' not found in model. Available: {available}")
return model[modality]
def _rank_k_svd_of_product(
X: np.ndarray, B: np.ndarray
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""SVD of (X @ B) for X of shape (J, K) and B of shape (K, N), rank ≤ K.
X @ B is rank-K by construction, so its non-trivial SVD is fully captured
by three small SVDs: thin SVD of X (J×K), thin SVD of B^T (N×K), and a
K×K SVD of the contracted middle. Cost O((J+N)·K²) rather than the full
O(J·N·min(J,N)) — for J=15 000, N=100 000, K=20 this is ~5e5× cheaper
and avoids ever materialising the J×N product or its (J×J) U / (J×N) Vt
factors that would otherwise blow out host RAM.
Returns ``(U, d, V)`` of shapes ``(J, K)``, ``(K,)``, ``(N, K)`` such
that ``X @ B ≈ U @ diag(d) @ V.T``.
The three thin SVDs run under ``threadpoolctl.threadpool_limits(1)``.
When this function is reached inside ``tl.multigedi(use_gpu=True)`` —
where multigedi's CUDA backend and torch are loaded in the same
process — the multi-threaded host BLAS pool can deadlock, hanging the
SVD indefinitely on inputs that complete in well under a second in a
fresh process. Forcing single-thread BLAS for these (tiny,
``O((J+N)·K²)``) decompositions sidesteps the deadlock at no measurable
cost. See issue #10.
"""
# threadpoolctl is already an install-time dependency (via scikit-learn);
# imported at module scope -- see the top of this file.
with threadpool_limits(limits=1):
Ux, sx, Vxt = np.linalg.svd(X, full_matrices=False) # (J×K, K, K×K)
Ub, sb, Vbt = np.linalg.svd(B.T, full_matrices=False) # (N×K, K, K×K)
# X @ B = Ux · diag(sx) · Vxt · Vbt.T · diag(sb) · Ub.T = Ux · M · Ub.T
M = (sx[:, None] * Vxt) @ Vbt.T * sb[None, :] # (K, K)
Um, s, Vmt = np.linalg.svd(M, full_matrices=False)
return Ux @ Um, s, Ub @ Vmt.T
def _compute_multimodal_svd(
model: dict[str, Any],
modality: str = "joint",
) -> dict[str, np.ndarray]:
"""Compute SVD for multi-modal GEDI results.
For ``modality="joint"``:
ZDB = vstack([Z_m @ diag(D_m) @ B_shared for each modality m])
For a named modality:
ZDB = Z_m @ diag(D_m) @ B_m (identical to single-modal formula)
The product ZDB is rank-K by construction, so this routes through
:func:`_rank_k_svd_of_product` and never materialises the full J×N
matrix or its full SVD factors.
Parameters
----------
model
The ``adata.uns[key]["model"]`` dict, containing ``"joint"`` and
per-modality sub-dicts.
modality
``"joint"`` (default) or a specific modality name.
Returns
-------
dict with keys ``"d"``, ``"u"``, ``"v"``.
"""
mod_names = [k for k in model if k != "joint"]
K = len(model[mod_names[0]]["D"])
if modality == "joint":
joint = model.get("joint", {})
if not joint.get("Bi"):
raise ValueError(
"joint Bi not found. Run tl.multigedi() to populate joint model params."
)
B = np.hstack(joint["Bi"]) # (K × N_total)
ZD_blocks = []
for mn in mod_names:
mp = model[mn]
Z = np.asarray(mp["Z"]) # (J_m × K)
D = np.asarray(mp["D"]) # (K,)
ZD_blocks.append(Z * D[None, :])
X = np.vstack(ZD_blocks) # (sum(J_m) × K)
else:
mp = model[modality]
Z = np.asarray(mp["Z"])
D = np.asarray(mp["D"])
Bi = mp["Bi"]
X = Z * D[None, :] # (J × K)
B = np.hstack([np.asarray(b) for b in Bi])
U, d, V = _rank_k_svd_of_product(X, B)
return {"d": d[:K], "u": U[:, :K], "v": V[:, :K]}
[docs]
def svd(
adata: AnnData,
*,
key: str = "gedi",
modality: str | None = None,
copy: bool = False,
) -> dict[str, np.ndarray] | None:
r"""Compute factorized SVD from GEDI decomposition.
Computes SVD while preserving GEDI's factorized structure:
``SVD(Z) × SVD(middle) × SVD(DB)``. This maintains biological
interpretability by respecting the decomposition structure.
Parameters
----------
adata
Annotated data matrix with GEDI results in ``.uns[key]``.
key
Key in ``adata.uns`` where GEDI results are stored.
copy
If ``True``, return the SVD result as a dict instead of
storing in ``adata``.
Returns
-------
If ``copy=True``, returns dict with keys ``'d'``, ``'u'``, ``'v'``.
Otherwise, stores results in ``adata.uns[key]['svd']`` and returns ``None``.
The SVD components are:
- ``d``: Singular values (K,)
- ``u``: Left singular vectors (n_genes, K) - gene loadings
- ``v``: Right singular vectors (n_cells, K) - cell embeddings
Examples
--------
>>> import multigedi as gd
>>> gd.tl.gedi(adata, batch_key="sample", n_latent=10)
>>> gd.tl.svd(adata)
>>> adata.uns["gedi"]["svd"]["d"] # singular values
"""
if key not in adata.uns:
raise ValueError(
f"No GEDI results found at adata.uns['{key}']. "
f"Run gd.tl.gedi() or gd.tl.multigedi() first."
)
gedi_data = adata.uns[key]
model = gedi_data.get("model", {})
if modality is not None:
# Multi-modal path
debug(f"Computing factorized SVD (modality={modality!r})")
svd_result = _compute_multimodal_svd(model, modality)
if copy:
return svd_result
gedi_data.setdefault("svd", {})[modality] = svd_result
info(f"Added SVD results to adata.uns['{key}']['svd']['{modality}']")
return None
# Single-modal path (backward-compatible with gedi2py)
Z = model.get("Z")
D = model.get("D")
Bi_list = model.get("Bi")
if Z is None or D is None or Bi_list is None:
raise ValueError(
"Missing model parameters for SVD. "
"Ensure gd.tl.gedi() completed successfully, "
"or pass modality= for multi-modal results."
)
debug("Computing factorized SVD")
svd_result = _compute_svd_factorized(Z, D, Bi_list)
if copy:
return svd_result
gedi_data["svd"] = svd_result
info(f"Added SVD results to adata.uns['{key}']['svd']")
return None
def _compute_svd_factorized(
Z: np.ndarray,
D: np.ndarray,
Bi_list: list[np.ndarray],
) -> dict[str, np.ndarray]:
"""Compute rank-K SVD of (Z·diag(D)·B) without materialising the J×N product.
Routes through :func:`_rank_k_svd_of_product`. Result is mathematically
identical (up to ULP) to ``np.linalg.svd(Z @ diag(D) @ B)`` truncated to
the top K, but ``O((J+N)·K²)`` instead of ``O(J·N·min(J,N))``.
"""
K = len(D)
X = np.asarray(Z) * np.asarray(D)[None, :] # (J × K)
B = np.hstack([np.asarray(b) for b in Bi_list]) # (K × N)
U, d, V = _rank_k_svd_of_product(X, B)
return {"d": d[:K], "u": U[:, :K], "v": V[:, :K]}
[docs]
def pca(
adata: AnnData,
*,
n_components: int | None = None,
key: str = "gedi",
modality: str | None = None,
key_added: str | None = None,
copy: bool = False,
) -> AnnData | np.ndarray | None:
r"""Compute PCA coordinates from GEDI decomposition.
PCA coordinates are computed as ``V @ diag(d)`` from the factorized SVD,
where V are the right singular vectors (cell embeddings).
Parameters
----------
adata
Annotated data matrix with GEDI results in ``.uns[key]``.
n_components
Number of PCs to compute. If ``None``, uses all K latent factors.
key
Key in ``adata.uns`` where GEDI results are stored.
modality
``None`` — single-modal (gedi2py-compatible).
``"joint"`` — multi-modal joint embedding using shared B.
``"<name>"`` — per-modality embedding for a named modality.
key_added
Key to store PCA in ``adata.obsm``.
Defaults to ``X_{key}_pca`` for joint/None, ``X_{key}_{modality}_pca`` otherwise.
copy
If ``True``, return the PCA coordinates instead of storing in ``adata``.
Returns
-------
If ``copy=True``, returns PCA coordinates as numpy array (n_cells, n_components).
Otherwise, stores in ``adata.obsm[key_added]`` and returns ``None``.
Examples
--------
>>> import multigedi as gd
>>> gd.tl.gedi(adata, batch_key="sample", n_latent=10)
>>> gd.tl.pca(adata, n_components=20)
>>> adata.obsm["X_gedi_pca"]
"""
if key not in adata.uns:
raise ValueError(
f"No GEDI results found at adata.uns['{key}']. "
f"Run gd.tl.gedi() or gd.tl.multigedi() first."
)
gedi_data = adata.uns[key]
if modality is not None:
# Multi-modal path: ensure SVD for this modality is computed
svd_store = gedi_data.setdefault("svd", {})
if modality not in svd_store:
svd(adata, key=key, modality=modality)
svd_result = gedi_data["svd"][modality]
else:
# Single-modal path (backward-compatible)
if (
"svd" not in gedi_data
or isinstance(gedi_data["svd"], dict)
and "d" not in gedi_data["svd"]
):
svd(adata, key=key)
svd_result = gedi_data["svd"]
d = svd_result["d"]
v = svd_result["v"]
K = len(d)
if n_components is None:
n_components = K
n_components = min(n_components, K)
debug(f"Computing PCA with {n_components} components (modality={modality!r})")
# PCA = V @ diag(d) (element-wise multiply is equivalent and faster)
pca_coords = v[:, :n_components] * d[:n_components]
if copy:
return pca_coords
# Key convention: joint uses same key as single-modal (primary embedding)
if key_added is None:
if modality is None or modality == "joint":
key_added = f"X_{key}_pca"
else:
key_added = f"X_{key}_{modality}_pca"
adata.obsm[key_added] = pca_coords
info(f"Added PCA to adata.obsm['{key_added}'] ({pca_coords.shape[1]} components)")
return None
[docs]
def umap(
adata: AnnData,
*,
n_neighbors: int = 15,
min_dist: float = 0.1,
n_components: int = 2,
metric: str = "euclidean",
input_key: Literal["pca", "db", "zdb"] = "pca",
key: str = "gedi",
modality: str | None = None,
key_added: str | None = None,
random_state: int | None = None,
copy: bool = False,
) -> AnnData | np.ndarray | None:
r"""Compute UMAP embedding from GEDI results.
Parameters
----------
adata
Annotated data matrix with GEDI results.
n_neighbors
Size of local neighborhood for UMAP.
min_dist
Minimum distance between points in the embedding.
n_components
Dimensionality of the UMAP embedding.
metric
Distance metric for neighbor search.
input_key
Which GEDI representation to use as input:
- ``"pca"``: PCA coordinates (default)
- ``"db"``: DB latent factor embedding
- ``"zdb"``: ZDB shared manifold projection
key
Key in ``adata.uns`` where GEDI results are stored.
modality
``None`` — single-modal (gedi2py-compatible).
``"joint"`` — use joint PCA from ``X_{key}_pca``.
``"<name>"`` — use per-modality PCA from ``X_{key}_{name}_pca``.
key_added
Key to store UMAP in ``adata.obsm``. Defaults to ``X_{key}_umap``
for joint/None, or ``X_{key}_{modality}_umap`` for named modalities.
random_state
Random seed for reproducibility. If ``None``, uses ``settings.random_state``.
copy
If ``True``, return UMAP coordinates instead of storing in ``adata``.
Returns
-------
If ``copy=True``, returns UMAP coordinates as numpy array (n_cells, n_components).
Otherwise, stores in ``adata.obsm[key_added]`` and returns ``None``.
Examples
--------
>>> import multigedi as gd
>>> gd.tl.gedi(adata, batch_key="sample", n_latent=10)
>>> gd.tl.umap(adata, n_neighbors=30)
>>> gd.pl.embedding(adata, basis="X_gedi_umap", color="cell_type")
"""
try:
import umap as umap_module
except ImportError as err:
raise ImportError(
"umap-learn is required for UMAP computation. Install with: pip install umap-learn"
) from err
if random_state is None:
random_state = settings.random_state
# Resolve which obsm key to use as UMAP input
if input_key == "pca":
if modality is None or modality == "joint":
pca_key = f"X_{key}_pca"
else:
pca_key = f"X_{key}_{modality}_pca"
if pca_key not in adata.obsm:
pca(adata, key=key, modality=modality)
X = adata.obsm[pca_key]
debug(f"Using PCA coordinates as UMAP input ({pca_key!r})")
elif input_key == "db":
from ._projections import get_projection
X = get_projection(adata, "db", key=key, modality=modality, copy=True)
debug("Using DB projection as UMAP input")
elif input_key == "zdb":
from ._projections import get_projection
X = get_projection(adata, "zdb", key=key, modality=modality, copy=True)
debug("Using ZDB projection as UMAP input")
else:
raise ValueError(f"Invalid input_key: {input_key}. Must be one of: 'pca', 'db', 'zdb'.")
info(f"Computing UMAP (n_neighbors={n_neighbors}, min_dist={min_dist})")
reducer = umap_module.UMAP(
n_neighbors=n_neighbors,
min_dist=min_dist,
n_components=n_components,
metric=metric,
random_state=random_state,
)
umap_coords = reducer.fit_transform(X)
if copy:
return umap_coords
if key_added is None:
if modality is None or modality == "joint":
key_added = f"X_{key}_umap"
else:
key_added = f"X_{key}_{modality}_umap"
# Store in adata.obsm
if key_added is None:
key_added = f"X_{key}_umap"
adata.obsm[key_added] = umap_coords
info(f"Added UMAP to adata.obsm['{key_added}']")
return None