matrix_voting.py

import bisect
from collections import Counter, defaultdict
from copy import deepcopy
from itertools import chain, combinations
from pprint import pprint

import numpy
import numpy as np
import json
import matplotlib.pyplot as plt
from matplotlib import image as mpimg
from matplotlib.colors import LinearSegmentedColormap


# Function to calculate f1 index
def calc_f1(set1, set2):
    fp = fn = tp = 0
    for el in set1:
        if el in set2:
            tp += 1
        if el not in set2:
            fp += 1
    for el in set2:
        if el not in set1:
            fn += 1
    F1 = tp / (tp + 0.5 * (fp + fn))
    precision = 0
    if tp + fp > 0:
        precision = tp / (tp + fp)
    recall = tp / (tp + fn)
    return F1


def make_f1_matrix(file_paths: dict, possible_llm_names: list) -> numpy.array:
    # Initialize the proteins_complexes dictionary
    proteins_complexes = {}

    indexed_llm_names = {possible_llm_names[i]: i for i in range(len(possible_llm_names))}

    for model, path in file_paths.items():
        with open(path, 'r') as f:
            proteins_complexes[model] = json.load(f)

    with open("../sources/verified_complexes.json") as f:
        verified_proteins_complexes = json.load(f)

    best_votes = [0] * (len(possible_llm_names) + 1)

    # Keep just one option

    allsubsets = chain.from_iterable(combinations(possible_llm_names, r) for r in range(1, len(possible_llm_names) + 1))

    allsubsets = list(allsubsets)

    result_matrix = [[0] * len(possible_llm_names) for _ in range(len(possible_llm_names))]
    i = 0
    highest_score = 0

    print(possible_llm_names)
    for llm_names in allsubsets:
        if llm_names != i:
            highest_score = 0
        i = len(llm_names)
        j = max([indexed_llm_names[llm] for llm in llm_names])

        # print(llm_names)
        num_llms = len(llm_names)
        votes_success = [0] * (len(possible_llm_names) + 1)

        # calculate avg f1 index for each pair of LLMs
        complexes_counters = defaultdict(lambda: Counter())
        for macromolecular_complex in verified_proteins_complexes:
            for i in range(num_llms):
                llm = llm_names[i]
                try:
                    for prot in set(proteins_complexes[llm][macromolecular_complex]):
                        complexes_counters[macromolecular_complex][prot] += 1
                except KeyError as e:
                    pass
            # print(e)
            # print(llm)
            # print(macromolecular_complex)
            # print(macromolecular_complex, )

        for votes_threshold in range(1, num_llms + 1):
            for macromolecular_complex in sorted(list(verified_proteins_complexes)):
                decided_complex = set()
                for prot in sorted(complexes_counters[macromolecular_complex]):
                    if complexes_counters[macromolecular_complex][prot] >= votes_threshold:
                        decided_complex.add(prot)
                # print(decided_complex)

                curated_complex = verified_proteins_complexes[macromolecular_complex]
                votes_success[votes_threshold] += calc_f1(decided_complex, curated_complex)

                # print(macromolecular_complex)
                # print(sorted(list(decided_complex)))

        votes_success = list(
            map(lambda x: round(x / len(verified_proteins_complexes), 3), votes_success))
        # print(votes_success)

        if max(votes_success) > max(best_votes):
            best_votes = votes_success
            # print(llm_names)
            # print(votes_success)
        highest_score = max(highest_score, max(votes_success))
        print(i, j)
        print(llm_names)
        print(votes_success)
        result_matrix[i][j] = max(highest_score, result_matrix[i][j])
        print(result_matrix)

    pprint(result_matrix)
    for i in range(len(possible_llm_names)):
        for j in range(len(possible_llm_names)):
            result_matrix[i][j] = round(100 * result_matrix[i][j], 1)

    return numpy.array(result_matrix)


def graph_matrix(matrix, x_labels, y_labels, filename, save=0):
    """
    Graphs a NumPy matrix with gradient color and cell values,
    including row and column names, adjustable bottom spacing,
    and x-axis labels on top.

    Args:
      matrix: The NumPy matrix to graph.
      x_labels: A list of strings representing the names of the columns (x-axis).
      y_labels: A list of strings representing the names of the rows (y-axis).
      x_label: The label for the x-axis (default: "X Axis").
      y_label: The label for the y-axis (default: "Y Axis").
    """

    # Create a custom colormap from red to yellow to green
    cmap = LinearSegmentedColormap.from_list('mycmap', ['red', 'yellow', 'green'])
    cmap.set_under('white')

    # Create the plot
    fig, ax = plt.subplots()
    im = ax.imshow(matrix, cmap=cmap, vmin=0.1, vmax=100)

    # Add colorbar
    fig.colorbar(im, shrink=0.85)

    # Add text annotations for each cell
    for i in range(matrix.shape[0]):
        for j in range(matrix.shape[1]):
            text = ax.text(j, i, f"{matrix[i, j]:.1f}",
                           ha="center", va="center", color="black")

    # Set axis labels
    ax.set_xlabel("LLMs")
    ax.set_ylabel("LLMs")

    # Set ticks and tick labels
    ax.set_xticks(np.arange(len(x_labels)))
    ax.set_yticks(np.arange(len(y_labels)))
    ax.set_xticklabels(x_labels, fontsize=8)
    ax.set_yticklabels(y_labels, fontsize=8)

    # Move x-axis labels to the top
    ax.xaxis.tick_top()

    # Rotate x-axis labels if needed
    plt.setp(ax.get_xticklabels(), rotation=45, ha="left",
             rotation_mode="anchor")

    # Adjust subplot spacing (optional)
    plt.subplots_adjust(bottom=-0.05)
    plt.subplots_adjust(left=0.18)

    # Show the plot

    if save:
        plt.savefig('../plots/f1/' + filename)

        plt.clf()
        img = mpimg.imread('../plots/f1/' + filename)
        plt.imshow(img)
        plt.axis('off')  # Hide axes

    plt.show()


# Example usage
if __name__ == '__main__':
    file_paths = {
        "ChatGPT": "../sources/chatgpto1preview_human_complex_prots.json",
        "Perplexity": "../sources/perplexity_context_3blocks_proteins.json",
        "Claude": "../sources/claude_context_4blocks_proteins.json",
        "Deepseek": "../sources/deepseek_context_4blocks_prots.json",
        "Llama": "../sources/llama3.1_general.json",
        "Gemini API": "../sources/gemini1.5_API_prots.json"
    }

    llm_names = list(file_paths.keys())

    output_filename = "voting_incremental_matrix_integrated_brs.png"
    votes_success_matrix = make_f1_matrix(file_paths, llm_names)
    number_range = list(range(1, len(llm_names) + 1))
    graph_matrix(votes_success_matrix, llm_names, number_range, output_filename, 1)