File size: 2,133 Bytes
e14f793 c781acb 08f6d28 e14f793 8d7f260 b93d2a5 e14f793 2ef1426 4749b40 f032ab3 e14f793 ec53b97 d4d7a40 e14f793 244de08 e14f793 613c843 244de08 613c843 e14f793 aaccffd b54f868 59c706f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | import pandas as pd
import numpy as np
import seaborn as sns
from turftopic import load_model
import matplotlib.pyplot as plt
def prepare_alignment_data(cfg):
sample = pd.read_csv(f"{cfg.output_dir}/{cfg.output_file}")
sample.columns = sample.columns.astype(str)
model = load_model(cfg.data[cfg.level].topic_model)
classification_cols = ['message_id', 'query', 'class']
topic_cols = [str(c) for c in range(len(model.get_topics()))]
print(classification_cols+topic_cols)
print(cfg)
return sample[classification_cols+topic_cols], model
def plot_heatmap(cfg):
alignment_df, model = prepare_alignment_data(cfg)
topic_cols = [str(c) for c in range(len(model.get_topics()))]
classes = alignment_df['class'].unique()
# Data matrix must have shape (len(y_edges)-1, len(x_edges)-1) -> (2, 3)
data = np.array([(alignment_df[alignment_df['class']==cl][topic_cols]!=0).sum().values for cl in classes])
col_labels = [f'topic_{k}' for k in range(len(model.get_topics()))]
# 2. Configure the plot layout
fig, ax = plt.subplots(figsize=(5, 3))
sns.set_theme(font_scale=1.5)
# 3. Create the heatmap
sns.heatmap(
data,
annot=True, # Write data numbers inside each cell
annot_kws={"size": 6},
fmt=".2f", # Format labels to 2 decimal places
cmap="coolwarm", # Use a blue-to-red color palette
xticklabels=col_labels, # Set X axis labels
yticklabels=classes, # Set Y axis labels
linewidths=0.5, # Add thin lines between cells
square=True, # Force cells to be perfect squares
)
# 2. Adjust axis titles
# plt.xlabel("X Axis", fontsize=6)
# plt.ylabel("Y Axis", fontsize=6)
# # 3. Adjust axis tick labels (numbers on the grid edge)
# ax.tick_params(labelsize=12)
# 4. Adjust colorbar tick font size
ax.figure.axes[-1].tick_params(labelsize=6)
# 4. Show the visual output
ax.set_title("Alignment Heatmap", fontsize=5, pad=6)
ax.tick_params(axis='x', labelrotation=45, labelsize=6)
ax.tick_params(axis='y', labelsize=6)
return fig |