Matplotlib - How to plot streamlines in polar coordinates?

问题

I have been trying to plot streamlines on a polar axis in matplotlib 1.4.3. The streamplot function has been around since 1.2.0 and is considered functional and stable by the documentation. Here is a little test script:

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)

This script fails with the following traceback:

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

Apparently streamplot is iterating forever and has to stop at some point. I have also tried a set of regularly-spaced cartesian points applied to the polar axis, but that fails in the same way. Making a polar plot using cartesian axes is not an option as I need a polar grid, but such a grid is not regularly-spaced in cartesian coordinates, and streamplot requires regularly-spaced points.

Does anybody know how to get matplotlib to plot streamlines in polar coordinates?

回答1:

You simply need to switch the radial and azimuthal coordinates. The following code makes this plot; note the division by r in the third argument to streamplot(), which converts linear to angular velocity:

import math
import numpy
import matplotlib
from matplotlib import pyplot

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

r = numpy.linspace(0, 1, 11)
phi = numpy.linspace(-1, 1, 361) * math.pi
r, phi = numpy.meshgrid(r, phi)

ones = numpy.ones_like(r)
zeros = numpy.zeros_like(r)
pyplot.streamplot(
  phi.transpose(), r.transpose(),
  (ones/r).transpose(), zeros.transpose(),
  color='red')
pyplot.streamplot(
  phi.transpose(), r.transpose(),
  (zeros/r).transpose(), ones.transpose(),
  color='blue')
pyplot.ylim(0, 1)
pyplot.annotate(matplotlib.__version__,
  (0, 0), (1, 0), 'axes fraction', 'axes fraction',
  ha='left', va='top')
pyplot.savefig('stream.png')

来源：https://stackoverflow.com/questions/30601865/matplotlib-how-to-plot-streamlines-in-polar-coordinates

标签

python

matplotlib

polar-coordinates