如何生成每行含30个氨基酸的注意力分数热力图？现有代码输出不符预期

问题描述 #

需要可视化模型对氨基酸序列的注意力分数，生成每行展示30个氨基酸及其对应注意力分数的热力图，但现有代码生成的结果不符合预期。

原代码问题分析 #

1. 氨基酸标签不完整：仅为x轴设置了前30个氨基酸标签，未关联后续行对应位置的氨基酸
2. 单元格信息缺失：每个单元格仅显示注意力分数，未展示对应的氨基酸
3. 视觉适配不足：文字颜色未根据背景色调整，部分场景下可读性差；网格线缺失，单元格边界不清晰

修正后的代码 #

import numpy as np
import matplotlib.pyplot as plt

amino_acids = ['M', 'L', 'I', 'V', 'N', 'Y', 'L', 'G', 'D', 'L', 'W', 'Q', 'V', 'T', 'Q', 'T', 'T', 'N', 'S', 'L', 'E', 'K', 'R', 'E', 'Q', 'F', 'R', 'G', 'L', 'Y', 'L', 'N', 'K', 'I', 'S', 'E', 'Y', 'S', 'G', 'C', 'I', 'T', 'K', 'S', 'L', 'D', 'E', 'K', 'L', 'L', 'G', 'P', 'I', 'L', 'C', 'E', 'S', 'F', 'F', 'V', 'N', 'G', 'L', 'A', 'Q', 'L', 'Y', 'D', 'K', 'K', 'Q', 'S', 'Q', 'H', 'Q', 'V', 'A', 'E', 'A', 'K', 'R', 'V', 'M', 'Y', 'I', 'V', 'A', 'K', 'E', 'F', 'N', 'A', 'D', 'Y', 'D']
attention_scores = [0.66119576, 0.8003918, 0.7251454, 0.72951543, 0.65185624, 0.70199203, 0.80630046, 0.6887356, 0.6589288, 0.7341603, 0.7876512, 0.7288329, 0.69848996, 0.71069247, 0.67025244, 0.73884994, 0.5722088, 0.8040074, 0.5754876, 0.71131456, 0.6967427, 0.7382109, 0.81139785, 0.8124091, 0.68958503, 0.7973248, 0.81046563, 0.6905091, 0.7415714, 0.70773214, 0.7040749, 0.66361755, 0.75876045, 0.7386744, 0.7828561, 0.7754259, 0.58837545, 0.7422827, 0.8162603, 0.7287266, 0.735986, 0.7265348, 0.7996047, 0.6885884, 0.7867183, 0.7414569, 0.5857099, 0.70292705, 0.76017207, 0.73468393, 0.73548526, 0.7142082, 0.691998, 0.7387566, 0.70642775, 0.7064969, 0.71976453, 0.71234685, 0.6858974, 0.632145, 0.6200939, 0.80994266, 0.70530456, 0.7354963, 0.8044978, 0.80209994, 0.7402193, 0.6048866, 0.73206097, 0.6948871, 0.7328906, 0.8089224, 0.79542226, 0.701279, 0.7086161, 0.80316234, 0.737327, 0.6821824, 0.7155741, 0.7198679, 0.81061196, 0.61442053, 0.8037533, 0.63040495, 0.7192761, 0.76812285, 0.7024652, 0.70812845, 0.80918, 0.64278257, 0.7862615, 0.81110525, 0.77845424, 0.7100564, 0.78630733, 0.8025117]

# 计算行数，取整处理前N*30个数据
num_rows = len(attention_scores) // 30
total_residues = num_rows * 30
# 分割注意力分数和氨基酸序列
score_matrix = np.array(attention_scores[:total_residues]).reshape(num_rows, 30)
aa_matrix = np.array(amino_acids[:total_residues]).reshape(num_rows, 30)

# 创建画布
plt.figure(figsize=(15, 6))
# 绘制热力图，设置颜色映射和数值范围
im = plt.imshow(score_matrix, cmap='hot', vmin=0, vmax=1, aspect='auto')

# 为每个单元格添加氨基酸和分数文本
for i in range(num_rows):
    for j in range(30):
        aa = aa_matrix[i, j]
        score = score_matrix[i, j]
        # 根据背景色调整文字颜色，提升可读性
        text_color = 'white' if score < 0.6 else 'black'
        plt.text(j, i, f'{aa}\n{score:.2f}', ha='center', va='center', color=text_color, fontsize=8)

# 自定义图表元素
plt.xlabel('Position in Row', fontsize=10)
plt.ylabel('Row Number', fontsize=10)
plt.title('Amino Acid Attention Score Heatmap', fontsize=12)

# 设置刻度
plt.xticks(range(30), [f'{x+1}' for x in range(30)], fontsize=8)
plt.yticks(range(num_rows), [f'Row {x+1}' for x in range(num_rows)], fontsize=8)

# 添加颜色条并设置标签
cbar = plt.colorbar(im)
cbar.set_label('Attention Score', fontsize=10)

# 添加网格线，清晰划分单元格
plt.grid(which='both', color='gray', linestyle='-', linewidth=0.5)
# 调整布局，避免标签重叠
plt.tight_layout()

plt.show()

代码优化说明 #

1. 完整关联氨基酸：将氨基酸序列按行分割，每个单元格对应显示所在位置的氨基酸
2. 双信息展示：每个单元格同时显示氨基酸缩写和注意力分数，匹配预期样式
3. 可读性优化：根据注意力分数动态调整文字颜色，暗背景用白色文字，亮背景用黑色文字
4. 视觉规整：添加网格线划分单元格，调整图表尺寸和字体大小，避免元素重叠
5. 刻度清晰：x轴标注行内位置，y轴标注行号，提升图表辨识度

内容的提问来源于stack exchange，提问作者Jhngir Khn