如何在python中绘制圆柱体周围的流线图？

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

# 创建圆柱体的数据
radius = 1.0
height = 2.0
theta = np.linspace(0, 2*np.pi, 100)
z = np.linspace(0, height, 10)
theta, z = np.meshgrid(theta, z)
x = radius * np.cos(theta)
y = radius * np.sin(theta)

# 创建流线图的数据
start_points = np.vstack([x.flatten(), y.flatten(), z.flatten()]).T
end_points = np.vstack([x.flatten(), y.flatten(), z.flatten()+1]).T

# 绘制3D图形
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# 绘制圆柱体表面
ax.plot_surface(x, y, z, alpha=0.5)

# 绘制流线图
ax.quiver(start_points[:, 0], start_points[:, 1], start_points[:, 2],
          end_points[:, 0]-start_points[:, 0], end_points[:, 1]-start_points[:, 1], end_points[:, 2]-start_points[:, 2],
          length=0.5, normalize=True, color='r')

# 设置坐标轴范围
ax.set_xlim(-radius, radius)
ax.set_ylim(-radius, radius)
ax.set_zlim(0, height)

# 设置坐标轴标签
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

# 显示图形
plt.show()