Posted on 23 February 2013
A long while ago, I wrote a script for
matplotlib that later (after lots of work by
Tony Yu and others) became the matplotlib
streamplot
function. This function generates plots that are constructed out of
streamlines of a flow, in such a way as to fill the domain uniformly
with such streamlines. I wrote the script to plot winds in the
tropical upper troposphere/lower stratosphere where there is a lot of
divergence and convergence (principally caused by convection over the
western Pacific and subsidence over the eastern Pacific.)
The streamplot
algorithm uses an adaptive 1st/2nd order
Runge-Kutta
scheme and linear interpolation of the velocity field to integrate the
streamlines. This was implemented in pure python, which I originally
chose to do because the portion of the algorithm responsible for
ensuring an even domain-filling distribution of streamlines imposes a
very strange termination condition on the integral that does not fit
well with the integrators in scipy
. In addition, matplotlib does not
depend on scipy
.
I initially expected a pure python algorithm to be incredibly slow, so
I used a very simple interpolation scheme. Any interpolation scheme
interpolating some data D[i]
defined at points x[i]
on to the
target point x0
is composed of two steps; first, the algorithm must
identify the nearest grid points to the target point. Second, the
interpolator must apply the appropriate interpolator to those
neighbouring points.
The second step is usually trivial to implement, but the first step is
non-trivial in general. The first step becomes a lot simpler when we
restrict ourselves to a regular grid with spacing dx
as we can then
find the index i
of the point just above x0
by rescaling and
casting to an integer
I thought this was probably the fastest way of implementing an
interpolator and so used this technique to implement the interpolator
in streamplot
. However, it added some complexity to the algorithm
and, most importantly, imposed a restriction that only regular grids
(i.e. grids with evenly spaced axes) could be used. For my original
purpose, this was not a problem but it is clearly a rather serious
limitation in general.
Recently I have been hitting this limitation and decided over the last
couple of evenings to try to free streamplot
of this limitation. To
generalise the first interpolation step to work with any x[i]
, a
search is needed to find the smallest i
such that x[i] > x0
. Any
pure python search algorithm would be too slow. Using argmax
to do a
linear search is an option:
However, this does not scale well with the length of each dimension and is too slow in this context.
Looking through the standard library, I came across the bisect
module which implements the
interval bisection method
for searching ordered lists and arrays – exactly what I needed
to do. Implementing this search is really simple; it’s just
This method is really fast – way faster than I expected. After re-implementing the streamplot integrator and interpolator to use this type of interpolation, typical plotting times were barely increased.
Here is an example of the new script in action with a very non-uniform vertical grid spacing:
(This plot shows the boreal winter average zonal and vertical winds in the inner tropics.)
The new script is available here. Feel free to try it out if you need to do streamplots with uneven grids. Please note that this is a work-in-progress and so use with caution :-)