model_corr_plot¶

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']
    }

def model_corr_plot(
    uniprot_id: str,
    model1: str = "AlphaMissense",
    model2: str = "GEMME",
    cache_dir: str = ".cache",
    figsize: Tuple[float, float] = (10, 8),
    bins: int = 60,
    cmap: str = 'viridis'
) -> Figure:
    """
    Compare prediction scores for two ProteinGym models using Spearman correlation.

    Creates a 2D density plot (hexbin) with marginal distributions showing the
    relationship between predicted scores from two models.

    Args:
        uniprot_id: Valid UniProt accession identifier
        model1: First model to plot (default: "AlphaMissense")
        model2: Second model to plot (default: "GEMME")
        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 GEMME)
        >>> fig = model_corr_plot(uniprot_id="Q9NV35")
        >>> plt.show()

        >>> # Compare specific models
        >>> fig = model_corr_plot(
        ...     uniprot_id="P04637",
        ...     model1="Kermut",
        ...     model2="EVE_single"
        ... )
        >>> 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 models
    valid_models = (
        available_zero_shot_models() +
        available_supervised_models() +
        ["AlphaMissense"]
    )

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

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

    # Load respective data for uniprot_id
    print(f"Loading data for {model1}...")
    model_df1 = _get_model_dataframe(model1, uniprot_id, cache_dir)

    print(f"Loading data for {model2}...")
    model_df2 = _get_model_dataframe(model2, uniprot_id, cache_dir)

    # Merge tables
    print("Merging model predictions...")
    merged_table = _merge_model_tables(
        model_df1=model_df1,
        model_df2=model_df2,
        model1=model1,
        model2=model2
    )

    # Check if merged table is empty
    if len(merged_table) == 0:
        raise ValueError(
            f"No common mutants between chosen models 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 = merged_table['mean_model1'].values
    y = merged_table['mean_model2'].values

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

    # Style main plot
    ax_main.set_xlabel(f'{model1} score', fontsize=14)
    ax_main.set_ylabel(f'{model2} score', fontsize=14)
    ax_main.tick_params(labelsize=12)

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

    # Top marginal (histogram + KDE for model1)
    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=10)
    ax_top.tick_params(labelbottom=False, labelsize=10)
    ax_top.spines['top'].set_visible(False)
    ax_top.spines['right'].set_visible(False)

    # Right marginal (histogram + KDE for model2)
    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=10)
    ax_right.tick_params(labelleft=False, labelsize=10)
    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