chatspatial 1.1.0__py3-none-any.whl

chatspatial/tools/visualization/integration.py

@@ -0,0 +1,205 @@
+"""
+Batch integration visualization functions.
+
+This module contains:
+- Batch integration quality assessment visualization
+"""
+
+from typing import TYPE_CHECKING, Optional
+
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.stats import entropy
+
+from ...models.data import VisualizationParameters
+from ...utils.adata_utils import get_spatial_key, validate_obs_column
+
+if TYPE_CHECKING:
+    import anndata as ad
+
+    from ...spatial_mcp_adapter import ToolContext
+
+
+# =============================================================================
+# Batch Integration Visualization
+# =============================================================================
+
+
+async def create_batch_integration_visualization(
+    adata: "ad.AnnData",
+    params: VisualizationParameters,
+    context: Optional["ToolContext"] = None,
+) -> plt.Figure:
+    """Create multi-panel visualization to assess batch integration quality.
+
+    This visualization is specifically for evaluating the quality of batch correction
+    after integrating multiple samples. It requires proper batch information.
+
+    Args:
+        adata: AnnData object with integrated samples
+        params: Visualization parameters (batch_key required)
+        context: MCP context for logging
+
+    Returns:
+        matplotlib Figure object
+
+    Raises:
+        DataNotFoundError: If batch information not found
+    """
+    if context:
+        await context.info("Creating batch integration quality visualization")
+
+    # Validate batch key exists
+    batch_key = params.batch_key
+    validate_obs_column(adata, batch_key, "Batch")
+
+    # Create multi-panel figure (2x2 layout)
+    figsize = params.figure_size if params.figure_size else (16, 12)
+    fig, axes = plt.subplots(2, 2, figsize=figsize)
+
+    batch_values = adata.obs[batch_key]
+    unique_batches = batch_values.unique()
+    colors = plt.get_cmap("Set3")(np.linspace(0, 1, len(unique_batches)))
+
+    # Panel 1: UMAP colored by batch (shows mixing)
+    if "X_umap" in adata.obsm:
+        umap_coords = adata.obsm["X_umap"]
+
+        for i, batch in enumerate(unique_batches):
+            mask = batch_values == batch
+            axes[0, 0].scatter(
+                umap_coords[mask, 0],
+                umap_coords[mask, 1],
+                c=[colors[i]],
+                label=f"{batch}",
+                s=5,
+                alpha=0.7,
+            )
+
+        axes[0, 0].set_title(
+            "UMAP colored by batch\n(Good integration = mixed colors)", fontsize=12
+        )
+        axes[0, 0].set_xlabel("UMAP 1")
+        axes[0, 0].set_ylabel("UMAP 2")
+        axes[0, 0].legend(bbox_to_anchor=(1.05, 1), loc="upper left")
+    else:
+        if context:
+            await context.warning(
+                "UMAP coordinates not available. "
+                "Run preprocessing with UMAP computation for complete visualization."
+            )
+        axes[0, 0].text(
+            0.5, 0.5, "UMAP coordinates not available", ha="center", va="center"
+        )
+        axes[0, 0].set_title("UMAP (Not Available)", fontsize=12)
+
+    # Panel 2: Spatial plot colored by batch (if spatial data available)
+    spatial_key = get_spatial_key(adata)
+    if spatial_key:
+        spatial_coords = adata.obsm[spatial_key]
+
+        for i, batch in enumerate(unique_batches):
+            mask = batch_values == batch
+            axes[0, 1].scatter(
+                spatial_coords[mask, 0],
+                spatial_coords[mask, 1],
+                c=[colors[i]],
+                label=f"{batch}",
+                s=10,
+                alpha=0.7,
+            )
+
+        axes[0, 1].set_title("Spatial coordinates colored by batch", fontsize=12)
+        axes[0, 1].set_xlabel("Spatial X")
+        axes[0, 1].set_ylabel("Spatial Y")
+        axes[0, 1].set_aspect("equal")
+        axes[0, 1].legend(bbox_to_anchor=(1.05, 1), loc="upper left")
+    else:
+        if context:
+            await context.info(
+                "Spatial coordinates not available. "
+                "This is expected for non-spatial datasets."
+            )
+        axes[0, 1].text(
+            0.5, 0.5, "Spatial coordinates not available", ha="center", va="center"
+        )
+        axes[0, 1].set_title("Spatial (Not Available)", fontsize=12)
+
+    # Panel 3: Batch composition bar plot
+    batch_counts = adata.obs[batch_key].value_counts()
+    axes[1, 0].bar(
+        range(len(batch_counts)),
+        batch_counts.values,
+        color=colors[: len(batch_counts)],
+    )
+    axes[1, 0].set_xticks(range(len(batch_counts)))
+    axes[1, 0].set_xticklabels(batch_counts.index, rotation=45, ha="right")
+    axes[1, 0].set_title("Cell counts per batch", fontsize=12)
+    axes[1, 0].set_ylabel("Number of cells")
+
+    # Panel 4: Integration quality metrics (if available)
+    axes[1, 1].text(
+        0.1,
+        0.9,
+        "Integration Quality Assessment:",
+        fontsize=14,
+        fontweight="bold",
+        transform=axes[1, 1].transAxes,
+    )
+
+    metrics_text = f"Total cells: {adata.n_obs:,}\n"
+    metrics_text += f"Total genes: {adata.n_vars:,}\n"
+    metrics_text += (
+        f"Batches: {len(unique_batches)} ({', '.join(map(str, unique_batches))})\n\n"
+    )
+
+    if params.integration_method:
+        metrics_text += f"Integration method: {params.integration_method}\n"
+
+    # Add basic mixing metrics
+    if "X_umap" in adata.obsm:
+        # Calculate simple mixing metric (entropy)
+        umap_coords = adata.obsm["X_umap"]
+        x_bins = np.linspace(umap_coords[:, 0].min(), umap_coords[:, 0].max(), 10)
+        y_bins = np.linspace(umap_coords[:, 1].min(), umap_coords[:, 1].max(), 10)
+
+        entropies = []
+        for i in range(len(x_bins) - 1):
+            for j in range(len(y_bins) - 1):
+                mask = (
+                    (umap_coords[:, 0] >= x_bins[i])
+                    & (umap_coords[:, 0] < x_bins[i + 1])
+                    & (umap_coords[:, 1] >= y_bins[j])
+                    & (umap_coords[:, 1] < y_bins[j + 1])
+                )
+                if mask.sum() > 10:  # Only consider regions with enough cells
+                    batch_props = adata.obs[batch_key][mask].value_counts(
+                        normalize=True
+                    )
+                    entropies.append(entropy(batch_props))
+
+        if entropies:
+            avg_entropy = np.mean(entropies)
+            max_entropy = np.log(len(unique_batches))  # Perfect mixing entropy
+            mixing_score = avg_entropy / max_entropy if max_entropy > 0 else 0
+            metrics_text += (
+                f"Mixing score: {mixing_score:.3f} (0=segregated, 1=perfectly mixed)\n"
+            )
+
+    axes[1, 1].text(
+        0.1,
+        0.7,
+        metrics_text,
+        fontsize=10,
+        transform=axes[1, 1].transAxes,
+        verticalalignment="top",
+        fontfamily="monospace",
+    )
+    axes[1, 1].set_xlim(0, 1)
+    axes[1, 1].set_ylim(0, 1)
+    axes[1, 1].set_xticks([])
+    axes[1, 1].set_yticks([])
+    axes[1, 1].set_title("Integration Metrics", fontsize=12)
+
+    plt.tight_layout()
+    return fig

