在 3D MatplotLib 散点图中绘制超平面

时间:2021-05-20 20:32:44

标签: python matplotlib 3d plotly plane

我有以下代码。它可以绘制 3D 散点图(没有平面)。我也想添加飞机。我有 Z 轴值,但是当我尝试使用这些值时,出现以下错误 shape mismatch: objects cannot be broadcast to a single shape

如何绘制平面的 Z 轴?

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

data = {
        'x': [30, 48, 15, 25, 19, 11, 19, 13, 11, 43, 37, 30, 20, 41, 18, 39, 14, 48, 26, 44, 36, 24, 20, 14, 19, 28, 13, 30, 37, 17, 20, 45, 46, 40, 36, 17],
        'y': [38, 35, 25, 31, 20, 31, 37, 27, 32, 34, 36, 35, 38, 24, 27, 38, 32, 39, 28, 40, 26, 22, 24, 40, 21, 30, 32, 22, 39, 39, 24, 37, 28, 23, 34, 34],
        'z': [1.57, 1.6, 1.64, 1.67, 1.75, 1.79, 1.87, 1.95, 1.96, 1.99, 2.01, 2.34, 2.45, 2.47, 2.54, 2.67, 2.85, 2.92, 2.94, 3.29, 3.4, 3.43, 3.52, 3.54, 3.63, 3.69, 3.69, 3.82, 3.99, 3.99, 4.35, 4.37, 4.92, 4.94, 5, 5],
        'plane': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5]
        }

df = pd.DataFrame(data)

planes = [1, 2, 3, 4, 5]

cdict = {1: 'blue',
         2: 'yellow',
         3: 'orange',
         4: 'red',
         5: 'green'
        }

fig = plt.figure(figsize=(12, 12))

ax = fig.add_subplot(111, projection='3d')

for i in planes:
    df2 = df[df['plane']==i]
    x = df2['x'].values.reshape(-1, 1)
    y = df2['y'].values.reshape(-1, 1)
    z = df2['z'].values.reshape(-1, 1)
    
    ax.scatter(xs = x,
               ys = y,
               zs = z, 
               label = i, 
               c = cdict[i]
              )
    
    xx , yy = np.meshgrid(np.arange(min(x), max(x), 1), np.arange(min(y), max(y), 1))
    zz = np.array(np.meshgrid(np.arange(min(z), max(z), 0.01)) * len(z)) #Multiply by count to have the same array size as xx and yy
    ax.plot_surface(xx, yy, zz, alpha=0.2)
    
ax.legend(loc="best")

plt.show()

0 个答案:

没有答案