XGBoost案例之红酒品质分类-数据预处理

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split

# 1. 加载数据
print("=== 加载红酒数据 ===")
data = pd.read_csv('data/winequality.csv')
print(f"数据形状: {data.shape}")
print(f"\n前5行:")
print(data.head())
print(f"\n列名: {data.columns.tolist()}")

# 2. 提取特征和标签
print("\n=== 提取特征和标签 ===")
# 最后一列是标签，其他是特征
X = data.iloc[:, :-1]
y = data.iloc[:, -1]

print(f"特征形状: {X.shape}")
print(f"标签形状: {y.shape}")
print(f"\n标签分布:")
print(y.value_counts().sort_index())

# 3. 标签处理（减3，转为0-5）
print("\n=== 标签处理 ===")
print(f"原始标签: {sorted(y.unique())}")
y = y - 3
print(f"处理后标签: {sorted(y.unique())}")
print(f"\n处理后标签分布:")
print(y.value_counts().sort_index())

# 4. 划分数据集
print("\n=== 划分数据集 ===")
X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,
    random_state=22,
    stratify=y  # 分层抽样，参考Y轴比例
)

print(f"训练集: {X_train.shape}")
print(f"测试集: {X_test.shape}")

# 5. 查看分层效果
print("\n=== 分层抽样效果 ===")
print("训练集标签分布:")
print(y_train.value_counts().sort_index())
print("\n测试集标签分布:")
print(y_test.value_counts().sort_index())

# 6. 保存数据集
print("\n=== 保存数据集 ===")

# 训练集
train_data = X_train.copy()
train_data['label'] = y_train
train_data.to_csv('data/wine_train.csv', index=False)
print(f"训练集已保存: {train_data.shape}")

# 测试集
test_data = X_test.copy()
test_data['label'] = y_test
test_data.to_csv('data/wine_test.csv', index=False)
print(f"测试集已保存: {test_data.shape}")

# 7. 数据探索
print("\n=== 数据探索 ===")
print(f"特征统计:")
print(X.describe())

# 8. 检查缺失值
print("\n=== 检查缺失值 ===")
print(f"特征缺失值:")
print(X.isnull().sum())

# 9. 特征相关性
print("\n=== 特征相关性 ===")
correlation = X.corr()
print("特征相关性矩阵（前5x5）:")
print(correlation.iloc[:5, :5])

# 10. 可视化准备
print("\n=== 数据可视化准备 ===")
import matplotlib.pyplot as plt

# 标签分布
fig, axes = plt.subplots(1, 2, figsize=(12, 5))

# 原始标签分布
axes[0].bar(y_train.value_counts().sort_index().index, 
            y_train.value_counts().sort_index().values)
axes[0].set_title('训练集标签分布')
axes[0].set_xlabel('类别')
axes[0].set_ylabel('数量')

# 测试集标签分布
axes[1].bar(y_test.value_counts().sort_index().index, 
            y_test.value_counts().sort_index().values)
axes[1].set_title('测试集标签分布')
axes[1].set_xlabel('类别')
axes[1].set_ylabel('数量')

plt.tight_layout()
plt.savefig('data/label_distribution.png', dpi=100)
print("标签分布图已保存")

# 11. 完整流程函数
def preprocess_wine_data(input_path, output_dir='data'):
    """红酒数据预处理完整流程"""
    print("=" * 50)
    print("红酒品质分类 - 数据预处理")
    print("=" * 50)
    
    # 1. 加载数据
    data = pd.read_csv(input_path)
    print(f"1. 加载数据: {data.shape}")
    
    # 2. 提取特征和标签
    X = data.iloc[:, :-1]
    y = data.iloc[:, -1]
    print(f"2. 特征: {X.shape}, 标签: {y.shape}")
    
    # 3. 标签处理
    y = y - 3  # 转为0-5
    print(f"3. 标签处理: {sorted(y.unique())}")
    
    # 4. 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y,
        test_size=0.2,
        random_state=22,
        stratify=y
    )
    print(f"4. 划分数据集: 训练{X_train.shape}, 测试{X_test.shape}")
    
    # 5. 保存
    train_data = X_train.copy()
    train_data['label'] = y_train
    train_data.to_csv(f'{output_dir}/wine_train.csv', index=False)
    
    test_data = X_test.copy()
    test_data['label'] = y_test
    test_data.to_csv(f'{output_dir}/wine_test.csv', index=False)
    
    print(f"5. 数据已保存")
    print("\n预处理完成!")
    
    return X_train, X_test, y_train, y_test

# 运行
# X_train, X_test, y_train, y_test = preprocess_wine_data('data/winequality.csv')