wordcloud.py

from wordcloud import WordCloud, STOPWORDS
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import image as mpimg
from matplotlib.cm import ScalarMappable
import json
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import image as mpimg
from scipy.stats import spearmanr

verified_complexes_js: dict


def stats(a: list):
    a.sort()
    return sum(a) / len(a), a[len(a) // 4], a[len(a) // 2], a[3 * len(a) // 4]


def calc_f1_ext(set1, set2) -> (float, list, list, list):
    tp = []
    fp = []
    fn = []
    for el in set1:
        if el in set2:
            tp.append(el)
        if el not in set2:
            fp.append(el)
    for el in set2:
        if el not in set1:
            fn.append(el)
    F1 = len(tp) / (len(tp) + 0.5 * (len(fp) + len(fn)))
    return F1, tp, fp, fn


import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.patches as mpatches

# 1. Read the CSV file
df = pd.read_csv('../results/results_05-24-2025-00-49-34.csv', sep=';')

with open("../sources/human_molc_mach_list.json") as f:
    verified_complexes_js = json.load(f)

# 2. Keep only rows where "integration method" == "Raw"
# 3. Keep only rows that match ANY of the specified conditions
conditions_to_keep = (
        (((df['LLM'] == 'Perplexity') & (df['prompt'] == '3 blocks')) |
         ((df['LLM'] == 'Claude') & (df['prompt'] == '4 blocks')) |
         ((df['LLM'] == 'ChatGPTo1-prev') & (df['prompt'] == 'General')) |
         ((df['LLM'] == 'Gemini2.5') & (df['prompt'] == '3 blocks')) |
         ((df['LLM'] == 'Llama4') & (df['prompt'] == '4 blocks')) |
         ((df['LLM'] == 'Deepseek') & (df['prompt'] == '4 blocks'))) &
        (df['integration method'] == 'Raw')
)

df_final = df[conditions_to_keep]

llms_order = ["ChatGPTo1-prev", "Perplexity", "Claude", "Llama4", "Gemini2.5", "Deepseek"]

# Dictionary to store scatter plot data per LLM.
results_by_complex = {}

complex_names = []
graph_densities = []
protein_counts = []
# 5. Prepare data for scatter plot
for complex_name in verified_complexes_js:
    complex_names.append(complex_name)
    subset = df_final[df_final['complex_name'] == complex_name]
    graph_densities.append(np.mean(subset['graph_density'].values))
    protein_counts.append(np.mean(subset['N proteins'].values))
    print(complex_name)
    print(subset['graph_density'].values)
    print(subset['N proteins'].values)

# Create dictionaries for mapping
name_to_f1 = dict(zip(complex_names, graph_densities))
name_to_protein_count = dict(zip(complex_names, protein_counts))

# Create a colormap
cmap = plt.cm.viridis

# Normalize F1 scores to the range [0, 255] for color mapping
norm = plt.Normalize(min(graph_densities), max(graph_densities))


# Create a color function
def color_func(word, **kwargs):
    try:
        f1 = name_to_f1[word]
        normalized_f1 = int(norm(f1) * 255)
        R, G, B, A = cmap(normalized_f1)
        R = int(R * 255)
        G = int(G * 255)
        B = int(B * 255)
        print(R, G, B, A)
        return R, G, B
    except KeyError:
        return 'grey'


# Create the word cloud
wordcloud = WordCloud(
    width=2000, height=1000,
    background_color=(255, 255, 255),
    stopwords=STOPWORDS,
    color_func=color_func
).generate_from_frequencies(name_to_protein_count)  # Use frequencies for word size

# Display the word cloud
fig, ax = plt.subplots(figsize=(20, 10))

ax.imshow(wordcloud, interpolation='bilinear')
ax.axis('off')

# ------------------------------------------------------------------
# create a dummy mappable and attach some data
# ------------------------------------------------------------------
sm = ScalarMappable(norm=norm, cmap=cmap)
sm.set_array([])  # or   sm.set_array(graph_densities)

# ------------------------------------------------------------------
# add the colour-bar; tell it which axes to shrink
# ------------------------------------------------------------------
cbar = fig.colorbar(sm,
             ax=ax,  # <- the important part
             orientation='horizontal',
             label='Avg. Graph Density',
             pad=0.05)

cbar.set_label('Avg. Graph Density', fontsize=24)
cbar.ax.tick_params(labelsize=20)

filename = "wordcloud_graph_density.png"

save = 1
if save:
    plt.savefig('../plots/supplemental/' + filename)

    plt.clf()
    img = mpimg.imread('../plots/supplemental/' + filename)
    plt.imshow(img)
    plt.axis('off')

plt.show()