matplotlib - 3D scatter plot

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  #3D图表需要使用“mpl_toolkits”模块
 
#①构造数据
np.random.seed(42)  #随机数种子
n_samples = 500  # 采样个数为500

dim = 3
#生成一组3维正态分布数据，数据方向完全随机
samples = np.random.multivariate_normal(np.zeros(dim), np.eye(dim), n_samples)  #500x3
#通过把每个样本到原点距离和均匀分布吻合得到球体内均匀分布的样本
for i in range(samples.shape[0]):
    r = np.power(np.random.random(), 1.0/3.0)
    samples[i] *= r/np.linalg.norm(samples[i])

    
#②将数据分为两组，以3x+2y-z=1作为判别平面
upper_samples = []
lower_samples = []
for x, y, z in samples: 
    if z > 3*x+2*y-1:
        upper_samples.append((x, y, z))
    else:
        lower_samples.append((x, y, z))

#③制作散点图
fig = plt.figure('3D scatter plot')
ax = fig.add_subplot(111, projection='3d')  #3d图需要加projection='3d'
uppers = np.array(upper_samples)
lowers = np.array(lower_samples)
 
# 用不同颜色、形状的图标区分两组数据
ax.scatter(uppers[:,0], uppers[:,1], uppers[:,2], c='r', marker='o')  #上半部分为红色圆点
ax.scatter(lowers[:,0], lowers[:,1], lowers[:,2], c='g', marker='^')  #下半部分为绿色三角

plt.savefig('3D scatter plot.png')  #保存为图片
plt.show()