vis/auroc_2.py

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import roc_curve, precision_recall_curve, roc_auc_score, average_precision_score

# 定义子图标题及每个子图需要绘制的列名
titles = {
    'R2': 'r2_case',
    # 'MHB': 'test', 
    # '12.5% FBS': '125fbs', 
    # '25% FBS': '25fbs',
}

methods = {
    'LSTM': "#665f88", 
    'GRU': "#658f5f",
    'MHA': "#5f858f",
    'Mamba': "#AC9E53", 
    'GAT': "#15A9B1",
    'GATv2': "#0EA488",
    # 'ADAPT': "#FF5733",
    'ADAPT': "#FF0000",
}

folders = {
    'LSTM': '../run-cls/lstm-diff-256-ce-32-0.001-50', 
    'GRU': '../run-cls/gru-diff-256-ce-32-0.001-50',
    'MHA': '../run-cls/mha-diff-256-ce-32-0.001-50',
    'Mamba': '../run-cls/mamba-diff-256-ce-32-0.001-50',
    'GAT': '../../DAminoMuta_graph/run-cls/gat-diff-256-ce-32-0.001-35',
    'GATv2': '../../DAminoMuta_graph/run-cls/gatv2-diff-256-ce-32-0.001-35',
    # 'ADAPT': '../run-cls/rn18-diff-16-768-ce-32-0.001-50',
    'ADAPT': '../run-cls/rn18-diff-16-mamba-pcs-768-ce-32-0.001-50/uda_r2',
}

# 创建三个子图，调整图形尺寸适合展示
fig, axes = plt.subplots(nrows=1, ncols=2*len(titles), figsize=(10.66*len(titles), 5))

# 遍历每个子图与对应的标题、方法列数据
for ax, (title, alias) in zip(axes[0::2], titles.items()):
    for method, color in methods.items():
        # 获取预测概率及真实标签
        df = pd.read_csv(f'{folders[method]}/preds_{alias}.csv')
        try:
            y_score = np.concat([df[f'model_{i}_test'].values for i in range(5)], axis=0)
            y_true = np.concat([df[f'gt'].values for _ in range(5)], axis=0)
        except:
            y_score = df[f'model_uda_teacher'].values
            y_true = df[f'gt'].values
        # 计算 ROC 曲线以及 AUROC 值
        fpr, tpr, thresholds = roc_curve(y_true, y_score)
        auroc = roc_auc_score(y_true, y_score)
        
        # 绘制 ROC 曲线，使用已分配的颜色，学术展示的曲线线宽合适
        ax.plot(fpr, tpr, label=f"{method} AUC: {auroc:.2f}", 
                color=color, lw=2, alpha=0.7)
    
    # 设置子图标题
    ax.set_title(f'{title} ROC Curve', fontsize=18, weight='bold')
    # 设置横纵坐标的标签
    ax.set_xlabel("False Positive Rate", fontsize=14)
    ax.set_ylabel("True Positive Rate", fontsize=14)
    ax.yaxis.set_label_position("right")
    
    # 仅保留下边框和右边框, 去掉上边框和左边框
    ax.spines['top'].set_visible(False)
    ax.spines['left'].set_visible(False)
    # 仅在底部和右侧显示刻度
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('right')
    
    # 设置图例，放置在右下角
    ax.legend(loc='lower right', fontsize=10, frameon=False, alignment='right', markerfirst=False)

# 遍历每个子图与对应的标题、方法列数据
for ax, (title, alias) in zip(axes[1::2], titles.items()):
    for method, color in methods.items():
        # 获取预测概率及真实标签
        df = pd.read_csv(f'{folders[method]}/preds_{alias}.csv')
        try:
            y_score = np.concat([df[f'model_{i}_test'].values for i in range(5)], axis=0)
            y_true = np.concat([df[f'gt'].values for _ in range(5)], axis=0)
        except:
            y_score = df[f'model_uda_teacher'].values
            y_true = df[f'gt'].values
        # 计算 ROC 曲线以及 AUPRC 值
        pr, rc, thresholds = precision_recall_curve(y_true, y_score, drop_intermediate=True)
        auprc = average_precision_score(y_true, y_score)
        
        # 绘制 PRC 曲线，使用已分配的颜色，学术展示的曲线线宽合适
        ax.plot(rc, pr, label=f"{method} AUC: {auprc:.2f}", 
                color=color, lw=2, alpha=0.7)
    
    # 设置子图标题
    ax.set_title(f'{title} PR Curve', fontsize=18, weight='bold')
    # 设置横纵坐标的标签
    ax.set_xlabel("Recall", fontsize=14)
    ax.set_ylabel("Precision", fontsize=14)
    ax.yaxis.set_label_position("left")
    
    # 仅保留下边框和右边框, 去掉上边框和左边框
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    # 仅在底部和右侧显示刻度
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    
    # 设置图例，放置在右下角
    ax.legend(loc='upper right', fontsize=10, frameon=False, alignment='right', markerfirst=False)

# 调整整个图形的布局，避免子图之间重叠
plt.tight_layout(w_pad=5)
plt.savefig('auroc_2.svg')
plt.show()