chatspatial/tools/visualization/main.py

@@ -0,0 +1,164 @@
+"""
+Main visualization entry point.
+
+This module contains the main visualize_data function that dispatches
+to appropriate visualization handlers based on plot_type.
+"""
+
+import traceback
+from typing import TYPE_CHECKING, Union
+
+import matplotlib.pyplot as plt
+import scanpy as sc
+from mcp.server.fastmcp.utilities.types import ImageContent
+from mcp.types import EmbeddedResource
+
+from ...models.data import VisualizationParameters
+from ...utils.exceptions import (
+    DataCompatibilityError,
+    DataNotFoundError,
+    ParameterError,
+    ProcessingError,
+)
+from ...utils.image_utils import optimize_fig_to_image_with_cache
+
+# Import all visualization handlers
+from .basic import (
+    create_dotplot_visualization,
+    create_heatmap_visualization,
+    create_spatial_visualization,
+    create_umap_visualization,
+    create_violin_visualization,
+)
+from .cell_comm import create_cell_communication_visualization
+from .cnv import create_cnv_heatmap_visualization, create_spatial_cnv_visualization
+from .deconvolution import (
+    create_card_imputation_visualization,
+    create_deconvolution_visualization,
+)
+from .enrichment import create_pathway_enrichment_visualization
+from .integration import create_batch_integration_visualization
+from .multi_gene import (
+    create_gene_correlation_visualization,
+    create_lr_pairs_visualization,
+    create_multi_gene_visualization,
+    create_spatial_interaction_visualization,
+)
+from .spatial_stats import create_spatial_statistics_visualization
+from .trajectory import create_trajectory_visualization
+from .velocity import create_rna_velocity_visualization
+
+if TYPE_CHECKING:
+    from ...spatial_mcp_adapter import ToolContext
+
+
+# Handler registry for dispatch - defined here to avoid circular imports
+PLOT_HANDLERS = {
+    # Basic plots
+    "spatial": create_spatial_visualization,
+    "umap": create_umap_visualization,
+    "heatmap": create_heatmap_visualization,
+    "violin": create_violin_visualization,
+    "dotplot": create_dotplot_visualization,
+    # Analysis-specific plots
+    "deconvolution": create_deconvolution_visualization,
+    "cell_communication": create_cell_communication_visualization,
+    "rna_velocity": create_rna_velocity_visualization,
+    "trajectory": create_trajectory_visualization,
+    "spatial_statistics": create_spatial_statistics_visualization,
+    "pathway_enrichment": create_pathway_enrichment_visualization,
+    # CNV plots
+    "card_imputation": create_card_imputation_visualization,
+    "spatial_cnv": create_spatial_cnv_visualization,
+    "cnv_heatmap": create_cnv_heatmap_visualization,
+    # Integration plots
+    "batch_integration": create_batch_integration_visualization,
+    # Multi-gene plots
+    "multi_gene": create_multi_gene_visualization,
+    "lr_pairs": create_lr_pairs_visualization,
+    "gene_correlation": create_gene_correlation_visualization,
+    "spatial_interaction": create_spatial_interaction_visualization,
+}
+
+
+async def visualize_data(
+    data_id: str,
+    ctx: "ToolContext",
+    params: VisualizationParameters = VisualizationParameters(),  # type: ignore[call-arg]
+) -> Union[ImageContent, tuple[ImageContent, EmbeddedResource]]:
+    """Visualize spatial transcriptomics data.
+
+    Args:
+        data_id: Dataset ID
+        ctx: ToolContext for unified data access and logging
+        params: Visualization parameters
+
+    Returns:
+        Union[ImageContent, Tuple[ImageContent, EmbeddedResource]]:
+            - Small images (<100KB): ImageContent object
+            - Large images (>=100KB): Tuple[Preview ImageContent, High-quality Resource]
+
+    Raises:
+        DataNotFoundError: If the dataset is not found
+        ParameterError: If parameters are invalid
+        DataCompatibilityError: If data is not compatible with the visualization
+        ProcessingError: If processing fails
+    """
+    # Validate parameters - use PLOT_HANDLERS as single source of truth
+    if params.plot_type not in PLOT_HANDLERS:
+        raise ParameterError(
+            f"Invalid plot_type: {params.plot_type}. "
+            f"Must be one of {list(PLOT_HANDLERS)}"
+        )
+
+    try:
+        # Retrieve the AnnData object via ToolContext
+        adata = await ctx.get_adata(data_id)
+
+        # Validate AnnData object - basic validation
+        if adata.n_obs < 5:
+            raise DataNotFoundError("Dataset has too few cells (minimum 5 required)")
+        if adata.n_vars < 5:
+            raise DataNotFoundError("Dataset has too few genes (minimum 5 required)")
+
+        # Set matplotlib style for better visualizations
+        sc.settings.set_figure_params(dpi=params.dpi or 100, facecolor="white")
+
+        # Dispatch to appropriate handler
+        handler = PLOT_HANDLERS[params.plot_type]
+        fig = await handler(adata, params, ctx._mcp_context)
+
+        # Generate plot_type_key with subtype if applicable (for cache consistency)
+        subtype = params.subtype
+        plot_type_key = f"{params.plot_type}_{subtype}" if subtype else params.plot_type
+
+        # Use the optimized conversion function
+        return await optimize_fig_to_image_with_cache(
+            fig,
+            params,
+            ctx._mcp_context,
+            data_id=data_id,
+            plot_type=plot_type_key,
+            mode="auto",
+        )
+
+    except Exception as e:
+        # Make sure to close any open figures in case of error
+        plt.close("all")
+
+        # For image conversion errors, return error message as string
+        if "fig_to_image" in str(e) or "convert" in str(e).lower():
+            error_details = traceback.format_exc()
+            return (
+                f"Error in {params.plot_type} visualization:\n\n"
+                f"{e}\n\n"
+                f"Technical details:\n{error_details}"
+            )
+
+        # Wrap the error in a more informative exception
+        if isinstance(e, (DataNotFoundError, ParameterError, DataCompatibilityError)):
+            raise
+        else:
+            raise ProcessingError(
+                f"Failed to create {params.plot_type} visualization: {e}"
+            ) from e