dms_corr_plot¶

def dms_corr_plot(
    uniprot_id: str,
    model: str = "AlphaMissense",
    dms_table: Optional[Dict[str, pd.DataFrame]] = None,
    cache_dir: str = ".cache",
    figsize: Tuple[float, float] = (10, 8),
    bins: int = 60,
    cmap: str = 'viridis'
) -> Figure:
    """
    Compare DMS experimental scores with model predictions using Spearman correlation.

    Creates a 2D density plot (hexbin) with marginal distributions showing the
    relationship between experimental DMS scores and predicted model scores.

    Args:
        uniprot_id: Valid UniProt accession identifier
        model: Model to plot (default: "AlphaMissense")
        dms_table: Optional dict of DMS DataFrames (if None, loads default)
        cache_dir: Directory to cache downloaded files
        figsize: Figure size as (width, height)
        bins: Number of bins for hexbin plot
        cmap: Colormap for hexbin plot

    Returns:
        matplotlib Figure object

    Examples:
        >>> # Use defaults (AlphaMissense vs DMS scores)
        >>> fig = dms_corr_plot(uniprot_id="Q9NV35")
        >>> plt.show()

        >>> # Compare specific model with DMS
        >>> fig = dms_corr_plot(
        ...     uniprot_id="P04637",
        ...     model="Kermut"
        ... )
        >>> plt.show()

    Notes:
        - Requires scipy for correlation analysis
        - Requires matplotlib for visualization
        - seaborn is optional but recommended for better styling
        - Model names can be from zero-shot, supervised, or AlphaMissense
        - Use available_zero_shot_models() and available_supervised_models()
          to see available models

    """
    # Check dependencies
    if stats is None:
        raise ImportError(
            "scipy is required for correlation analysis. "
            "Install with: pip install scipy"
        )

    # Validate uniprot_id
    if not isinstance(uniprot_id, str):
        raise ValueError("uniprot_id must be a string")

    # Validate model
    valid_models = (
        available_zero_shot_models() +
        available_supervised_models() +
        ["AlphaMissense"]
    )

    if model not in valid_models:
        raise ValueError(f"Invalid model specified: {model}")

    # Load model data for uniprot_id
    print(f"Loading data for {model}...")
    model_df = _get_model_dataframe(model, uniprot_id, cache_dir)

    # Load DMS data for uniprot_id
    print(f"Loading DMS scores for {uniprot_id}...")
    dms_df = _filter_dms_uniprot(dms_table, uniprot_id, cache_dir)

    # Merge tables
    print("Merging model predictions with DMS scores...")
    merged_table = _merge_model_dms_tables(
        model_df=model_df,
        dms_df=dms_df,
        model=model
    )

    # Check if merged table is empty
    if len(merged_table) == 0:
        raise ValueError(
            f"No common mutants between {model} and DMS scores for accession '{uniprot_id}'"
        )

    print(f"Found {len(merged_table)} common mutants")

    # Calculate correlation
    correlation, pvalue = _calculate_spearman_correlation(merged_table)

    print(f"Spearman r = {correlation:.3f}, p-value = {pvalue:.2e}")

    # Create figure with GridSpec for layout
    fig = plt.figure(figsize=figsize)
    gs = GridSpec(
        4, 3,
        figure=fig,
        hspace=0.15,
        wspace=0.05,
        width_ratios=[1, 4, 0.3],
        height_ratios=[1, 4, 0.3, 0.4]
    )

    # Main plot (center)
    ax_main = fig.add_subplot(gs[1, 1])

    # Marginal plots
    ax_top = fig.add_subplot(gs[0, 1], sharex=ax_main)
    ax_right = fig.add_subplot(gs[1, 2], sharey=ax_main)

    # Colorbar axis (row 3, below spacer row)
    ax_cbar = fig.add_subplot(gs[3, 1])

    # Create hexbin plot (x=DMS, y=model - matching R version)
    x = merged_table['mean_dms'].values
    y = merged_table['mean_model'].values

    hexbin = ax_main.hexbin(
        x, y,
        gridsize=bins,
        cmap=cmap,
        mincnt=1,
        edgecolors='none'
    )

    # Style main plot
    ax_main.set_xlabel('DMS score', fontsize=16)
    ax_main.set_ylabel(f'{model} score', fontsize=16)
    ax_main.tick_params(labelsize=16)

    # Add colorbar
    cbar = plt.colorbar(hexbin, cax=ax_cbar, orientation='horizontal')
    cbar.set_label('Count', fontsize=16)
    cbar.ax.tick_params(labelsize=16)

    # Top marginal (histogram + KDE for DMS scores)
    ax_top.hist(x, bins=30, color='#B0C4DE', edgecolor='black', alpha=0.7, density=True)
    if sns is not None:
        try:
            from scipy.stats import gaussian_kde
            kde = gaussian_kde(x)
            x_range = np.linspace(x.min(), x.max(), 100)
            ax_top.plot(x_range, kde(x_range), 'k-', linewidth=1.5)
        except:
            pass  # Skip KDE if it fails
    ax_top.set_ylabel('Density', fontsize=16)
    ax_top.tick_params(labelbottom=False, labelsize=16)
    ax_top.spines['top'].set_visible(False)
    ax_top.spines['right'].set_visible(False)

    # Right marginal (histogram + KDE for model scores)
    ax_right.hist(y, bins=30, color='#B0C4DE', edgecolor='black',
                  alpha=0.7, orientation='horizontal', density=True)
    if sns is not None:
        try:
            from scipy.stats import gaussian_kde
            kde = gaussian_kde(y)
            y_range = np.linspace(y.min(), y.max(), 100)
            ax_right.plot(kde(y_range), y_range, 'k-', linewidth=1.5)
        except:
            pass  # Skip KDE if it fails
    ax_right.set_xlabel('Density', fontsize=16)
    ax_right.tick_params(labelleft=False, labelsize=16)
    ax_right.spines['top'].set_visible(False)
    ax_right.spines['right'].set_visible(False)

    # Add title with correlation info
    title = (
        f"UniProt ID: {uniprot_id}\n"
        f"Spearman r = {correlation:.2f}, p-value = {pvalue:.2e}"
    )
    fig.suptitle(title, fontsize=14, y=0.98)

    return fig

def get_available_models() -> Dict[str, list]:
    """
    Get dictionary of all available models in ProteinGym.

    Returns:
        Dictionary with keys 'zero_shot', 'supervised', and 'other'
    """
    return {
        'zero_shot': available_zero_shot_models(),
        'supervised': available_supervised_models(),
        'other': ['AlphaMissense']
    }