Matplotlib - How to plot streamlines in polar coordinates?

Question

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?

Answer 1

You simply need to switch the radial and azimuthal coordinates. Consider the following figure, and the code used to generate it. Note the division by radius in the third argument to streamplot(), which converts linear velocity 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)

# 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')