vis/dataset_vis.py

import sys 
sys.path.append("..") 
from dataset import PeptidePairDataset
from torch.utils.data import Subset
from sklearn.model_selection import KFold
import matplotlib.pyplot as plt
import numpy as np

def extract_label(dataset: PeptidePairDataset):
    labels = []
    for _, label in dataset:
        labels.append(label)
    return labels

def extract_labels_cls(dataset: PeptidePairDataset):
    labels = []
    for _, label in dataset:
        if label >= 0.5:
            labels.append(2)
        elif label <= -0.5:
            labels.append(1)
        else:
            labels.append(0)
    return labels

all_set = PeptidePairDataset(mode='train', pad_length=30, task='reg')
test_set = PeptidePairDataset(mode='r2_case', pad_length=30, task='reg', one_way=True)

# kf = KFold(n_splits=5, shuffle=True, random_state=42)
# train_idx, val_idx = next(kf.split(all_set))
# train_set = Subset(all_set, train_idx)
# valid_set = Subset(all_set, val_idx)

bins = np.arange(-6.5, 6.6, 0.5)

def plot_gradient_bar_histogram(ax, data, bins, cmap_name='Blues'):
    """使用渐变配色方案绘制直方图，每个柱子颜色不同但连续，并在柱子上添加数值标注"""
    # 计算直方图数据
    counts, bin_edges = np.histogram(data, bins=bins)
    n_bins = len(counts)
    # 获取 colormap
    cmap = plt.get_cmap(cmap_name)
    width = np.diff(bin_edges)[0]  # 本例中每个柱子的宽度均为固定值
    # 绘制每个柱子，并添加数值标注
    for i in range(n_bins):
        # 归一化的颜色比例
        if n_bins > 1:
            ratio = (i / (n_bins - 1)) * 0.5 + 0.3
        else:
            ratio = 0.5
        color = cmap(ratio)
        # 绘制柱子
        ax.bar(bin_edges[i], counts[i], width=width, align='edge',
               color=color, zorder=2)
        # 为非空柱子添加数值标注
        if counts[i] != 0:
            x_pos = bin_edges[i] + width / 2
            y_pos = counts[i]
            ax.text(
                x_pos, y_pos, f'{counts[i]}', 
                ha='center', va='bottom', fontweight='bold', fontsize=9
                )
    return counts, bin_edges

def plot_gradient_bar_histogram_cls(ax, data, cmap_name='Blues'):
    """将类别数据统计后绘制直方图，并添加数值标注，类别固定为 [0, 1, 2]"""
    # 固定类别顺序
    classes = np.array([0, 1, 2])
    counts = np.array([data.count(cls) for cls in classes])
    num_classes = len(classes)
    width = 0.6
    cmap = plt.get_cmap(cmap_name)
    for i, cls in enumerate(classes):
        # 渐变配色比例
        ratio = (i / (num_classes - 1)) * 0.5 + 0.3 if num_classes > 1 else 0.5
        color = cmap(ratio)
        ax.bar(cls, counts[i], width=width, color=color, zorder=2)
        if counts[i] != 0:
            ax.text(cls, counts[i], f'{counts[i]}', ha='center', va='bottom', fontweight='bold')
    ax.set_xticks(classes)
    ax.set_xticklabels(['Unchanged', 'Lower', 'Higher'])
    return counts

# # 设置全局字体大小
# plt.rcParams.update({
#     'font.size': 10,
#     'axes.titlesize': 14,
#     'axes.labelsize': 14,
#     'xtick.labelsize': 13,
#     'ytick.labelsize': 13,
#     'legend.fontsize': 14,
#     'figure.titlesize': 16
# })
plt.style.use('seaborn-v0_8-whitegrid')

# 创建 3 行 2 列子图，第一列为回归数据的可视化，第二列为分类数据的可视化
# fig, axes = plt.subplots(3, 2, figsize=(12, 9), width_ratios=[2.5, 1])
fig, axes = plt.subplots(1, 2, figsize=(11.4, 3))

# 数据提取：回归数据
reg_train = extract_label(all_set)
# reg_val = extract_label(valid_set)
reg_test = extract_label(test_set)

# 数据提取：分类数据
# cls_train = extract_labels_cls(all_set)
# cls_val = extract_labels_cls(valid_set)
# cls_test = extract_labels_cls(test_set)

