spatialcheckpoint 0.1.0__py3-none-any.whl

spatialcheckpoint/__init__.py

@@ -0,0 +1,39 @@
+"""SpatialCheckpoint: Spatial heterogeneity profiling of immune checkpoints."""
+__version__ = "0.1.0"
+
+from spatialcheckpoint.data.preprocess import SpatialDataPreprocessor
+from spatialcheckpoint.data.loader import SpatialDataLoader
+from spatialcheckpoint.analysis.spatial_expression import SpatialCheckpointProfiler
+from spatialcheckpoint.analysis.spatial_features import SpatialFeatureEngineer
+from spatialcheckpoint.analysis.colocalization import CheckpointColocalizationAnalyzer
+from spatialcheckpoint.model.archetype_discovery import SpatialArchetypeDiscovery
+from spatialcheckpoint.model.classifier import SpatialArchetypeClassifier
+from spatialcheckpoint.model.trainer import ArchetypeModelTrainer
+from spatialcheckpoint.model.explainer import ArchetypeExplainer
+from spatialcheckpoint.utils.gene_sets import (
+    get_all_checkpoint_genes,
+    get_category_genes,
+    get_immune_cell_markers,
+    get_ligand_receptor_pairs,
+)
+
+__all__ = [
+    "__version__",
+    # Data
+    "SpatialDataPreprocessor",
+    "SpatialDataLoader",
+    # Analysis
+    "SpatialCheckpointProfiler",
+    "SpatialFeatureEngineer",
+    "CheckpointColocalizationAnalyzer",
+    # Model
+    "SpatialArchetypeDiscovery",
+    "SpatialArchetypeClassifier",
+    "ArchetypeModelTrainer",
+    "ArchetypeExplainer",
+    # Gene sets
+    "get_all_checkpoint_genes",
+    "get_category_genes",
+    "get_immune_cell_markers",
+    "get_ligand_receptor_pairs",
+]

spatialcheckpoint/analysis/__init__.py

@@ -0,0 +1 @@
+"""Analysis module."""

spatialcheckpoint/analysis/colocalization.py

