chatspatial 1.1.0__py3-none-any.whl

chatspatial/tools/deconvolution/base.py

@@ -0,0 +1,318 @@
+"""
+Base utilities for deconvolution methods.
+
+Design Philosophy:
+- Immutable data container (frozen dataclass) for prepared data
+- Single function API for the common case
+- Hook pattern for method-specific preprocessing (e.g., cell2location)
+"""
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Optional,
+)
+
+import anndata as ad
+import numpy as np
+import pandas as pd
+from numpy.typing import NDArray
+
+if TYPE_CHECKING:
+    from ...spatial_mcp_adapter import ToolContext
+
+from ...utils.adata_utils import (
+    find_common_genes,
+    get_raw_data_source,
+    get_spatial_key,
+    to_dense,
+    validate_gene_overlap,
+    validate_obs_column,
+)
+from ...utils.exceptions import DataError
+
+# Type alias for preprocess hook
+PreprocessHook = Callable[
+    [ad.AnnData, ad.AnnData, "ToolContext"],
+    Awaitable[tuple[ad.AnnData, ad.AnnData]],
+]
+
+
+# =============================================================================
+# Immutable Data Container
+# =============================================================================
+
+
+@dataclass(frozen=True)
+class PreparedDeconvolutionData:
+    """Immutable container for prepared deconvolution data.
+
+    All fields are populated by prepare_deconvolution() and cannot be modified.
+    This eliminates state machine complexity and makes data flow explicit.
+
+    Attributes:
+        spatial: Spatial AnnData subset to common genes (raw counts)
+        reference: Reference AnnData subset to common genes (raw counts)
+        cell_type_key: Column name for cell types in reference
+        cell_types: List of unique cell types
+        common_genes: List of genes present in both datasets
+        spatial_coords: Spatial coordinates array (n_spots, 2) or None
+        ctx: ToolContext for logging/warnings
+
+    Usage:
+        data = await prepare_deconvolution(spatial, ref, "cell_type", ctx)
+        proportions = run_method(data.spatial, data.reference, data.cell_types)
+    """
+
+    spatial: ad.AnnData
+    reference: ad.AnnData
+    cell_type_key: str
+    cell_types: list[str]
+    common_genes: list[str]
+    spatial_coords: Optional[NDArray[np.floating]]
+    ctx: "ToolContext"
+
+    @property
+    def n_spots(self) -> int:
+        """Number of spatial spots."""
+        return self.spatial.n_obs
+
+    @property
+    def n_cell_types(self) -> int:
+        """Number of cell types."""
+        return len(self.cell_types)
+
+    @property
+    def n_genes(self) -> int:
+        """Number of common genes."""
+        return len(self.common_genes)
+
+
+# =============================================================================
+# Single Entry Point
+# =============================================================================
+
+
+async def prepare_deconvolution(
+    spatial_adata: ad.AnnData,
+    reference_adata: ad.AnnData,
+    cell_type_key: str,
+    ctx: "ToolContext",
+    require_int_dtype: bool = False,
+    min_common_genes: int = 100,
+    preprocess: Optional[PreprocessHook] = None,
+) -> PreparedDeconvolutionData:
+    """Prepare data for deconvolution in a single function call.
+
+    This is the primary API for deconvolution data preparation. It handles:
+    1. Validation of cell type key
+    2. Raw count restoration for both datasets
+    3. Optional method-specific preprocessing (via hook)
+    4. Common gene identification and validation
+    5. Subsetting to common genes
+
+    Args:
+        spatial_adata: Spatial transcriptomics AnnData
+        reference_adata: Single-cell reference AnnData
+        cell_type_key: Column in reference.obs containing cell type labels
+        ctx: ToolContext for logging
+        require_int_dtype: Convert to int32 (required for R-based methods)
+        min_common_genes: Minimum required gene overlap
+        preprocess: Optional async hook for method-specific preprocessing.
+                   Signature: async (spatial, reference, ctx) -> (spatial, reference)
+                   Called after raw count restoration, before gene finding.
+
+    Returns:
+        PreparedDeconvolutionData with all fields populated
+
+    Examples:
+        # Standard usage (most methods)
+        data = await prepare_deconvolution(spatial, ref, "cell_type", ctx)
+
+        # With custom preprocessing (e.g., cell2location)
+        async def custom_filter(sp, ref, ctx):
+            sp = await apply_filtering(sp, ctx)
+            ref = await apply_filtering(ref, ctx)
+            return sp, ref
+
+        data = await prepare_deconvolution(
+            spatial, ref, "cell_type", ctx,
+            preprocess=custom_filter
+        )
+    """
+    # 1. Extract spatial coordinates from original data (before any processing)
+    spatial_coords: Optional[NDArray[np.floating]] = None
+    spatial_key = get_spatial_key(spatial_adata)
+    if spatial_key:
+        spatial_coords = np.asarray(spatial_adata.obsm[spatial_key], dtype=np.float64)
+
+    # 2. Validate cell type key
+    validate_obs_column(reference_adata, cell_type_key, "Cell type")
+
+    # 3. Extract cell types
+    cell_types = list(reference_adata.obs[cell_type_key].unique())
+    if len(cell_types) < 2:
+        raise DataError(
+            f"Reference data must have at least 2 cell types, found {len(cell_types)}"
+        )
+
+    # 4. Restore raw counts
+    spatial_prep = await _prepare_counts(
+        spatial_adata, "Spatial", ctx, require_int_dtype
+    )
+    reference_prep = await _prepare_counts(
+        reference_adata, "Reference", ctx, require_int_dtype
+    )
+
+    # 5. Optional method-specific preprocessing
+    if preprocess is not None:
+        spatial_prep, reference_prep = await preprocess(
+            spatial_prep, reference_prep, ctx
+        )
+
+    # 6. Find common genes
+    common_genes = find_common_genes(
+        spatial_prep.var_names,
+        reference_prep.var_names,
+    )
+
+    # 7. Validate gene overlap
+    validate_gene_overlap(
+        common_genes,
+        spatial_prep.n_vars,
+        reference_prep.n_vars,
+        min_genes=min_common_genes,
+        source_name="spatial",
+        target_name="reference",
+    )
+
+    # 8. Return immutable result with data subset to common genes
+    return PreparedDeconvolutionData(
+        spatial=spatial_prep[:, common_genes].copy(),
+        reference=reference_prep[:, common_genes].copy(),
+        cell_type_key=cell_type_key,
+        cell_types=cell_types,
+        common_genes=common_genes,
+        spatial_coords=spatial_coords,
+        ctx=ctx,
+    )
+
+
+async def _prepare_counts(
+    adata: ad.AnnData,
+    label: str,
+    ctx: "ToolContext",
+    require_int_dtype: bool,
+) -> ad.AnnData:
+    """Prepare AnnData by restoring raw counts."""
+    result = get_raw_data_source(
+        adata,
+        prefer_complete_genes=True,
+        require_integer_counts=True,
+        sample_size=100,
+    )
+
+    if not result.is_integer_counts:
+        await ctx.warning(
+            f"{label}: Using normalized data (no raw counts available). "
+            f"This may be acceptable for some reference datasets."
+        )
+
+    # Construct copy based on source
+    if result.source == "raw":
+        adata_copy = adata.raw.to_adata()
+        # Preserve obsm from original (raw.to_adata() doesn't include it)
+        for key in adata.obsm:
+            adata_copy.obsm[key] = adata.obsm[key].copy()
+    elif result.source == "counts_layer":
+        adata_copy = adata.copy()
+        adata_copy.X = adata_copy.layers["counts"]
+    else:
+        adata_copy = adata.copy()
+
+    # Convert to int32 if required (R-based methods)
+    if require_int_dtype and result.is_integer_counts:
+        dense = to_dense(adata_copy.X)
+        adata_copy.X = (
+            dense.astype(np.int32, copy=False) if dense.dtype != np.int32 else dense
+        )
+
+    return adata_copy
+
+
+# =============================================================================
+# Statistics Helper
+# =============================================================================
+
+
+def create_deconvolution_stats(
+    proportions: pd.DataFrame,
+    common_genes: list[str],
+    method: str,
+    device: str = "CPU",
+    **method_specific_params,
+) -> dict[str, Any]:
+    """Create standardized statistics dictionary for deconvolution results."""
+    cell_types = list(proportions.columns)
+    stats = {
+        "method": method,
+        "device": device,
+        "n_spots": len(proportions),
+        "n_cell_types": len(cell_types),
+        "cell_types": cell_types,
+        "genes_used": len(common_genes),
+        "common_genes": len(common_genes),
+        "mean_proportions": proportions.mean().to_dict(),
+        "dominant_types": proportions.idxmax(axis=1).value_counts().to_dict(),
+    }
+    stats.update(method_specific_params)
+    return stats
+
+
+# =============================================================================
+# Convergence Checking
+# =============================================================================
+
+
+def check_model_convergence(
+    model,
+    model_name: str,
+    convergence_threshold: float = 0.001,
+    convergence_window: int = 50,
+) -> tuple[bool, Optional[str]]:
+    """Check if a scvi-tools model has converged based on ELBO history."""
+    if not hasattr(model, "history") or model.history is None:
+        return True, None
+
+    history = model.history
+    elbo_keys = ["elbo_train", "elbo_validation", "train_loss_epoch"]
+    elbo_history = None
+
+    for key in elbo_keys:
+        if key in history and len(history[key]) > 0:
+            elbo_history = history[key]
+            break
+
+    if elbo_history is None or len(elbo_history) < convergence_window:
+        return True, None
+
+    elbo_arr = np.asarray(elbo_history).ravel()
+    recent_elbo = elbo_arr[-convergence_window:]
+    elbo_changes = np.abs(np.diff(recent_elbo))
+
+    mean_value = np.abs(np.mean(recent_elbo))
+    if mean_value > 0:
+        relative_changes = elbo_changes / mean_value
+        mean_relative_change = np.mean(relative_changes)
+
+        if mean_relative_change > convergence_threshold:
+            return False, (
+                f"{model_name} may not have fully converged. "
+                f"Mean relative ELBO change: {mean_relative_change:.4f} "
+                f"(threshold: {convergence_threshold}). "
+                "Consider increasing training epochs."
+            )
+
+    return True, None