# 第一行：训练集
counts, _ = plot_gradient_bar_histogram(axes[0], reg_train, bins, cmap_name='Blues')
axes[0].set_xlabel('log2 (MIC ratio)', size=14)
axes[0].set_xlim(-6.5, 6.5)
axes[0].set_ylabel('Counts', size=14)
axes[0].set_ylim(0, max(counts)*1.15)
axes[0].set_title(f'Train Set ({len(reg_train)})', weight='bold', size=16)

# counts_cls = plot_gradient_bar_histogram_cls(axes[0, 1], cls_train, cmap_name='Blues')
# axes[0, 1].set_xlabel('Class')
# axes[0, 1].set_xlim(-0.5, 2.5)
# axes[0, 1].set_ylabel('Counts')
# axes[0, 1].set_ylim(0, max(counts_cls)*1.15)
# axes[0, 1].set_title(f'Train Set (Cls.) ({len(cls_train)})', weight='bold')

# 第二行：验证集
counts, _ = plot_gradient_bar_histogram(axes[1], reg_test, bins, cmap_name='Blues')
axes[1].set_xlabel('log2 (MIC ratio)', size=14)
axes[1].set_xlim(-6.5, 6.5)
axes[1].set_ylabel('Counts', size=14)
axes[1].set_ylim(0, max(counts)*1.15)
axes[1].set_title(f'R2 Set ({len(reg_test)})', weight='bold', size=16)

# counts_cls = plot_gradient_bar_histogram_cls(axes[1, 1], cls_val, cmap_name='Blues')
# axes[1, 1].set_xlabel('Class')
# axes[1, 1].set_xlim(-0.5, 2.5)
# axes[1, 1].set_ylabel('Counts')
# axes[1, 1].set_ylim(0, max(counts_cls)*1.15)
# axes[1, 1].set_title(f'Validation Set (Cls.) ({len(cls_val)})', weight='bold')

# # 第三行：测试集
# counts, _ = plot_gradient_bar_histogram(axes[2, 0], reg_test, bins, cmap_name='Blues')
# axes[2, 0].set_xlabel('log2 (MIC ratio)')
# axes[2, 0].set_xlim(-6.5, 6.5)
# axes[2, 0].set_ylabel('Counts')
# axes[2, 0].set_ylim(0, max(counts)*1.15)
# axes[2, 0].set_title(f'Test Set (Reg.) ({len(reg_test)})', weight='bold')

# counts_cls = plot_gradient_bar_histogram_cls(axes[2, 1], cls_test, cmap_name='Blues')
# axes[2, 1].set_xlabel('Class')
# axes[2, 1].set_xlim(-0.5, 2.5)
# axes[2, 1].set_ylabel('Counts')
# axes[2, 1].set_ylim(0, max(counts_cls)*1.15)
# axes[2, 1].set_title(f'Test Set (Cls.) ({len(cls_test)})', weight='bold')

plt.tight_layout(w_pad=3.2)
plt.savefig('dataset.svg')
plt.show()

# 输出数据集统计信息（回归数据）
print("\n训练集回归统计信息:")
print(f"样本总数: {len(reg_train)}")
print(f"范围: {min(reg_train)} - {max(reg_train)}")

# print("\n验证集回归统计信息:")
# print(f"样本总数: {len(reg_val)}")
# print(f"范围: {min(reg_val)} - {max(reg_val)}")

print("\n测试集回归统计信息:")
print(f"样本总数: {len(reg_test)}")
print(f"范围: {min(reg_test)} - {max(reg_test)}")

# 输出数据集统计信息（分类数据）
# print("\n训练集分类统计信息:")
# print(f"样本总数: {len(cls_train)}")
# print(f"各类别统计: 0类: {cls_train.count(0)}, 1类: {cls_train.count(1)}, 2类: {cls_train.count(2)}")

# print("\n验证集分类统计信息:")
# print(f"样本总数: {len(cls_val)}")
# print(f"各类别统计: 0类: {cls_val.count(0)}, 1类: {cls_val.count(1)}, 2类: {cls_val.count(2)}")

# print("\n测试集分类统计信息:")
# print(f"样本总数: {len(cls_test)}")
# print(f"各类别统计: 0类: {cls_test.count(0)}, 1类: {cls_test.count(1)}, 2类: {cls_test.count(2)}")