用Python进行3D绘图

Question

我正在尝试在Python中绘制一个表面图。我有一个N乘N的数值表格。我创建了两个向量X和Y，每个向量都有N个元素。当我尝试绘制这个图时，出现了一个错误：

ValueError: total size of new array must be unchanged

我查看了一些例子，发现对于Z的N个元素，X和Y也需要有N个元素。

这让我感到困惑。为什么我需要N个元素，而不是N乘N的元素呢？

下面是一个示例代码：

import random
import math

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np

bignum = 100

mat = []
X = []
Y = []

for x in range(0,bignum):
    mat.append([])
    X.append(x);
    for y in range (0,bignum):
        mat[x].append(random.random())
        Y.append(y)

fig = plt.figure(figsize=plt.figaspect(2.))
ax = fig.add_subplot(1,1,1, projection='3d')
surf = ax.plot_surface(X,Y,mat)

Answer 1

首先，千万不要这样做：

mat = []
X = []
Y = []

for x in range(0,bignum):
    mat.append([])
    X.append(x);
    for y in range (0,bignum):
        mat[x].append(random.random())
        Y.append(y)

这相当于：

mat = np.random.random((bignum, bignum))
X, Y = np.mgrid[:bignum, :bignum]

...但这样做要快得多，而且占用的内存也少得多，相比于先用列表再转换成数组。

不过，你的例子是完全可以工作的。

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np

bignum = 100
mat = np.random.random((bignum, bignum))
X, Y = np.mgrid[:bignum, :bignum]

fig = plt.figure()
ax = fig.add_subplot(1,1,1, projection='3d')
surf = ax.plot_surface(X,Y,mat)
plt.show()

如果你查看一下plot_surface的文档，会发现它明确说明X、Y和Z应该是二维数组。

这样做是为了让你能够绘制更复杂的表面（比如球体），通过定义点之间的连接关系来实现。（例如，可以参考这个来自matplotlib画廊的例子：http://matplotlib.sourceforge.net/examples/mplot3d/surface3d_demo2.html）

如果你有一维的X和Y数组，并且想从二维网格中生成一个简单的表面，那么可以使用numpy.meshgrid或numpy.mgrid来生成合适的X和Y二维数组。

补充：为了说明mgrid和meshgrid的作用，我们来看看它们的输出：

print np.mgrid[:5, :5]

结果是：

array([[[0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1],
        [2, 2, 2, 2, 2],
        [3, 3, 3, 3, 3],
        [4, 4, 4, 4, 4]],

       [[0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4]]])

所以，它返回一个形状为2x5x5的三维数组，但更容易理解为两个二维数组。一个表示5x5网格上任意点的i坐标，另一个表示j坐标。

由于Python的解包方式，我们可以直接写：

xx, yy = np.mgrid[:5, :5]

Python并不关心mgrid具体返回了什么，它只会尝试将其解包为两个项目。因为numpy数组会遍历它们第一个轴的切片，所以如果我们解包一个形状为(2x5x5)的数组，就会得到两个5x5的数组。同样，我们也可以做类似的事情：

xx, yy, zz = np.mgrid[:5, :5, :5]

...并得到三个形状为5x5x5的三维数组的索引。如果我们用不同的范围切片（例如xx, yy = np.mgrid[10:15, 3:8]），它会返回从10到14（包括10和14）和从3到7（包括3和7）的索引。

mgrid还有更多功能（它可以接受复杂的步长参数来模拟linspace，例如xx, yy = np.mgrid[0:1:10j, 0:5:5j]会返回两个10x5的数组，分别在0到1和0到5之间递增），但我们先跳到meshgrid。

meshgrid接受两个数组，并以类似mgrid的方式将它们拼接在一起。举个例子：

x = np.arange(5)
y = np.arange(5)
xx, yy = np.meshgrid(x, y)
print xx, yy

结果是：

(array([[0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4]]), 

array([[0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1],
       [2, 2, 2, 2, 2],
       [3, 3, 3, 3, 3],
       [4, 4, 4, 4, 4]]))

meshgrid实际上返回的是两个5x5的二维数组的元组，但这个区别并不重要。关键的不同在于索引不需要朝特定方向递增。它只是拼接了给定的数组。例如：

x = [0.1, 2.4, -5, 19]
y = [-4.3, 2, -1, 18.4]
xx, yy = np.meshgrid(x, y)

结果是：

(array([[  0.1,   2.4,  -5. ,  19. ],
       [  0.1,   2.4,  -5. ,  19. ],
       [  0.1,   2.4,  -5. ,  19. ],
       [  0.1,   2.4,  -5. ,  19. ]]),
 array([[ -4.3,  -4.3,  -4.3,  -4.3],
       [  2. ,   2. ,   2. ,   2. ],
       [ -1. ,  -1. ,  -1. ,  -1. ],
       [ 18.4,  18.4,  18.4,  18.4]]))

你会注意到，它只是拼接了我们给定的值。

基本上，当你需要处理与输入网格形状相同的索引时，就可以使用这些工具。它们主要在你想在网格值上评估一个函数时很有用。

例如：

import numpy as np
import matplotlib.pyplot as plt

x, y = np.mgrid[-10:10, -10:10]
dist = np.hypot(x, y) # Linear distance from point 0, 0
z = np.cos(2 * dist / np.pi)

plt.title(r'$\cos(\frac{2*\sqrt{x^2 + y^2}}{\pi})$', size=20)
plt.imshow(z, origin='lower', interpolation='bicubic',
          extent=(x.min(), x.max(), y.min(), y.max()))
plt.colorbar()
plt.show()