chatspatial/tools/deconvolution/card.py

@@ -0,0 +1,244 @@
+"""
+CARD (Conditional AutoRegressive-based Deconvolution) method.
+
+CARD models spatial correlation in cell type composition using a
+CAR (Conditional AutoRegressive) model. Unique features:
+- Spatial correlation modeling
+- Optional high-resolution imputation
+"""
+
+from typing import TYPE_CHECKING, Any, Optional
+
+import numpy as np
+import pandas as pd
+
+if TYPE_CHECKING:
+    pass
+
+from ...utils.dependency_manager import validate_r_package
+from ...utils.exceptions import ProcessingError
+from .base import PreparedDeconvolutionData, create_deconvolution_stats
+
+
+def deconvolve(
+    data: PreparedDeconvolutionData,
+    sample_key: Optional[str] = None,
+    minCountGene: int = 100,
+    minCountSpot: int = 5,
+    imputation: bool = False,
+    NumGrids: int = 2000,
+    ineibor: int = 10,
+) -> tuple[pd.DataFrame, dict[str, Any]]:
+    """Deconvolve spatial data using CARD R package.
+
+    Args:
+        data: Prepared deconvolution data (immutable, includes spatial coordinates)
+        sample_key: Optional sample/batch key in reference data
+        minCountGene: Include genes with at least this many counts
+        minCountSpot: Include genes expressed in at least this many spots
+        imputation: Whether to perform spatial imputation
+        NumGrids: Number of grids for imputation
+        ineibor: Number of neighbors for imputation
+
+    Returns:
+        Tuple of (proportions DataFrame, statistics dictionary)
+    """
+    import anndata2ri
+    import rpy2.robjects as ro
+    from rpy2.robjects import numpy2ri, pandas2ri
+    from rpy2.robjects.conversion import localconverter
+
+    ctx = data.ctx
+
+    # Validate R package
+    validate_r_package(
+        "CARD",
+        ctx,
+        install_cmd="devtools::install_github('YingMa0107/CARD')",
+    )
+
+    try:
+        # Load CARD
+        with localconverter(ro.default_converter + pandas2ri.converter):
+            ro.r("library(CARD)")
+
+        # Data already copied in prepare_deconvolution
+        spatial_data = data.spatial
+        reference_data = data.reference
+
+        # Get spatial coordinates from prepared data
+        if data.spatial_coords is not None:
+            spatial_location = pd.DataFrame(
+                data.spatial_coords[:, :2],
+                index=spatial_data.obs_names,
+                columns=["x", "y"],
+            )
+        else:
+            spatial_location = pd.DataFrame(
+                {"x": range(spatial_data.n_obs), "y": [0] * spatial_data.n_obs},
+                index=spatial_data.obs_names,
+            )
+
+        # Prepare metadata
+        sc_meta = reference_data.obs[[data.cell_type_key]].copy()
+        sc_meta.columns = ["cellType"]
+
+        if sample_key and sample_key in reference_data.obs:
+            sc_meta["sampleInfo"] = reference_data.obs[sample_key]
+        else:
+            sc_meta["sampleInfo"] = "sample1"
+
+        # Transfer matrices to R
+        with localconverter(ro.default_converter + anndata2ri.converter):
+            ro.globalenv["sc_count"] = reference_data.X.T
+            ro.globalenv["spatial_count"] = spatial_data.X.T
+
+            ro.globalenv["gene_names_ref"] = ro.StrVector(reference_data.var_names)
+            ro.globalenv["cell_names"] = ro.StrVector(reference_data.obs_names)
+            ro.globalenv["gene_names_spatial"] = ro.StrVector(spatial_data.var_names)
+            ro.globalenv["spot_names"] = ro.StrVector(spatial_data.obs_names)
+
+            ro.r(
+                """
+                rownames(sc_count) <- gene_names_ref
+                colnames(sc_count) <- cell_names
+                rownames(spatial_count) <- gene_names_spatial
+                colnames(spatial_count) <- spot_names
+            """
+            )
+
+        # Transfer metadata
+        with localconverter(ro.default_converter + pandas2ri.converter):
+            ro.globalenv["sc_meta"] = ro.conversion.py2rpy(sc_meta)
+            ro.globalenv["spatial_location"] = ro.conversion.py2rpy(spatial_location)
+            ro.globalenv["minCountGene"] = minCountGene
+            ro.globalenv["minCountSpot"] = minCountSpot
+
+        # Create CARD object and run deconvolution
+        ro.r(
+            """
+            capture.output(
+                CARD_obj <- createCARDObject(
+                    sc_count = sc_count,
+                    sc_meta = sc_meta,
+                    spatial_count = spatial_count,
+                    spatial_location = spatial_location,
+                    ct.varname = "cellType",
+                    ct.select = unique(sc_meta$cellType),
+                    sample.varname = "sampleInfo",
+                    minCountGene = minCountGene,
+                    minCountSpot = minCountSpot
+                ),
+                file = "/dev/null"
+            )
+            capture.output(
+                CARD_obj <- CARD_deconvolution(CARD_object = CARD_obj),
+                file = "/dev/null"
+            )
+        """
+        )
+
+        # Extract results
+        with localconverter(
+            ro.default_converter + pandas2ri.converter + numpy2ri.converter
+        ):
+            row_names = list(ro.r("rownames(CARD_obj@Proportion_CARD)"))
+            col_names = list(ro.r("colnames(CARD_obj@Proportion_CARD)"))
+            proportions_r = ro.r("CARD_obj@Proportion_CARD")
+            proportions_array = np.array(proportions_r)
+
+            proportions = pd.DataFrame(
+                proportions_array, index=row_names, columns=col_names
+            )
+
+        # Optional imputation
+        imputed_proportions = None
+        imputed_coordinates = None
+
+        if imputation:
+            ro.r(
+                f"""
+                capture.output(
+                    CARD_impute <- CARD.imputation(
+                        CARD_object = CARD_obj,
+                        NumGrids = {NumGrids},
+                        ineibor = {ineibor}
+                    ),
+                    file = "/dev/null"
+                )
+            """
+            )
+
+            with localconverter(ro.default_converter + pandas2ri.converter):
+                imputed_row_names = list(ro.r("rownames(CARD_impute@refined_prop)"))
+                imputed_col_names = list(ro.r("colnames(CARD_impute@refined_prop)"))
+                imputed_proportions_r = ro.r("CARD_impute@refined_prop")
+                imputed_proportions_array = np.array(imputed_proportions_r)
+
+                # Parse coordinates from rownames
+                coords_list = []
+                for name in imputed_row_names:
+                    parts = name.split("x")
+                    coords_list.append([float(parts[0]), float(parts[1])])
+
+                imputed_proportions = pd.DataFrame(
+                    imputed_proportions_array,
+                    index=imputed_row_names,
+                    columns=imputed_col_names,
+                )
+                imputed_coordinates = pd.DataFrame(
+                    coords_list, index=imputed_row_names, columns=["x", "y"]
+                )
+
+        # Create statistics
+        stats = create_deconvolution_stats(
+            proportions,
+            data.common_genes,
+            method="CARD",
+            device="CPU",
+            minCountGene=minCountGene,
+            minCountSpot=minCountSpot,
+        )
+
+        if imputation and imputed_proportions is not None:
+            stats["imputation"] = {
+                "enabled": True,
+                "n_imputed_locations": len(imputed_proportions),
+                "resolution_increase": (
+                    f"{len(imputed_proportions) / len(row_names):.1f}x"
+                ),
+                "imputed_proportions": imputed_proportions,
+                "imputed_coordinates": imputed_coordinates,
+            }
+
+        # Clean up R global environment
+        cleanup_vars = [
+            "sc_count",
+            "spatial_count",
+            "gene_names_ref",
+            "cell_names",
+            "gene_names_spatial",
+            "spot_names",
+            "sc_meta",
+            "spatial_location",
+            "minCountGene",
+            "minCountSpot",
+            "CARD_obj",
+        ]
+        if imputation:
+            cleanup_vars.append("CARD_impute")
+
+        ro.r(
+            f"""
+            rm(list = c({', '.join(f'"{v}"' for v in cleanup_vars)}),
+               envir = .GlobalEnv)
+            gc()
+        """
+        )
+
+        return proportions, stats
+
+    except Exception as e:
+        if isinstance(e, ProcessingError):
+            raise
+        raise ProcessingError(f"CARD deconvolution failed: {e}") from e