""" Computational and Statistical Analysis Tools Tools for data preprocessing, integration, clustering, statistical analysis, and pattern summarization using algorithmic and computational methods. All tools are standalone functions following Biomni pattern. """ # Lightweight imports - keep at module level import os import json import pickle import warnings import numpy as np import pandas as pd from typing import Annotated from os.path import exists from glob import glob from langchain_core.tools import tool from langchain_core.prompts import ChatPromptTemplate from langchain_core.output_parsers import StrOutputParser from pydantic import Field warnings.simplefilter(action="ignore", category=FutureWarning) warnings.filterwarnings('ignore') # Module-level config (set via configure_analytics_tools) _config = { "save_path": "./experiments", } def configure_analytics_tools(save_path: str = "./experiments"): """Configure paths for analytics tools. Call this before using the tools.""" _config["save_path"] = save_path # Default model for subagent LLM calls (fallback if agent model not set) DEFAULT_SUBAGENT_MODEL = "claude-sonnet-4-5-20250929" def _get_subagent_model() -> str: """Get the model to use for subagent calls. Uses the main agent's model if available, otherwise falls back to DEFAULT_SUBAGENT_MODEL. """ try: from ..agent import get_agent_model model = get_agent_model() return model if model else DEFAULT_SUBAGENT_MODEL except ImportError: return DEFAULT_SUBAGENT_MODEL # Heavy imports moved inside tools: # - scanpy (2-3s import time) # - scanpy.external (sce) # - scipy # - sklearn.neural_network.MLPClassifier # - tqdm # ============================================================================= # Tool 1: Preprocess # ============================================================================= @tool def preprocess_spatial_data( adata_path: Annotated[str, Field(description="Path to raw spatial transcriptomics h5ad file")], save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Preprocess spatial transcriptomics data using Scanpy pipeline.""" save_path = save_path or _config["save_path"] # Heavy imports - only load when this tool is called import scanpy as sc output_path = f"{save_path}/preprocessed.h5ad" if exists(output_path): msg = f"Preprocessed data already exists at {output_path}" print(msg) return msg print(f"[preprocess_spatial_data] Loading data from {adata_path}...") # Load and preprocess adata = sc.read_h5ad(adata_path) adata.var.index = adata.var.index.str.upper() print(f"[preprocess_spatial_data] Loaded {adata.n_obs} cells, {adata.n_vars} genes. Running QC...") # QC filtering sc.pp.filter_cells(adata, min_genes=5) sc.pp.filter_cells(adata, min_counts=10) sc.pp.filter_genes(adata, min_cells=5) print(f"[preprocess_spatial_data] After QC: {adata.n_obs} cells, {adata.n_vars} genes. Normalizing...") # Normalization sc.pp.normalize_per_cell(adata) sc.pp.log1p(adata) print("[preprocess_spatial_data] Running PCA and UMAP...") # Dimensionality reduction sc.pp.pca(adata) sc.pp.neighbors(adata) sc.tl.umap(adata) # Save adata.write(output_path, compression="gzip") msg = f"Successfully preprocessed data: {adata.n_obs} cells, {adata.n_vars} genes. Saved to {output_path}" print(msg) return msg # ============================================================================= # Tool 2: Harmony Transfer # ============================================================================= @tool def harmony_transfer_labels( adata_path: Annotated[str, Field(description="Path to preprocessed spatial data")], ref_path: Annotated[str, Field(description="Path to CZI reference scRNA data")], save_path: Annotated[str, Field(description="Directory to save results")] = None, czi_index: Annotated[int, Field(ge=0, description="Index for naming output files")] = 0, ) -> str: """Transfer cell type labels from CZI reference to spatial data using Harmony integration.""" save_path = save_path or _config["save_path"] # Heavy imports - only load when this tool is called import scanpy as sc import scanpy.external as sce from sklearn.neural_network import MLPClassifier output_path = f"{save_path}/celltype-transferred_{czi_index}.h5ad" if exists(output_path): msg = f"Harmony results already exist at {output_path}" print(msg) return msg print(f"[harmony_transfer_labels] Loading spatial and reference data...") # Load data adata_sp = sc.read_h5ad(adata_path) adata_sc = sc.read_h5ad(ref_path) print(f"[harmony_transfer_labels] Spatial: {adata_sp.n_obs} cells, Reference: {adata_sc.n_obs} cells") # Subset reference if larger if adata_sp.shape[0] < adata_sc.shape[0]: # Sample proportionally from each cell type cell_type_counts = adata_sc.obs["cell_type"].value_counts() scale = adata_sp.shape[0] / cell_type_counts.sum() samples_per_type = (cell_type_counts * scale).astype(int) sampled_indices = [] for cell_type, count in samples_per_type.items(): if count > 0: mask = adata_sc.obs["cell_type"] == cell_type indices = adata_sc.obs.index[mask] sampled = np.random.RandomState(42).choice(indices, count, replace=False) sampled_indices.extend(sampled) adata_sc = adata_sc[sampled_indices] # Harmonize gene names adata_sc.var.index = adata_sc.var["feature_name"].str.upper() adata_sc.var_names_make_unique() # Select common genes common_genes = adata_sp.var.index.intersection(adata_sc.var.index) adata_sp, adata_sc = adata_sp[:, common_genes], adata_sc[:, common_genes] # Scale sc.pp.scale(adata_sp) sc.pp.normalize_per_cell(adata_sc) sc.pp.log1p(adata_sc) sc.pp.scale(adata_sc) # Combine and integrate combined = adata_sp.concatenate(adata_sc, batch_key="dataset", batch_categories=["st", "scrna"]) sc.pp.pca(combined, n_comps=30) sce.pp.harmony_integrate(combined, "dataset") # Save combined.write_h5ad(output_path, compression="gzip") # Transfer labels using MLP ad_sp = combined[combined.obs["dataset"] == "st", :] ad_sc = combined[combined.obs["dataset"] == "scrna", :] X_train = ad_sc.obsm["X_pca_harmony"] y_train = ad_sc.obs["cell_type"] X_test = ad_sp.obsm["X_pca_harmony"] clf = MLPClassifier(hidden_layer_sizes=(100, 50), max_iter=500, random_state=42) clf.fit(X_train, y_train) predictions = clf.predict(X_test) # Save predictions (strip batch suffix added by concatenate) cell_ids = ad_sp.obs.index.str.replace(r'-st$', '', regex=True) result_df = pd.DataFrame({"predicted_celltype": predictions}, index=cell_ids) csv_path = f"{save_path}/celltype_transferred.csv" result_df.to_csv(csv_path) msg = f"Successfully transferred labels using Harmony. Predictions for {len(predictions)} cells saved to {csv_path}" print(msg) return msg # ============================================================================= # Tool 3: UTAG (Spatial Clustering) # ============================================================================= def _estimate_max_dist(adata, slide_key=None, target_neighbors=3): """Estimate optimal max_dist for UTAG based on spatial density.""" import numpy as np # Use first sample if slide_key provided if slide_key and slide_key in adata.obs.columns: sample = adata.obs[slide_key].unique()[0] sample_data = adata[adata.obs[slide_key] == sample] else: sample_data = adata # Sample random cells for efficiency n_cells = min(10000, sample_data.shape[0]) idx = np.random.permutation(sample_data.shape[0])[:n_cells] coords = sample_data.obsm["spatial"][idx] # Find distance where each cell has ~target_neighbors neighbors for distance in range(10, 500, 5): distances = np.sqrt(np.sum((coords[:, np.newaxis, :] - coords[np.newaxis, :, :]) ** 2, axis=2)) avg_neighbors = np.mean(np.sum(distances < distance, axis=1) - 1) if avg_neighbors > target_neighbors: return distance return 50 # Default fallback def _remove_small_clusters(adata, label_key, slide_key=None, min_cells=100): """Remove small clusters by assigning cells to nearest larger cluster. IMPORTANT: Filtering is done PER-SAMPLE (not globally) to ensure each sample has no small clusters. A cluster might have many cells globally but only a few in a specific sample, which causes issues in per-sample analysis (e.g., statistical tests like Wilcoxon require >= 2 samples per group). """ from sklearn.neighbors import NearestNeighbors import numpy as np # Convert to string to avoid categorical issues adata.obs[label_key] = adata.obs[label_key].astype(str) batches = adata.obs[slide_key].unique() if slide_key else [None] for batch in batches: if batch is not None: batch_mask = adata.obs[slide_key] == batch batch_data = adata[batch_mask] else: batch_data = adata batch_mask = np.ones(adata.shape[0], dtype=bool) # Identify small clusters WITHIN THIS BATCH (not globally) batch_counts = batch_data.obs[label_key].value_counts() small_clusters_in_batch = list(batch_counts[batch_counts < min_cells].index) if not small_clusters_in_batch: continue cells_to_assign = batch_data.obs[label_key].isin(small_clusters_in_batch) if cells_to_assign.sum() == 0: continue reference_mask = ~batch_data.obs[label_key].isin(small_clusters_in_batch) if reference_mask.sum() == 0: # All clusters are small in this batch - skip (can't reassign) print(f"Warning: All clusters in batch '{batch}' have < {min_cells} cells, skipping merge") continue # Find nearest neighbor from larger clusters nn = NearestNeighbors(n_neighbors=1) nn.fit(batch_data[reference_mask].obsm["spatial"]) _, indices = nn.kneighbors(batch_data[cells_to_assign].obsm["spatial"]) # Assign to nearest larger cluster new_labels = batch_data[reference_mask].obs[label_key].values[indices.flatten()] cell_indices = batch_data[cells_to_assign].obs.index adata.obs.loc[cell_indices, label_key] = new_labels if small_clusters_in_batch: print(f"Batch '{batch}': merged {len(small_clusters_in_batch)} small clusters " f"({cells_to_assign.sum()} cells) to nearest neighbors") return adata @tool def run_utag_clustering( adata_path: Annotated[str, Field(description="Path to annotated spatial data")], save_path: Annotated[str, Field(description="Directory to save results")] = None, slide_key: Annotated[str, Field(description="Column for sample/slide ID to run UTAG per sample (e.g., 'batch', 'sample_id')")] = None, max_dist: Annotated[float, Field(description="Max distance for neighbors. Use 0 for auto-estimation.")] = 0, min_cluster_size: Annotated[int, Field(description="Min cells per cluster (smaller merged to nearest)")] = 100, resolutions: Annotated[list, Field(description="Clustering resolutions to try")] = [0.05, 0.1, 0.3], min_niches: Annotated[int, Field(description="Minimum number of niches required")] = 5, ) -> str: """Run UTAG spatial clustering to identify tissue niches. UTAG combines gene expression with spatial information to identify tissue domains/niches. It uses message passing on a spatial graph. When slide_key is provided, UTAG runs separately on each sample/slide to account for batch effects and sample-specific spatial patterns. Features: - Auto-estimates max_dist if set to 0 - Removes small clusters and reassigns cells to nearest larger cluster - Tries multiple resolutions and picks one with enough niches """ save_path = save_path or _config["save_path"] import scanpy as sc import matplotlib.pyplot as plt from utag import utag output_path = f"{save_path}/utag_main_result.csv" utag_h5ad_path = f"{save_path}/utag_clustered.h5ad" if exists(output_path): msg = f"UTAG results already exist at {output_path}" print(msg) return msg print(f"[run_utag_clustering] Loading data and running UTAG...") # Load data adata = sc.read_h5ad(adata_path) # Validate slide_key if provided if slide_key and slide_key not in adata.obs.columns: msg = f"ERROR: slide_key '{slide_key}' not found. Available: {list(adata.obs.columns)}" print(msg) return msg # Auto-estimate max_dist if not provided if max_dist <= 0: max_dist = _estimate_max_dist(adata, slide_key) print(f"Auto-estimated max_dist: {max_dist}") # Try resolutions until we find one with enough niches utag_results = None best_label_key = None for resolution in resolutions: print(f"Trying UTAG with resolution={resolution}...") utag_results = utag( adata, slide_key=slide_key, max_dist=max_dist, normalization_mode="l1_norm", apply_umap=True, apply_clustering=True, clustering_method=["leiden"], resolutions=[resolution], ) label_key = f"UTAG Label_leiden_{resolution}" n_niches = utag_results.obs[label_key].nunique() print(f" Found {n_niches} niches") if n_niches >= min_niches: best_label_key = label_key break if utag_results is None: msg = "ERROR: UTAG clustering failed" print(msg) return msg if best_label_key is None: # Use last resolution if none met threshold best_label_key = f"UTAG Label_leiden_{resolutions[-1]}" # Remove small clusters utag_results = _remove_small_clusters(utag_results, best_label_key, slide_key, min_cluster_size) # Add unified 'utag' column utag_results.obs["utag"] = utag_results.obs[best_label_key].astype("category") n_final = utag_results.obs["utag"].nunique() # Save results utag_results.write_h5ad(utag_h5ad_path, compression="gzip") pd.DataFrame(utag_results.obs).to_csv(output_path) # Generate per-sample plots if slide_key provided if slide_key: plt.ioff() for sample in utag_results.obs[slide_key].unique(): sample_data = utag_results[utag_results.obs[slide_key] == sample] fig, ax = plt.subplots(figsize=(6, 6)) sc.pl.embedding(sample_data, basis="spatial", color="utag", palette="tab20", size=3, ax=ax, show=False) ax.set_title(f"UTAG Niches - {sample}") plt.savefig(f"{save_path}/utag_niche_{sample}.png", dpi=150, bbox_inches="tight") plt.close() sample_info = f" (per sample: {slide_key})" if slide_key else "" msg = f"Successfully ran UTAG clustering{sample_info}. Found {n_final} niches using {best_label_key}. Saved to {output_path}" print(msg) return msg # ============================================================================= # Tool 4: Gene Voting (LLM-based) # ============================================================================= @tool def aggregate_gene_voting( adata_path: Annotated[str, Field(description="Path to annotated spatial data")], group_by: Annotated[str, Field(description="Column to group by (e.g., 'celltype', 'niche')")], save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Aggregate marker genes across cells/niches using LLM-based voting.""" save_path = save_path or _config["save_path"] # Heavy imports import scanpy as sc from ..agent import make_llm # Create LLM instance llm = make_llm(_get_subagent_model(), stop_sequences=[]) output_path = f"{save_path}/gene_voting_results.csv" if exists(output_path): msg = f"Gene voting results already exist at {output_path}" print(msg) return msg print(f"[aggregate_gene_voting] Loading data and aggregating genes by {group_by}...") # Load data adata = sc.read_h5ad(adata_path) # Group by specified column and aggregate results = [] for group_name in adata.obs[group_by].unique(): mask = adata.obs[group_by] == group_name subset = adata[mask] # Get top marker genes sc.tl.rank_genes_groups(subset, groupby=group_by) markers = subset.uns['rank_genes_groups']['names'][:20] results.append({ "group": group_name, "marker_genes": markers.tolist() }) # Save pd.DataFrame(results).to_csv(output_path) msg = f"Successfully aggregated genes for {len(results)} groups. Saved to {output_path}" print(msg) return msg # ============================================================================= # Tool 5: Cell-Cell Interactions (Full LIANA + Cell2Cell + TensorLy) # ============================================================================= @tool def liana_tensor( adata_path: Annotated[str, Field(description="Path to AnnData h5ad file")], sample_key: Annotated[str, Field(description="Sample/batch column name")], condition_key: Annotated[str, Field(description="Condition column name")], cell_type_key: Annotated[str, Field(description="Cell type column name")], organism: Annotated[str, Field(description="'human', 'mouse', or 'auto'")] = "auto", save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Run LIANA + Tensor-Cell2Cell for multi-sample interaction analysis. Performs tensor decomposition to identify context-specific communication patterns: 1. LIANA rank_aggregate per sample 2. Build interaction tensor (samples x interactions x senders x receivers) 3. Non-negative tensor factorization into latent factors """ save_path = save_path or _config["save_path"] # Heavy imports - only load when this tool is called import scanpy as sc from tqdm.auto import tqdm import liana as li output_path = f"{save_path}/cci_analysis/factor.pkl" if exists(output_path): msg = f"CCI analysis results already exist at {output_path}" print(msg) return msg print(f"[liana_tensor] Running LIANA + Tensor-Cell2Cell analysis...") # Create output directory os.makedirs(f"{save_path}/cci_analysis", exist_ok=True) # Load data adata = sc.read_h5ad(adata_path) # Auto-detect organism based on gene name format if organism == "auto": # Mouse genes: first letter uppercase, rest lowercase (e.g., Actb) # Human genes: all uppercase (e.g., ACTB) sample_genes = adata.var_names[:100].tolist() n_mouse_format = sum(1 for g in sample_genes if g[0].isupper() and len(g) > 1 and g[1:].islower()) organism = "mouse" if n_mouse_format > 50 else "human" print(f"Auto-detected organism: {organism}") # Select resource based on organism resource_name = "mouseconsensus" if organism == "mouse" else "consensus" # Step 1: Run LIANA per sample to get ligand-receptor interactions print("Running LIANA per sample...") li.mt.rank_aggregate.by_sample( adata, sample_key=sample_key, groupby=cell_type_key, resource_name=resource_name, use_raw=False, verbose=True, key_added='liana_res' ) # Get unique LR pairs liana_res = adata.uns['liana_res'] lr_pairs = liana_res[['ligand_complex', 'receptor_complex']].drop_duplicates() print(f"Found {len(lr_pairs)} unique ligand-receptor pairs") # Step 2: Build tensor using LIANA's built-in function print("Building interaction tensor...") tensor = li.multi.to_tensor_c2c( adata, sample_key=sample_key, score_key='magnitude_rank', non_negative=True ) # Step 3: Tensor factorization using PreBuiltTensor's built-in method print("Running tensor factorization...") try: tensor.compute_tensor_factorization( rank=3, init='random', random_state=42 ) except Exception as e: # If factorization fails, return partial results print(f"Warning: Tensor factorization failed ({e}). Returning LIANA results only.") liana_res.to_csv(f"{save_path}/cci_analysis/liana_results.csv", index=False) msg = f"LIANA analysis complete for {len(adata.obs[sample_key].unique())} samples, {len(lr_pairs)} LR pairs. Factorization skipped (insufficient data). Saved to {save_path}/cci_analysis/" print(msg) return msg # Get context (condition) mapping samples = adata.obs[sample_key].unique() context_dict = dict(zip(samples, adata.obs.groupby(sample_key)[condition_key].first())) # Save results results = { 'tensor': tensor, # Contains factorization results 'context_dict': context_dict, 'liana_res': liana_res, 'lr_pairs': lr_pairs.values.tolist(), 'samples': list(samples) } with open(output_path, 'wb') as f: pickle.dump(results, f) # Also save LIANA results as CSV liana_res.to_csv(f"{save_path}/cci_analysis/liana_results.csv", index=False) msg = f"Successfully computed cell-cell interactions for {len(samples)} samples. Found {len(lr_pairs)} LR pairs. Saved to {output_path}" print(msg) return msg # ============================================================================= # Tool 6: Dynamics (DEG Analysis) # ============================================================================= @tool def infer_dynamics( adata_path: Annotated[str, Field(description="Path to spatial data with condition labels")], condition_column: Annotated[str, Field(description="Column containing condition labels")], condition1: Annotated[str, Field(description="First condition (e.g., 'control')")], condition2: Annotated[str, Field(description="Second condition (e.g., 'disease')")], save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Compare conditions with differential expression gene (DEG) analysis.""" save_path = save_path or _config["save_path"] # Heavy imports import scanpy as sc output_path = f"{save_path}/deg_results.csv" if exists(output_path): msg = f"DEG results already exist at {output_path}" print(msg) return msg print(f"[infer_dynamics] Comparing {condition1} vs {condition2}...") # Load data adata = sc.read_h5ad(adata_path) # Filter to specified conditions mask = adata.obs[condition_column].isin([condition1, condition2]) adata_subset = adata[mask] # Run differential expression sc.tl.rank_genes_groups( adata_subset, groupby=condition_column, groups=[condition2], reference=condition1, method='wilcoxon' ) # Extract results result = sc.get.rank_genes_groups_df(adata_subset, group=condition2) result.to_csv(output_path) msg = f"Successfully analyzed dynamics: {condition1} vs {condition2}. Found {len(result)} DEGs. Saved to {output_path}" print(msg) return msg # ============================================================================= # Tool 7: Summarize Conditions # ============================================================================= @tool def summarize_conditions( adata_path: Annotated[str, Field(description="Path to spatial data with condition labels")], condition_key: Annotated[str, Field(description="Column name for condition labels")], cell_type_key: Annotated[str, Field(description="Column name for cell type annotations")] = "cell_type", save_path: Annotated[str, Field(description="Directory to save summary")] = None, ) -> str: """Summarize cell type distributions across different conditions.""" save_path = save_path or _config["save_path"] # Heavy imports import scanpy as sc output_path = f"{save_path}/condition_summary.txt" if exists(output_path): msg = f"Condition summary already exists at {output_path}" print(msg) return msg print(f"[summarize_conditions] Analyzing conditions by {condition_key}...") # Load data adata = sc.read_h5ad(adata_path) if condition_key not in adata.obs.columns: msg = f"Error: Condition column '{condition_key}' not found. Available columns: {list(adata.obs.columns)}" print(msg) return msg if cell_type_key not in adata.obs.columns: msg = f"Error: Cell type column '{cell_type_key}' not found. Available columns: {list(adata.obs.columns)}" print(msg) return msg # Get condition distribution condition_counts = adata.obs[condition_key].value_counts() total_cells = len(adata) # Build summary lines = [] lines.append("# Condition Summary") lines.append(f"\nTotal cells: {total_cells}") lines.append(f"Conditions: {len(condition_counts)}\n") lines.append("## Condition Overview\n") lines.append("| Condition | Cells | % Total |") lines.append("|-----------|-------|---------|") for condition in condition_counts.index: count = condition_counts[condition] pct = 100 * count / total_cells lines.append(f"| {condition} | {count} | {pct:.1f}% |") # Cell type distribution per condition lines.append("\n## Cell Type Distribution by Condition\n") for condition in condition_counts.index: mask = adata.obs[condition_key] == condition subset = adata[mask] celltype_counts = subset.obs[cell_type_key].value_counts() lines.append(f"### {condition} ({len(subset)} cells)\n") for ct, count in celltype_counts.head(5).items(): ct_pct = 100 * count / len(subset) lines.append(f"- {ct}: {count} ({ct_pct:.1f}%)") lines.append("") # Save summary_text = "\n".join(lines) with open(output_path, 'w') as f: f.write(summary_text) # Also save cross-tabulation as CSV csv_path = f"{save_path}/condition_celltype_matrix.csv" crosstab = pd.crosstab(adata.obs[condition_key], adata.obs[cell_type_key]) crosstab.to_csv(csv_path) msg = f"Condition summary:\n{summary_text}\n\nSaved to {output_path} and {csv_path}" print(msg) return msg # ============================================================================= # Tool 8: Summarize Cell Types (LLM-based) # ============================================================================= @tool def summarize_celltypes( adata_path: Annotated[str, Field(description="Path to annotated spatial data")], cell_type_key: Annotated[str, Field(description="Column name for cell type annotations")] = "cell_type", save_path: Annotated[str, Field(description="Directory to save summary")] = None, ) -> str: """Summarize cell type distributions and marker genes in the dataset.""" save_path = save_path or _config["save_path"] # Heavy imports import scanpy as sc output_path = f"{save_path}/celltype_summary.txt" if exists(output_path): msg = f"Cell type summary already exists at {output_path}" print(msg) return msg print(f"[summarize_celltypes] Analyzing cell types...") # Load data adata = sc.read_h5ad(adata_path) if cell_type_key not in adata.obs.columns: msg = f"Error: Column '{cell_type_key}' not found. Available columns: {list(adata.obs.columns)}" print(msg) return msg # Get cell type distribution celltype_counts = adata.obs[cell_type_key].value_counts() total_cells = len(adata) # Run differential expression to find markers for each cell type print(f"Running marker gene analysis for {len(celltype_counts)} cell types...") sc.tl.rank_genes_groups(adata, groupby=cell_type_key, method='wilcoxon') # Build summary lines = [] lines.append("# Cell Type Summary") lines.append(f"\nTotal cells: {total_cells}") lines.append(f"Cell types: {len(celltype_counts)}\n") lines.append("## Cell Type Distribution\n") lines.append("| Cell Type | Count | Percentage | Top Markers |") lines.append("|-----------|-------|------------|-------------|") for celltype in celltype_counts.index: count = celltype_counts[celltype] pct = 100 * count / total_cells # Get marker genes for this cell type try: markers = adata.uns['rank_genes_groups']['names'][celltype][:5].tolist() marker_str = ", ".join(markers) except Exception: marker_str = "N/A" lines.append(f"| {celltype} | {count} | {pct:.1f}% | {marker_str} |") # Save summary_text = "\n".join(lines) with open(output_path, 'w') as f: f.write(summary_text) # Also save as CSV for easy access csv_path = f"{save_path}/celltype_summary.csv" summary_df = pd.DataFrame({ 'cell_type': celltype_counts.index, 'count': celltype_counts.values, 'percentage': [100 * c / total_cells for c in celltype_counts.values] }) summary_df.to_csv(csv_path, index=False) msg = f"Cell type summary:\n{summary_text}\n\nSaved to {output_path} and {csv_path}" print(msg) return msg # ============================================================================= # Tool 9: Summarize Tissue Regions (LLM-based) # ============================================================================= @tool def summarize_tissue_regions( adata_path: Annotated[str, Field(description="Path to spatial data with region annotations")], region_key: Annotated[str, Field(description="Column name for region/niche annotations")] = "spatial_cluster", cell_type_key: Annotated[str, Field(description="Column name for cell type annotations")] = "cell_type", save_path: Annotated[str, Field(description="Directory to save summary")] = None, ) -> str: """Summarize tissue regions and their cell type compositions.""" save_path = save_path or _config["save_path"] # Heavy imports import scanpy as sc output_path = f"{save_path}/tissue_region_summary.txt" if exists(output_path): msg = f"Tissue region summary already exists at {output_path}" print(msg) return msg print(f"[summarize_tissue_regions] Analyzing tissue regions by {region_key}...") # Load data adata = sc.read_h5ad(adata_path) if region_key not in adata.obs.columns: msg = f"Error: Region column '{region_key}' not found. Available columns: {list(adata.obs.columns)}" print(msg) return msg if cell_type_key not in adata.obs.columns: msg = f"Error: Cell type column '{cell_type_key}' not found. Available columns: {list(adata.obs.columns)}" print(msg) return msg # Get region counts region_counts = adata.obs[region_key].value_counts() total_cells = len(adata) # Build summary lines = [] lines.append("# Tissue Region Summary") lines.append(f"\nTotal cells: {total_cells}") lines.append(f"Regions: {len(region_counts)}\n") # Build composition matrix composition_data = [] for region in region_counts.index: mask = adata.obs[region_key] == region subset = adata[mask] region_pct = 100 * len(subset) / total_cells # Get cell type composition celltype_counts = subset.obs[cell_type_key].value_counts() top_celltypes = [] for ct, count in celltype_counts.head(3).items(): ct_pct = 100 * count / len(subset) top_celltypes.append(f"{ct} ({ct_pct:.0f}%)") composition_data.append({ 'region': region, 'count': len(subset), 'percentage': region_pct, 'top_celltypes': ", ".join(top_celltypes) }) lines.append("## Region Overview\n") lines.append("| Region | Cells | % Total | Top Cell Types |") lines.append("|--------|-------|---------|----------------|") for row in composition_data: lines.append(f"| {row['region']} | {row['count']} | {row['percentage']:.1f}% | {row['top_celltypes']} |") # Save summary_text = "\n".join(lines) with open(output_path, 'w') as f: f.write(summary_text) # Also save composition matrix as CSV csv_path = f"{save_path}/tissue_region_composition.csv" # Create cross-tabulation crosstab = pd.crosstab(adata.obs[region_key], adata.obs[cell_type_key], normalize='index') * 100 crosstab.to_csv(csv_path) msg = f"Tissue region summary:\n{summary_text}\n\nSaved to {output_path} and {csv_path}" print(msg) return msg # ============================================================================= # Tangram Tools for Spatial Mapping # ============================================================================= @tool def tangram_preprocess( adata_sc_path: Annotated[str, Field(description="Path to single-cell RNA-seq data (h5ad)")], adata_sp_path: Annotated[str, Field(description="Path to spatial transcriptomics data (h5ad)")], marker_genes: Annotated[str, Field(description="Comma-separated genes, or 'auto' to compute from scRNA-seq")] = "auto", cell_type_key: Annotated[str, Field(description="Cell type column for auto marker detection")] = "cell_type", n_markers: Annotated[int, Field(description="Top markers per cell type if auto")] = 100, save_path: Annotated[str, Field(description="Directory to save preprocessed data")] = None, ) -> str: """Preprocess scRNA-seq and spatial data for Tangram mapping. Finds shared genes, removes zero-valued genes, and computes density priors. Must run before tangram_map_cells. """ save_path = save_path or _config["save_path"] import scanpy as sc import tangram as tg os.makedirs(save_path, exist_ok=True) sc_out = f"{save_path}/tangram_sc_prep.h5ad" sp_out = f"{save_path}/tangram_sp_prep.h5ad" adata_sc = sc.read_h5ad(adata_sc_path) adata_sp = sc.read_h5ad(adata_sp_path) if marker_genes == "auto": sc.tl.rank_genes_groups(adata_sc, groupby=cell_type_key, use_raw=False) markers_df = pd.DataFrame(adata_sc.uns["rank_genes_groups"]["names"]).iloc[:n_markers, :] markers = list(np.unique(markers_df.melt().value.values)) else: from .utils import parse_list_string markers = parse_list_string(marker_genes) tg.pp_adatas(adata_sc, adata_sp, genes=markers) n_train = len(adata_sc.uns.get("training_genes", [])) adata_sc.write_h5ad(sc_out, compression="gzip") adata_sp.write_h5ad(sp_out, compression="gzip") msg = f"Preprocessed for Tangram: {n_train} training genes. Saved to {sc_out} and {sp_out}" print(msg) return msg @tool def tangram_map_cells( adata_sc_path: Annotated[str, Field(description="Path to preprocessed scRNA-seq (from tangram_preprocess)")], adata_sp_path: Annotated[str, Field(description="Path to preprocessed spatial data")], mode: Annotated[str, Field(description="'cells' (single-cell, GPU recommended) or 'clusters' (faster)")] = "cells", cluster_label: Annotated[str, Field(description="Cell type column for 'clusters' mode")] = "cell_type", device: Annotated[str, Field(description="'cpu' or 'cuda:0'")] = "cpu", num_epochs: Annotated[int, Field(description="Training epochs")] = 500, save_path: Annotated[str, Field(description="Directory to save mapping")] = None, ) -> str: """Map single cells onto spatial locations using Tangram. Creates cell-by-spot probability matrix by optimizing gene expression similarity. """ save_path = save_path or _config["save_path"] import scanpy as sc import tangram as tg out_path = f"{save_path}/tangram_mapping.h5ad" if exists(out_path): msg = f"Mapping exists at {out_path}" print(msg) return msg print(f"[tangram_map_cells] Running Tangram mapping in {mode} mode...") adata_sc = sc.read_h5ad(adata_sc_path) adata_sp = sc.read_h5ad(adata_sp_path) if mode == "clusters": ad_map = tg.map_cells_to_space(adata_sc, adata_sp, mode="clusters", cluster_label=cluster_label, density_prior='rna_count_based', num_epochs=num_epochs, device=device) else: ad_map = tg.map_cells_to_space(adata_sc, adata_sp, mode="cells", density_prior='rna_count_based', num_epochs=num_epochs, device=device) ad_map.write_h5ad(out_path, compression="gzip") score = ad_map.uns.get("train_genes_df", pd.DataFrame())["train_score"].mean() if "train_genes_df" in ad_map.uns else 0 msg = f"Mapped {ad_map.shape[0]} cells to {ad_map.shape[1]} spots. Avg score: {score:.3f}. Saved to {out_path}" print(msg) return msg @tool def tangram_project_annotations( adata_map_path: Annotated[str, Field(description="Path to Tangram mapping")], adata_sp_path: Annotated[str, Field(description="Path to spatial data")], annotation: Annotated[str, Field(description="Annotation to project (e.g., 'cell_type')")] = "cell_type", save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Project cell type annotations onto spatial data. Transfers cell type probabilities to spatial spots based on mapping. """ save_path = save_path or _config["save_path"] import scanpy as sc import tangram as tg ad_map = sc.read_h5ad(adata_map_path) adata_sp = sc.read_h5ad(adata_sp_path) tg.project_cell_annotations(ad_map, adata_sp, annotation=annotation) ct_pred = adata_sp.obsm.get("tangram_ct_pred") if ct_pred is not None: ct_pred.to_csv(f"{save_path}/tangram_celltype_probs.csv") adata_sp.obs["tangram_celltype"] = ct_pred.idxmax(axis=1) adata_sp.write_h5ad(f"{save_path}/tangram_annotated.h5ad", compression="gzip") n_types = ct_pred.shape[1] if ct_pred is not None else 0 msg = f"Projected {n_types} cell types. Saved to {save_path}/tangram_annotated.h5ad" print(msg) return msg @tool def tangram_project_genes( adata_map_path: Annotated[str, Field(description="Path to Tangram mapping")], adata_sc_path: Annotated[str, Field(description="Path to scRNA-seq data")], cluster_label: Annotated[str, Field(description="Cluster label if 'clusters' mode used")] = "", save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Project gene expression from scRNA-seq onto spatial locations. Creates imputed spatial gene expression for the entire transcriptome. """ save_path = save_path or _config["save_path"] import scanpy as sc import tangram as tg ad_map = sc.read_h5ad(adata_map_path) adata_sc = sc.read_h5ad(adata_sc_path) if cluster_label and cluster_label in adata_sc.obs.columns: ad_ge = tg.project_genes(ad_map, adata_sc, cluster_label=cluster_label) else: ad_ge = tg.project_genes(ad_map, adata_sc) out_path = f"{save_path}/tangram_projected.h5ad" ad_ge.write_h5ad(out_path, compression="gzip") msg = f"Projected {ad_ge.shape[1]} genes to {ad_ge.shape[0]} spots. Saved to {out_path}" print(msg) return msg @tool def tangram_evaluate( adata_ge_path: Annotated[str, Field(description="Path to projected genes (from tangram_project_genes)")], adata_sp_path: Annotated[str, Field(description="Path to spatial data (original or preprocessed)")], save_path: Annotated[str, Field(description="Directory to save evaluation")] = None, ) -> str: """Evaluate Tangram mapping quality. Computes per-gene cosine similarity between predicted and measured expression. """ save_path = save_path or _config["save_path"] import scanpy as sc import matplotlib.pyplot as plt from scipy import sparse os.makedirs(save_path, exist_ok=True) ad_ge = sc.read_h5ad(adata_ge_path) adata_sp = sc.read_h5ad(adata_sp_path) # Find overlapping genes (case-insensitive) ge_genes = set(g.lower() for g in ad_ge.var_names) sp_genes = set(g.lower() for g in adata_sp.var_names) overlap = sorted(ge_genes & sp_genes) if len(overlap) < 5: msg = f"ERROR: Only {len(overlap)} overlapping genes. Cannot evaluate." print(msg) return msg # Get expression matrices for overlapping genes ge_idx = [i for i, g in enumerate(ad_ge.var_names) if g.lower() in overlap] sp_idx = [i for i, g in enumerate(adata_sp.var_names) if g.lower() in overlap] X_ge = ad_ge.X[:, ge_idx] X_sp = adata_sp.X[:, sp_idx] if sparse.issparse(X_ge): X_ge = X_ge.toarray() if sparse.issparse(X_sp): X_sp = X_sp.toarray() # Compute per-gene cosine similarity scores = [] gene_names = [ad_ge.var_names[i] for i in ge_idx] for i in range(len(overlap)): v1, v2 = X_ge[:, i], X_sp[:, i] norm = np.linalg.norm(v1) * np.linalg.norm(v2) score = (v1 @ v2) / norm if norm > 0 else 0 scores.append(score) df = pd.DataFrame({'gene': gene_names, 'score': scores}) df = df.sort_values('score', ascending=False) df.to_csv(f"{save_path}/tangram_scores.csv", index=False) avg_score = np.nanmean(scores) median_score = np.nanmedian(scores) # Plot score distribution plt.ioff() fig, axes = plt.subplots(1, 2, figsize=(10, 4)) axes[0].hist(scores, bins=20, color='coral', alpha=0.7, edgecolor='black') axes[0].axvline(avg_score, color='red', linestyle='--', label=f'Mean: {avg_score:.3f}') axes[0].set_xlabel('Cosine Similarity') axes[0].set_ylabel('Count') axes[0].set_title('Gene Score Distribution') axes[0].legend() axes[1].scatter(range(len(scores)), sorted(scores, reverse=True), s=10, alpha=0.6) axes[1].set_xlabel('Gene Rank') axes[1].set_ylabel('Score') axes[1].set_title('Ranked Gene Scores') plt.tight_layout() plt.savefig(f"{save_path}/tangram_scores.png", dpi=150, bbox_inches="tight") plt.close() top5 = ', '.join(df.head(5)['gene'].tolist()) msg = f"Evaluation: Mean={avg_score:.3f}, Median={median_score:.3f} ({len(overlap)} genes). Top: {top5}. Saved to {save_path}/tangram_scores.csv" print(msg) return msg # ============================================================================= # CellPhoneDB Tools - Cell-cell communication analysis # ============================================================================= @tool def cellphonedb_prepare( adata_path: Annotated[str, Field(description="Path to AnnData h5ad file")], cell_type_key: Annotated[str, Field(description="Column in obs for cell type annotations")] = "cell_type", layer: Annotated[str, Field(description="Layer to use for counts (empty for .X)")] = "", save_path: Annotated[str, Field(description="Directory to save prepared files")] = None, ) -> str: """Prepare AnnData for CellPhoneDB analysis. Extracts counts matrix and metadata files required by CellPhoneDB methods. Converts gene symbols to human format if needed. """ save_path = save_path or _config["save_path"] import scanpy as sc os.makedirs(save_path, exist_ok=True) adata = sc.read_h5ad(adata_path) # Extract metadata meta = pd.DataFrame({ 'Cell': adata.obs_names, 'cell_type': adata.obs[cell_type_key].values }) meta_path = f"{save_path}/cellphonedb_meta.txt" meta.to_csv(meta_path, sep='\t', index=False) # Extract counts - CellPhoneDB expects normalized counts counts_path = f"{save_path}/cellphonedb_counts.h5ad" if layer and layer in adata.layers: adata_out = adata.copy() adata_out.X = adata.layers[layer] else: adata_out = adata adata_out.write_h5ad(counts_path, compression="gzip") n_cells = adata.n_obs n_genes = adata.n_vars n_types = adata.obs[cell_type_key].nunique() msg = f"Prepared CellPhoneDB inputs: {n_cells} cells, {n_genes} genes, {n_types} cell types. Saved to {meta_path} and {counts_path}" print(msg) return msg @tool def cellphonedb_analysis( counts_path: Annotated[str, Field(description="Path to counts h5ad file")], meta_path: Annotated[str, Field(description="Path to metadata txt file")], iterations: Annotated[int, Field(description="Permutation iterations (0 for simple/no-stats method)")] = 1000, threshold: Annotated[float, Field(description="Min fraction of cells expressing gene")] = 0.1, microenvs_path: Annotated[str, Field(description="Path to microenvironments file (optional, restricts to colocalized cells)")] = "", score_interactions: Annotated[bool, Field(description="Score interactions by specificity (0-10 scale)")] = False, threads: Annotated[int, Field(description="Number of threads")] = 4, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Run CellPhoneDB analysis (statistical or simple). METHOD 2 (Statistical): Permutation-based p-values for enriched interactions. Set iterations=0 for METHOD 1 (Simple): Mean expression only, no statistics. Outputs: means.txt, pvalues.txt, significant_means.txt, deconvoluted.txt """ save_path = save_path or _config["save_path"] os.makedirs(save_path, exist_ok=True) # Get CellPhoneDB database path cpdb_dir = os.path.expanduser("~/.cpdb") cpdb_file_path = os.path.join(cpdb_dir, "cellphonedb.zip") # Prepare optional parameters microenvs = microenvs_path if microenvs_path and os.path.exists(microenvs_path) else None if iterations == 0: # METHOD 1: Simple analysis (no statistics) from cellphonedb.src.core.methods import cpdb_analysis_method out_path = f"{save_path}/cellphonedb_simple" os.makedirs(out_path, exist_ok=True) results = cpdb_analysis_method.call( cpdb_file_path=cpdb_file_path, meta_file_path=meta_path, counts_file_path=counts_path, counts_data='hgnc_symbol', output_path=out_path, threshold=threshold, microenvs_file_path=microenvs, score_interactions=score_interactions, ) method_name = "simple" else: # METHOD 2: Statistical analysis from cellphonedb.src.core.methods import cpdb_statistical_analysis_method out_path = f"{save_path}/cellphonedb_statistical" os.makedirs(out_path, exist_ok=True) results = cpdb_statistical_analysis_method.call( cpdb_file_path=cpdb_file_path, meta_file_path=meta_path, counts_file_path=counts_path, counts_data='hgnc_symbol', output_path=out_path, threshold=threshold, iterations=iterations, threads=threads, microenvs_file_path=microenvs, score_interactions=score_interactions, ) method_name = "statistical" # Count results means = results.get('means', pd.DataFrame()) pvals = results.get('pvalues', pd.DataFrame()) sig_means = results.get('significant_means', pd.DataFrame()) n_interactions = len(means) if len(means) > 0 else 0 n_significant = (sig_means.notna().sum(axis=1) > 0).sum() if len(sig_means) > 0 else 0 msg = f"CellPhoneDB {method_name} analysis complete: {n_interactions} interactions tested, {n_significant} significant. Results saved to {out_path}" print(msg) return msg @tool def cellphonedb_degs_analysis( counts_path: Annotated[str, Field(description="Path to counts h5ad file")], meta_path: Annotated[str, Field(description="Path to metadata txt file")], degs_path: Annotated[str, Field(description="Path to DEGs file (two columns: Cell, Gene)")], threshold: Annotated[float, Field(description="Min fraction of cells expressing gene")] = 0.1, microenvs_path: Annotated[str, Field(description="Path to microenvironments file (optional, restricts to colocalized cells)")] = "", score_interactions: Annotated[bool, Field(description="Score interactions by specificity (0-10 scale)")] = False, threads: Annotated[int, Field(description="Number of threads")] = 4, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Run CellPhoneDB DEG-based analysis (METHOD 3). Retrieves interactions where at least one gene is differentially expressed. More targeted than statistical method for specific comparisons. DEGs file format: Two columns - 'Cell' (cell type) and 'Gene' (gene symbol). Only interactions with at least one DEG are returned as relevant. """ save_path = save_path or _config["save_path"] from cellphonedb.src.core.methods import cpdb_degs_analysis_method os.makedirs(save_path, exist_ok=True) out_path = f"{save_path}/cellphonedb_degs" os.makedirs(out_path, exist_ok=True) # Get CellPhoneDB database path cpdb_dir = os.path.expanduser("~/.cpdb") cpdb_file_path = os.path.join(cpdb_dir, "cellphonedb.zip") # Prepare optional parameters microenvs = microenvs_path if microenvs_path and os.path.exists(microenvs_path) else None results = cpdb_degs_analysis_method.call( cpdb_file_path=cpdb_file_path, meta_file_path=meta_path, counts_file_path=counts_path, degs_file_path=degs_path, counts_data='hgnc_symbol', output_path=out_path, threshold=threshold, microenvs_file_path=microenvs, score_interactions=score_interactions, threads=threads, ) relevant = results.get('relevant_interactions', pd.DataFrame()) sig_means = results.get('significant_means', pd.DataFrame()) n_relevant = (relevant == 1).sum().sum() if len(relevant) > 0 else 0 n_interactions = len(sig_means) if len(sig_means) > 0 else 0 msg = f"CellPhoneDB DEG analysis complete (METHOD 3): {n_relevant} relevant interactions from {n_interactions} tested. Results saved to {out_path}" print(msg) return msg @tool def cellphonedb_filter( results_path: Annotated[str, Field(description="Path to CellPhoneDB results directory")], cell_types: Annotated[str, Field(description="Comma-separated cell types to filter (e.g., 'T cell,B cell')")] = "", genes: Annotated[str, Field(description="Comma-separated genes to filter")] = "", min_mean: Annotated[float, Field(description="Minimum mean expression threshold")] = 0.0, save_path: Annotated[str, Field(description="Directory to save filtered results")] = None, ) -> str: """Search and filter CellPhoneDB results. Filters significant interactions by cell types, genes, or expression levels. """ save_path = save_path or _config["save_path"] os.makedirs(save_path, exist_ok=True) # Load results sig_means_path = f"{results_path}/significant_means.csv" if not exists(sig_means_path): sig_means_path = f"{results_path}/significant_means.txt" if not exists(sig_means_path): msg = f"ERROR: No significant_means file found in {results_path}" print(msg) return msg sig_means = pd.read_csv(sig_means_path, sep='\t' if sig_means_path.endswith('.txt') else ',') # Get cell type pair columns (exclude metadata columns) meta_cols = ['id_cp_interaction', 'interacting_pair', 'partner_a', 'partner_b', 'gene_a', 'gene_b', 'secreted', 'receptor_a', 'receptor_b', 'annotation_strategy', 'is_integrin', 'rank'] pair_cols = [c for c in sig_means.columns if c not in meta_cols and '|' in c] filtered = sig_means.copy() # Filter by cell types if cell_types: ct_list = [ct.strip() for ct in cell_types.split(',')] matching_cols = [c for c in pair_cols if any(ct in c for ct in ct_list)] if matching_cols: keep_cols = [c for c in sig_means.columns if c not in pair_cols] + matching_cols filtered = filtered[keep_cols] pair_cols = matching_cols # Filter by genes if genes: from .utils import parse_list_string gene_list = parse_list_string(genes, uppercase=True) mask = filtered['interacting_pair'].str.upper().apply( lambda x: any(g in x for g in gene_list) ) filtered = filtered[mask] # Filter by minimum mean expression if min_mean > 0 and pair_cols: mask = filtered[pair_cols].max(axis=1) >= min_mean filtered = filtered[mask] out_path = f"{save_path}/cellphonedb_filtered.csv" filtered.to_csv(out_path, index=False) msg = f"Filtered results: {len(filtered)} interactions. Saved to {out_path}" print(msg) return msg @tool def cellphonedb_plot( results_path: Annotated[str, Field(description="Path to CellPhoneDB results directory")], plot_type: Annotated[str, Field(description="'dotplot', 'heatmap', or 'chord'")] = "dotplot", cell_types: Annotated[str, Field(description="Comma-separated cell types to include (empty for all)")] = "", top_n: Annotated[int, Field(description="Number of top interactions to plot")] = 30, save_path: Annotated[str, Field(description="Directory to save plots")] = None, ) -> str: """Visualize CellPhoneDB results. Creates dot plots, heatmaps, or chord diagrams of significant interactions. """ save_path = save_path or _config["save_path"] import matplotlib.pyplot as plt import seaborn as sns os.makedirs(save_path, exist_ok=True) # Load results sig_means_path = f"{results_path}/significant_means.csv" pvals_path = f"{results_path}/pvalues.csv" if not exists(sig_means_path): sig_means_path = f"{results_path}/significant_means.txt" pvals_path = f"{results_path}/pvalues.txt" if not exists(sig_means_path): msg = f"ERROR: No results found in {results_path}" print(msg) return msg sep = '\t' if sig_means_path.endswith('.txt') else ',' sig_means = pd.read_csv(sig_means_path, sep=sep) pvals = pd.read_csv(pvals_path, sep=sep) if exists(pvals_path) else None # Get cell type pair columns meta_cols = ['id_cp_interaction', 'interacting_pair', 'partner_a', 'partner_b', 'gene_a', 'gene_b', 'secreted', 'receptor_a', 'receptor_b', 'annotation_strategy', 'is_integrin', 'rank'] pair_cols = [c for c in sig_means.columns if c not in meta_cols and '|' in c] # Filter by cell types if specified if cell_types: ct_list = [ct.strip() for ct in cell_types.split(',')] pair_cols = [c for c in pair_cols if any(ct in c for ct in ct_list)] if not pair_cols: msg = "ERROR: No cell type pairs found after filtering" print(msg) return msg # Get top interactions by mean expression means_data = sig_means[['interacting_pair'] + pair_cols].copy() means_data['max_mean'] = means_data[pair_cols].max(axis=1) means_data = means_data.nlargest(top_n, 'max_mean') plt.ioff() if plot_type == "dotplot": # Prepare data for dot plot plot_data = means_data.melt( id_vars=['interacting_pair'], value_vars=pair_cols, var_name='cell_pair', value_name='mean' ).dropna() if pvals is not None: pval_data = pvals[['interacting_pair'] + pair_cols].melt( id_vars=['interacting_pair'], value_vars=pair_cols, var_name='cell_pair', value_name='pvalue' ) plot_data = plot_data.merge(pval_data, on=['interacting_pair', 'cell_pair']) plot_data['-log10(pval)'] = -np.log10(plot_data['pvalue'] + 1e-10) fig, ax = plt.subplots(figsize=(max(12, len(pair_cols)*0.8), max(8, top_n*0.3))) # Create pivot for heatmap-style dot plot pivot = plot_data.pivot(index='interacting_pair', columns='cell_pair', values='mean') sns.heatmap(pivot, cmap='Reds', ax=ax, cbar_kws={'label': 'Mean Expression'}) ax.set_title(f'Top {top_n} Interactions') plt.xticks(rotation=45, ha='right') elif plot_type == "heatmap": # Interaction count heatmap per cell type pair counts = (sig_means[pair_cols].notna() & (sig_means[pair_cols] > 0)).sum() count_df = pd.DataFrame({'pair': counts.index, 'count': counts.values}) count_df[['source', 'target']] = count_df['pair'].str.split('|', expand=True) pivot = count_df.pivot(index='source', columns='target', values='count').fillna(0) fig, ax = plt.subplots(figsize=(10, 8)) sns.heatmap(pivot, cmap='YlOrRd', annot=True, fmt='.0f', ax=ax) ax.set_title('Significant Interactions per Cell Type Pair') elif plot_type == "chord": # Simplified chord-like visualization as bar plot counts = (sig_means[pair_cols].notna() & (sig_means[pair_cols] > 0)).sum() counts = counts.sort_values(ascending=False).head(20) fig, ax = plt.subplots(figsize=(12, 6)) counts.plot(kind='bar', ax=ax, color='coral', edgecolor='black') ax.set_ylabel('Number of Interactions') ax.set_title('Interactions per Cell Type Pair') plt.xticks(rotation=45, ha='right') else: msg = f"ERROR: Unknown plot type '{plot_type}'. Use 'dotplot', 'heatmap', or 'chord'" print(msg) return msg plt.tight_layout() out_path = f"{save_path}/cellphonedb_{plot_type}.png" plt.savefig(out_path, dpi=150, bbox_inches='tight') plt.close() msg = f"Created {plot_type} visualization. Saved to {out_path}" print(msg) return msg # ============================================================================= # LIANA Tools - Multi-method cell-cell communication analysis # ============================================================================= @tool def liana_inference( adata_path: Annotated[str, Field(description="Path to AnnData h5ad file")], cell_type_key: Annotated[str, Field(description="Column in obs for cell type annotations")] = "cell_type", sample_key: Annotated[str, Field(description="Column for sample IDs (empty for single-sample)")] = "", organism: Annotated[str, Field(description="'human', 'mouse', or 'auto'")] = "auto", expr_prop: Annotated[float, Field(description="Min proportion of cells expressing gene")] = 0.1, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Run LIANA rank aggregate for ligand-receptor inference. Combines multiple LR methods (CellPhoneDB, CellChat, NATMI, etc.) into consensus rankings. Works for single-sample or multi-sample analysis. """ save_path = save_path or _config["save_path"] import scanpy as sc import liana as li os.makedirs(save_path, exist_ok=True) adata = sc.read_h5ad(adata_path) # Auto-detect organism if organism == "auto": sample_genes = adata.var_names[:100].tolist() n_mouse = sum(1 for g in sample_genes if len(g) > 1 and g[0].isupper() and g[1:].islower()) organism = "mouse" if n_mouse > 50 else "human" resource = "mouseconsensus" if organism == "mouse" else "consensus" if sample_key and sample_key in adata.obs.columns: # Multi-sample analysis li.mt.rank_aggregate.by_sample( adata, sample_key=sample_key, groupby=cell_type_key, resource_name=resource, expr_prop=expr_prop, use_raw=False, verbose=True, key_added='liana_res' ) n_samples = adata.obs[sample_key].nunique() mode = f"multi-sample ({n_samples} samples)" else: # Single-sample analysis li.mt.rank_aggregate( adata, groupby=cell_type_key, resource_name=resource, expr_prop=expr_prop, use_raw=False, verbose=True, key_added='liana_res' ) mode = "single-sample" # Save results liana_res = adata.uns['liana_res'] out_csv = f"{save_path}/liana_results.csv" liana_res.to_csv(out_csv, index=False) # Save updated adata out_h5ad = f"{save_path}/liana_adata.h5ad" adata.write_h5ad(out_h5ad, compression="gzip") n_interactions = len(liana_res) n_pairs = liana_res[['source', 'target']].drop_duplicates().shape[0] msg = f"LIANA {mode} analysis complete: {n_interactions} interactions across {n_pairs} cell type pairs. Saved to {out_csv}" print(msg) return msg @tool def liana_spatial( adata_path: Annotated[str, Field(description="Path to spatial AnnData h5ad file")], local_metric: Annotated[str, Field(description="Local metric: 'cosine', 'pearson', 'spearman', 'jaccard'")] = "cosine", global_metric: Annotated[str, Field(description="Global metric: 'morans' or 'lee'")] = "morans", bandwidth: Annotated[float, Field(description="Spatial bandwidth for neighbor weights")] = 200, n_perms: Annotated[int, Field(description="Permutations for p-value calculation")] = 100, organism: Annotated[str, Field(description="'human', 'mouse', or 'auto'")] = "auto", save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Run LIANA spatial bivariate analysis. Computes local and global spatial correlations between ligand-receptor pairs. Requires spatial coordinates in adata.obsm['spatial']. """ save_path = save_path or _config["save_path"] import scanpy as sc import liana as li os.makedirs(save_path, exist_ok=True) adata = sc.read_h5ad(adata_path) # Check for spatial coordinates if 'spatial' not in adata.obsm: return "ERROR: No spatial coordinates found in adata.obsm['spatial']" # Auto-detect organism if organism == "auto": sample_genes = adata.var_names[:100].tolist() n_mouse = sum(1 for g in sample_genes if len(g) > 1 and g[0].isupper() and g[1:].islower()) organism = "mouse" if n_mouse > 50 else "human" resource = "mouseconsensus" if organism == "mouse" else "consensus" # Build spatial neighbors li.ut.spatial_neighbors(adata, bandwidth=bandwidth, cutoff=0.1, kernel='gaussian', set_diag=True) # Run bivariate analysis lrdata = li.mt.bivariate( adata, resource_name=resource, local_name=local_metric, global_name=global_metric, n_perms=n_perms, mask_negatives=False, add_categories=True, nz_prop=0.1, use_raw=False, verbose=True ) # Save results out_h5ad = f"{save_path}/liana_spatial.h5ad" lrdata.write_h5ad(out_h5ad, compression="gzip") # Save summary statistics var_df = lrdata.var.copy() var_df.to_csv(f"{save_path}/liana_spatial_summary.csv") n_interactions = lrdata.n_vars top_by_global = var_df.sort_values(global_metric, ascending=False).head(5).index.tolist() msg = f"LIANA bivariate analysis complete: {n_interactions} interactions. Top by {global_metric}: {', '.join(top_by_global[:3])}. Saved to {out_h5ad}" print(msg) return msg @tool def liana_misty( adata_path: Annotated[str, Field(description="Path to spatial AnnData h5ad file")], target_key: Annotated[str, Field(description="obsm key for target features (e.g., 'compositions')")] = "", predictor_key: Annotated[str, Field(description="obsm key for predictor features (e.g., 'pathway_scores')")] = "", bandwidth: Annotated[float, Field(description="Spatial bandwidth for para view")] = 200, n_neighs: Annotated[int, Field(description="Number of neighbors for juxta view")] = 6, model_type: Annotated[str, Field(description="'linear' or 'rf' (random forest)")] = "linear", save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Run LIANA MISTy for learning spatial relationships. Models how features in different spatial contexts (intra, juxta, para) predict target features. Useful for understanding spatial dependencies between cell types, pathways, or gene programs. """ save_path = save_path or _config["save_path"] import scanpy as sc import liana as li from liana.method import genericMistyData from liana.method.sp import LinearModel, RandomForestModel os.makedirs(save_path, exist_ok=True) adata = sc.read_h5ad(adata_path) # Check for spatial coordinates if 'spatial' not in adata.obsm: return "ERROR: No spatial coordinates found in adata.obsm['spatial']" # Build spatial neighbors li.ut.spatial_neighbors(adata, bandwidth=bandwidth, cutoff=0.1, set_diag=False) # Extract target and predictor data if target_key and target_key in adata.obsm: intra = li.ut.obsm_to_adata(adata, target_key) else: # Use highly variable genes as default if 'highly_variable' not in adata.var.columns: sc.pp.highly_variable_genes(adata, n_top_genes=500) hvg = adata.var[adata.var['highly_variable']].index intra = adata[:, hvg].copy() if predictor_key and predictor_key in adata.obsm: extra = li.ut.obsm_to_adata(adata, predictor_key) else: extra = None # Create MISTy data if extra is not None: misty = genericMistyData( intra=intra, extra=extra, cutoff=0.05, bandwidth=bandwidth, n_neighs=n_neighs ) else: misty = genericMistyData( intra=intra, cutoff=0.05, bandwidth=bandwidth, n_neighs=n_neighs ) # Select model model = LinearModel if model_type == "linear" else RandomForestModel # Run MISTy misty(model=model, bypass_intra=False, verbose=True) # Extract results target_metrics = misty.uns.get('target_metrics', pd.DataFrame()) interactions = misty.uns.get('interactions', {}) # Save results as CSV (avoid pickling non-serializable objects) if len(target_metrics) > 0: target_metrics.to_csv(f"{save_path}/misty_metrics.csv", index=False) # Save interaction importances per view for view_name, view_df in interactions.items(): if isinstance(view_df, pd.DataFrame) and len(view_df) > 0: view_df.to_csv(f"{save_path}/misty_{view_name}.csv") if len(target_metrics) > 0: avg_r2 = target_metrics['multi_R2'].mean() if 'multi_R2' in target_metrics.columns else 0 avg_gain = target_metrics['gain_R2'].mean() if 'gain_R2' in target_metrics.columns else 0 msg = f"MISTy analysis complete: {len(target_metrics)} targets. Avg R²={avg_r2:.3f}, Avg gain={avg_gain:.3f}. Saved to {save_path}/misty_*.csv" print(msg) return msg msg = f"MISTy analysis complete. Saved to {save_path}/misty_*.csv" print(msg) return msg @tool def liana_plot( results_path: Annotated[str, Field(description="Path to LIANA results CSV or h5ad")], plot_type: Annotated[str, Field(description="'dotplot', 'tileplot', or 'source_target'")] = "dotplot", source_cells: Annotated[str, Field(description="Comma-separated source cell types to include")] = "", target_cells: Annotated[str, Field(description="Comma-separated target cell types to include")] = "", top_n: Annotated[int, Field(description="Number of top interactions to plot")] = 20, save_path: Annotated[str, Field(description="Directory to save plots")] = None, ) -> str: """Visualize LIANA results. Creates dotplots, tileplots, or source-target network plots of LR interactions. """ save_path = save_path or _config["save_path"] import scanpy as sc import matplotlib.pyplot as plt import seaborn as sns os.makedirs(save_path, exist_ok=True) # Load results if results_path.endswith('.csv'): liana_res = pd.read_csv(results_path) elif results_path.endswith('.h5ad'): adata = sc.read_h5ad(results_path) if 'liana_res' in adata.uns: liana_res = adata.uns['liana_res'] else: return "ERROR: No liana_res found in adata.uns" else: return "ERROR: Provide .csv or .h5ad file" # Filter by cell types if specified if source_cells: src_list = [s.strip() for s in source_cells.split(',')] liana_res = liana_res[liana_res['source'].isin(src_list)] if target_cells: tgt_list = [t.strip() for t in target_cells.split(',')] liana_res = liana_res[liana_res['target'].isin(tgt_list)] if len(liana_res) == 0: return "ERROR: No interactions found after filtering" # Sort by magnitude or specificity sort_col = 'magnitude_rank' if 'magnitude_rank' in liana_res.columns else 'lr_means' ascending = True if 'rank' in sort_col else False liana_res = liana_res.sort_values(sort_col, ascending=ascending) plt.ioff() if plot_type == "dotplot": # Create interaction label liana_res = liana_res.head(top_n * 5) # Get more for filtering liana_res['interaction'] = liana_res['ligand_complex'] + ' → ' + liana_res['receptor_complex'] liana_res['cell_pair'] = liana_res['source'] + ' → ' + liana_res['target'] # Get top interactions top_interactions = liana_res.groupby('interaction')[sort_col].mean().sort_values(ascending=ascending).head(top_n).index plot_df = liana_res[liana_res['interaction'].isin(top_interactions)] # Pivot for heatmap if 'lr_means' in plot_df.columns: pivot = plot_df.pivot_table(index='interaction', columns='cell_pair', values='lr_means', aggfunc='mean') else: pivot = plot_df.pivot_table(index='interaction', columns='cell_pair', values='magnitude_rank', aggfunc='mean') fig, ax = plt.subplots(figsize=(max(10, len(pivot.columns)*0.6), max(8, len(pivot)*0.4))) sns.heatmap(pivot, cmap='Reds', ax=ax, cbar_kws={'label': 'Expression'}) ax.set_title(f'Top {top_n} Ligand-Receptor Interactions') plt.xticks(rotation=45, ha='right') plt.yticks(rotation=0) elif plot_type == "tileplot": # Interaction counts per cell pair counts = liana_res.groupby(['source', 'target']).size().reset_index(name='n_interactions') pivot = counts.pivot(index='source', columns='target', values='n_interactions').fillna(0) fig, ax = plt.subplots(figsize=(10, 8)) sns.heatmap(pivot, cmap='YlOrRd', annot=True, fmt='.0f', ax=ax) ax.set_title('Number of Interactions per Cell Type Pair') elif plot_type == "source_target": # Bar plot of interactions per source/target fig, axes = plt.subplots(1, 2, figsize=(14, 6)) source_counts = liana_res['source'].value_counts().head(15) source_counts.plot(kind='barh', ax=axes[0], color='steelblue') axes[0].set_xlabel('Number of Interactions') axes[0].set_title('Interactions by Source Cell Type') axes[0].invert_yaxis() target_counts = liana_res['target'].value_counts().head(15) target_counts.plot(kind='barh', ax=axes[1], color='coral') axes[1].set_xlabel('Number of Interactions') axes[1].set_title('Interactions by Target Cell Type') axes[1].invert_yaxis() else: return f"ERROR: Unknown plot type '{plot_type}'. Use 'dotplot', 'tileplot', or 'source_target'" plt.tight_layout() out_path = f"{save_path}/liana_{plot_type}.png" plt.savefig(out_path, dpi=150, bbox_inches='tight') plt.close() msg = f"Created {plot_type} visualization. Saved to {out_path}" print(msg) return msg # ============================================================================= # Squidpy Tools - Spatial analysis and graph statistics # ============================================================================= @tool def squidpy_spatial_neighbors( adata_path: Annotated[str, Field(description="Path to AnnData h5ad file")], coord_type: Annotated[str, Field(description="'grid' (for Visium/hex), 'generic' (for other), or 'auto'")] = "auto", n_neighs: Annotated[int, Field(description="Number of neighbors (for generic coord_type)")] = 6, n_rings: Annotated[int, Field(description="Number of rings (for grid coord_type, e.g., Visium)")] = 1, delaunay: Annotated[bool, Field(description="Use Delaunay triangulation")] = False, radius: Annotated[float, Field(description="Radius for neighbor search (0 to disable)")] = 0, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Build spatial neighbors graph for downstream spatial analysis. Creates connectivity and distance matrices stored in adata.obsp. Required before running most Squidpy graph analyses. coord_type options: - 'auto': Auto-detect (grid if Visium-like, generic otherwise) - 'grid': For Visium and other grid-based spatial data - 'generic': For MERFISH, SeqFISH, Slide-seq, and other non-grid data """ save_path = save_path or _config["save_path"] import scanpy as sc import squidpy as sq os.makedirs(save_path, exist_ok=True) adata = sc.read_h5ad(adata_path) # Check for spatial coordinates if 'spatial' not in adata.obsm: return "ERROR: No spatial coordinates found in adata.obsm['spatial']" # Map coord_type (handle 'visium' as alias for 'grid', 'auto' as None) coord_type_map = {"visium": "grid", "auto": None} sq_coord_type = coord_type_map.get(coord_type, coord_type) # Build spatial neighbors based on coord_type if sq_coord_type == "grid": sq.gr.spatial_neighbors(adata, n_rings=n_rings, coord_type="grid") elif delaunay: sq.gr.spatial_neighbors(adata, delaunay=True, coord_type="generic") elif radius > 0: sq.gr.spatial_neighbors(adata, radius=radius, coord_type="generic") elif sq_coord_type is None: # Auto-detect sq.gr.spatial_neighbors(adata, n_neighs=n_neighs, coord_type=None) else: sq.gr.spatial_neighbors(adata, n_neighs=n_neighs, coord_type="generic") # Save results out_path = f"{save_path}/spatial_neighbors.h5ad" adata.write_h5ad(out_path, compression="gzip") n_cells = adata.n_obs n_edges = adata.obsp['spatial_connectivities'].nnz msg = f"Built spatial neighbors graph: {n_cells} cells, {n_edges} edges. Saved to {out_path}" print(msg) return msg @tool def squidpy_nhood_enrichment( adata_path: Annotated[str, Field(description="Path to AnnData h5ad file (with spatial neighbors)")], cluster_key: Annotated[str, Field(description="Column in obs for cluster/cell type annotations")] = "cell_type", n_perms: Annotated[int, Field(description="Number of permutations for significance")] = 1000, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Compute neighborhood enrichment between cell type clusters. Tests whether cell types are enriched or depleted as neighbors compared to random. Requires spatial neighbors graph (run squidpy_spatial_neighbors first). """ save_path = save_path or _config["save_path"] import scanpy as sc import squidpy as sq import matplotlib.pyplot as plt os.makedirs(save_path, exist_ok=True) adata = sc.read_h5ad(adata_path) # Check prerequisites if 'spatial_connectivities' not in adata.obsp: return "ERROR: No spatial neighbors found. Run squidpy_spatial_neighbors first." if cluster_key not in adata.obs.columns: return f"ERROR: Column '{cluster_key}' not found in adata.obs" # Compute neighborhood enrichment sq.gr.nhood_enrichment(adata, cluster_key=cluster_key, n_perms=n_perms) # Save plot plt.ioff() fig, ax = plt.subplots(figsize=(10, 8)) sq.pl.nhood_enrichment(adata, cluster_key=cluster_key, ax=ax) plt.tight_layout() plt.savefig(f"{save_path}/squidpy_nhood_enrichment.png", dpi=150, bbox_inches='tight') plt.close() # Save results out_path = f"{save_path}/squidpy_nhood.h5ad" adata.write_h5ad(out_path, compression="gzip") # Extract z-scores summary zscore_key = f"{cluster_key}_nhood_enrichment" if zscore_key in adata.uns: zscores = adata.uns[zscore_key]['zscore'] max_enrich = zscores.max().max() min_enrich = zscores.min().min() msg = f"Neighborhood enrichment complete. Z-scores range: [{min_enrich:.2f}, {max_enrich:.2f}]. Saved to {out_path}" print(msg) return msg msg = f"Neighborhood enrichment complete. Saved to {out_path}" print(msg) return msg @tool def squidpy_co_occurrence( adata_path: Annotated[str, Field(description="Path to AnnData h5ad file")], cluster_key: Annotated[str, Field(description="Column in obs for cluster/cell type annotations")] = "cell_type", spatial_key: Annotated[str, Field(description="Key in obsm for spatial coordinates")] = "spatial", interval: Annotated[int, Field(description="Number of distance intervals")] = 50, n_splits: Annotated[int, Field(description="Number of spatial splits for computation")] = 2, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Compute co-occurrence probability between cell types across distances. Calculates conditional probability of observing one cell type given another at various spatial distances. """ save_path = save_path or _config["save_path"] import scanpy as sc import squidpy as sq import matplotlib.pyplot as plt os.makedirs(save_path, exist_ok=True) adata = sc.read_h5ad(adata_path) if cluster_key not in adata.obs.columns: return f"ERROR: Column '{cluster_key}' not found in adata.obs" # Compute co-occurrence sq.gr.co_occurrence(adata, cluster_key=cluster_key, spatial_key=spatial_key, interval=interval, n_splits=n_splits) # Save plot for first few clusters plt.ioff() try: clusters = adata.obs[cluster_key].unique()[:3].tolist() sq.pl.co_occurrence(adata, cluster_key=cluster_key, clusters=clusters, figsize=(12, 4)) plt.savefig(f"{save_path}/co_occurrence.png", dpi=150, bbox_inches='tight') plt.close() except Exception as e: # Plotting may fail due to version issues, continue with data export pass # Save results out_path = f"{save_path}/squidpy_cooccur.h5ad" adata.write_h5ad(out_path, compression="gzip") msg = f"Co-occurrence analysis complete for {len(adata.obs[cluster_key].unique())} clusters. Saved to {out_path}" print(msg) return msg @tool def squidpy_spatial_autocorr( adata_path: Annotated[str, Field(description="Path to AnnData h5ad file (with spatial neighbors)")], mode: Annotated[str, Field(description="'moran' for Moran's I or 'geary' for Geary's C")] = "moran", genes: Annotated[str, Field(description="Comma-separated genes, or 'hvg' for highly variable, or 'all'")] = "hvg", n_perms: Annotated[int, Field(description="Number of permutations")] = 100, n_jobs: Annotated[int, Field(description="Number of parallel jobs")] = 4, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Compute spatial autocorrelation (Moran's I or Geary's C) for genes. Identifies genes with significant spatial patterns. Requires spatial neighbors graph (run squidpy_spatial_neighbors first). """ save_path = save_path or _config["save_path"] import scanpy as sc import squidpy as sq os.makedirs(save_path, exist_ok=True) adata = sc.read_h5ad(adata_path) # Check prerequisites if 'spatial_connectivities' not in adata.obsp: return "ERROR: No spatial neighbors found. Run squidpy_spatial_neighbors first." # Select genes if genes in ["hvg", "highly_variable"]: if 'highly_variable' not in adata.var.columns: # Compute HVGs if not present sc.pp.highly_variable_genes(adata, n_top_genes=min(500, adata.n_vars)) gene_list = adata.var[adata.var['highly_variable']].index.tolist() elif genes == "all": gene_list = adata.var_names.tolist()[:1000] # Limit to avoid memory issues else: from .utils import parse_list_string gene_list = parse_list_string(genes) gene_list = [g for g in gene_list if g in adata.var_names] if len(gene_list) == 0: # Fallback to top variable genes gene_list = adata.var_names.tolist()[:100] if len(gene_list) == 0: return "ERROR: No genes available in dataset" # Compute spatial autocorrelation sq.gr.spatial_autocorr(adata, mode=mode, genes=gene_list, n_perms=n_perms, n_jobs=n_jobs) # Get results result_key = "moranI" if mode == "moran" else "gearyC" results_df = adata.uns[result_key] results_df.to_csv(f"{save_path}/squidpy_{mode}.csv") # Save adata out_path = f"{save_path}/squidpy_autocorr.h5ad" adata.write_h5ad(out_path, compression="gzip") # Summary sig_genes = results_df[results_df['pval_norm'] < 0.05] top_genes = results_df.head(5).index.tolist() msg = f"Spatial autocorrelation ({mode}): {len(gene_list)} genes tested, {len(sig_genes)} significant. Top: {', '.join(top_genes)}. Saved to {save_path}/squidpy_{mode}.csv" print(msg) return msg @tool def squidpy_ripley( adata_path: Annotated[str, Field(description="Path to AnnData h5ad file")], cluster_key: Annotated[str, Field(description="Column in obs for cluster annotations")] = "cell_type", mode: Annotated[str, Field(description="'L' (Ripley's L), 'F', or 'G'")] = "L", save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Compute Ripley's statistics for spatial point pattern analysis. Determines whether cell types show clustered, random, or dispersed spatial patterns. """ save_path = save_path or _config["save_path"] import scanpy as sc import squidpy as sq import matplotlib.pyplot as plt os.makedirs(save_path, exist_ok=True) adata = sc.read_h5ad(adata_path) if cluster_key not in adata.obs.columns: return f"ERROR: Column '{cluster_key}' not found in adata.obs" # Compute Ripley's statistic sq.gr.ripley(adata, cluster_key=cluster_key, mode=mode) # Save plot plt.ioff() fig, ax = plt.subplots(figsize=(10, 6)) sq.pl.ripley(adata, cluster_key=cluster_key, mode=mode, ax=ax) plt.tight_layout() plt.savefig(f"{save_path}/squidpy_ripley_{mode}.png", dpi=150, bbox_inches='tight') plt.close() # Save results out_path = f"{save_path}/squidpy_ripley.h5ad" adata.write_h5ad(out_path, compression="gzip") n_clusters = adata.obs[cluster_key].nunique() msg = f"Ripley's {mode} analysis complete for {n_clusters} clusters. Saved to {out_path}" print(msg) return msg @tool def squidpy_centrality( adata_path: Annotated[str, Field(description="Path to AnnData h5ad file (with spatial neighbors)")], cluster_key: Annotated[str, Field(description="Column in obs for cluster annotations")] = "cell_type", save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Compute centrality scores for cell type clusters in spatial graph. Calculates closeness centrality, clustering coefficient, and degree centrality. Requires spatial neighbors graph (run squidpy_spatial_neighbors first). """ save_path = save_path or _config["save_path"] import scanpy as sc import squidpy as sq import matplotlib.pyplot as plt os.makedirs(save_path, exist_ok=True) adata = sc.read_h5ad(adata_path) # Check prerequisites if 'spatial_connectivities' not in adata.obsp: return "ERROR: No spatial neighbors found. Run squidpy_spatial_neighbors first." if cluster_key not in adata.obs.columns: return f"ERROR: Column '{cluster_key}' not found in adata.obs" # Compute centrality scores sq.gr.centrality_scores(adata, cluster_key=cluster_key) # Save plot plt.ioff() try: sq.pl.centrality_scores(adata, cluster_key=cluster_key, figsize=(12, 6)) plt.savefig(f"{save_path}/centrality.png", dpi=150, bbox_inches='tight') plt.close() except Exception: # Plotting may fail due to version issues pass # Save results out_path = f"{save_path}/squidpy_centrality.h5ad" adata.write_h5ad(out_path, compression="gzip") # Extract centrality results cent_key = f"{cluster_key}_centrality_scores" if cent_key in adata.uns: cent_df = adata.uns[cent_key] cent_df.to_csv(f"{save_path}/squidpy_centrality.csv") msg = f"Centrality scores computed for {adata.obs[cluster_key].nunique()} clusters. Saved to {out_path}" print(msg) return msg @tool def squidpy_interaction_matrix( adata_path: Annotated[str, Field(description="Path to AnnData h5ad file (with spatial neighbors)")], cluster_key: Annotated[str, Field(description="Column in obs for cluster annotations")] = "cell_type", normalized: Annotated[bool, Field(description="Row-normalize the interaction matrix")] = True, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Compute interaction matrix between cell type clusters. Quantifies the number of edges between cell types in the spatial graph. Requires spatial neighbors graph (run squidpy_spatial_neighbors first). """ save_path = save_path or _config["save_path"] import scanpy as sc import squidpy as sq import matplotlib.pyplot as plt os.makedirs(save_path, exist_ok=True) adata = sc.read_h5ad(adata_path) # Check prerequisites if 'spatial_connectivities' not in adata.obsp: return "ERROR: No spatial neighbors found. Run squidpy_spatial_neighbors first." if cluster_key not in adata.obs.columns: return f"ERROR: Column '{cluster_key}' not found in adata.obs" # Compute interaction matrix sq.gr.interaction_matrix(adata, cluster_key=cluster_key, normalized=normalized) # Save plot plt.ioff() fig, ax = plt.subplots(figsize=(10, 8)) sq.pl.interaction_matrix(adata, cluster_key=cluster_key, ax=ax) plt.tight_layout() plt.savefig(f"{save_path}/squidpy_interaction_matrix.png", dpi=150, bbox_inches='tight') plt.close() # Save results out_path = f"{save_path}/squidpy_interaction.h5ad" adata.write_h5ad(out_path, compression="gzip") # Save matrix as CSV int_key = f"{cluster_key}_interactions" if int_key in adata.uns: int_df = pd.DataFrame(adata.uns[int_key]) int_df.to_csv(f"{save_path}/squidpy_interaction_matrix.csv") msg = f"Interaction matrix computed for {adata.obs[cluster_key].nunique()} clusters. Saved to {out_path}" print(msg) return msg @tool def squidpy_ligrec( adata_path: Annotated[str, Field(description="Path to AnnData h5ad file")], cluster_key: Annotated[str, Field(description="Column in obs for cluster annotations")] = "cell_type", n_perms: Annotated[int, Field(description="Number of permutations")] = 1000, threshold: Annotated[float, Field(description="Min fraction of cells expressing gene")] = 0.1, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Run receptor-ligand analysis using Squidpy and OmniPath database. Identifies significant ligand-receptor interactions between cell types. """ save_path = save_path or _config["save_path"] import scanpy as sc import squidpy as sq os.makedirs(save_path, exist_ok=True) adata = sc.read_h5ad(adata_path) if cluster_key not in adata.obs.columns: return f"ERROR: Column '{cluster_key}' not found in adata.obs" # Run ligand-receptor analysis res = sq.gr.ligrec( adata, n_perms=n_perms, cluster_key=cluster_key, copy=True, use_raw=False, threshold=threshold, transmitter_params={"categories": "ligand"}, receiver_params={"categories": "receptor"}, ) # Save results means_df = res['means'] pvals_df = res['pvalues'] means_df.to_csv(f"{save_path}/squidpy_ligrec_means.csv") pvals_df.to_csv(f"{save_path}/squidpy_ligrec_pvalues.csv") # Count significant interactions n_significant = (pvals_df < 0.05).sum().sum() n_total = pvals_df.notna().sum().sum() msg = f"Ligand-receptor analysis complete: {n_significant} significant interactions out of {n_total} tested. Saved to {save_path}/squidpy_ligrec_*.csv" print(msg) return msg # ============================================================================= # scvi-tools Spatial Deconvolution Tools # ============================================================================= @tool def destvi_deconvolution( sc_adata_path: Annotated[str, Field(description="Path to single-cell reference h5ad file")], st_adata_path: Annotated[str, Field(description="Path to spatial transcriptomics h5ad file")], cell_type_key: Annotated[str, Field(description="Column in sc_adata.obs for cell type labels")] = "cell_type", layer: Annotated[str, Field(description="Layer to use for counts (empty for .X)")] = "", sc_max_epochs: Annotated[int, Field(description="Max epochs for single-cell model")] = 300, st_max_epochs: Annotated[int, Field(description="Max epochs for spatial model")] = 2500, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Run DestVI for multi-resolution spatial deconvolution. DestVI decomposes spatial spots into cell type proportions while capturing intra-cell-type variation (gamma). Requires a single-cell reference with cell type annotations. Outputs: Cell type proportions per spot, gamma latent space, trained models. """ save_path = save_path or _config["save_path"] import scanpy as sc import scvi from scvi.model import CondSCVI, DestVI os.makedirs(save_path, exist_ok=True) # Load data sc_adata = sc.read_h5ad(sc_adata_path) st_adata = sc.read_h5ad(st_adata_path) # Validate cell type column if cell_type_key not in sc_adata.obs.columns: return f"ERROR: Column '{cell_type_key}' not found in sc_adata.obs. Available: {list(sc_adata.obs.columns)}" # Find common genes common_genes = list(set(sc_adata.var_names) & set(st_adata.var_names)) if len(common_genes) < 100: return f"ERROR: Only {len(common_genes)} common genes found. Need at least 100." sc_adata = sc_adata[:, common_genes].copy() st_adata = st_adata[:, common_genes].copy() # Setup single-cell model if layer and layer in sc_adata.layers: CondSCVI.setup_anndata(sc_adata, layer=layer, labels_key=cell_type_key) else: CondSCVI.setup_anndata(sc_adata, labels_key=cell_type_key) # Train single-cell model sc_model = CondSCVI(sc_adata, weight_obs=False) sc_model.train(max_epochs=sc_max_epochs, accelerator="auto") # Setup spatial model if layer and layer in st_adata.layers: DestVI.setup_anndata(st_adata, layer=layer) else: DestVI.setup_anndata(st_adata) # Train spatial model st_model = DestVI.from_rna_model(st_adata, sc_model) st_model.train(max_epochs=st_max_epochs, accelerator="auto") # Get proportions proportions = st_model.get_proportions() st_adata.obsm["destvi_proportions"] = proportions # Save results proportions.to_csv(f"{save_path}/destvi_proportions.csv") st_adata.write_h5ad(f"{save_path}/destvi_spatial.h5ad", compression="gzip") # Save models sc_model.save(f"{save_path}/destvi_sc_model", overwrite=True) st_model.save(f"{save_path}/destvi_st_model", overwrite=True) n_spots = st_adata.n_obs n_celltypes = proportions.shape[1] msg = f"DestVI deconvolution complete: {n_spots} spots, {n_celltypes} cell types, {len(common_genes)} genes. Results saved to {save_path}" print(msg) return msg @tool def cell2location_mapping( sc_adata_path: Annotated[str, Field(description="Path to single-cell reference h5ad file")], st_adata_path: Annotated[str, Field(description="Path to spatial transcriptomics h5ad file")], cell_type_key: Annotated[str, Field(description="Column in sc_adata.obs for cell type labels")] = "cell_type", batch_key: Annotated[str, Field(description="Column for batch information (optional)")] = "", n_cells_per_location: Annotated[int, Field(description="Expected cells per spot (tissue-dependent)")] = 30, detection_alpha: Annotated[float, Field(description="Detection sensitivity (200 default, 20 for high variation)")] = 200, sc_max_epochs: Annotated[int, Field(description="Max epochs for reference model")] = 250, st_max_epochs: Annotated[int, Field(description="Max epochs for spatial model")] = 30000, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Run Cell2location for spatial cell type mapping. Cell2location uses Bayesian inference to estimate cell type abundance at each spatial location. Integrates scRNA-seq reference with Visium/spatial data. Parameters: - n_cells_per_location: ~30 for lymph node, ~8 for brain, adjust per tissue - detection_alpha: Lower (20) for high within-batch variation Outputs: Cell abundance per spot (q05 quantile = confident estimates). """ save_path = save_path or _config["save_path"] import scanpy as sc from cell2location.models import Cell2location, RegressionModel os.makedirs(save_path, exist_ok=True) # Load data sc_adata = sc.read_h5ad(sc_adata_path) st_adata = sc.read_h5ad(st_adata_path) # Validate cell type column if cell_type_key not in sc_adata.obs.columns: return f"ERROR: Column '{cell_type_key}' not found in sc_adata.obs. Available: {list(sc_adata.obs.columns)}" # Find common genes common_genes = list(set(sc_adata.var_names) & set(st_adata.var_names)) if len(common_genes) < 100: return f"ERROR: Only {len(common_genes)} common genes found. Need at least 100." sc_adata = sc_adata[:, common_genes].copy() st_adata = st_adata[:, common_genes].copy() # Setup reference model setup_kwargs = {"labels_key": cell_type_key} if batch_key and batch_key in sc_adata.obs.columns: setup_kwargs["batch_key"] = batch_key RegressionModel.setup_anndata(sc_adata, **setup_kwargs) # Train reference model ref_model = RegressionModel(sc_adata) ref_model.train(max_epochs=sc_max_epochs, accelerator="auto") # Export signatures sc_adata = ref_model.export_posterior( sc_adata, sample_kwargs={"num_samples": 1000, "batch_size": 2500} ) # Get signatures for spatial model if "means_per_cluster_mu_fg" in sc_adata.varm: inf_aver = sc_adata.varm["means_per_cluster_mu_fg"].T.copy() else: return "ERROR: Failed to extract cell type signatures from reference model" # Setup spatial model Cell2location.setup_anndata(st_adata) # Train spatial model st_model = Cell2location( st_adata, cell_state_df=inf_aver, N_cells_per_location=n_cells_per_location, detection_alpha=detection_alpha, ) st_model.train(max_epochs=st_max_epochs, accelerator="auto") # Export results st_adata = st_model.export_posterior( st_adata, sample_kwargs={"num_samples": 1000, "batch_size": st_adata.n_obs} ) # Save results if "q05_cell_abundance_w_sf" in st_adata.obsm: abundance = pd.DataFrame( st_adata.obsm["q05_cell_abundance_w_sf"], index=st_adata.obs_names, columns=inf_aver.index ) abundance.to_csv(f"{save_path}/cell2location_abundance.csv") st_adata.write_h5ad(f"{save_path}/cell2location_spatial.h5ad", compression="gzip") # Save models ref_model.save(f"{save_path}/cell2location_ref_model", overwrite=True) st_model.save(f"{save_path}/cell2location_st_model", overwrite=True) n_spots = st_adata.n_obs n_celltypes = inf_aver.shape[0] msg = f"Cell2location mapping complete: {n_spots} spots, {n_celltypes} cell types, {len(common_genes)} genes. Results saved to {save_path}" print(msg) return msg @tool def stereoscope_deconvolution( sc_adata_path: Annotated[str, Field(description="Path to single-cell reference h5ad file")], st_adata_path: Annotated[str, Field(description="Path to spatial transcriptomics h5ad file")], cell_type_key: Annotated[str, Field(description="Column in sc_adata.obs for cell type labels")] = "cell_type", layer: Annotated[str, Field(description="Layer to use for counts (empty for .X)")] = "", sc_max_epochs: Annotated[int, Field(description="Max epochs for single-cell model")] = 100, st_max_epochs: Annotated[int, Field(description="Max epochs for spatial model")] = 2000, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Run Stereoscope for spatial cell type deconvolution. Stereoscope learns cell-type-specific gene expression from scRNA-seq and infers cell type proportions in spatial spots. Outputs: Cell type proportions per spot, trained models. """ save_path = save_path or _config["save_path"] import scanpy as sc from scvi.external import RNAStereoscope, SpatialStereoscope os.makedirs(save_path, exist_ok=True) # Load data sc_adata = sc.read_h5ad(sc_adata_path) st_adata = sc.read_h5ad(st_adata_path) # Validate cell type column if cell_type_key not in sc_adata.obs.columns: return f"ERROR: Column '{cell_type_key}' not found in sc_adata.obs. Available: {list(sc_adata.obs.columns)}" # Find common genes common_genes = list(set(sc_adata.var_names) & set(st_adata.var_names)) if len(common_genes) < 100: return f"ERROR: Only {len(common_genes)} common genes found. Need at least 100." sc_adata = sc_adata[:, common_genes].copy() st_adata = st_adata[:, common_genes].copy() # Setup single-cell model if layer and layer in sc_adata.layers: RNAStereoscope.setup_anndata(sc_adata, layer=layer, labels_key=cell_type_key) else: RNAStereoscope.setup_anndata(sc_adata, labels_key=cell_type_key) # Train single-cell model sc_model = RNAStereoscope(sc_adata) sc_model.train(max_epochs=sc_max_epochs, accelerator="auto") # Setup spatial model if layer and layer in st_adata.layers: SpatialStereoscope.setup_anndata(st_adata, layer=layer) else: SpatialStereoscope.setup_anndata(st_adata) # Train spatial model st_model = SpatialStereoscope.from_rna_model(st_adata, sc_model) st_model.train(max_epochs=st_max_epochs, accelerator="auto") # Get proportions proportions = st_model.get_proportions() st_adata.obsm["stereoscope_proportions"] = proportions # Save results proportions.to_csv(f"{save_path}/stereoscope_proportions.csv") st_adata.write_h5ad(f"{save_path}/stereoscope_spatial.h5ad", compression="gzip") # Save models sc_model.save(f"{save_path}/stereoscope_sc_model", overwrite=True) st_model.save(f"{save_path}/stereoscope_st_model", overwrite=True) n_spots = st_adata.n_obs n_celltypes = proportions.shape[1] msg = f"Stereoscope deconvolution complete: {n_spots} spots, {n_celltypes} cell types, {len(common_genes)} genes. Results saved to {save_path}" print(msg) return msg @tool def gimvi_imputation( sc_adata_path: Annotated[str, Field(description="Path to single-cell reference h5ad file")], st_adata_path: Annotated[str, Field(description="Path to spatial transcriptomics h5ad file")], genes_to_impute: Annotated[str, Field(description="Comma-separated genes to impute (empty for all missing)")] = "", layer: Annotated[str, Field(description="Layer to use for counts (empty for .X)")] = "", max_epochs: Annotated[int, Field(description="Max training epochs")] = 200, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Run gimVI for gene imputation in spatial data. gimVI learns a joint model of scRNA-seq and spatial data to impute genes missing from the spatial measurements using the single-cell reference. Useful for spatial technologies with limited gene panels (FISH-based). Outputs: Imputed gene expression matrix, imputed spatial AnnData. """ save_path = save_path or _config["save_path"] import scanpy as sc from scvi.external import GIMVI os.makedirs(save_path, exist_ok=True) # Load data sc_adata = sc.read_h5ad(sc_adata_path) st_adata = sc.read_h5ad(st_adata_path) # Determine genes to impute sc_genes = set(sc_adata.var_names) st_genes = set(st_adata.var_names) common_genes = list(sc_genes & st_genes) if len(common_genes) < 50: return f"ERROR: Only {len(common_genes)} common genes. Need at least 50 for training." if genes_to_impute: from .utils import parse_list_string target_genes = parse_list_string(genes_to_impute) missing_genes = [g for g in target_genes if g in sc_genes and g not in st_genes] if not missing_genes: return f"ERROR: No valid genes to impute. Genes must be in scRNA-seq but not in spatial data." else: missing_genes = list(sc_genes - st_genes) if not missing_genes: return "No genes to impute - all scRNA-seq genes already in spatial data." # Subset to common genes for training sc_adata_train = sc_adata[:, common_genes].copy() st_adata_train = st_adata[:, common_genes].copy() # Setup gimVI if layer and layer in sc_adata_train.layers: GIMVI.setup_anndata(sc_adata_train, layer=layer) GIMVI.setup_anndata(st_adata_train, layer=layer) else: GIMVI.setup_anndata(sc_adata_train) GIMVI.setup_anndata(st_adata_train) # Train model model = GIMVI(sc_adata_train, st_adata_train) model.train(max_epochs=max_epochs, accelerator="auto") # Get imputed values for spatial data # gimVI returns (seq_imputed, spatial_imputed) tuple _, imputed_spatial = model.get_imputed_values(normalized=True) # Create imputed expression matrix # gimVI imputes expression for common genes using the joint latent space imputed_df = pd.DataFrame( imputed_spatial, index=st_adata_train.obs_names, columns=common_genes ) # Save results imputed_df.to_csv(f"{save_path}/gimvi_imputed.csv") # Add imputed values to spatial AnnData st_adata_train.obsm["gimvi_imputed"] = imputed_spatial st_adata_train.write_h5ad(f"{save_path}/gimvi_spatial.h5ad", compression="gzip") # Save model model.save(f"{save_path}/gimvi_model", overwrite=True) n_spots = st_adata_train.n_obs n_imputed = len(common_genes) msg = f"gimVI imputation complete: {n_spots} spots, {n_imputed} genes imputed. Results saved to {save_path}" print(msg) return msg # ============================================================================= # Spatial Domain Detection Tools (SpaGCN, GraphST) # ============================================================================= @tool def spagcn_clustering( adata_path: Annotated[str, Field(description="Path to spatial transcriptomics h5ad file")], n_clusters: Annotated[int, Field(description="Target number of spatial domains")] = 7, histology_path: Annotated[str, Field(description="Path to histology image (optional, tif/png/jpg)")] = "", x_pixel_col: Annotated[str, Field(description="Column in obs for x pixel coordinates")] = "x_pixel", y_pixel_col: Annotated[str, Field(description="Column in obs for y pixel coordinates")] = "y_pixel", x_array_col: Annotated[str, Field(description="Column in obs for x array coordinates")] = "x_array", y_array_col: Annotated[str, Field(description="Column in obs for y array coordinates")] = "y_array", p: Annotated[float, Field(description="Percentage of expression from spatial neighbors")] = 0.5, alpha: Annotated[float, Field(description="Histology weight (higher = more weight)")] = 1, beta: Annotated[int, Field(description="Spot area parameter for histology")] = 49, refine: Annotated[bool, Field(description="Whether to refine clusters using spatial adjacency")] = True, shape: Annotated[str, Field(description="Spot shape: 'hexagon' for Visium, 'square' for ST")] = "hexagon", max_epochs: Annotated[int, Field(description="Maximum training epochs")] = 200, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Run SpaGCN for spatial domain detection. SpaGCN uses graph convolutional networks to identify spatial domains by integrating gene expression with spatial location and histology images. Outputs: Cluster assignments, refined clusters (optional), trained model. """ save_path = save_path or _config["save_path"] import scanpy as sc import SpaGCN as spg import random import torch os.makedirs(save_path, exist_ok=True) # Load data adata = sc.read_h5ad(adata_path) adata.var_names_make_unique() # Get coordinates if x_pixel_col in adata.obs.columns and y_pixel_col in adata.obs.columns: x_pixel = adata.obs[x_pixel_col].values y_pixel = adata.obs[y_pixel_col].values elif 'spatial' in adata.obsm: x_pixel = adata.obsm['spatial'][:, 0] y_pixel = adata.obsm['spatial'][:, 1] adata.obs['x_pixel'] = x_pixel adata.obs['y_pixel'] = y_pixel else: return f"ERROR: No spatial coordinates found. Need '{x_pixel_col}'/'{y_pixel_col}' in obs or 'spatial' in obsm." # Get array coordinates for refinement if x_array_col in adata.obs.columns and y_array_col in adata.obs.columns: x_array = adata.obs[x_array_col].values y_array = adata.obs[y_array_col].values else: x_array = x_pixel y_array = y_pixel # Calculate adjacency matrix if histology_path and os.path.exists(histology_path): import cv2 img = cv2.imread(histology_path) adj = spg.calculate_adj_matrix( x=x_pixel, y=y_pixel, x_pixel=x_pixel, y_pixel=y_pixel, image=img, beta=beta, alpha=alpha, histology=True ) else: adj = spg.calculate_adj_matrix(x=x_pixel, y=y_pixel, histology=False) # Preprocess expression data spg.prefilter_genes(adata, min_cells=3) spg.prefilter_specialgenes(adata) sc.pp.normalize_per_cell(adata) sc.pp.log1p(adata) # Search for l parameter l = spg.search_l(p, adj, start=0.01, end=1000, tol=0.01, max_run=100) # Set random seeds r_seed = t_seed = n_seed = 100 random.seed(r_seed) torch.manual_seed(t_seed) np.random.seed(n_seed) # Search for resolution res = spg.search_res( adata, adj, l, n_clusters, start=0.7, step=0.1, tol=5e-3, lr=0.05, max_epochs=20, r_seed=r_seed, t_seed=t_seed, n_seed=n_seed ) # Train SpaGCN clf = spg.SpaGCN() clf.set_l(l) clf.train( adata, adj, init_spa=True, init="leiden", res=res, tol=5e-3, lr=0.05, max_epochs=max_epochs ) # Get predictions y_pred, prob = clf.predict() adata.obs["spagcn_pred"] = y_pred adata.obs["spagcn_pred"] = adata.obs["spagcn_pred"].astype('category') # Refine predictions if refine: adj_2d = spg.calculate_adj_matrix(x=x_array, y=y_array, histology=False) refined_pred = spg.refine( sample_id=adata.obs.index.tolist(), pred=adata.obs["spagcn_pred"].tolist(), dis=adj_2d, shape=shape ) adata.obs["spagcn_refined"] = refined_pred adata.obs["spagcn_refined"] = adata.obs["spagcn_refined"].astype('category') # Save results adata.write_h5ad(f"{save_path}/spagcn_result.h5ad", compression="gzip") # Save cluster assignments cluster_col = "spagcn_refined" if refine else "spagcn_pred" clusters_df = adata.obs[[cluster_col]].copy() clusters_df.to_csv(f"{save_path}/spagcn_clusters.csv") n_spots = adata.n_obs n_domains = adata.obs[cluster_col].nunique() msg = f"SpaGCN clustering complete: {n_spots} spots, {n_domains} spatial domains detected. Results saved to {save_path}" print(msg) return msg @tool def graphst_clustering( adata_path: Annotated[str, Field(description="Path to spatial transcriptomics h5ad file")], n_clusters: Annotated[int, Field(description="Target number of spatial domains")] = 7, cluster_method: Annotated[str, Field(description="Clustering method: 'leiden'")] = "leiden", n_pcs: Annotated[int, Field(description="Number of principal components")] = 30, n_neighbors: Annotated[int, Field(description="Number of neighbors for graph construction")] = 10, random_seed: Annotated[int, Field(description="Random seed for reproducibility")] = 42, device: Annotated[str, Field(description="Device: 'cuda' or 'cpu'")] = "cuda", save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Run GraphST for spatial domain detection. GraphST uses graph self-supervised contrastive learning to identify spatial domains by integrating gene expression with spatial information. Clustering methods: - leiden: Community detection (recommended) Outputs: Cluster assignments, spatial embeddings, trained model. """ save_path = save_path or _config["save_path"] import scanpy as sc import torch from GraphST.GraphST import GraphST from GraphST import clustering os.makedirs(save_path, exist_ok=True) # Set device if device == "cuda" and torch.cuda.is_available(): device_obj = torch.device('cuda') else: device_obj = torch.device('cpu') # Load data adata = sc.read_h5ad(adata_path) adata.var_names_make_unique() # Check for spatial coordinates if 'spatial' not in adata.obsm: if 'x_pixel' in adata.obs.columns and 'y_pixel' in adata.obs.columns: adata.obsm['spatial'] = np.array([ adata.obs['x_pixel'].values, adata.obs['y_pixel'].values ]).T else: return "ERROR: No spatial coordinates found. Need 'spatial' in obsm or 'x_pixel'/'y_pixel' in obs." # Preprocess (normalize and select HVGs) sc.pp.highly_variable_genes(adata, flavor="seurat_v3", n_top_genes=3000) sc.pp.normalize_total(adata, target_sum=1e4) sc.pp.log1p(adata) # Build and train GraphST model model = GraphST(adata, device=device_obj, random_seed=random_seed) adata = model.train() # Perform clustering on the learned embeddings if cluster_method == 'leiden': clustering(adata, n_clusters=n_clusters, method=cluster_method, start=0.1, end=2.0, increment=0.01) else: return f"ERROR: Unknown clustering method '{cluster_method}'. Use 'leiden'." # Rename cluster column for consistency if 'domain' in adata.obs.columns: adata.obs['graphst_cluster'] = adata.obs['domain'].astype('category') # Save results adata.write_h5ad(f"{save_path}/graphst_result.h5ad", compression="gzip") # Save cluster assignments if 'graphst_cluster' in adata.obs.columns: clusters_df = adata.obs[['graphst_cluster']].copy() clusters_df.to_csv(f"{save_path}/graphst_clusters.csv") n_domains = adata.obs['graphst_cluster'].nunique() else: n_domains = "unknown" # Save embeddings if 'emb' in adata.obsm: emb_df = pd.DataFrame( adata.obsm['emb'], index=adata.obs_names, columns=[f'GraphST_{i}' for i in range(adata.obsm['emb'].shape[1])] ) emb_df.to_csv(f"{save_path}/graphst_embeddings.csv") n_spots = adata.n_obs msg = f"GraphST clustering complete: {n_spots} spots, {n_domains} spatial domains detected. Results saved to {save_path}" print(msg) return msg # ============================================================================= # Scanpy Tools: Gene Scoring # ============================================================================= @tool def scanpy_score_genes( adata_path: Annotated[str, Field(description="Path to h5ad file")], gene_list: Annotated[list[str], Field(description="List of genes to score (e.g., pathway genes, signature genes)")], score_name: Annotated[str, Field(description="Name for the score column in adata.obs")] = "gene_score", ctrl_size: Annotated[int, Field(description="Number of control genes per binned expression level")] = 50, n_bins: Annotated[int, Field(description="Number of expression bins for control gene selection")] = 25, use_raw: Annotated[bool, Field(description="Use raw expression data if available")] = False, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Score cells/spots based on expression of a gene signature using Scanpy. Computes a score for each cell by comparing the average expression of the input gene set to the average expression of a reference set of control genes. The control genes are randomly selected from genes binned by their average expression, so that the control gene set has a similar expression level distribution as the input genes. Common use cases: - Score cells for pathway activity (e.g., hypoxia, cell cycle, EMT) - Score cells for cell type markers - Score cells for custom gene signatures from literature """ save_path = save_path or _config["save_path"] import scanpy as sc adata = sc.read_h5ad(adata_path) # Ensure gene names are uppercase for matching adata.var_names_make_unique() gene_list_upper = [g.upper() for g in gene_list] adata.var.index = adata.var.index.str.upper() # Check which genes are present genes_found = [g for g in gene_list_upper if g in adata.var_names] genes_missing = [g for g in gene_list_upper if g not in adata.var_names] if len(genes_found) == 0: return f"ERROR: None of the {len(gene_list)} genes found in the data. First few: {gene_list[:5]}" # Score genes sc.tl.score_genes( adata, gene_list=genes_found, score_name=score_name, ctrl_size=ctrl_size, n_bins=n_bins, use_raw=use_raw, ) # Save results os.makedirs(save_path, exist_ok=True) adata.write_h5ad(f"{save_path}/scored.h5ad", compression="gzip") # Save scores as CSV scores_df = adata.obs[[score_name]].copy() scores_df.to_csv(f"{save_path}/{score_name}.csv") # Basic stats score_mean = adata.obs[score_name].mean() score_std = adata.obs[score_name].std() score_min = adata.obs[score_name].min() score_max = adata.obs[score_name].max() result = f"Gene scoring complete for '{score_name}':\n" result += f" - {len(genes_found)}/{len(gene_list)} genes found in data\n" if genes_missing and len(genes_missing) <= 10: result += f" - Missing genes: {genes_missing}\n" elif genes_missing: result += f" - {len(genes_missing)} genes missing (first 5: {genes_missing[:5]})\n" result += f" - Score stats: mean={score_mean:.3f}, std={score_std:.3f}, range=[{score_min:.3f}, {score_max:.3f}]\n" result += f" - Saved to {save_path}/scored.h5ad and {save_path}/{score_name}.csv" print(result) return result # ============================================================================= # Scanpy Tools: Integration - Ingest # ============================================================================= @tool def scanpy_ingest( adata_query_path: Annotated[str, Field(description="Path to query h5ad file (data to annotate)")], adata_ref_path: Annotated[str, Field(description="Path to reference h5ad file (annotated data)")], obs_keys: Annotated[list[str], Field(description="Observation keys to transfer from reference (e.g., ['cell_type', 'cluster'])")], embedding: Annotated[str, Field(description="Embedding to use for mapping: 'pca' or 'umap'")] = "umap", use_rep: Annotated[str, Field(description="Key in adata.obsm for representation (e.g., 'X_pca')")] = "X_pca", save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Transfer annotations from a reference dataset to a query dataset using Scanpy ingest. Ingest maps cells from a query dataset to a reference dataset that has been preprocessed and annotated. It projects query cells onto the PCA/UMAP of the reference and transfers labels based on nearest neighbors. Requirements: - Reference must have computed PCA (and UMAP if embedding='umap') - Reference must have the annotation columns specified in obs_keys - Both datasets should have overlapping genes Common use cases: - Transfer cell type labels from annotated scRNA-seq to spatial data - Annotate new samples using a well-curated reference atlas """ save_path = save_path or _config["save_path"] import scanpy as sc adata_query = sc.read_h5ad(adata_query_path) adata_ref = sc.read_h5ad(adata_ref_path) # Ensure gene names are uppercase adata_query.var.index = adata_query.var.index.str.upper() adata_ref.var.index = adata_ref.var.index.str.upper() # Find common genes common_genes = list(set(adata_query.var_names) & set(adata_ref.var_names)) if len(common_genes) < 100: return f"ERROR: Only {len(common_genes)} common genes found between query and reference. Need at least 100." # Subset to common genes adata_query = adata_query[:, common_genes].copy() adata_ref = adata_ref[:, common_genes].copy() # Check if reference has required representations if use_rep not in adata_ref.obsm: return f"ERROR: Reference missing '{use_rep}' in obsm. Run PCA on reference first." if embedding == 'umap' and 'X_umap' not in adata_ref.obsm: return f"ERROR: Reference missing 'X_umap' in obsm. Run UMAP on reference first or use embedding='pca'." # Check if obs_keys exist in reference missing_keys = [k for k in obs_keys if k not in adata_ref.obs.columns] if missing_keys: return f"ERROR: Reference missing annotation keys: {missing_keys}. Available: {list(adata_ref.obs.columns)}" # Run ingest sc.tl.ingest( adata_query, adata_ref, obs=obs_keys, embedding_method=embedding, ) # Save results os.makedirs(save_path, exist_ok=True) adata_query.write_h5ad(f"{save_path}/ingest_result.h5ad", compression="gzip") # Save transferred annotations as CSV transferred_df = adata_query.obs[obs_keys].copy() transferred_df.to_csv(f"{save_path}/ingest_annotations.csv") result = f"Ingest complete:\n" result += f" - Query: {adata_query.n_obs} cells\n" result += f" - Reference: {adata_ref.n_obs} cells\n" result += f" - Common genes: {len(common_genes)}\n" result += f" - Transferred annotations: {obs_keys}\n" for key in obs_keys: n_categories = adata_query.obs[key].nunique() result += f" - {key}: {n_categories} unique values\n" result += f" - Saved to {save_path}/ingest_result.h5ad" print(result) return result # ============================================================================= # Scanpy Tools: Integration - BBKNN # ============================================================================= @tool def scanpy_bbknn( adata_path: Annotated[str, Field(description="Path to h5ad file with multiple batches")], batch_key: Annotated[str, Field(description="Column in adata.obs containing batch information")], n_pcs: Annotated[int, Field(description="Number of principal components to use")] = 50, neighbors_within_batch: Annotated[int, Field(description="Number of neighbors per batch")] = 3, trim: Annotated[int, Field(description="Trim KNN graph to this many neighbors per cell (0=no trim)")] = 0, use_rep: Annotated[str, Field(description="Key in adata.obsm for representation")] = "X_pca", run_umap: Annotated[bool, Field(description="Compute UMAP after batch correction")] = True, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Perform batch-balanced KNN (BBKNN) integration for multi-batch data. BBKNN modifies the neighborhood graph to balance connections across batches. Instead of finding k-nearest neighbors globally, it finds k neighbors from each batch, effectively removing batch effects in the graph structure. Requirements: - Data must have PCA computed - Batch key must exist in adata.obs Common use cases: - Integrate multiple spatial tissue sections - Integrate spatial data with scRNA-seq reference - Remove technical batch effects while preserving biological variation """ save_path = save_path or _config["save_path"] import scanpy as sc try: import bbknn except ImportError: return "ERROR: bbknn not installed. Install with: pip install bbknn" adata = sc.read_h5ad(adata_path) # Check batch key exists if batch_key not in adata.obs.columns: return f"ERROR: Batch key '{batch_key}' not found in adata.obs. Available: {list(adata.obs.columns)}" # Check PCA exists if use_rep not in adata.obsm: # Run PCA if not present if adata.n_vars > 2000: sc.pp.highly_variable_genes(adata, n_top_genes=2000) sc.tl.pca(adata, use_highly_variable=True) else: sc.tl.pca(adata) # Get batch statistics batch_counts = adata.obs[batch_key].value_counts() n_batches = len(batch_counts) # Run BBKNN bbknn.bbknn( adata, batch_key=batch_key, n_pcs=min(n_pcs, adata.obsm[use_rep].shape[1]), neighbors_within_batch=neighbors_within_batch, trim=trim if trim > 0 else None, ) # Optionally run UMAP if run_umap: sc.tl.umap(adata) # Save results os.makedirs(save_path, exist_ok=True) adata.write_h5ad(f"{save_path}/bbknn_result.h5ad", compression="gzip") result = f"BBKNN integration complete:\n" result += f" - {adata.n_obs} cells across {n_batches} batches\n" result += f" - Batch distribution:\n" for batch, count in batch_counts.head(5).items(): result += f" {batch}: {count} cells\n" if n_batches > 5: result += f" ... and {n_batches - 5} more batches\n" result += f" - UMAP computed: {run_umap}\n" result += f" - Saved to {save_path}/bbknn_result.h5ad" print(result) return result # ============================================================================= # Trajectory Inference: scVelo # ============================================================================= @tool def scvelo_velocity( adata_path: Annotated[str, Field(description="Path to h5ad file with spliced/unspliced counts")], mode: Annotated[str, Field(description="Velocity mode: 'deterministic', 'stochastic', or 'dynamical'")] = "stochastic", n_top_genes: Annotated[int, Field(description="Number of highly variable genes to use")] = 2000, n_pcs: Annotated[int, Field(description="Number of PCs for neighbors")] = 30, n_neighbors: Annotated[int, Field(description="Number of neighbors")] = 30, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Compute RNA velocity using scVelo. RNA velocity estimates the rate of gene expression change by comparing spliced and unspliced mRNA counts. This reveals the future state of cells and can infer developmental trajectories. Requirements: - Data must have spliced ('spliced') and unspliced ('unspliced') layers - These are typically from velocyto, kallisto bustools, or STARsolo Modes: - deterministic: Fast, assumes steady-state (original velocyto model) - stochastic: Accounts for transcriptional stochasticity (recommended) - dynamical: Full dynamical model, most accurate but slower """ save_path = save_path or _config["save_path"] import scvelo as scv import scanpy as sc adata = sc.read_h5ad(adata_path) # Check for required layers if 'spliced' not in adata.layers or 'unspliced' not in adata.layers: return "ERROR: Data must have 'spliced' and 'unspliced' layers. Run velocyto or kallisto bustools first." # Filter and normalize scv.pp.filter_and_normalize(adata, min_shared_counts=20, n_top_genes=n_top_genes) # Compute moments scv.pp.moments(adata, n_pcs=n_pcs, n_neighbors=n_neighbors) # Compute velocity if mode == 'dynamical': scv.tl.recover_dynamics(adata) scv.tl.velocity(adata, mode='dynamical') else: scv.tl.velocity(adata, mode=mode) # Compute velocity graph scv.tl.velocity_graph(adata) # Save results os.makedirs(save_path, exist_ok=True) adata.write_h5ad(f"{save_path}/velocity_result.h5ad", compression="gzip") # Get velocity statistics n_cells = adata.n_obs n_genes_velocity = adata.var['velocity_genes'].sum() if 'velocity_genes' in adata.var else 'N/A' result = f"RNA velocity computed successfully:\n" result += f" - Mode: {mode}\n" result += f" - Cells: {n_cells}\n" result += f" - Velocity genes: {n_genes_velocity}\n" result += f" - Saved to {save_path}/velocity_result.h5ad\n" result += f"\nNext steps: Use scvelo_velocity_embedding for visualization" print(result) return result @tool def scvelo_velocity_embedding( adata_path: Annotated[str, Field(description="Path to h5ad file with computed velocity")], basis: Annotated[str, Field(description="Embedding to use: 'umap', 'tsne', or 'pca'")] = "umap", color_by: Annotated[str, Field(description="Column in adata.obs to color by")] = None, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Project RNA velocity onto embedding and create velocity stream plot. Creates a velocity stream plot showing the direction and magnitude of cellular state changes on a low-dimensional embedding. Requirements: - Data must have velocity computed (from scvelo_velocity) - Data must have the specified embedding (e.g., X_umap) """ save_path = save_path or _config["save_path"] import scvelo as scv import scanpy as sc import matplotlib.pyplot as plt adata = sc.read_h5ad(adata_path) # Check velocity exists if 'velocity' not in adata.layers: return "ERROR: Velocity not computed. Run scvelo_velocity first." # Check embedding exists embed_key = f"X_{basis}" if embed_key not in adata.obsm: # Try to compute it import scanpy as sc if basis == 'umap': sc.tl.umap(adata) elif basis == 'tsne': sc.tl.tsne(adata) else: return f"ERROR: Embedding '{basis}' not found. Available: {list(adata.obsm.keys())}" os.makedirs(save_path, exist_ok=True) # Create velocity embedding stream plot fig, ax = plt.subplots(figsize=(10, 8)) scv.pl.velocity_embedding_stream(adata, basis=basis, color=color_by, ax=ax, show=False) plt.savefig(f"{save_path}/velocity_stream_{basis}.png", dpi=150, bbox_inches='tight') plt.close() # Create velocity embedding grid plot fig, ax = plt.subplots(figsize=(10, 8)) scv.pl.velocity_embedding_grid(adata, basis=basis, color=color_by, ax=ax, show=False) plt.savefig(f"{save_path}/velocity_grid_{basis}.png", dpi=150, bbox_inches='tight') plt.close() # Save updated adata adata.write_h5ad(f"{save_path}/velocity_embedding.h5ad", compression="gzip") result = f"Velocity embedding created:\n" result += f" - Basis: {basis}\n" result += f" - Stream plot: {save_path}/velocity_stream_{basis}.png\n" result += f" - Grid plot: {save_path}/velocity_grid_{basis}.png\n" result += f" - Data saved to: {save_path}/velocity_embedding.h5ad" print(result) return result # ============================================================================= # Trajectory Inference: CellRank # ============================================================================= @tool def cellrank_terminal_states( adata_path: Annotated[str, Field(description="Path to h5ad file with velocity or connectivity")], cluster_key: Annotated[str, Field(description="Column in adata.obs for cell clusters")] = None, n_states: Annotated[int, Field(description="Number of terminal states to find (0=auto)")] = 0, use_velocity: Annotated[bool, Field(description="Use RNA velocity (requires velocity layer)")] = True, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Identify terminal/macrostates using CellRank. CellRank uses Markov chain analysis to identify terminal states (endpoints) and initial states of cellular differentiation trajectories. Requirements: - If use_velocity=True: Data must have velocity (from scvelo_velocity) - If use_velocity=False: Data must have neighbors graph Outputs: - Terminal states (macrostates) identified in adata.obs - Transition probabilities stored in adata """ save_path = save_path or _config["save_path"] import cellrank as cr import scanpy as sc adata = sc.read_h5ad(adata_path) # Create kernel based on data type if use_velocity: if 'velocity' not in adata.layers: return "ERROR: Velocity not found. Run scvelo_velocity first or set use_velocity=False." # Velocity kernel vk = cr.kernels.VelocityKernel(adata) vk.compute_transition_matrix() kernel = vk else: # Connectivity kernel (based on neighbors graph) if 'neighbors' not in adata.uns: sc.pp.neighbors(adata) ck = cr.kernels.ConnectivityKernel(adata) ck.compute_transition_matrix() kernel = ck # Create GPCCA estimator for terminal states g = cr.estimators.GPCCA(kernel) # Compute Schur decomposition g.compute_schur(n_components=20) # Compute macrostates if n_states > 0: g.compute_macrostates(n_states=n_states, cluster_key=cluster_key) else: g.compute_macrostates(cluster_key=cluster_key) # Identify terminal states g.predict_terminal_states() # Save results os.makedirs(save_path, exist_ok=True) adata.write_h5ad(f"{save_path}/cellrank_terminal.h5ad", compression="gzip") # Get terminal state info terminal_states = adata.obs['terminal_states'].dropna().unique().tolist() if 'terminal_states' in adata.obs else [] result = f"CellRank terminal state analysis complete:\n" result += f" - Method: {'Velocity kernel' if use_velocity else 'Connectivity kernel'}\n" result += f" - Terminal states found: {len(terminal_states)}\n" if terminal_states: result += f" - States: {terminal_states}\n" result += f" - Saved to {save_path}/cellrank_terminal.h5ad\n" result += f"\nNext steps: Use cellrank_fate_probabilities to compute fate maps" print(result) return result @tool def cellrank_fate_probabilities( adata_path: Annotated[str, Field(description="Path to h5ad file with terminal states")], save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Compute fate probabilities for each cell towards terminal states. After identifying terminal states with cellrank_terminal_states, this computes the probability of each cell reaching each terminal state. Requirements: - Data must have terminal states computed (from cellrank_terminal_states) Outputs: - Fate probabilities in adata.obsm['to_terminal_states'] - Lineage drivers (genes correlated with fate) """ save_path = save_path or _config["save_path"] import cellrank as cr import scanpy as sc import matplotlib.pyplot as plt adata = sc.read_h5ad(adata_path) # Check terminal states exist if 'terminal_states' not in adata.obs: return "ERROR: Terminal states not found. Run cellrank_terminal_states first." # Recreate estimator if 'velocity' in adata.layers: vk = cr.kernels.VelocityKernel(adata) vk.compute_transition_matrix() kernel = vk else: ck = cr.kernels.ConnectivityKernel(adata) ck.compute_transition_matrix() kernel = ck g = cr.estimators.GPCCA(kernel) g.compute_schur(n_components=20) # Set terminal states from previous analysis terminal_states = adata.obs['terminal_states'].dropna().unique().tolist() g.set_terminal_states(terminal_states) # Compute fate probabilities g.compute_fate_probabilities() os.makedirs(save_path, exist_ok=True) # Save fate probabilities if hasattr(g, 'fate_probabilities') and g.fate_probabilities is not None: fate_df = g.fate_probabilities.to_frame() fate_df.to_csv(f"{save_path}/fate_probabilities.csv") # Save updated adata adata.write_h5ad(f"{save_path}/cellrank_fate.h5ad", compression="gzip") result = f"Fate probabilities computed:\n" result += f" - Terminal states: {terminal_states}\n" result += f" - Fate probabilities saved to: {save_path}/fate_probabilities.csv\n" result += f" - Data saved to: {save_path}/cellrank_fate.h5ad" print(result) return result @tool def paga_trajectory( adata_path: Annotated[str, Field(description="Path to preprocessed h5ad file")], groups_key: Annotated[str, Field(description="Column in adata.obs for cell groups (e.g., 'leiden', 'cell_type')")], threshold: Annotated[float, Field(description="Threshold for PAGA edges (0-1)")] = 0.05, root_group: Annotated[str, Field(description="Name of root group for pseudotime (optional)")] = None, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Compute PAGA (Partition-based Graph Abstraction) trajectory. PAGA computes a coarse-grained graph of cell groups showing connectivity and potential differentiation paths. It can also compute diffusion pseudotime. Requirements: - Data must have neighbors computed - Data must have the specified groups column """ save_path = save_path or _config["save_path"] import scanpy as sc import matplotlib.pyplot as plt adata = sc.read_h5ad(adata_path) # Check groups exist if groups_key not in adata.obs: return f"ERROR: Groups key '{groups_key}' not found. Available: {list(adata.obs.columns)}" # Compute neighbors if not present if 'neighbors' not in adata.uns: sc.pp.pca(adata) sc.pp.neighbors(adata) # Compute PAGA sc.tl.paga(adata, groups=groups_key) os.makedirs(save_path, exist_ok=True) # Plot PAGA fig, ax = plt.subplots(figsize=(10, 8)) sc.pl.paga(adata, threshold=threshold, ax=ax, show=False) plt.savefig(f"{save_path}/paga_graph.png", dpi=150, bbox_inches='tight') plt.close() # Compute diffusion pseudotime if root specified if root_group: # Find a cell in root group root_cells = adata.obs[adata.obs[groups_key] == root_group].index if len(root_cells) == 0: return f"ERROR: Root group '{root_group}' not found. Available: {adata.obs[groups_key].unique().tolist()}" adata.uns['iroot'] = adata.obs_names.get_loc(root_cells[0]) sc.tl.diffmap(adata) sc.tl.dpt(adata) # Plot pseudotime if 'X_umap' in adata.obsm: fig, ax = plt.subplots(figsize=(10, 8)) sc.pl.umap(adata, color='dpt_pseudotime', ax=ax, show=False) plt.savefig(f"{save_path}/pseudotime_umap.png", dpi=150, bbox_inches='tight') plt.close() # Save PAGA adjacency paga_adj = adata.uns['paga']['connectivities'].toarray() groups = adata.obs[groups_key].cat.categories.tolist() paga_df = pd.DataFrame(paga_adj, index=groups, columns=groups) paga_df.to_csv(f"{save_path}/paga_adjacency.csv") # Save results adata.write_h5ad(f"{save_path}/paga_result.h5ad", compression="gzip") result = f"PAGA trajectory analysis complete:\n" result += f" - Groups: {groups_key} ({len(groups)} groups)\n" result += f" - PAGA graph: {save_path}/paga_graph.png\n" result += f" - Adjacency matrix: {save_path}/paga_adjacency.csv\n" if root_group: result += f" - Pseudotime computed from root: {root_group}\n" result += f" - Pseudotime plot: {save_path}/pseudotime_umap.png\n" result += f" - Data saved to: {save_path}/paga_result.h5ad" print(result) return result # ============================================================================= # Multimodal Integration: scvi-tools # ============================================================================= @tool def totalvi_integration( adata_path: Annotated[str, Field(description="Path to h5ad file with RNA and protein (CITE-seq) data")], protein_layer: Annotated[str, Field(description="Key in adata.obsm for protein counts")] = "protein_expression", n_latent: Annotated[int, Field(description="Dimensionality of latent space")] = 20, max_epochs: Annotated[int, Field(description="Maximum training epochs")] = 400, batch_key: Annotated[str, Field(description="Batch column for integration (optional)")] = None, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Integrate RNA and protein data using totalVI (CITE-seq analysis). totalVI is a deep generative model for joint analysis of RNA and protein measurements from CITE-seq or similar multi-modal single-cell technologies. Requirements: - RNA counts in adata.X (raw counts, not normalized) - Protein counts in adata.obsm[protein_layer] Outputs: - Latent representation in adata.obsm['X_totalVI'] - Denoised protein values - Normalized RNA values """ save_path = save_path or _config["save_path"] import scvi import scanpy as sc adata = sc.read_h5ad(adata_path) # Check protein data exists if protein_layer not in adata.obsm: return f"ERROR: Protein layer '{protein_layer}' not found. Available obsm: {list(adata.obsm.keys())}" # Get protein names if available protein_names = None if f"{protein_layer}_names" in adata.uns: protein_names = adata.uns[f"{protein_layer}_names"] # Setup anndata for totalVI scvi.model.TOTALVI.setup_anndata( adata, protein_expression_obsm_key=protein_layer, batch_key=batch_key, ) # Create and train model model = scvi.model.TOTALVI(adata, n_latent=n_latent) model.train(max_epochs=max_epochs, early_stopping=True) # Get latent representation adata.obsm['X_totalVI'] = model.get_latent_representation() # Get normalized values rna_norm, protein_norm = model.get_normalized_expression(return_mean=True) adata.layers['totalvi_rna_normalized'] = rna_norm # Convert protein_norm to numpy array to avoid column name issues if hasattr(protein_norm, 'values'): adata.obsm['totalvi_protein_normalized'] = protein_norm.values else: adata.obsm['totalvi_protein_normalized'] = protein_norm # Compute UMAP on totalVI latent space sc.pp.neighbors(adata, use_rep='X_totalVI') sc.tl.umap(adata) os.makedirs(save_path, exist_ok=True) # Save model model.save(f"{save_path}/totalvi_model", overwrite=True) # Save results adata.write_h5ad(f"{save_path}/totalvi_result.h5ad", compression="gzip") result = f"totalVI integration complete:\n" result += f" - Cells: {adata.n_obs}\n" result += f" - RNA genes: {adata.n_vars}\n" result += f" - Proteins: {adata.obsm[protein_layer].shape[1]}\n" result += f" - Latent dims: {n_latent}\n" result += f" - Batch integration: {batch_key if batch_key else 'None'}\n" result += f" - Model saved to: {save_path}/totalvi_model\n" result += f" - Data saved to: {save_path}/totalvi_result.h5ad" print(result) return result @tool def multivi_integration( adata_path: Annotated[str, Field(description="Path to h5ad file with RNA and ATAC data")], atac_layer: Annotated[str, Field(description="Key in adata.obsm or layer for ATAC peaks")] = None, n_latent: Annotated[int, Field(description="Dimensionality of latent space")] = 20, max_epochs: Annotated[int, Field(description="Maximum training epochs")] = 500, batch_key: Annotated[str, Field(description="Batch column for integration (optional)")] = None, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Integrate RNA and ATAC data using MultiVI (multiome analysis). MultiVI is a deep generative model for joint analysis of gene expression and chromatin accessibility from 10x Multiome or similar technologies. Requirements: - Data with both modalities, or MuData object - ATAC peaks either as separate vars or in a layer Outputs: - Latent representation in adata.obsm['X_multivi'] - Imputed accessibility for RNA-only cells - Imputed expression for ATAC-only cells """ save_path = save_path or _config["save_path"] import scvi import scanpy as sc adata = sc.read_h5ad(adata_path) # For MultiVI, we need to properly organize the data # This is a simplified version - real usage may need MuData # Setup anndata scvi.model.MULTIVI.setup_anndata( adata, batch_key=batch_key, ) # Create and train model model = scvi.model.MULTIVI(adata, n_latent=n_latent) model.train(max_epochs=max_epochs, early_stopping=True) # Get latent representation adata.obsm['X_multivi'] = model.get_latent_representation() # Compute UMAP on MultiVI latent space sc.pp.neighbors(adata, use_rep='X_multivi') sc.tl.umap(adata) os.makedirs(save_path, exist_ok=True) # Save model model.save(f"{save_path}/multivi_model", overwrite=True) # Save results adata.write_h5ad(f"{save_path}/multivi_result.h5ad", compression="gzip") result = f"MultiVI integration complete:\n" result += f" - Cells: {adata.n_obs}\n" result += f" - Features: {adata.n_vars}\n" result += f" - Latent dims: {n_latent}\n" result += f" - Model saved to: {save_path}/multivi_model\n" result += f" - Data saved to: {save_path}/multivi_result.h5ad" print(result) return result @tool def mofa_integration( adata_path: Annotated[str, Field(description="Path to h5ad file OR comma-separated paths for multiple modalities")], n_factors: Annotated[int, Field(description="Number of latent factors to learn")] = 10, modality_key: Annotated[str, Field(description="Column in adata.var indicating modality (if single file)")] = None, groups_key: Annotated[str, Field(description="Column in adata.obs for sample/group structure")] = None, max_iterations: Annotated[int, Field(description="Maximum training iterations")] = 1000, save_path: Annotated[str, Field(description="Directory to save results")] = None, ) -> str: """Perform multi-omics factor analysis using MOFA+. MOFA+ identifies latent factors that explain variation across multiple data modalities (e.g., RNA, protein, ATAC, methylation). Requirements: - Data with multiple modalities, organized by: 1. Single adata with modality_key in var, OR 2. Comma-separated paths to multiple h5ad files Outputs: - Latent factors in adata.obsm['X_mofa'] - Factor loadings (weights) per modality - Variance explained per factor per modality """ save_path = save_path or _config["save_path"] from mofapy2.run.entry_point import entry_point import scanpy as sc # Check if multiple files if ',' in adata_path: paths = [p.strip() for p in adata_path.split(',')] adatas = [sc.read_h5ad(p) for p in paths] modality_names = [f"modality_{i}" for i in range(len(adatas))] else: adata = sc.read_h5ad(adata_path) if modality_key and modality_key in adata.var: # Split by modality modalities = adata.var[modality_key].unique() adatas = [adata[:, adata.var[modality_key] == m].copy() for m in modalities] modality_names = list(modalities) else: # Single modality adatas = [adata] modality_names = ["expression"] # Prepare data for MOFA # MOFA expects: data[modality][group] = matrix data = {} for mod_name, ad in zip(modality_names, adatas): if groups_key and groups_key in ad.obs: data[mod_name] = {} for group in ad.obs[groups_key].unique(): mask = ad.obs[groups_key] == group data[mod_name][group] = ad[mask].X.toarray() if hasattr(ad.X, 'toarray') else ad.X else: data[mod_name] = {"group1": ad.X.toarray() if hasattr(ad.X, 'toarray') else ad.X} # Initialize MOFA ent = entry_point() ent.set_data_options(scale_views=True) ent.set_data_matrix(data) ent.set_model_options(factors=n_factors) ent.set_train_options(iter=max_iterations, convergence_mode="slow", seed=42) # Train ent.build() ent.run() os.makedirs(save_path, exist_ok=True) # Save MOFA model ent.save(f"{save_path}/mofa_model.hdf5") # Get factors and add to first adata factors = ent.model.nodes["Z"].getExpectation() if isinstance(factors, dict): # Multiple groups - concatenate all_factors = np.vstack(list(factors.values())) else: all_factors = factors # Get weights for each modality weights = ent.model.nodes["W"].getExpectation() # Get variance explained var_explained = ent.model.calculate_variance_explained() # Save factors factors_df = pd.DataFrame( all_factors, columns=[f'Factor_{i+1}' for i in range(n_factors)] ) factors_df.to_csv(f"{save_path}/mofa_factors.csv") # Save variance explained (handle 3D output: groups x views x factors) try: if var_explained.ndim == 3: # Reshape to 2D (views x factors) for single group var_explained_2d = var_explained[0] # Take first group var_df = pd.DataFrame( var_explained_2d, index=[f'view_{i}' for i in range(var_explained_2d.shape[0])], columns=[f'Factor_{i+1}' for i in range(var_explained_2d.shape[1])] ) else: var_df = pd.DataFrame(var_explained) var_df.to_csv(f"{save_path}/mofa_variance_explained.csv") except Exception as e: # If variance explained fails, just save factors pass result = f"MOFA+ analysis complete:\n" result += f" - Modalities: {modality_names}\n" result += f" - Factors: {n_factors}\n" result += f" - Samples: {all_factors.shape[0]}\n" result += f" - Model saved to: {save_path}/mofa_model.hdf5\n" result += f" - Factors saved to: {save_path}/mofa_factors.csv\n" result += f" - Variance explained: {save_path}/mofa_variance_explained.csv" print(result) return result