plot_PSVS_Z_scores.py

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns

# Load the updated CSV file
df = pd.read_csv('./psvs_stats.csv')

fontsize_t = 16
# List of columns to plot for histograms
columns_to_plot_nmr = ["NMR_Procheck_phi_psi_Z-score", "NMR_Procheck_Dihe_z_score", 'NMR_MolProbity_clashscore', "NMR_Most_favoured", "NMR_Allowed", "NMR_Disallowed"]
columns_to_plot_af2 = ["AF2_Procheck_phi_psi_Z-score", "AF2_Procheck_Dihe_z_score", 'AF2_MolProbity_clashscore', "AF2_Most_favoured", "AF2_Allowed", "AF2_Disallowed"]
columns_to_plot = columns_to_plot_nmr + columns_to_plot_af2
suffix = '_trim'

for i in range(0, len(columns_to_plot_nmr)):

    # Setting up the figure and subplots in a 1x2 grid with additional axes for histograms
    fig = plt.figure(figsize=(15, 7))

    # Scatter plot
    ax_scatter = plt.subplot2grid((4, 4), (1, 0), rowspan=3, colspan=3)
    # Histograms
    ax_hist_x = plt.subplot2grid((4, 4), (0, 0), colspan=3, sharex=ax_scatter)
    ax_hist_y = plt.subplot2grid((4, 4), (1, 3), rowspan=3, sharey=ax_scatter)

    # Scatter plot comparing AF2 and NMR data
    sns.scatterplot(ax=ax_scatter, x=df[columns_to_plot_nmr[i]+suffix], y=df[columns_to_plot_af2[i]+suffix])
    ax_scatter.set_xlabel(columns_to_plot_nmr[i].replace('NMR', 'PDB').replace('_',' '), fontsize=fontsize_t)
    ax_scatter.set_ylabel(columns_to_plot_af2[i].replace('_', ' '), fontsize=fontsize_t)
    
    # Set the limits of the scatter plot to the min and max values of the data
    ax_scatter.set_xlim(df[columns_to_plot_nmr[i]+suffix].min()-1, df[columns_to_plot_nmr[i]+suffix].max()+1)
    ax_scatter.set_ylim(df[columns_to_plot_af2[i]+suffix].min()-1, df[columns_to_plot_af2[i]+suffix].max()+1)

    # Add a red line to indicate the y=x line
    ax_scatter.plot([df[columns_to_plot_nmr[i]+suffix].min()-1, df[columns_to_plot_nmr[i]+suffix].max()+1],
                    [df[columns_to_plot_nmr[i]+suffix].min()-1, df[columns_to_plot_nmr[i]+suffix].max()+1], color='red')
    
    # Label the y=x line
    ax_scatter.text(df[columns_to_plot_nmr[i]+suffix].max(), df[columns_to_plot_nmr[i]+suffix].max(), 'y=x', color='red', fontsize=12, verticalalignment='bottom')

    # Histograms
    ax_hist_x.hist(df[columns_to_plot_nmr[i]+suffix].dropna(), bins=20, color='blue', alpha=0.7)
    ax_hist_y.hist(df[columns_to_plot_af2[i]+suffix].dropna(), bins=20, color='blue', alpha=0.7, orientation='horizontal')

    # Hide x labels and tick labels for top plots and y ticks for right plots.
    plt.setp(ax_hist_x.get_xticklabels(), visible=False)
    plt.setp(ax_hist_y.get_yticklabels(), visible=False)

    # Adjust the layout to add more space between the plots
    plt.tight_layout()
    
    # Save the plot to the Figures directory with a reasonable file name
    plt.savefig(f'./Figures/plot_{columns_to_plot_nmr[i]}_vs_{columns_to_plot_af2[i]}.png')
    # plt.show()