plot_utils.py

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import scipy.interpolate


def plot_hotpotQA_knowledge_precision_log_matrix(result_root, precision_log_matrix, N):
    plt.figure()
    plt.imshow(precision_log_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.power(10,np.linspace(-3, -1, N)), decimals=3), rotation=45)
    plt.yticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Knowledge Precision Log Matrix')
    plt.colorbar()
    plt.savefig(result_root+'knowledge_precision_log_matrix.pdf')

def plot_hotpotQA_knowledge_recall_log_matrix(result_root, recall_log_matrix, N):
    plt.figure()
    plt.imshow(recall_log_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.power(10,np.linspace(-3, -1, N)), decimals=3), rotation=45)
    plt.yticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Knowledge Recall Log Matrix')
    plt.colorbar()
    plt.savefig(result_root+'knowledge_recall_log_matrix.pdf')

def plot_hotpotQA_knowledge_f1_log_matrix(result_root, f1_log_matrix, N):
    plt.figure()
    plt.imshow(f1_log_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.power(10,np.linspace(-3, -1, N)), decimals=3), rotation=45)
    plt.yticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Knowledge F1 Log Matrix')
    plt.colorbar()
    plt.savefig(result_root+'knowledge_f1_log_matrix.pdf')


def plot_hotpotQA_answer_f1_log_matrix(result_root, log_answer_f1_matrix, N):
    plt.figure()
    plt.imshow(log_answer_f1_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.power(10,np.linspace(-3, -1, N)), decimals=3), rotation=45)
    plt.yticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Hotpot QA Answer F1 Matrix')
    plt.colorbar()
    plt.savefig(result_root+'answer_f1_log_matrix.pdf')

