如何用matplotlib绘制图像分类任务的类别特征空间分布

import matplotlib.pyplot as plt
import numpy as np
from sklearn.decomposition import PCA
from sklearn.datasets import load_iris
from mpl_toolkits.mplot3d import Axes3D

# 加载示例数据（Iris 数据集）
data = load_iris()
X = data.data
y = data.target
categories = data.target_names

# 使用PCA将数据降维到二维
pca_2d = PCA(n_components=2)
X_reduced_2d = pca_2d.fit_transform(X)

# 绘制二维特征空间分布
plt.figure(figsize=(10, 8))
for i, category in enumerate(categories):
    plt.scatter(X_reduced_2d[y == i, 0], X_reduced_2d[y == i, 1], label=category)
plt.xlabel('PCA 1')
plt.ylabel('PCA 2')
plt.title('类别特征空间分布（二维）')
plt.legend()
plt.savefig('2d_feature_space.png')
plt.show()

# 使用PCA将数据降维到三维
pca_3d = PCA(n_components=3)
X_reduced_3d = pca_3d.fit_transform(X)

# 绘制三维特征空间分布
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(111, projection='3d')
for i, category in enumerate(categories):
    ax.scatter(X_reduced_3d[y == i, 0], X_reduced_3d[y == i, 1], X_reduced_3d[y == i, 2], label=category)
ax.set_xlabel('PCA 1')
ax.set_ylabel('PCA 2')
ax.set_zlabel('PCA 3')
ax.set_title('类别特征空间分布（三维）')
ax.legend()
plt.savefig('3d_feature_space.png')
plt.show()