@@ -0,0 +1,516 @@
+"""Spatial co-localization analysis of immune checkpoints."""
+from __future__ import annotations
+
+from typing import List, Optional
+import warnings
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+from scipy import stats
+from scipy.spatial import cKDTree
+import scanpy as sc
+
+try:
+    import squidpy as sq
+    HAS_SQUIDPY = True
+except ImportError:
+    HAS_SQUIDPY = False
+
+
+class CheckpointColocalizationAnalyzer:
+    """Spatial co-localization analysis for immune checkpoints in Visium data."""
+
+    # Known ligand-receptor pairs for checkpoints
+    LR_PAIRS = [
+        ('CD274', 'PDCD1'),     # PD-L1 / PD-1
+        ('PDCD1LG2', 'PDCD1'),  # PD-L2 / PD-1
+        ('LGALS9', 'HAVCR2'),   # Galectin-9 / TIM-3
+        ('PVR', 'TIGIT'),       # CD155 / TIGIT
+        ('NECTIN2', 'TIGIT'),   # CD112 / TIGIT
+        ('FGL1', 'LAG3'),       # FGL1 / LAG-3
+        ('CEACAM1', 'LAG3'),    # CEACAM1 / LAG-3
+        ('CD47', 'SIRPA'),      # CD47 / SIRPα
+    ]
+
+    def __init__(self, adata: sc.AnnData, checkpoint_genes: List[str],
+                 immune_markers: Optional[dict] = None):
+        self.adata = adata
+        self.cp_genes = [g for g in checkpoint_genes if g in adata.var_names]
+        self.immune_markers = immune_markers or {}
+        self._coords = None
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _get_coords(self) -> np.ndarray:
+        """Get spot spatial coordinates (n_spots, 2)."""
+        if self._coords is None:
+            self._coords = np.array(self.adata.obsm['spatial'], dtype=float)
+        return self._coords
+
+    def _get_expr(self, genes: list[str]) -> np.ndarray:
+        """Get dense expression for given genes (n_spots, n_genes).
+        Handles sparse adata.X."""
+        indices = [self.adata.var_names.get_loc(g) for g in genes]
+        X = self.adata.X[:, indices]
+        if sp.issparse(X):
+            X = X.toarray()
+        return np.array(X, dtype=float)
+
+    # ------------------------------------------------------------------
+    # Public methods
+    # ------------------------------------------------------------------
+
+    def pairwise_cooccurrence(self, interval: int = 50,
+                              n_intervals: int = 10) -> dict:
+        """Pairwise checkpoint spatial co-occurrence across distance intervals.
+
+        For each pair of checkpoint genes (g1, g2):
+        - Define high-expression spots: top 25% expression
+        - For distances d = 0, interval, 2*interval, ... n_intervals*interval:
+          compute fraction of g2-high spots within distance d of g1-high spots
+          vs random expectation
+        - co_occurrence_ratio[d] = observed / expected
+
+        If squidpy available, use sq.gr.co_occurrence instead.
+
+        Returns: dict mapping (gene1, gene2) -> np.ndarray of co-occurrence ratios
+        Only compute for the top-20 most variable checkpoint gene pairs
+        (to keep computation tractable).
+        """
+        if len(self.cp_genes) < 2:
+            return {}
+
+        if HAS_SQUIDPY:
+            return self._pairwise_cooccurrence_squidpy(interval, n_intervals)
+        return self._pairwise_cooccurrence_fallback(interval, n_intervals)
+
+    def _pairwise_cooccurrence_squidpy(self, interval: int, n_intervals: int) -> dict:
+        """Co-occurrence via squidpy."""
+        top_genes = self._select_top_variable_genes(n=20)
+        if len(top_genes) < 2:
+            return {}
+
+        # Build interval array: squidpy expects an array of distances
+        distances = np.arange(0, (n_intervals + 1) * interval, interval, dtype=float)
+        distances[0] = 1e-6  # avoid zero
+
+        adata_sub = self.adata[:, top_genes].copy()
+        try:
+            sq.gr.co_occurrence(adata_sub, cluster_key=None, genes=top_genes,
+                                spatial_key='spatial', interval=distances)
+            result = {}
+            co_occ = adata_sub.uns.get('co_occurrence', {})
+            for g1 in top_genes:
+                for g2 in top_genes:
+                    if g1 == g2:
+                        continue
+                    key = (g1, g2)
+                    if key in co_occ:
+                        result[key] = np.array(co_occ[key], dtype=float)
+            return result
+        except Exception:
+            warnings.warn(
+                "squidpy co_occurrence failed; falling back to manual computation.",
+                RuntimeWarning,
+            )
+            return self._pairwise_cooccurrence_fallback(interval, n_intervals)
+
+    def _pairwise_cooccurrence_fallback(self, interval: int, n_intervals: int) -> dict:
+        """Manual co-occurrence computation."""
+        top_genes = self._select_top_variable_genes(n=20)
+        if len(top_genes) < 2:
+            return {}
+
+        coords = self._get_coords()
+        expr = self._get_expr(top_genes)  # (n_spots, n_top)
+        gene_index = {g: i for i, g in enumerate(top_genes)}
+        n_spots = coords.shape[0]
+
+        distances = np.arange(0, (n_intervals + 1) * interval, interval, dtype=float)
+
+        # Build a KDTree for all spots
+        tree = cKDTree(coords)
+
+        # Global fraction of high spots for each gene (baseline / expected)
+        threshold_75 = np.percentile(expr, 75, axis=0)  # (n_top,)
+        high_masks = expr > threshold_75[np.newaxis, :]  # (n_spots, n_top)
+        global_frac = high_masks.mean(axis=0)  # (n_top,)
+
+        result = {}
+        for i, g1 in enumerate(top_genes):
+            g1_high = np.where(high_masks[:, i])[0]
+            if len(g1_high) == 0:
+                continue
+            for j, g2 in enumerate(top_genes):
+                if i == j:
+                    continue
+                g2_high_set = set(np.where(high_masks[:, j])[0])
+                ratios = np.zeros(len(distances))
+                for d_idx, d in enumerate(distances):
+                    if d < 1e-6:
+                        # Distance = 0: only the spot itself
+                        observed = np.mean([1.0 if k in g2_high_set else 0.0
+                                            for k in g1_high])
+                    else:
+                        # Spots within distance d for each g1-high spot
+                        neighbors_list = tree.query_ball_point(coords[g1_high], r=d)
+                        obs_fracs = []
+                        for spot_idx, nbrs in zip(g1_high, neighbors_list):
+                            nbrs_excl = [n for n in nbrs if n != spot_idx]
+                            if len(nbrs_excl) == 0:
+                                obs_fracs.append(0.0)
+                            else:
+                                frac = sum(1 for n in nbrs_excl if n in g2_high_set) / len(nbrs_excl)
+                                obs_fracs.append(frac)
+                        observed = float(np.mean(obs_fracs)) if obs_fracs else 0.0
+
+                    expected = global_frac[j]
+                    ratios[d_idx] = observed / (expected + 1e-9)
+
+                result[(g1, g2)] = ratios
+
+        return result
+
+    def _select_top_variable_genes(self, n: int = 20) -> list:
+        """Select top-n checkpoint genes by expression variance."""
+        if len(self.cp_genes) <= n:
+            return list(self.cp_genes)
+        expr = self._get_expr(self.cp_genes)
+        variances = expr.var(axis=0)
+        top_idx = np.argsort(variances)[::-1][:n]
+        return [self.cp_genes[i] for i in top_idx]
+
+    # ------------------------------------------------------------------
+
+    def checkpoint_immune_proximity(self, n_permutations: int = 200,
+                                    top_pct: float = 0.25) -> pd.DataFrame:
+        """Distance from checkpoint-high spots to nearest immune cell spots.
+
+        For each checkpoint gene:
+        1. Identify checkpoint-high spots (top top_pct expression)
+        2. For each immune cell type (from immune_markers):
+           a. Identify immune spots: ≥1 immune marker expressed
+           b. Use KDTree to compute mean distance: checkpoint-high → nearest immune
+           c. Permutation test (n_permutations): shuffle checkpoint-high labels,
+              compute null distribution of mean distances
+           d. z_score = (observed - null_mean) / null_std
+           e. p_value = fraction of permutations with smaller distance
+
+        Returns DataFrame: columns=[gene, immune_cell, mean_dist, expected_dist,
+                                    z_score, p_value]
+        Negative z_score = closer than random (co-localized).
+        """
+        if not self.cp_genes or not self.immune_markers:
+            return pd.DataFrame(columns=['gene', 'immune_cell', 'mean_dist',
+                                         'expected_dist', 'z_score', 'p_value'])
+
+        coords = self._get_coords()
+        n_spots = coords.shape[0]
+        rng = np.random.default_rng(42)
+
+        # Precompute checkpoint expressions
+        cp_expr = self._get_expr(self.cp_genes)  # (n_spots, n_cp)
+        cp_thresholds = np.quantile(cp_expr, 1.0 - top_pct, axis=0)  # (n_cp,)
+
+        # Precompute immune spot masks
+        immune_spot_masks = {}
+        for cell_type, markers in self.immune_markers.items():
+            avail = [m for m in markers if m in self.adata.var_names]
+            if not avail:
+                continue
+            marker_expr = self._get_expr(avail)  # (n_spots, n_markers)
+            # A spot is "immune" if ≥1 marker expressed (> 0)
+            immune_mask = (marker_expr > 0).any(axis=1)
+            immune_spot_masks[cell_type] = immune_mask
+
+        records = []
+        for g_idx, gene in enumerate(self.cp_genes):
+            cp_high_mask = cp_expr[:, g_idx] > cp_thresholds[g_idx]
+            cp_high_spots = np.where(cp_high_mask)[0]
+            if len(cp_high_spots) == 0:
+                continue
+
+            for cell_type, immune_mask in immune_spot_masks.items():
+                immune_spots = np.where(immune_mask)[0]
+                if len(immune_spots) == 0:
+                    continue
+
+                immune_tree = cKDTree(coords[immune_spots])
+
+                # Observed mean distance: checkpoint-high → nearest immune
+                dists, _ = immune_tree.query(coords[cp_high_spots], k=1)
+                observed_mean = float(np.mean(dists))
+
+                # Permutation test: shuffle which spots are "checkpoint-high"
+                n_cp_high = len(cp_high_spots)
+                null_means = np.empty(n_permutations)
+                for perm_i in range(n_permutations):
+                    perm_indices = rng.choice(n_spots, size=n_cp_high, replace=False)
+                    perm_dists, _ = immune_tree.query(coords[perm_indices], k=1)
+                    null_means[perm_i] = float(np.mean(perm_dists))
+
+                null_mean = float(np.mean(null_means))
+                null_std = float(np.std(null_means))
+                z_score = (observed_mean - null_mean) / (null_std + 1e-9)
+                # p_value = fraction of permutations with smaller (more co-localized) distance
+                p_value = float(np.mean(null_means <= observed_mean))
+
+                records.append({
+                    'gene': gene,
+                    'immune_cell': cell_type,
+                    'mean_dist': observed_mean,
+                    'expected_dist': null_mean,
+                    'z_score': z_score,
+                    'p_value': p_value,
+                })
+
+        return pd.DataFrame(records, columns=['gene', 'immune_cell', 'mean_dist',
+                                               'expected_dist', 'z_score', 'p_value'])
+
+    # ------------------------------------------------------------------
+
+    def ligand_receptor_spatial_analysis(self) -> pd.DataFrame:
+        """Spatial proximity analysis for known checkpoint LR pairs.
+
+        For each (ligand, receptor) pair that has both genes in dataset:
+        1. Get ligand expression per spot
+        2. For each ligand-high spot (top 25%), compute mean receptor expression
+           in k=10 nearest neighbors vs k=10 neighbors of random spots
+        3. Log2 fold change: LR_score = log2(neighbor_receptor / non_neighbor_receptor + 1e-6)
+        4. Permutation test for significance (n=500 permutations)
+
+        Returns DataFrame: columns=[ligand, receptor, alias, lr_score, pvalue,
+                                    ligand_in_dataset, receptor_in_dataset]
+        """
+        var_names_set = set(self.adata.var_names)
+        coords = self._get_coords()
+        n_spots = coords.shape[0]
+        k = min(10, n_spots - 1)
+        n_permutations = 500
+        rng = np.random.default_rng(42)
+
+        # Build alias map for LR pairs
+        alias_map = {
+            ('CD274', 'PDCD1'): 'PD-L1/PD-1',
+            ('PDCD1LG2', 'PDCD1'): 'PD-L2/PD-1',
+            ('LGALS9', 'HAVCR2'): 'Galectin-9/TIM-3',
+            ('PVR', 'TIGIT'): 'CD155/TIGIT',
+            ('NECTIN2', 'TIGIT'): 'CD112/TIGIT',
+            ('FGL1', 'LAG3'): 'FGL1/LAG-3',
+            ('CEACAM1', 'LAG3'): 'CEACAM1/LAG-3',
+            ('CD47', 'SIRPA'): 'CD47/SIRPα',
+        }
+
+        # Build KDTree once for all spots
+        tree = cKDTree(coords)
+
+        records = []
+        for ligand, receptor in self.LR_PAIRS:
+            lig_in = ligand in var_names_set
+            rec_in = receptor in var_names_set
+            alias = alias_map.get((ligand, receptor), f'{ligand}/{receptor}')
+
+            if not (lig_in and rec_in):
+                records.append({
+                    'ligand': ligand,
+                    'receptor': receptor,
+                    'alias': alias,
+                    'lr_score': np.nan,
+                    'pvalue': np.nan,
+                    'ligand_in_dataset': lig_in,
+                    'receptor_in_dataset': rec_in,
+                })
+                continue
+
+            lig_expr = self._get_expr([ligand])[:, 0]   # (n_spots,)
+            rec_expr = self._get_expr([receptor])[:, 0]  # (n_spots,)
+
+            # Ligand-high spots: top 25%
+            lig_threshold = np.percentile(lig_expr, 75)
+            lig_high = np.where(lig_expr > lig_threshold)[0]
+            if len(lig_high) == 0:
+                records.append({
+                    'ligand': ligand, 'receptor': receptor, 'alias': alias,
+                    'lr_score': np.nan, 'pvalue': np.nan,
+                    'ligand_in_dataset': lig_in, 'receptor_in_dataset': rec_in,
+                })
+                continue
+
+            # For each ligand-high spot, find k nearest neighbors
+            _, nbr_indices = tree.query(coords[lig_high], k=k + 1)
+            # Exclude the query spot itself (first column will be the spot itself if it's in tree)
+            # tree.query returns the spot and its neighbors; exclude self
+            nbr_indices = nbr_indices[:, 1:]  # (n_lig_high, k)
+
+            # Mean receptor expression in neighbors of ligand-high spots
+            neighbor_rec = rec_expr[nbr_indices].mean()  # scalar
+
+            # Mean receptor expression in neighbors of random (non-ligand-high) spots
+            non_lig_high = np.where(lig_expr <= lig_threshold)[0]
+            if len(non_lig_high) >= len(lig_high):
+                random_spots = rng.choice(non_lig_high, size=len(lig_high), replace=False)
+            else:
+                random_spots = rng.choice(n_spots, size=len(lig_high), replace=False)
+
+            _, rand_nbr_indices = tree.query(coords[random_spots], k=k + 1)
+            rand_nbr_indices = rand_nbr_indices[:, 1:]
+            random_rec = rec_expr[rand_nbr_indices].mean()
+
+            # LR score: log2 fold change
+            lr_score = float(np.log2((neighbor_rec + 1e-6) / (random_rec + 1e-6)))
+
+            # Permutation test: shuffle ligand-high labels
+            null_scores = np.empty(n_permutations)
+            for perm_i in range(n_permutations):
+                perm_lig_high = rng.choice(n_spots, size=len(lig_high), replace=False)
+                _, perm_nbr_idx = tree.query(coords[perm_lig_high], k=k + 1)
+                perm_nbr_idx = perm_nbr_idx[:, 1:]
+                perm_neighbor_rec = rec_expr[perm_nbr_idx].mean()
+
+                perm_rand = rng.choice(n_spots, size=len(lig_high), replace=False)
+                _, perm_rand_idx = tree.query(coords[perm_rand], k=k + 1)
+                perm_rand_idx = perm_rand_idx[:, 1:]
+                perm_random_rec = rec_expr[perm_rand_idx].mean()
+
+                null_scores[perm_i] = float(
+                    np.log2((perm_neighbor_rec + 1e-6) / (perm_random_rec + 1e-6))
+                )
+
+            # Two-sided p-value: fraction of permutations with |score| >= |observed|
+            pvalue = float(np.mean(np.abs(null_scores) >= np.abs(lr_score)))
+
+            records.append({
+                'ligand': ligand,
+                'receptor': receptor,
+                'alias': alias,
+                'lr_score': lr_score,
+                'pvalue': pvalue,
+                'ligand_in_dataset': lig_in,
+                'receptor_in_dataset': rec_in,
+            })
+
+        return pd.DataFrame(records, columns=['ligand', 'receptor', 'alias',
+                                               'lr_score', 'pvalue',
+                                               'ligand_in_dataset', 'receptor_in_dataset'])
+
+    # ------------------------------------------------------------------
+
+    def spatial_niche_analysis(self, n_niches: int = 5,
+                               k_neighbors: int = 15,
+                               random_state: int = 42) -> pd.DataFrame:
+        """Discover spatial niches based on local checkpoint + immune expression.
+
+        Algorithm:
+        1. For each spot, compute neighborhood profile:
+           - k=15 nearest neighbor mean expression for all cp_genes + immune markers
+        2. Cluster neighborhood profiles using k-means (n_niches clusters)
+           with fixed random_state for reproducibility
+        3. Store niche labels in adata.obs['spatial_niche']
+        4. Characterize each niche: top checkpoint and immune marker genes
+
+        Returns DataFrame: index=niche_id, columns:
+          - top_checkpoints: top 3 checkpoint genes (by mean expression in niche)
+          - top_immune: top 3 immune markers
+          - n_spots: number of spots in niche
+          - pct_tumor_core, pct_stroma, pct_immune_enriched: region composition
+        """
+        from sklearn.cluster import KMeans
+
+        coords = self._get_coords()
+        n_spots = coords.shape[0]
+        k = min(k_neighbors, n_spots - 1)
+
+        # Collect all genes: checkpoint + all immune marker genes
+        all_immune_genes = []
+        for markers in self.immune_markers.values():
+            all_immune_genes.extend(markers)
+        all_immune_genes = list(dict.fromkeys(all_immune_genes))  # unique, order-preserving
+        avail_immune = [g for g in all_immune_genes if g in self.adata.var_names]
+
+        feature_genes = list(self.cp_genes) + avail_immune
+        if not feature_genes:
+            warnings.warn("No feature genes available for niche analysis.", RuntimeWarning)
+            self.adata.obs['spatial_niche'] = 0
+            return pd.DataFrame()
+
+        expr = self._get_expr(feature_genes)  # (n_spots, n_features)
+
+        # Build KDTree and compute neighborhood mean profiles
+        tree = cKDTree(coords)
+        # Query k+1 neighbors (includes the spot itself)
+        _, nbr_indices = tree.query(coords, k=k + 1)  # (n_spots, k+1)
+        # Exclude self (first index is always the query point in cKDTree)
+        nbr_indices = nbr_indices[:, 1:]  # (n_spots, k)
+
+        # Neighborhood mean profile: for each spot, mean expression over k neighbors
+        nbr_profiles = np.zeros((n_spots, len(feature_genes)), dtype=float)
+        for i in range(n_spots):
+            nbr_profiles[i] = expr[nbr_indices[i]].mean(axis=0)
+
+        # K-means clustering on neighborhood profiles
+        n_niches_actual = min(n_niches, n_spots)
+        kmeans = KMeans(n_clusters=n_niches_actual, random_state=random_state, n_init=10)
+        niche_labels = kmeans.fit_predict(nbr_profiles)
+
+        # Store in adata.obs
+        self.adata.obs['spatial_niche'] = niche_labels.astype(str)
+
+        # Characterize each niche
+        n_cp = len(self.cp_genes)
+        cp_feature_indices = list(range(n_cp))
+        immune_feature_indices = list(range(n_cp, len(feature_genes)))
+        immune_feature_names = feature_genes[n_cp:]
+
+        has_region = 'region_type' in self.adata.obs.columns
+        region_labels = self.adata.obs['region_type'].values if has_region else None
+
+        records = []
+        for niche_id in range(n_niches_actual):
+            niche_mask = niche_labels == niche_id
+            n_spots_niche = int(niche_mask.sum())
+            if n_spots_niche == 0:
+                continue
+
+            niche_expr = expr[niche_mask]  # (n_niche_spots, n_features)
+
+            # Top 3 checkpoint genes by mean expression in niche
+            if self.cp_genes:
+                cp_means = niche_expr[:, cp_feature_indices].mean(axis=0)
+                top_cp_idx = np.argsort(cp_means)[::-1][:3]
+                top_checkpoints = ', '.join([self.cp_genes[i] for i in top_cp_idx])
+            else:
+                top_checkpoints = ''
+
+            # Top 3 immune marker genes by mean expression in niche
+            if immune_feature_indices and immune_feature_names:
+                immune_means = niche_expr[:, immune_feature_indices].mean(axis=0)
+                top_imm_idx = np.argsort(immune_means)[::-1][:3]
+                top_immune = ', '.join([immune_feature_names[i] for i in top_imm_idx])
+            else:
+                top_immune = ''
+
+            # Region composition
+            pct_tumor_core = 0.0
+            pct_stroma = 0.0
+            pct_immune_enriched = 0.0
+            if has_region and region_labels is not None:
+                niche_regions = region_labels[niche_mask]
+                pct_tumor_core = float(np.mean(niche_regions == 'tumor_core'))
+                pct_stroma = float(np.mean(niche_regions == 'stroma'))
+                pct_immune_enriched = float(np.mean(niche_regions == 'immune_enriched'))
+
+            records.append({
+                'niche_id': niche_id,
+                'top_checkpoints': top_checkpoints,
+                'top_immune': top_immune,
+                'n_spots': n_spots_niche,
+                'pct_tumor_core': pct_tumor_core,
+                'pct_stroma': pct_stroma,
+                'pct_immune_enriched': pct_immune_enriched,
+            })
+
+        df = pd.DataFrame(records).set_index('niche_id')
+        return df