def plot_hotpotQA_answer_f1_matrix(result_root, answer_f1_matrix, N):
    plt.figure()
    plt.imshow(answer_f1_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.yticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Hotpot QA Answer F1 Matrix')
    plt.colorbar()
    plt.savefig(result_root+'answer_f1_matrix.pdf')
    #plt.show()

def plot_hotpotQA_knowledge_matrices(result_root, precision_matrix, recall_matrix, f1_matrix, N):
    plt.figure()
    plt.imshow(precision_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.yticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Knowledge Precision Matrix')
    plt.colorbar()
    plt.savefig(result_root+'knowledge_precision_matrix.pdf')
    #plt.show()

    plt.figure()
    plt.imshow(recall_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.yticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Knowledge Recall Matrix')
    plt.colorbar()
    plt.savefig(result_root+'knowledge_recall_matrix.pdf')
    #plt.show()

    plt.figure()
    plt.imshow(f1_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.yticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Knowledge F1 Matrix')
    plt.colorbar()
    plt.savefig(result_root+'knowledge_f1_matrix.pdf')
    #plt.show()

def plot_hotpotQA_knowledge_matrices_local_log(result_root, local_log_precision_matrix, local_log_recall_matrix, local_log_f1_matrix, M, N):
    plt.figure()
    plt.imshow(local_log_precision_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.power(10,np.linspace(-3, -1, N)), decimals=3), rotation=45)
    plt.yticks(np.arange(0, M), np.around(np.linspace(0.8, 1, M), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Knowledge Precision Matrix')
    plt.colorbar(orientation='horizontal')
    plt.savefig(result_root+'knowledge_precision_local_log_matrix.pdf')

    plt.figure()
    plt.imshow(local_log_recall_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.power(10,np.linspace(-3, -1, N)), decimals=3), rotation=45)
    plt.yticks(np.arange(0, M), np.around(np.linspace(0.8, 1, M), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Knowledge Recall Matrix')
    plt.colorbar(orientation='horizontal')
    plt.savefig(result_root+'knowledge_recall_local_log_matrix.pdf')

    plt.figure()
    plt.imshow(local_log_f1_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.power(10,np.linspace(-3, -1, N)), decimals=3), rotation=45)
    plt.yticks(np.arange(0, M), np.around(np.linspace(0.8, 1, M), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Knowledge F1 Matrix')
    plt.colorbar(orientation='horizontal')
    plt.savefig(result_root+'knowledge_f1_local_log_matrix.pdf')

def plot_hotpotQA_area_matrix(result_root, answer_f1_matrix, no_knowledge_answer_f1, full_knowledge_answer_f1, N):
    higher_than_full_knowledge = answer_f1_matrix > full_knowledge_answer_f1
    lower_than_no_knowledge = answer_f1_matrix < no_knowledge_answer_f1
    area_matrix = np.zeros((N, N)) + 0.5
    area_matrix[higher_than_full_knowledge] += 0.5
    area_matrix[lower_than_no_knowledge] -= 0.5
    plt.figure()
    plt.imshow(area_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.yticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Hotpot QA Knowledge Selector Recommendation Matrix')
    #plt.colorbar()
    plt.savefig(result_root+'area_matrix.pdf')

def plot_hotpotQA_area_log_matrix(result_root, log_answer_f1_matrix, no_knowledge_answer_f1, full_knowledge_answer_f1, N):
    higher_than_full_knowledge = log_answer_f1_matrix > full_knowledge_answer_f1
    lower_than_no_knowledge = log_answer_f1_matrix < no_knowledge_answer_f1
    area_matrix = np.zeros((N, N)) + 0.5
    area_matrix[higher_than_full_knowledge] += 0.5
    area_matrix[lower_than_no_knowledge] -= 0.5
    plt.figure()
    plt.imshow(area_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.power(10,np.linspace(-3, -1, N)), decimals=3), rotation=45)
    plt.yticks(np.arange(0, N, N//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Hotpot QA Knowledge Selector Recommendation Log Matrix')
    #plt.colorbar()
    plt.savefig(result_root+'log_area_matrix.pdf')

def plot_hotpotQA_local_log_answer_f1_matrix(result_root, local_log_answer_f1_matrix, M, N):
    plt.figure()
    plt.imshow(local_log_answer_f1_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.power(10,np.linspace(-3, -1, N)), decimals=3), rotation=45)
    plt.yticks(np.arange(0, M), np.around(np.linspace(0.8, 1, M), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Hotpot QA Knowledge Selector Recommendation Matrix')
    plt.colorbar(orientation='horizontal')
    plt.savefig(result_root+'local_log_answer_f1_matrix.pdf')

def plot_hotpotQA_local_log_area_matrix(result_root, local_log_answer_f1_matrix, no_knowledge_answer_f1, full_knowledge_answer_f1, M, N):
    higher_than_full_knowledge = local_log_answer_f1_matrix > full_knowledge_answer_f1
    lower_than_no_knowledge = local_log_answer_f1_matrix < no_knowledge_answer_f1
    area_matrix = np.zeros((M, N)) + 0.5
    area_matrix[higher_than_full_knowledge] += 0.5
    area_matrix[lower_than_no_knowledge] -= 0.5
    plt.figure()
    plt.imshow(area_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.power(10,np.linspace(-3, -1, N)), decimals=3), rotation=45)
    plt.yticks(np.arange(0, M), np.around(np.linspace(0.8, 1, M), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Hotpot QA Knowledge Selector Recommendation Log Matrix')
    #plt.colorbar()
    plt.savefig(result_root+'local_log_area_matrix.pdf')

def plot_hotpotQA_fixed_recall_precision_vs_answer_f1(result_root, precision_matrix, precision_log_matrix, local_precision_matrix, local_log_precision_matrix, answer_f1_matrix, log_answer_f1_matrix, local_answer_f1_matrix, local_log_answer_f1_matrix, no_knowledge_answer_f1, full_knowledge_answer_f1, N, M, figure_title):
    fig, ax = plt.subplots()
    for i, gold_prob in enumerate(np.linspace(0, 1, N)):
        if i > 0:
            precision_curve = np.append(precision_matrix[i,:], precision_log_matrix[i,:])
            answer_f1_curve = np.append(answer_f1_matrix[i,:], log_answer_f1_matrix[i,:])
            sorted_id = np.argsort(precision_curve)
            plt.plot(precision_curve[sorted_id], answer_f1_curve[sorted_id], color=plt.cm.coolwarm(gold_prob))
    for i, gold_prob in enumerate(np.linspace(0.8, 1, M)):
        precision_curve = np.append(local_precision_matrix[i,:], local_log_precision_matrix[i,:])
        answer_f1_curve = np.append(local_answer_f1_matrix[i,:], local_log_answer_f1_matrix[i,:])
        sorted_id = np.argsort(precision_curve)
        plt.plot(precision_curve[sorted_id], answer_f1_curve[sorted_id], color=plt.cm.coolwarm(gold_prob))
        #plt.legend([f'Gold Prob: {gold_prob:.2f}' for gold_prob in np.linspace(0, 1, N)], title='Gold Probability')

    sm = plt.cm.ScalarMappable(cmap="coolwarm", norm=plt.Normalize(vmin=0, vmax=1))
    sm.set_array([])
    plt.colorbar(sm, ax=ax, label='Knowledge Recall')

    plt.xlabel('Knowledge Precision')
    plt.ylabel('Answer F1')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    #plt.title('Knowledge Precision vs Answer F1 by Fixed Recall')

    plt.axhline(y=full_knowledge_answer_f1, color='red', linewidth=1, linestyle='--')
    plt.axhline(y=no_knowledge_answer_f1, color='black', linewidth=1, linestyle='--')
    # BEGIN: Add legend for dashed lines
    handles = [
        plt.Line2D([0], [0], color='red', linewidth=1, linestyle='--', label='Full Knowledge Answer F1'),
        plt.Line2D([0], [0], color='black', linewidth=1, linestyle='--', label='No Knowledge Answer F1'),
    ]
    ax.legend(handles=handles, loc='lower right', framealpha=0.5)
    plt.title(figure_title)
    plt.savefig(result_root+'fixed_recall_precision_vs_answer_f1.pdf')
    #plt.show()

def make_selection_vs_answer_performance(precision_matrix, recall_matrix, f1_matrix, precision_log_matrix, recall_log_matrix, f1_log_matrix, local_precision_matrix, local_recall_matrix, local_f1_matrix, local_log_precision_matrix, local_log_recall_matrix, local_log_f1_matrix, answer_f1_matrix, log_answer_f1_matrix, local_answer_f1_matrix, local_log_answer_f1_matrix, M, N):
    selection_f1_vs_answer_f1 = []
    selection_prec_vs_answer_f1 = []
    selection_recall_vs_answer_f1 = []
    for i, gold_prob in enumerate(np.linspace(0, 1, N)):
        for j, distract_prob in enumerate(np.linspace(0, 1, N)):
            selection_config = (precision_matrix[i,j], recall_matrix[i,j])
            #if selection_config not in selection_configs:
            #   selection_configs.add(selection_config)
            selection_f1_vs_answer_f1.append([f1_matrix[i,j], answer_f1_matrix[i,j]])
            selection_prec_vs_answer_f1.append([precision_matrix[i,j], answer_f1_matrix[i,j]])
            selection_recall_vs_answer_f1.append([recall_matrix[i,j], answer_f1_matrix[i,j]])

    for i, gold_prob in enumerate(np.linspace(0, 1, N)):
        for j, pw in enumerate(np.linspace(-3, -1, N)):
            distract_prob = np.power(10,pw)
            selection_config = (precision_log_matrix[i,j], recall_log_matrix[i,j])
            #if selection_config not in selection_configs:
            #    selection_configs.add(selection_config)
            selection_f1_vs_answer_f1.append([f1_log_matrix[i,j], log_answer_f1_matrix[i,j]])
            selection_prec_vs_answer_f1.append([precision_log_matrix[i,j], log_answer_f1_matrix[i,j]])
            selection_recall_vs_answer_f1.append([recall_log_matrix[i,j], log_answer_f1_matrix[i,j]])

    for i, gold_prob in enumerate(np.linspace(0.8, 1, M)):
        for j, distract_prob in enumerate(np.linspace(0, 1, N)):
            selection_config = (local_precision_matrix[i,j], local_recall_matrix[i,j])
            #if selection_config not in selection_configs:
            #    selection_configs.add(selection_config)
            selection_f1_vs_answer_f1.append([local_f1_matrix[i,j], local_answer_f1_matrix[i,j]])
            selection_prec_vs_answer_f1.append([local_precision_matrix[i,j], local_answer_f1_matrix[i,j]])
            selection_recall_vs_answer_f1.append([local_recall_matrix[i,j], local_answer_f1_matrix[i,j]])

    for i, gold_prob in enumerate(np.linspace(0.8, 1, M)):
        for j, pw in enumerate(np.linspace(-3, -1, N)):
            distract_prob = np.power(10,pw)
            selection_config = (local_log_precision_matrix[i,j], local_log_recall_matrix[i,j])
            #if selection_config not in selection_configs:
            #    selection_configs.add(selection_config)
            selection_f1_vs_answer_f1.append([local_log_f1_matrix[i,j], local_log_answer_f1_matrix[i,j]])
            selection_prec_vs_answer_f1.append([local_log_precision_matrix[i,j], local_log_answer_f1_matrix[i,j]])
            selection_recall_vs_answer_f1.append([local_log_recall_matrix[i,j], local_log_answer_f1_matrix[i,j]])
    selection_f1_vs_answer_f1 = np.array(selection_f1_vs_answer_f1)
    selection_prec_vs_answer_f1 = np.array(selection_prec_vs_answer_f1)
    selection_recall_vs_answer_f1 = np.array(selection_recall_vs_answer_f1)
    return selection_f1_vs_answer_f1, selection_prec_vs_answer_f1, selection_recall_vs_answer_f1


def determine_performance_area(answer_f1, no_knowledge_answer_f1, full_knowledge_answer_f1):
    if answer_f1 > full_knowledge_answer_f1:
        return "red"
    elif answer_f1 < no_knowledge_answer_f1:
        return "black"
    else:
        return "blue"

def determine_performance_area_color(answer_f1, no_knowledge_answer_f1, full_knowledge_answer_f1):
    if answer_f1 > full_knowledge_answer_f1:
        return "orange"
    elif answer_f1 < no_knowledge_answer_f1:
        return "cyan"
    else:
        return "white"
    
def determine_performance_area_alpha(answer_f1, no_knowledge_answer_f1, full_knowledge_answer_f1):
    if answer_f1 > full_knowledge_answer_f1:
        return 0.5
    elif answer_f1 < no_knowledge_answer_f1:
        return 0.5
    else:
        return 0.8
    
def plot_hotpotQA_interpolation(result_root, selection_prec_vs_answer_f1, selection_recall_vs_answer_f1):
    x = selection_prec_vs_answer_f1[:,0]
    y = selection_recall_vs_answer_f1[:,0]
    z = selection_prec_vs_answer_f1[:,1]

    # Define grid
    xi = np.linspace(0, 1, 100)
    yi = np.linspace(0, 1, 100)
    xi, yi = np.meshgrid(xi, yi)

    # Interpolate data onto grid
    zi = scipy.interpolate.griddata((x, y), z, (xi, yi), method='linear')

    # Create a heatmap
    fig, ax = plt.subplots()
    heatmap = plt.contourf(xi, yi, zi, levels=100, cmap='coolwarm')
    plt.colorbar(heatmap)

    # Display the heatmap
    plt.xlabel('Knowledge Precision')
    plt.ylabel('Knowledge Recall')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    plt.show()

def plot_hotpotQA_knowledge_selection_precision_recall_vs_answer_f1(result_root, selection_prec_vs_answer_f1, selection_recall_vs_answer_f1, selection_f1_vs_answer_f1, figure_title):
    fig, ax = plt.subplots()
    scatter = ax.scatter(selection_prec_vs_answer_f1[:, 0], selection_recall_vs_answer_f1[:, 0], 
                        c=selection_f1_vs_answer_f1[:, 1], cmap="coolwarm", alpha=0.8, s=50)
    plt.colorbar(scatter, ax=ax)
    plt.xlabel('Knowledge Precision')
    plt.ylabel('Knowledge Recall')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    plt.title(figure_title)
    plt.savefig(result_root+'knowledge_selection_precision_recall_vs_answer_f1.pdf')

def plot_hotpotQA_knowledge_selection_precision_recall_vs_answer_f1_area(result_root, selection_prec_vs_answer_f1, selection_recall_vs_answer_f1, selection_f1_vs_answer_f1, no_knowledge_answer_f1, full_knowledge_answer_f1, figure_title):
    fig, ax = plt.subplots()
    for prec_pair, recall_pair in zip(selection_prec_vs_answer_f1, selection_recall_vs_answer_f1):
        plt.scatter(prec_pair[0], recall_pair[0], color=determine_performance_area_color(prec_pair[1], no_knowledge_answer_f1, full_knowledge_answer_f1), alpha=determine_performance_area_alpha(prec_pair[1], no_knowledge_answer_f1, full_knowledge_answer_f1), s=1200)
    scatter = ax.scatter(selection_prec_vs_answer_f1[:, 0], selection_recall_vs_answer_f1[:, 0], 
                        c=selection_f1_vs_answer_f1[:, 1], cmap="coolwarm", alpha=0.8, s=50)
    plt.colorbar(scatter, ax=ax)
    plt.xlabel('Knowledge Precision')
    plt.ylabel('Knowledge Recall')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    plt.title(figure_title)
    plt.savefig(result_root+'knowledge_selection_precision_recall_vs_answer_f1_area.pdf')

def plot_knowledge_selection_precision_recall_by_area(result_root, selection_prec_vs_answer_f1, selection_recall_vs_answer_f1, no_knowledge_answer_f1, full_knowledge_answer_f1, figure_title):
    fig, ax = plt.subplots()
    for prec_pair, recall_pair in zip(selection_prec_vs_answer_f1, selection_recall_vs_answer_f1):
        plt.scatter(prec_pair[0], recall_pair[0], color=determine_performance_area(prec_pair[1], no_knowledge_answer_f1, full_knowledge_answer_f1), alpha=0.25, s=50)
    plt.xlabel('Knowledge Precision')
    plt.ylabel('Knowledge Recall')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    # BEGIN: Add legend
    handles = [
        plt.Line2D([0], [0], marker='o', color='w', label='> Full Knowledge', 
                markerfacecolor='red', markersize=10, alpha=0.5),
        plt.Line2D([0], [0], marker='o', color='w', label='In Between', 
                markerfacecolor='blue', markersize=10, alpha=0.5),
        plt.Line2D([0], [0], marker='o', color='w', label='< No Knowledge', 
                markerfacecolor='black', markersize=10, alpha=0.5),
    ]
    ax.legend(handles=handles, loc='lower right', framealpha=0.5)
    # END: Add legend
    plt.title(figure_title)
    plt.savefig(result_root+'knowledge_selection_precision_recall_by_area.pdf')
    #plt.show()

def plot_selection_f1_vs_answer_f1(result_root, selection_f1_vs_answer_f1, no_knowledge_answer_f1, full_knowledge_answer_f1, figure_title):
    fig, ax = plt.subplots()
    plt.scatter(selection_f1_vs_answer_f1[:,0], selection_f1_vs_answer_f1[:,1], alpha=0.5, s=10)
    plt.xlabel('Knowledge F1')
    plt.ylabel('Answer F1')
    plt.axhline(y=full_knowledge_answer_f1, color='red', linewidth=1, linestyle='--')
    plt.axhline(y=no_knowledge_answer_f1, color='black', linewidth=1, linestyle='--')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    # BEGIN: Add legend for dashed lines
    handles = [
        plt.Line2D([0], [0], color='red', linewidth=1, linestyle='--', label='Full Knowledge Answer F1'),
        plt.Line2D([0], [0], color='black', linewidth=1, linestyle='--', label='No Knowledge Answer F1'),
    ]
    ax.legend(handles=handles, loc='lower right', framealpha=0.5)
    # END: Add legend for dashed lines
    #plt.show()
    plt.title(figure_title)
    plt.savefig(result_root+'selection_f1_vs_answer_f1.pdf')

def plot_selection_prec_vs_answer_f1(result_root, selection_prec_vs_answer_f1, no_knowledge_answer_f1, full_knowledge_answer_f1, figure_title):
    fig, ax = plt.subplots()
    plt.scatter(selection_prec_vs_answer_f1[:,0], selection_prec_vs_answer_f1[:,1], alpha=0.5, s=10)
    plt.xlabel('Knowledge Precision')
    plt.ylabel('Answer F1')

    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')

    plt.axhline(y=full_knowledge_answer_f1, color='red', linewidth=1, linestyle='--')
    plt.axhline(y=no_knowledge_answer_f1, color='black', linewidth=1, linestyle='--')
    # BEGIN: Add legend for dashed lines
    handles = [
        plt.Line2D([0], [0], color='red', linewidth=1, linestyle='--', label='Full Knowledge Answer F1'),
        plt.Line2D([0], [0], color='black', linewidth=1, linestyle='--', label='No Knowledge Answer F1'),
    ]
    ax.legend(handles=handles, loc='lower right', framealpha=0.5)
    # END: Add legend for dashed lines
    #plt.show()
    plt.title(figure_title)
    plt.savefig(result_root+'selection_prec_vs_answer_f1.pdf')

def plot_selection_recall_vs_answer_f1(result_root, selection_recall_vs_answer_f1, no_knowledge_answer_f1, full_knowledge_answer_f1, figure_title):
    fig, ax = plt.subplots()
    plt.scatter(selection_recall_vs_answer_f1[:,0], selection_recall_vs_answer_f1[:,1], alpha=0.5, s=10)
    plt.xlabel('Knowledge Recall')
    plt.ylabel('Answer F1')
    plt.axhline(y=full_knowledge_answer_f1, color='red', linewidth=1, linestyle='--')
    plt.axhline(y=no_knowledge_answer_f1, color='black', linewidth=1, linestyle='--')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    # BEGIN: Add legend for dashed lines
    handles = [
        plt.Line2D([0], [0], color='red', linewidth=1, linestyle='--', label='Full Knowledge Answer F1'),
        plt.Line2D([0], [0], color='black', linewidth=1, linestyle='--', label='No Knowledge Answer F1'),
    ]
    ax.legend(handles=handles, loc='lower right', framealpha=0.5)
    # END: Add legend for dashed lines
    #plt.show()
    plt.title(figure_title)
    plt.savefig(result_root+'selection_recall_vs_answer_f1.pdf')

def plot_wow_knowledge_matrices(result_root, precision_matrix, recall_matrix, f1_matrix, N, M):
    plt.figure()
    plt.imshow(precision_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.power(10, np.linspace(-4, 0, 11)), decimals=4), rotation=45)
    plt.yticks(np.arange(0, M, M//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Knowledge Precision Matrix')
    plt.colorbar()
    plt.savefig(result_root+'knowledge_precision_matrix.pdf')
    #plt.show()

    plt.figure()
    plt.imshow(recall_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.power(10, np.linspace(-4, 0, 11)), decimals=4), rotation=45)
    plt.yticks(np.arange(0, M, M//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Knowledge Recall Matrix')
    plt.colorbar()
    plt.savefig(result_root+'knowledge_recall_matrix.pdf')
    #plt.show()

    plt.figure()
    plt.imshow(f1_matrix, cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(0, N, N//10), np.around(np.power(10, np.linspace(-4, 0, 11)), decimals=4), rotation=45)
    plt.yticks(np.arange(0, M, M//10), np.around(np.linspace(0, 1, 11), decimals=2))
    plt.ylabel('Prob of sampling gold knowledge')
    plt.xlabel('Prob of sampling noise knowledge')
    plt.title('Knowledge F1 Matrix')
    plt.colorbar()
    plt.savefig(result_root+'knowledge_f1_matrix.pdf')
    #plt.show()

def determine_wow_performance_area_alpha(answer_f1, no_knowledge_answer_f1, full_knowledge_answer_f1):
    if answer_f1 > full_knowledge_answer_f1:
        return 0.5
    elif answer_f1 < no_knowledge_answer_f1:
        return 0.5
    else:
        return 0.5
    
def plot_wow_knowledge_selection_precision_recall_vs_answer_f1_area(result_root, selection_prec_vs_answer_f1, selection_recall_vs_answer_f1, selection_f1_vs_answer_f1, no_knowledge_answer_f1, full_knowledge_answer_f1, figure_title):
    fig, ax = plt.subplots()
    for prec_pair, recall_pair in zip(selection_prec_vs_answer_f1, selection_recall_vs_answer_f1):
        plt.scatter(prec_pair[0], recall_pair[0], color=determine_performance_area_color(prec_pair[1], no_knowledge_answer_f1, full_knowledge_answer_f1), alpha=determine_wow_performance_area_alpha(prec_pair[1], no_knowledge_answer_f1, full_knowledge_answer_f1), s=1200)
    scatter = ax.scatter(selection_prec_vs_answer_f1[:, 0], selection_recall_vs_answer_f1[:, 0], 
                        c=selection_f1_vs_answer_f1[:, 1], cmap="coolwarm", alpha=0.8, s=50)
    plt.colorbar(scatter, ax=ax)
    plt.xlabel('Knowledge Precision')
    plt.ylabel('Knowledge Recall')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    plt.title(figure_title)
    plt.savefig(result_root+'knowledge_selection_precision_recall_vs_answer_f1_area.pdf')

def plot_wow_knowledge_selection_precision_recall_vs_rougeL(result_root, selection_prec_vs_rouge, selection_recall_vs_rouge, selection_f1_vs_rouge, figure_title):
    fig, ax = plt.subplots()
    scatter = ax.scatter(selection_prec_vs_rouge[:, 0], selection_recall_vs_rouge[:, 0], 
                        c=selection_f1_vs_rouge[:, 1], cmap="coolwarm", alpha=0.8, s=50)
    plt.colorbar(scatter, ax=ax)
    plt.xlabel('Knowledge Precision')
    plt.ylabel('Knowledge Recall')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    plt.title(figure_title)
    plt.savefig(result_root+'knowledge_selection_precision_recall_vs_rougeL.pdf')

def plot_wow_knowledge_selection_precision_recall_vs_answer_f1(result_root, selection_prec_vs_answer_f1, selection_recall_vs_answer_f1, selection_f1_vs_answer_f1, figure_title):

    fig, ax = plt.subplots()
    scatter = ax.scatter(selection_prec_vs_answer_f1[:, 0], selection_recall_vs_answer_f1[:, 0], 
                        c=selection_f1_vs_answer_f1[:, 1], cmap="coolwarm", alpha=0.8, s=50)
    plt.colorbar(scatter, ax=ax)
    plt.xlabel('Knowledge Precision')
    plt.ylabel('Knowledge Recall')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    plt.title(figure_title)
    plt.savefig(result_root+'knowledge_selection_precision_recall_vs_answer_f1.pdf')

def determine_rouge_performance_area(answer_rouge, no_knowledge_rouge_L, full_knowledge_rouge_L):
    if answer_rouge > full_knowledge_rouge_L:
        return "red"
    elif answer_rouge < no_knowledge_rouge_L:
        return "black"
    else:
        return "blue"
    
def plot_wow_knowledge_selection_precision_recall_by_area_rouge(result_root, selection_prec_vs_rouge, selection_recall_vs_rouge, no_knowledge_rouge_L, full_knowledge_rouge_L, figure_title):
    
    fig, ax = plt.subplots()
    for prec_pair, recall_pair in zip(selection_prec_vs_rouge, selection_recall_vs_rouge):
        plt.scatter(prec_pair[0], recall_pair[0], color=determine_rouge_performance_area(prec_pair[1], no_knowledge_rouge_L, full_knowledge_rouge_L), alpha=0.25, s=50)
    plt.xlabel('Knowledge Precision')
    plt.ylabel('Knowledge Recall')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    # BEGIN: Add legend
    handles = [
        plt.Line2D([0], [0], marker='o', color='w', label='> Full Knowledge', 
                markerfacecolor='red', markersize=10, alpha=0.5),
        plt.Line2D([0], [0], marker='o', color='w', label='In Between', 
                markerfacecolor='blue', markersize=10, alpha=0.5),
        plt.Line2D([0], [0], marker='o', color='w', label='< No Knowledge', 
                markerfacecolor='black', markersize=10, alpha=0.5),
    ]
    ax.legend(handles=handles, loc='lower right', framealpha=0.5)
    # END: Add legend
    plt.title(figure_title)
    plt.savefig(result_root+'knowledge_selection_precision_recall_by_area_rouge.pdf')
#plt.show()

def plot_wow_fixed_recall_precision_vs_answer_f1(result_root, precision_matrix, answer_f1_matrix, no_knowledge_answer_f1, full_knowledge_answer_f1, M, figure_title):
    fig, ax = plt.subplots()
    for i, gold_prob in enumerate(np.linspace(0, 1, M)):
        if i > 0:
            precision_curve = precision_matrix[i,:]
            answer_f1_curve = answer_f1_matrix[i,:]
            sorted_id = np.argsort(precision_curve)
            plt.plot(precision_curve[sorted_id], answer_f1_curve[sorted_id], color=plt.cm.coolwarm(gold_prob))

    sm = plt.cm.ScalarMappable(cmap="coolwarm", norm=plt.Normalize(vmin=0, vmax=1))
    sm.set_array([])
    plt.colorbar(sm, ax=ax, label='Knowledge Recall')

    plt.xlabel('Knowledge Precision')
    plt.ylabel('Answer F1')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    #plt.title('Knowledge Precision vs Answer F1 by Fixed Recall')

    plt.axhline(y=full_knowledge_answer_f1, color='red', linewidth=1, linestyle='--')
    plt.axhline(y=no_knowledge_answer_f1, color='black', linewidth=1, linestyle='--')
    # BEGIN: Add legend for dashed lines
    handles = [
        plt.Line2D([0], [0], color='red', linewidth=1, linestyle='--', label='Full Knowledge Answer F1'),
        plt.Line2D([0], [0], color='black', linewidth=1, linestyle='--', label='No Knowledge Answer F1'),
    ]
    ax.legend(handles=handles, loc='lower right', framealpha=0.5)
    plt.title(figure_title)
    plt.savefig(result_root+'fixed_recall_precision_vs_answer_f1.pdf')
    #plt.show()