Matplotlib - 如何在极坐标中绘制流线？

Question

我一直在尝试在 matplotlib 1.4.3 的极轴上绘制流线。 streamplot 函数从 1.2.0 开始就已经存在，并且被文档认为是功能性和稳定的。 这是一个小测试脚本：

from matplotlib import pyplot as plt
import numpy as np

# Define polar grid
r = np.arange(0,2001,50)
theta = np.arange(-np.pi, np.pi+np.pi/180, 2*np.pi/180)
r2D, theta2D = np.meshgrid(r, theta)
# Define some data
u = -np.sin(theta2D)
v = np.cos(theta2D)
# Set up axes
fig = plt.figure()
ax = fig.add_axes([0.1, 0.1, 0.8, 0.8], polar=True)
# Plot streamlines
ax.streamplot(r, theta, u, v, color='k', density=1, linewidth=1)

此脚本失败并显示以下回溯：

Traceback (most recent call last):
  File "streamline_test.py", line 15, in <module>
    ax.streamplot(r, theta, u, v, color='k', density=1, linewidth=1)
  File "python2.7/site-packages/matplotlib/axes/_axes.py", line 4204, in streamplot
    zorder=zorder)
  File "python2.7/site-packages/matplotlib/streamplot.py", line 167, in streamplot
    axes.add_patch(p)
  File "python2.7/site-packages/matplotlib/axes/_base.py", line 1568, in add_patch
    self._update_patch_limits(p)
  File "python2.7/site-packages/matplotlib/axes/_base.py", line 1586, in _update_patch_limits
    vertices = patch.get_path().vertices
  File "python2.7/site-packages/matplotlib/patches.py", line 4033, in get_path
    _path, fillable = self.get_path_in_displaycoord()
  File "python2.7/site-packages/matplotlib/patches.py", line 4054, in get_path_in_displaycoord
    shrinkB=self.shrinkB * dpi_cor
  File "python2.7/site-packages/matplotlib/patches.py", line 2613, in __call__
    shrinked_path = self._shrink(clipped_path, shrinkA, shrinkB)
  File "python2.7/site-packages/matplotlib/patches.py", line 2586, in _shrink
    left, right = split_path_inout(path, insideA)
  File "python2.7/site-packages/matplotlib/bezier.py", line 246, in split_path_inout
    ctl_points, command = next(path_iter)
StopIteration

显然 streamplot 永远迭代并且必须在某个时候停止。 我还尝试了一组应用于极轴的规则间隔的笛卡尔点，但以同样的方式失败。 使用笛卡尔轴制作极坐标图不是一种选择，因为我需要一个极坐标网格，但这样的网格在笛卡尔坐标中不是规则间隔的，并且流图需要规则间隔的点。

有人知道如何让 matplotlib 在极坐标中绘制流线吗？

Answer 1

您只需要切换径向和方位角坐标。 考虑下图，以及用于生成它的代码。 注意streamplot()的第三个参数中的streamplot()半径，它将线速度转换为角速度：

import math
import numpy
import matplotlib
from matplotlib import pyplot

pyplot.gcf().add_axes([0.1, 0.1, 0.8, 0.8], polar=True)

# coordinates
r = numpy.linspace(0, 1, 11)
p = numpy.linspace(-math.pi, math.pi, 361)
rg, pg = numpy.meshgrid(r, p)

def repeat(x):
  return numpy.full_like(r, x)

epsilon = 1e-8

# cylindrical components of horizontal unit vector
xr =  numpy.cos(pg)
xp = -numpy.sin(pg)
# cylindrical components of vertical unit vector
yr =  numpy.sin(pg)
yp =  numpy.cos(pg)

# starting points of streamlines
sx = numpy.transpose([
  numpy.hstack([repeat(-math.pi/2), repeat(math.pi/2)]),
  numpy.hstack([r, r])
])
sy = numpy.transpose([
  numpy.hstack([repeat(-math.pi+epsilon), repeat(0), repeat(math.pi-epsilon)]),
  numpy.hstack([r, r, r])
])

# streamlines
pyplot.streamplot(
  pg.transpose(), rg.transpose(), (xp/rg).transpose(), xr.transpose(),
  color='red', start_points=sx)
pyplot.streamplot(
  pg.transpose(), rg.transpose(), (yp/rg).transpose(), yr.transpose(),
  color='blue', start_points=sy)

pyplot.ylim(0, 1)
pyplot.annotate(
  matplotlib.__version__,
  (0, 0), (1, 0), 'axes fraction', 'axes fraction',
  ha='left', va='top')
pyplot.savefig('stream.png')