harshameghadri/medgemma-spatial
0
1"""2Uncertainty-Aware Spatial Analysis Pipeline3Stage 0: Pre-processing with annotation quality validation4Stage 1: Spatial pattern extraction with uncertainty quantification5"""6 7import scanpy as sc8import numpy as np9import pandas as pd10from scipy import stats11from typing import Dict, Tuple, List12import warnings13warnings.filterwarnings('ignore')14 15# ==============================================================================16# STAGE 0: ANNOTATION QUALITY LAYER (Doublet Detection)17# ==============================================================================18 19def detect_doublets_scrublet(adata, expected_doublet_rate=0.06):20 """21 Detect doublets using Scrublet algorithm.22 23 Returns:24 adata: Updated with doublet scores and predictions25 quality_metrics: Dict with doublet statistics26 """27 try:28 import scrublet as scr29 except ImportError:30 print("WARNING: Scrublet not installed. Skipping doublet detection.")31 print("Install with: pip install scrublet")32 adata.obs['doublet_score'] = 0.033 adata.obs['predicted_doublet'] = False34 return adata, {35 'doublet_detection_available': False,36 'n_predicted_doublets': 0,37 'doublet_rate': 0.038 }39 40 scrub = scr.Scrublet(adata.X, expected_doublet_rate=expected_doublet_rate)41 doublet_scores, predicted_doublets = scrub.scrub_doublets(min_counts=2, min_cells=3, min_gene_variability_pctl=85)42 43 adata.obs['doublet_score'] = doublet_scores44 45 # Handle case when Scrublet fails to auto-detect threshold46 if predicted_doublets is None:47 print(" WARNING: Scrublet failed to auto-detect threshold. Using manual threshold at 0.25")48 threshold = 0.2549 predicted_doublets = doublet_scores > threshold50 51 adata.obs['predicted_doublet'] = predicted_doublets52 53 n_doublets = predicted_doublets.sum()54 doublet_rate = n_doublets / len(adata)55 56 quality_metrics = {57 'doublet_detection_available': True,58 'n_predicted_doublets': int(n_doublets),59 'doublet_rate': float(doublet_rate),60 'mean_doublet_score': float(doublet_scores.mean()),61 'max_doublet_score': float(doublet_scores.max())62 }63 64 print(f"Doublet detection: {n_doublets}/{len(adata)} spots ({doublet_rate:.1%})")65 66 return adata, quality_metrics67 68 69def assess_annotation_confidence(adata, confidence_key='conf_score'):70 """71 Assess cell type annotation confidence and flag low-quality annotations.72 73 Returns:74 annotation_quality: Dict with confidence statistics75 """76 if confidence_key not in adata.obs.columns:77 return {78 'annotation_confidence_available': False,79 'mean_confidence': np.nan,80 'low_confidence_rate': np.nan81 }82 83 confidence = adata.obs[confidence_key]84 low_conf_threshold = 0.585 86 annotation_quality = {87 'annotation_confidence_available': True,88 'mean_confidence': float(confidence.mean()),89 'median_confidence': float(confidence.median()),90 'min_confidence': float(confidence.min()),91 'low_confidence_rate': float((confidence < low_conf_threshold).mean()),92 'n_low_confidence': int((confidence < low_conf_threshold).sum())93 }94 95 print(f"Annotation confidence: mean={annotation_quality['mean_confidence']:.3f}, "96 f"low_conf_rate={annotation_quality['low_confidence_rate']:.1%}")97 98 return annotation_quality99 100 101# ==============================================================================102# STAGE 1: UNCERTAINTY-AWARE SPATIAL STATISTICS103# ==============================================================================104 105def compute_morans_i_with_uncertainty(adata, n_genes=100, n_permutations=999):106 """107 Compute Moran's I with permutation-based p-values and confidence intervals.108 109 Returns:110 results: Dict with genes, Moran's I values, p-values, CIs, and signal strength111 """112 from sklearn.preprocessing import StandardScaler113 114 if 'spatial_connectivities' not in adata.obsp:115 print("ERROR: Spatial neighbors not computed. Run sq.gr.spatial_neighbors first.")116 return None117 118 W = adata.obsp['spatial_connectivities']119 120 # Select highly variable genes121 if 'highly_variable' not in adata.var.columns:122 sc.pp.highly_variable_genes(adata, n_top_genes=n_genes, flavor='seurat_v3')123 124 hvg = adata.var_names[adata.var['highly_variable']][:n_genes]125 126 results = {127 'genes': [],128 'morans_i': [],129 'p_value': [],130 'ci_lower': [],131 'ci_upper': [],132 'signal_strength': []133 }134 135 print(f"Computing Moran's I with {n_permutations} permutations for {len(hvg)} genes...")136 137 for gene in hvg:138 X = adata[:, gene].X.toarray().flatten()139 140 # Compute observed Moran's I141 X_std = StandardScaler().fit_transform(X.reshape(-1, 1)).flatten()142 n = len(X_std)143 144 numerator = (W.multiply(np.outer(X_std, X_std))).sum()145 denominator = (X_std ** 2).sum()146 147 I_obs = (n / W.sum()) * (numerator / denominator)148 149 # Permutation test150 I_perm = []151 for _ in range(n_permutations):152 X_perm = np.random.permutation(X_std)153 numerator_perm = (W.multiply(np.outer(X_perm, X_perm))).sum()154 I_perm.append((n / W.sum()) * (numerator_perm / denominator))155 156 I_perm = np.array(I_perm)157 158 # P-value (two-tailed)159 p_val = ((np.abs(I_perm) >= np.abs(I_obs)).sum() + 1) / (n_permutations + 1)160 161 # 95% Confidence interval from permutation distribution162 ci_lower, ci_upper = np.percentile(I_perm, [2.5, 97.5])163 164 # Signal strength classification165 if p_val < 0.001 and np.abs(I_obs) > 0.3:166 signal = "STRONG"167 elif p_val < 0.05 and np.abs(I_obs) > 0.1:168 signal = "MODERATE"169 elif p_val < 0.05:170 signal = "WEAK"171 else:172 signal = "NONE"173 174 results['genes'].append(gene)175 results['morans_i'].append(float(I_obs))176 results['p_value'].append(float(p_val))177 results['ci_lower'].append(float(ci_lower))178 results['ci_upper'].append(float(ci_upper))179 results['signal_strength'].append(signal)180 181 # Sort by signal strength and p-value182 df = pd.DataFrame(results)183 df = df.sort_values(['p_value', 'morans_i'], ascending=[True, False])184 185 return df.to_dict('list')186 187 188def compute_multiscale_neighborhood_enrichment(adata, radii=[1, 2, 3], n_permutations=999):189 """190 Compute neighborhood enrichment at multiple spatial scales.191 192 Returns:193 multiscale_results: Dict with enrichment at each radius + scale stability flags194 """195 if 'cell_type' not in adata.obs.columns:196 print("ERROR: Cell types not annotated. Run CellTypist first.")197 return None198 199 try:200 import squidpy as sq201 squidpy_available = True202 except ImportError:203 print("WARNING: Squidpy not available. Using simplified scanpy-based enrichment.")204 squidpy_available = False205 206 if not squidpy_available:207 # Simplified enrichment using existing spatial graph208 return {209 'radii': radii,210 'enrichment_by_radius': {},211 'scale_stable_pairs': [],212 'scale_dependent_pairs': [],213 'note': 'Squidpy not available - using single-scale analysis'214 }215 216 multiscale_results = {217 'radii': radii,218 'enrichment_by_radius': {},219 'scale_stable_pairs': [],220 'scale_dependent_pairs': []221 }222 223 print(f"Computing neighborhood enrichment at {len(radii)} spatial scales...")224 225 all_z_scores = {}226 227 for radius in radii:228 print(f" Radius={radius}...")229 230 # Rebuild spatial graph with current radius231 sq.gr.spatial_neighbors(adata, coord_type='generic', n_neighs=radius*6, radius=radius*55)232 233 # Compute neighborhood enrichment234 sq.gr.nhood_enrichment(adata, cluster_key='cell_type', n_perms=n_permutations)235 236 z_scores = adata.uns['cell_type_nhood_enrichment']['zscore']237 cell_types = adata.obs['cell_type'].cat.categories.tolist()238 239 all_z_scores[radius] = z_scores240 241 # Extract significant enrichments242 enrichment_pairs = []243 for i, ct1 in enumerate(cell_types):244 for j, ct2 in enumerate(cell_types):245 z = z_scores[i, j]246 if np.abs(z) > 2.0: # |z| > 2 is significant247 enrichment_pairs.append({248 'cell_type_1': ct1,249 'cell_type_2': ct2,250 'z_score': float(z),251 'enriched': z > 0252 })253 254 multiscale_results['enrichment_by_radius'][f'radius_{radius}'] = enrichment_pairs255 256 # Assess scale stability257 if len(radii) >= 2:258 z1 = all_z_scores[radii[0]]259 z2 = all_z_scores[radii[-1]]260 261 correlation = np.corrcoef(z1.flatten(), z2.flatten())[0, 1]262 263 cell_types = adata.obs['cell_type'].cat.categories.tolist()264 265 for i, ct1 in enumerate(cell_types):266 for j, ct2 in enumerate(cell_types):267 # Check if sign and magnitude consistent across scales268 z_small = z1[i, j]269 z_large = z2[i, j]270 271 if np.sign(z_small) == np.sign(z_large) and np.abs(z_small) > 2 and np.abs(z_large) > 2:272 multiscale_results['scale_stable_pairs'].append(f"{ct1}-{ct2}")273 elif np.abs(z_small) > 2 or np.abs(z_large) > 2:274 multiscale_results['scale_dependent_pairs'].append(f"{ct1}-{ct2}")275 276 multiscale_results['scale_stability_correlation'] = float(correlation)277 278 print(f" Scale stability: {len(multiscale_results['scale_stable_pairs'])} stable pairs, "279 f"{len(multiscale_results['scale_dependent_pairs'])} scale-dependent pairs")280 281 return multiscale_results282 283 284def compute_spatial_entropy_with_bootstrapping(adata, n_bootstrap=1000):285 """286 Compute spatial entropy with bootstrap confidence intervals.287 288 Returns:289 entropy_results: Dict with entropy values and uncertainty estimates290 """291 if 'cell_type' not in adata.obs.columns:292 print("ERROR: Cell types not annotated.")293 return None294 295 from scipy.stats import entropy296 297 # Compute entropy for each spot's neighborhood298 spatial_entropy = []299 300 for spot_idx in range(len(adata)):301 # Get neighbors302 neighbors = adata.obsp['spatial_connectivities'][spot_idx].toarray().flatten()303 neighbor_indices = np.where(neighbors > 0)[0]304 305 if len(neighbor_indices) == 0:306 spatial_entropy.append(0.0)307 continue308 309 # Get cell type distribution in neighborhood310 neighbor_types = adata.obs['cell_type'].iloc[neighbor_indices]311 type_counts = neighbor_types.value_counts(normalize=True)312 313 # Shannon entropy314 H = entropy(type_counts.values)315 spatial_entropy.append(H)316 317 spatial_entropy = np.array(spatial_entropy)318 adata.obs['spatial_entropy'] = spatial_entropy319 320 # Bootstrap confidence intervals321 bootstrap_means = []322 n_spots = len(adata)323 324 for _ in range(n_bootstrap):325 sample_indices = np.random.choice(n_spots, size=n_spots, replace=True)326 bootstrap_means.append(spatial_entropy[sample_indices].mean())327 328 ci_lower, ci_upper = np.percentile(bootstrap_means, [2.5, 97.5])329 330 entropy_results = {331 'mean_entropy': float(spatial_entropy.mean()),332 'median_entropy': float(np.median(spatial_entropy)),333 'std_entropy': float(spatial_entropy.std()),334 'ci_lower': float(ci_lower),335 'ci_upper': float(ci_upper),336 'bootstrap_samples': n_bootstrap337 }338 339 print(f"Spatial entropy: {entropy_results['mean_entropy']:.3f} "340 f"(95% CI: [{ci_lower:.3f}, {ci_upper:.3f}])")341 342 return entropy_results343 344 345# ==============================================================================346# STAGE STOPPING LOGIC347# ==============================================================================348 349def assess_signal_quality(morans_results, entropy_results, annotation_quality):350 """351 Determine if signal quality is sufficient to proceed to comparative analysis.352 353 Returns:354 decision: "PROCEED", "STOP_WEAK_SIGNAL", or "STOP_LOW_QUALITY"355 rationale: Explanation for decision356 """357 stop_conditions = []358 359 # Check 1: Annotation quality360 if annotation_quality.get('annotation_confidence_available'):361 if annotation_quality['mean_confidence'] < 0.4:362 stop_conditions.append(f"Low annotation confidence (mean={annotation_quality['mean_confidence']:.2f} < 0.4)")363 if annotation_quality['low_confidence_rate'] > 0.5:364 stop_conditions.append(f"High low-confidence rate ({annotation_quality['low_confidence_rate']:.1%} > 50%)")365 366 # Check 2: Spatial signal strength367 if morans_results:368 strong_signals = sum(1 for s in morans_results['signal_strength'] if s == "STRONG")369 moderate_signals = sum(1 for s in morans_results['signal_strength'] if s == "MODERATE")370 371 if strong_signals == 0 and moderate_signals < 3:372 stop_conditions.append(f"Weak spatial signal (0 STRONG, {moderate_signals} MODERATE genes)")373 374 # Check 3: Spatial heterogeneity375 if entropy_results:376 if entropy_results['mean_entropy'] < 0.2:377 stop_conditions.append(f"Low spatial heterogeneity (entropy={entropy_results['mean_entropy']:.2f} < 0.2)")378 379 # Decision logic380 if len(stop_conditions) >= 2:381 return "STOP_WEAK_SIGNAL", "; ".join(stop_conditions)382 elif len(stop_conditions) == 1 and "Low annotation confidence" in stop_conditions[0]:383 return "STOP_LOW_QUALITY", stop_conditions[0]384 else:385 return "PROCEED", "Signal quality sufficient for comparative analysis"386 387 388# ==============================================================================389# MAIN PIPELINE390# ==============================================================================391 392def run_uncertainty_aware_spatial_analysis(adata_path, output_path):393 """394 Execute full uncertainty-aware spatial analysis pipeline.395 """396 print("=" * 80)397 print("UNCERTAINTY-AWARE SPATIAL ANALYSIS PIPELINE")398 print("=" * 80)399 400 # Load data401 print("\n[1/6] Loading data...")402 adata = sc.read_h5ad(adata_path)403 print(f" Loaded: {adata.n_obs} spots, {adata.n_vars} genes")404 405 # Stage 0: Annotation quality406 print("\n[2/6] STAGE 0: Annotation Quality Assessment")407 adata, doublet_metrics = detect_doublets_scrublet(adata)408 annotation_quality = assess_annotation_confidence(adata)409 410 # Build spatial graph (needed for all subsequent stages)411 print("\n[3/6] Building spatial neighbor graph...")412 try:413 import squidpy as sq414 sq.gr.spatial_neighbors(adata, coord_type='generic', n_neighs=6)415 except ImportError:416 print(" WARNING: Squidpy not available, using scanpy spatial neighbors")417 sc.pp.neighbors(adata, use_rep='spatial')418 419 # Stage 1: Moran's I with uncertainty420 print("\n[4/6] STAGE 1: Spatial Autocorrelation (Moran's I)")421 morans_results = compute_morans_i_with_uncertainty(adata, n_genes=50, n_permutations=999)422 423 # Spatial entropy with uncertainty424 print("\n[5/6] Spatial Entropy (with bootstrap CI)")425 entropy_results = compute_spatial_entropy_with_bootstrapping(adata, n_bootstrap=1000)426 427 # Multi-scale analysis428 print("\n[6/6] Multi-Scale Neighborhood Enrichment")429 multiscale_results = compute_multiscale_neighborhood_enrichment(adata, radii=[1, 2, 3], n_permutations=999)430 431 # Stage stopping decision432 print("\n" + "=" * 80)433 print("STAGE STOPPING ASSESSMENT")434 print("=" * 80)435 decision, rationale = assess_signal_quality(morans_results, entropy_results, annotation_quality)436 437 print(f"\nDecision: {decision}")438 print(f"Rationale: {rationale}")439 440 # Compile results441 results = {442 'stage_0_annotation_quality': {443 'doublet_metrics': doublet_metrics,444 'annotation_confidence': annotation_quality445 },446 'stage_1_spatial_patterns': {447 'morans_i': morans_results,448 'spatial_entropy': entropy_results,449 'multiscale_enrichment': multiscale_results450 },451 'stage_stopping': {452 'decision': decision,453 'rationale': rationale454 }455 }456 457 # Save results458 import json459 with open(output_path, 'w') as f:460 json.dump(results, f, indent=2)461 462 print(f"\nResults saved to: {output_path}")463 print("=" * 80)464 465 return results466 467 468if __name__ == "__main__":469 import sys470 471 if len(sys.argv) < 3:472 print("Usage: python uncertainty_spatial_analysis.py <input.h5ad> <output.json>")473 sys.exit(1)474 475 adata_path = sys.argv[1]476 output_path = sys.argv[2]477 478 run_uncertainty_aware_spatial_analysis(adata_path, output_path)479