oncoordinate 0.1.7__py3-none-any.whl

oncoordinate/lt.py

@@ -0,0 +1,472 @@
+from __future__ import annotations
+
+import logging
+from importlib import resources
+from pathlib import Path
+from typing import Optional, Union, Sequence
+
+import anndata as ad
+import numpy as np
+import pandas as pd
+import scanpy as sc
+from scipy import sparse
+import scvi
+
+logger = logging.getLogger(__name__)
+
+def _load_reference_adata(
+    reference: Optional[Union[str, Path, ad.AnnData]] = None,
+) -> ad.AnnData:
+    if isinstance(reference, ad.AnnData):
+        return reference.copy()
+
+    if reference is not None:
+        path = Path(reference)
+        if not path.is_file():
+            raise FileNotFoundError(f"Reference AnnData not found: {path}")
+        return ad.read_h5ad(path)
+
+    try:
+        pkg_root = resources.files("oncoordinate")
+        res = pkg_root.joinpath("reference_sc.h5ad")
+        with resources.as_file(res) as p:
+            return ad.read_h5ad(p)
+    except Exception as e:
+        raise FileNotFoundError(
+            "No reference AnnData provided and default 'reference_sc.h5ad' "
+            "not found within the oncoordinate package."
+        ) from e
+
+def _load_spatial_adata(
+    spatial: Union[str, Path, ad.AnnData],
+) -> ad.AnnData:
+    if isinstance(spatial, ad.AnnData):
+        return spatial.copy()
+    path = Path(spatial)
+    if not path.is_file():
+        raise FileNotFoundError(f"Spatial AnnData not found: {path}")
+    return ad.read_h5ad(path)
+
+def _get_counts_layer(
+    adata: ad.AnnData,
+    counts_layer: Optional[str] = "counts",
+) -> np.ndarray:
+    if counts_layer is not None and counts_layer in adata.layers:
+        X = adata.layers[counts_layer]
+    else:
+        X = adata.X
+
+    if sparse.issparse(X):
+        X = X.tocsr(copy=False)
+        X.data = np.rint(np.clip(X.data, 0, None)).astype(np.int32)
+    else:
+        X = np.rint(np.clip(np.asarray(X), 0, None)).astype(np.int32)
+    return X
+
+def _make_pseudospots(
+    adata_ref: ad.AnnData,
+    *,
+    celltype_key: str,
+    stage_key: str,
+    pseudospot_size: int = 10,
+) -> ad.AnnData:
+    X = adata_ref.layers["counts"]
+    if sparse.issparse(X):
+        X = X.tocsr(copy=False)
+
+    obs = adata_ref.obs[[celltype_key, stage_key]].astype(str).copy()
+    obs["cell_state"] = obs[celltype_key] + "::" + obs[stage_key]
+
+    groups = obs.groupby("cell_state", observed=True).indices
+
+    X_ps_list = []
+    obs_ps_rows = []
+    rng = np.random.default_rng(0)
+
+    for cs, idxs in groups.items():
+        idxs = np.asarray(idxs)
+        if idxs.size == 0:
+            continue
+
+        rng.shuffle(idxs)
+        n_spots = max(1, idxs.size // pseudospot_size)
+
+        for k in range(n_spots):
+            start = k * pseudospot_size
+            end = (k + 1) * pseudospot_size
+            sel = idxs[start:end]
+            if sel.size == 0:
+                continue
+
+            v = X[sel].sum(axis=0)
+            v = np.asarray(v).ravel() if sparse.issparse(v) else np.asarray(v).ravel()
+
+            ct, st = cs.split("::", 1)
+            X_ps_list.append(v)
+            obs_ps_rows.append((ct, st, cs))
+
+    if not X_ps_list:
+        raise RuntimeError(
+            "No pseudospots could be created. Check celltype_key/stage_key "
+            "and pseudospot_size."
+        )
+
+    X_ps = np.vstack(X_ps_list).astype(np.int32)
+    obs_ps = pd.DataFrame(
+        obs_ps_rows,
+        columns=[celltype_key, stage_key, "cell_state"],
+        index=[f"ps_{i}" for i in range(len(X_ps))],
+    )
+    obs_ps["tech"] = "reference"
+
+    adata_ps = ad.AnnData(X=X_ps, obs=obs_ps, var=adata_ref.var.copy())
+    adata_ps.layers["counts"] = X_ps
+    return adata_ps
+
+
+def _match_library_sizes(
+    adata_ref_ps: ad.AnnData,
+    adata_sp: ad.AnnData,
+    *,
+    layer: str = "counts",
+) -> None:
+    
+    def _libsizes(a: ad.AnnData) -> np.ndarray:
+        Xc = a.layers[layer]
+        return (
+            np.asarray(Xc.sum(axis=1)).ravel()
+            if sparse.issparse(Xc)
+            else np.asarray(Xc.sum(axis=1)).ravel()
+        )
+
+    lib_ref = _libsizes(adata_ref_ps)
+    lib_sp = _libsizes(adata_sp)
+
+    med_ref = float(np.median(lib_ref)) if lib_ref.size else 0.0
+    med_sp = float(np.median(lib_sp)) if lib_sp.size else 0.0
+
+    if med_ref <= 0 or med_sp <= 0:
+        logger.warning("Unable to estimate library sizes; skipping scaling.")
+        return
+
+    ratio = med_sp / med_ref
+    Xc = adata_ref_ps.layers[layer]
+    if sparse.issparse(Xc):
+        Xc = Xc.tocsr(copy=False)
+        Xc.data = np.rint(Xc.data * ratio).astype(np.int32)
+    else:
+        Xc = np.rint(Xc * ratio).astype(np.int32)
+    adata_ref_ps.layers[layer] = Xc
+
+def _run_transfer_single(
+    adata_spatial_in: ad.AnnData,
+    adata_ref_in: ad.AnnData,
+    *,
+    celltype_key: str = "celltype",
+    stage_key: str = "oncoordinate_stage",
+    counts_layer: Optional[str] = "counts",
+    stage_order: Optional[Sequence[str]] = ("normal", "abnormal", "pre-malignant", "malignant"),
+    max_reference_cells: int = 900_000,
+    pseudospot_size: int = 10,
+    n_hvg: int = 2000,
+    scvi_max_epochs: int = 250,
+    scanvi_max_epochs: int = 250,
+    scvi_model_dir: Optional[Union[str, Path]] = None,
+    scanvi_model_dir: Optional[Union[str, Path]] = None,
+) -> ad.AnnData:
+    adata_ref = adata_ref_in.copy()
+    adata_spatial = adata_spatial_in.copy()
+
+    if celltype_key not in adata_ref.obs.columns:
+        raise KeyError(f"celltype_key '{celltype_key}' not found in reference.obs")
+    if stage_key not in adata_ref.obs.columns:
+        raise KeyError(f"stage_key '{stage_key}' not found in reference.obs")
+
+    for a in (adata_ref, adata_spatial):
+        a.obs.index = a.obs.index.astype(str)
+        a.var.index = a.var.index.astype(str)
+        a.var_names = a.var_names.astype(str).str.upper()
+        a.var_names_make_unique()
+
+    shared_genes = np.intersect1d(adata_ref.var_names, adata_spatial.var_names)
+    if shared_genes.size == 0:
+        raise ValueError("No overlapping genes between reference and spatial data.")
+    
+    adata_ref = adata_ref[:, shared_genes].copy()
+    adata_sp = adata_spatial[:, shared_genes].copy()
+    adata_ref.layers["counts"] = _get_counts_layer(adata_ref, counts_layer)
+    adata_sp.layers["counts"] = _get_counts_layer(adata_sp, counts_layer)
+
+    if adata_ref.n_obs > max_reference_cells:
+        logger.info(
+            f"Reference has {adata_ref.n_obs} cells; downsampling to ~{max_reference_cells}."
+        )
+        rng = np.random.default_rng(0)
+        groups = adata_ref.obs.groupby([celltype_key, stage_key], observed=True).indices
+        total = adata_ref.n_obs
+        keep_idx = []
+        for _, idxs in groups.items():
+            idxs = np.asarray(idxs)
+            if idxs.size == 0:
+                continue
+            n = min(int(idxs.size / total * max_reference_cells), idxs.size)
+            if n > 0:
+                keep_idx.append(rng.choice(idxs, size=n, replace=False))
+        if not keep_idx:
+            raise RuntimeError("Subsampling removed all reference cells.")
+        keep_idx = np.concatenate(keep_idx)
+        adata_ref = adata_ref[keep_idx].copy()
+
+    idx_multi = pd.MultiIndex.from_frame(
+        adata_ref.obs[[celltype_key, stage_key]].astype(str)
+    )
+    counts = idx_multi.value_counts()
+    valid_states = counts[counts >= 3].index
+    mask_valid = idx_multi.isin(valid_states)
+    adata_ref = adata_ref[mask_valid].copy()
+    if adata_ref.n_obs == 0:
+        raise RuntimeError(
+            "No reference cells left after filtering small (celltype, stage) groups."
+        )
+
+    adata_ps = _make_pseudospots(
+        adata_ref,
+        celltype_key=celltype_key,
+        stage_key=stage_key,
+        pseudospot_size=pseudospot_size,
+    )
+
+    adata_sp = adata_sp.copy()
+    adata_sp.obs["tech"] = "spatial"
+    adata_sp.obs["cell_state"] = "unlabeled"
+
+    _match_library_sizes(adata_ps, adata_sp, layer="counts")
+
+    adata_comb = ad.concat(
+        [adata_ps, adata_sp],
+        join="inner",
+        merge="same",
+        label="__source__",
+        keys=["reference", "spatial"],
+        index_unique=None,
+    )
+
+    sc.pp.highly_variable_genes(
+        adata_comb,
+        layer="counts",
+        n_top_genes=n_hvg,
+        flavor="seurat_v3",
+        batch_key="tech",
+        inplace=True,
+    )
+    adata_comb = adata_comb[:, adata_comb.var["highly_variable"].values].copy()
+
+    scvi.model.SCVI.setup_anndata(
+        adata_comb,
+        batch_key="tech",
+        labels_key="cell_state",
+        layer="counts",
+    )
+
+    scvi_dir = Path(scvi_model_dir) if scvi_model_dir is not None else None
+    if scvi_dir is not None and scvi_dir.exists():
+        logger.info(f"Loading SCVI model from {scvi_dir}")
+        scvi_model = scvi.model.SCVI.load(scvi_dir, adata=adata_comb)
+    else:
+        scvi_model = scvi.model.SCVI(adata_comb, n_layers=3, n_latent=32)
+        scvi_model.train(max_epochs=scvi_max_epochs, batch_size=512)
+        if scvi_dir is not None:
+            scvi_model.save(scvi_dir, overwrite=True)
+
+    adata_comb.obsm["X_oncoordinate_scvi"] = scvi_model.get_latent_representation()
+    scanvi_dir = Path(scanvi_model_dir) if scanvi_model_dir is not None else None
+    if scanvi_dir is not None and scanvi_dir.exists():
+        logger.info(f"Loading SCANVI model from {scanvi_dir}")
+        scanvi_model = scvi.model.SCANVI.load(scanvi_dir, adata=adata_comb)
+    else:
+        scanvi_model = scvi.model.SCANVI.from_scvi_model(
+            scvi_model,
+            unlabeled_category="unlabeled",
+        )
+        scanvi_model.train(
+            max_epochs=scanvi_max_epochs,
+            batch_size=256,
+            plan_kwargs={"lr": 5e-4},
+            gradient_clip_val=10.0,
+        )
+        if scanvi_dir is not None:
+            scanvi_model.save(scanvi_dir, overwrite=True)
+
+    adata_comb.obsm["X_oncoordinate_scanvi"] = scanvi_model.get_latent_representation(
+        adata_comb
+    )
+
+    state_pred = scanvi_model.predict(adata_comb)
+    adata_comb.obs["oncoordinate_tl_cell_state"] = state_pred.astype(str)
+    soft = scanvi_model.predict(adata_comb, soft=True)
+
+    if isinstance(soft, pd.DataFrame):
+        soft_df = soft.loc[adata_comb.obs_names]
+    else:
+        soft_arr = np.asarray(soft)
+        if soft_arr.ndim == 1:
+            soft_arr = soft_arr.reshape(-1, 1)
+        soft_df = pd.DataFrame(
+            soft_arr,
+            index=adata_comb.obs_names,
+        )
+
+    adata_comb.obsm["oncoordinate_tl_state_proba"] = soft_df.to_numpy()
+    adata_comb.obs["oncoordinate_tl_state_confidence"] = soft_df.max(axis=1)
+
+    def _split_state(s: str) -> tuple[str, str]:
+        return s.split("::", 1) if "::" in s else (s, "NA")
+
+    ct_pred, st_pred = zip(
+        *[_split_state(s) for s in adata_comb.obs["oncoordinate_tl_cell_state"].astype(str)]
+    )
+    adata_comb.obs["oncoordinate_tl_celltype"] = list(ct_pred)
+
+    if stage_order is None:
+        stage_order_use = sorted(set(st_pred))
+    else:
+        stage_order_use = list(stage_order)
+
+    adata_comb.obs["oncoordinate_tl_stage"] = pd.Categorical(
+        list(st_pred),
+        categories=stage_order_use,
+        ordered=True,
+    )
+
+    if isinstance(soft, pd.DataFrame):
+        for stage in stage_order_use:
+            cols_stage = [c for c in soft_df.columns if c.endswith(f"::{stage}")]
+            if cols_stage:
+                series = soft_df[cols_stage].sum(axis=1)
+            else:
+                series = pd.Series(0.0, index=soft_df.index)
+            adata_comb.obs[f"oncoordinate_tl_stage_proba_{stage}"] = series
+    else:
+        for stage in stage_order_use:
+            adata_comb.obs[f"oncoordinate_tl_stage_proba_{stage}"] = 0.0
+
+    is_spatial = adata_comb.obs["tech"].values == "spatial"
+    idx_spatial_comb = adata_comb.obs.index[is_spatial]
+
+    spatial_out = adata_spatial_in.copy()
+    pred_obs = adata_comb.obs.loc[idx_spatial_comb]
+    pred_obs = pred_obs.reindex(spatial_out.obs_names)
+
+    cols_to_copy = [
+        "oncoordinate_tl_cell_state",
+        "oncoordinate_tl_celltype",
+        "oncoordinate_tl_stage",
+        "oncoordinate_tl_state_confidence",
+    ] + [f"oncoordinate_tl_stage_proba_{stage}" for stage in stage_order_use]
+
+    for col in cols_to_copy:
+        spatial_out.obs[col] = pred_obs[col]
+
+    latent_all = adata_comb.obsm["X_oncoordinate_scanvi"]
+    latent_spatial = latent_all[is_spatial, :]
+    latent_df = pd.DataFrame(latent_spatial, index=idx_spatial_comb)
+    latent_df = latent_df.reindex(spatial_out.obs_names)
+    spatial_out.obsm["X_oncoordinate_scanvi"] = latent_df.to_numpy()
+
+    spatial_out.uns.setdefault("oncoordinate_tl_params", {})
+    spatial_out.uns["oncoordinate_tl_params"].update(
+        dict(
+            celltype_key=celltype_key,
+            stage_key=stage_key,
+            counts_layer=counts_layer,
+            max_reference_cells=int(max_reference_cells),
+            pseudospot_size=int(pseudospot_size),
+            n_hvg=int(n_hvg),
+            stage_order=list(stage_order_use),
+        )
+    )
+
+    return spatial_out
+
+def run_transfer(
+    spatial: Union[str, Path, ad.AnnData],
+    reference: Optional[Union[str, Path, ad.AnnData]] = None,
+    *,
+    celltype_key: str = "celltype",
+    stage_key: str = "oncoordinate_stage",
+    counts_layer: Optional[str] = "counts",
+    stage_order: Optional[Sequence[str]] = ("normal", "abnormal", "pre-malignant", "malignant"),
+    max_reference_cells: int = 900_000,
+    pseudospot_size: int = 10,
+    n_hvg: int = 2000,
+    scvi_max_epochs: int = 250,
+    scanvi_max_epochs: int = 250,
+    scvi_model_dir: Optional[Union[str, Path]] = None,
+    scanvi_model_dir: Optional[Union[str, Path]] = None,
+    per_sample: bool = True,
+    sample_key: str = "sample",
+) -> ad.AnnData:
+    adata_ref = _load_reference_adata(reference)
+    adata_spatial_in = _load_spatial_adata(spatial)
+
+    if (
+        per_sample
+        and sample_key is not None
+        and sample_key in adata_spatial_in.obs.columns
+    ):
+        groups = adata_spatial_in.obs.groupby(sample_key, observed=True).indices
+        if len(groups) > 1:
+            logger.info(
+                "Found %d spatial samples in obs['%s']; running transfer per sample.",
+                len(groups),
+                sample_key,
+            )
+            outs = []
+            for sample_name, idx in groups.items():
+                sub = adata_spatial_in[idx].copy()
+                logger.info(
+                    "Running transfer for sample '%s' with %d spots.",
+                    sample_name,
+                    sub.n_obs,
+                )
+                out_sub = _run_transfer_single(
+                    sub,
+                    adata_ref,
+                    celltype_key=celltype_key,
+                    stage_key=stage_key,
+                    counts_layer=counts_layer,
+                    stage_order=stage_order,
+                    max_reference_cells=max_reference_cells,
+                    pseudospot_size=pseudospot_size,
+                    n_hvg=n_hvg,
+                    scvi_max_epochs=scvi_max_epochs,
+                    scanvi_max_epochs=scanvi_max_epochs,
+                    scvi_model_dir=scvi_model_dir,
+                    scanvi_model_dir=scanvi_model_dir,
+                )
+                outs.append(out_sub)
+
+            combined = ad.concat(
+                outs,
+                join="outer",
+                merge="same",
+                label=None,    
+                index_unique=None,
+            )
+            return combined
+
+    return _run_transfer_single(
+        adata_spatial_in,
+        adata_ref,
+        celltype_key=celltype_key,
+        stage_key=stage_key,
+        counts_layer=counts_layer,
+        stage_order=stage_order,
+        max_reference_cells=max_reference_cells,
+        pseudospot_size=pseudospot_size,
+        n_hvg=n_hvg,
+        scvi_max_epochs=scvi_max_epochs,
+        scanvi_max_epochs=scanvi_max_epochs,
+        scvi_model_dir=scvi_model_dir,
+        scanvi_model_dir=scanvi_model_dir,
+    )