matplotlib (mplot3d) - 怎样拉伸/缩短坐标轴的长度

首先用python的matplotlib库画一个图像。
代码如下：

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

fig = plt.figure()
ax = Axes3D(fig)

theta = np.linspace(-4 * np.pi, 4* np.pi, 100)
z = np.linspace(-2, 2, 100)
r = z ** 2 + 1
x = r * np.sin(theta)
y = r * np.cos(theta)

ax.plot(x, y , z)
plt.show()

上面的代码作出的图如下所示：