Add `streamplot` axes method to plot streamlines #664

Merged
merged 23 commits into from Feb 27, 2012

Projects

None yet

8 participants

@tonysyu
Contributor
tonysyu commented Jan 8, 2012

Addresses Issue #366

tonysyu added some commits Jan 8, 2012
@tonysyu tonysyu Add `streamplot` axes method to plot streamlines 687af18
@tonysyu tonysyu Add improvements from Tom Flannaghan.
* Add faster adaptive-step algorithm (RK12) to replace RK45 and RK4 integrators.
* Fix arrow placement when streamline is only 2 points.
* Calculate stepsize based on grid size.
0657e73
@tonysyu tonysyu Fix bounds checking.
Checking that the coordinate less than the array Nx/Ny does not work when the coordinate is float (e.g. it can be between `Nx` and `Nx - 1`).
a88f621
@tonysyu tonysyu Remove `Terminated` exception. 28a9e27
@tonysyu tonysyu Remove OutOfBounds exception and factor out Euler step e6e977a
@tonysyu tonysyu Move comments into docstring. 3bf9f09
@tonysyu tonysyu Adjust exponent of step-size correction to match method-order. 3144f6d
@tonysyu tonysyu Fix error raised when trajectory lands on Nx or Ny. 5632418
@tonysyu tonysyu Remove integrator parameter for streamplot
* RK12 is 1.5--2x faster than RK4 and RK45 produces visually-similar results.
* Remove implementations of RK4 and RK45.
d60c9fc
@mdboom
Member
mdboom commented Jan 12, 2012

I haven't done a detailed review, but on first glance, I think it would be nice to have a pyplot method for this as well.

@tonysyu
Contributor
tonysyu commented Jan 12, 2012

I agree. I mentioned this in an email to the dev list when I posted the pull request. When I looked at the pyplot module, I noticed that functions had comments that they were autogenerated by boilerplate.py, and it wasn't clear to me how to do that.

@WeatherGod
Member

it means that you run boilerplate.py to regenerate the pyplot.py file and commit the changes to it that way.

@WeatherGod
Member

Also note that you should add an entry to the CHANGELOG file

@efiring efiring and 2 others commented on an outdated diff Jan 16, 2012
lib/matplotlib/streamplot.py
@@ -0,0 +1,523 @@
+"""
+Streamline plotting like Mathematica.
+Copyright (c) 2011 Tom Flannaghan.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
@efiring
efiring Jan 16, 2012 Member

I think the copyright notice should go in a separate string, not in the module docstring. The module docstring could refer to it.

@tonysyu
tonysyu Jan 17, 2012 Contributor

Sounds good, but I'm not sure exactly what you mean. Should I move the notice to a separate string (or comment?) in the same module, or should I point to a copyright notice that lives in another file? Also, are you referring to both the "Copyright (c)..." and "Permission is hereby..." parts, or just one of these?

@efiring
efiring Jan 17, 2012 Member

It looks like there is no single standard way to do it. pyparsing_py2.py puts it in a comment. There are a few other copyright notices sprinkled around other modules, but not many. mlab.py has a free-standing multi-line string acting as if it were a comment. You could ask Tom what his preference is. The two best candidates seem to me to be the comment block (as in pyparsing_py2.py), and a multi-line string that acts like a comment but is not part of the module docstring. In either case, the "Copyright" and "Permission" parts normally go together.

A nice convention might be to have a standard name for any copyright info that differs from the global mpl copyright and license info:

copyright = """Copyright 2012, Joe Smith
Permission...
"""

Then it would be easy to sweep through a package and find all copyright information. I am not aware of any such convention being already established, however. You could start one...

@tomflannaghan
tomflannaghan Jan 18, 2012

I think it should be removed. I'm not at all bothered about the copyright - I only added it because on the mailing list someone requested clarification for a previous script.

@tonysyu
tonysyu Jan 18, 2012 Contributor

Thanks for your input Tom. I've removed the copyright notice in the most recent commit.

@efiring
Member
efiring commented Jan 16, 2012

In boilerplate.py, you need to add your axes method to the _plotcommands list, and add an appropriate entry to the cmappable dictionary to set the current image appropriately. Then you run boilerplate.py and capture its stdout. Manually edit pyplot.py, deleting the part provided by the previous boilerplate.py, and substituting the new output. Test the result. Now you are ready for a commit. It's all very clumsy, but it doesn't need to be done often.

@efiring
Member
efiring commented Jan 16, 2012

A limitation of this code is that is has no support for masked arrays; everything has to be a complete quadrilateral grid. Since most of mpl does support masked arrays, this unfortunate limitation probably should be highlighted in the docstring. (Any thoughts on what it would take to support masking? It is critical for many applications in fluid mechanics, both for engineering and for oceanography.)

@tonysyu
Contributor
tonysyu commented Jan 16, 2012

@efiring Thanks for the detailed instructions for adding pyplot commands. I think this raises a possible issue with this streamplot implementation: The streamplot command doesn't currently return a ScalarMappable instance. A line collection and individual arrow patches are added separately to the axes. I'm guessing I can collect the FancyArrowPatches into a PatchCollection, but that still leaves me with two separate collections. I'm not really sure what to do here.

Is there any reason pyplot isn't implemented as a subpackage (instead of a single file)? If it were a subpackage the output of boilerplate.py could be piped directly to a file (instead of being pasted into a file), and that file could be imported by the subpackage.

@WeatherGod
Member

don't forget about the Legend containers. What should be plotted in the legend for this plot?

@efiring
Member
efiring commented Jan 16, 2012

Contouring is similar, in that it is color-mappable, but doesn't return a simple Collection. The streamplot could be handled in a similar manner, with considerable reworking. One option is to leave that out for now, label streamplot "experimental", note that its return value is likely to change, and that at present it is not a proper mappable. One of the consequences is that the colorbar command won't work with it.

@efiring
Member
efiring commented Jan 16, 2012

Regarding your pyplot-as-submodule suggestion: I don't think this is even required. The boilerplate output could go to _pyplot_auto.py, and "from _pyplot_auto import *" could be used in pyplot.py to pull it in. But this is getting off-topic, and more suitable for the devel mailing list.

@tonysyu tonysyu Add streamplot function to pyplot.
As a temporary fix, streamplot returns the last plotted streamline and arrow for use with colorbars and legends. Instead, all lines should instead be grouped into a single LineCollection, and all FancyArrowPatches should be grouped into a patch collection.
6f3053e
@tonysyu
Contributor
tonysyu commented Jan 16, 2012

@efiring For the color-mappable, I ended up returning the last-plotted streamline and arrow, and use only the streamline as the color-mappable. A future fix will collect all the lines into a single line collection so that a colorbar would be properly scaled.

@WeatherGod This change plots a color patch in the legend---not ideal, but this is the same behavior as quiver.

@tonysyu tonysyu commented on the diff Jan 16, 2012
lib/matplotlib/pyplot.py
@@ -2623,7 +2623,7 @@ def stem(x, y, linefmt='b-', markerfmt='bo', basefmt='r-', hold=None):
if hold is not None:
ax.hold(hold)
try:
- ret = ax.stem(x, y, linefmt, markerfmt, basefmt)
+ ret = ax.stem(x, y, linefmt, markerfmt, basefmt, bottom, label)
draw_if_interactive()
@tonysyu
tonysyu Jan 16, 2012 Contributor

These changes to stem were auto-generated when I ran boilerplate.py.

tonysyu added some commits Jan 17, 2012
@tonysyu tonysyu Fix streamplot to work with colorbar.
Streamplot no longer returns an arrow patch and, instead, only returns a dummy LineCollection.
72abc5d
@tonysyu tonysyu Fix selection of arrow's linewidth and color b884fc8
@tonysyu tonysyu Add support for masked arrays.
Note that masked arrays are converted to float and filled with NaNs. Filling integer masked arrays with zeros would also work, but would give incorrect results in some instances.
7d01052
@tonysyu
Contributor
tonysyu commented Jan 17, 2012

@efiring In the most recent commit, I added support for masked arrays and NaN values in arrays. I wasn't certain how to handle integer masked arrays, so I just convert to float and fill with NaNs. I'm not certain that's optimal.

@WeatherGod
Member

The SOP is to convert NaNs into masked elements and to design algorithms
around masked arrays. This will help keep compatibility with future
changes to masked arrays such as numpy's NA.

@tonysyu tonysyu Change handling of masked arrays.
Instead of changing masked values to NaNs, change NaNs to masked values.
40dce9c
@tonysyu
Contributor
tonysyu commented Jan 17, 2012

@WeatherGod Thanks for the advice. I changed the behavior of NaNs and masked arrays as suggested.

tonysyu added some commits Jan 18, 2012
@tonysyu tonysyu Remove copyright notice.
As Tom Flannaghan suggested in GitHub discussions, the copyright notice was removed from streamplot.py. Added attribution to CHANGELOG.
07553f0
@tonysyu tonysyu Return collection of streamlines from `streamplot`.
Previously, `streamplot` just returned a dummy LineCollection so that colorbars and legends were properly set. Bonus: moving the creation of LineCollection objects out of the loop boosted the speed by about 10%.
3259d0a
@tonysyu tonysyu Add example to plot streamplot with masked values. a98c543
@tonysyu
Contributor
tonysyu commented Jan 19, 2012

Update: I fixed the previous hack for getting colorbars and legends to work properly. (streamplot now returns the actual streamlines as a LineCollection). Also, I added an example for masked arrays.

@pelson pelson and 1 other commented on an outdated diff Feb 7, 2012
lib/matplotlib/streamplot.py
+ self.dx = x[1] - x[0]
+ self.dy = y[1] - y[0]
+
+ self.x_origin = x[0]
+ self.y_origin = y[0]
+
+ self.width = x[-1] - x[0]
+ self.height = y[-1] - y[0]
+
+ @property
+ def shape(self):
+ return self.ny, self.nx
+
+ def valid_index(self, xi, yi):
+ """Return True if point is a valid index of grid."""
+ return xi >= 0 and xi <= self.nx-1 and yi >= 0 and yi <= self.ny-1
@pelson
pelson Feb 7, 2012 Member

Think this could be simplified to 0 <= xi <= self.nx - 1 and 0 <= yi <= self.ny - 1.

@tonysyu
tonysyu Feb 7, 2012 Contributor

I thought the same thing, but since xi and yi can be float values, xi and yi can be larger the nx-1 and ny-1 but less than nx and ny. I think the use of this function changed at some point so the name valid_index is misleading (since an index should be an integer). I'll probably change this in a future commit.

@pelson pelson and 1 other commented on an outdated diff Feb 7, 2012
lib/matplotlib/streamplot.py
+ def valid_index(self, xi, yi):
+ """Return True if point is a valid index of grid."""
+ return xi >= 0 and xi <= self.nx-1 and yi >= 0 and yi <= self.ny-1
+
+
+class StreamMask(object):
+ """Mask to keep track of discrete regions crossed by streamlines.
+
+ The resolution of this grid determines the approximate spacing between
+ trajectories. Streamlines are only allowed to pass through zeroed cells:
+ When a streamline enters a cell, that cell is set to 1, and no new
+ streamlines are allowed to enter.
+ """
+
+ def __init__(self, density):
+ if type(density) == float or type(density) == int:
@pelson
pelson Feb 7, 2012 Member

Stylistically, PEP8 suggests this should be isinstance(density, (float, int)).

@jdh2358
jdh2358 Feb 26, 2012 Collaborator

or more in line with mpl duck-typing:

if cbook.is_numlike(density)

this is also True for complex, but in practice this is rarely a problem.

@pelson pelson commented on the diff Feb 7, 2012
lib/matplotlib/streamplot.py
+#========================
+
+def get_integrator(u, v, dmap, minlength):
+
+ # rescale velocity onto grid-coordinates for integrations.
+ u, v = dmap.data2grid(u, v)
+
+ # speed (path length) will be in axes-coordinates
+ u_ax = u / dmap.grid.nx
+ v_ax = v / dmap.grid.ny
+ speed = np.sqrt(u_ax**2 + v_ax**2)
+
+ def forward_time(xi, yi):
+ ds_dt = interpgrid(speed, xi, yi)
+ if ds_dt == 0:
+ raise TerminateTrajectory
@pelson
pelson Feb 7, 2012 Member

PEP8 suggests this should be an instance of TerminateTrajectory, i.e. raise TerminateTrajectory().

@pelson
Member
pelson commented Feb 7, 2012

I should add, I am not a matplotlib developer & my comments hold no authority, but I was sucked in by the pull request and have used this streamlines code to good effect in the past (before this pull request). Hope the devs don't mind :-)

@tonysyu
Contributor
tonysyu commented Feb 7, 2012

@PhilipElson Thanks for your comments and for reading through the code. There tend to be more Pull Requests than there is spare time to read them; so I'm sure the devs are happy for your help. I definitely appreciate it.

@tonysyu tonysyu Rename `valid_index` to `within_grid`.
`valid_index` suggests that inputs are integer, which is not true.
398e8ec
@WeatherGod WeatherGod and 1 other commented on an outdated diff Feb 22, 2012
lib/matplotlib/streamplot.py
@@ -0,0 +1,546 @@
+"""
+Streamline plotting for 2D vector fields.
+
+"""
+import numpy as np
+import matplotlib
+import matplotlib.patches as mpp
+
+
+__all__ = ['streamplot']
+
+
+def streamplot(axes, x, y, u, v, density=1, linewidth=1, color='k', cmap=None,
@WeatherGod
WeatherGod Feb 22, 2012 Member

For the call signature, I would like to avoid using explicit defaults for things like linewidth and color. Instead, I would rather that the default be set to None and if they are None, then grab the rcparam values. Don't know if there are defaults for arrowsize and arrowstyle, but maybe one should added?

Even better would be if the color (when None) is pulled from the existing color cycling mechanism! But I would be happy with just pulling the color from rcparams.

@tonysyu
tonysyu Feb 22, 2012 Contributor

Good call. I changed the defaults to None and use the rc parameters if None. I decided to use the color cycling mechanism (is there a better way than calling axes._get_lines.color_cycle.next()?) .

There doesn't appear to be arrowsize or arrowstyle defaults.

@WeatherGod WeatherGod and 1 other commented on an outdated diff Feb 22, 2012
lib/matplotlib/streamplot.py
+ if x == (Nx - 2): xn = x
+ else: xn = x + 1
+ if y == (Ny - 2): yn = y
+ else: yn = y + 1
+
+ a00 = a[y, x]
+ a01 = a[y, xn]
+ a10 = a[yn, x]
+ a11 = a[yn, xn]
+ xt = xi - x
+ yt = yi - y
+ a0 = a00 * (1 - xt) + a01 * xt
+ a1 = a10 * (1 - xt) + a11 * xt
+ ai = a0 * (1 - yt) + a1 * yt
+
+ if not type(xi) == np.ndarray:
@WeatherGod
WeatherGod Feb 22, 2012 Member

There has to be a better way to test for ndarray-like objects. At the very least a "isinstance()" test?

@tonysyu
tonysyu Feb 22, 2012 Contributor

Done.

@gzahl gzahl and 1 other commented on an outdated diff Feb 23, 2012
lib/matplotlib/streamplot.py
+"""
+import numpy as np
+import matplotlib
+import matplotlib.patches as mpp
+
+
+__all__ = ['streamplot']
+
+
+def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
+ cmap=None, arrowsize=1, arrowstyle='-|>', minlength=0.1):
+ """Draws streamlines of a vector flow.
+
+ Parameters
+ ----------
+ x, y : 1d arrays
@gzahl
gzahl Feb 23, 2012 Contributor

Why is it only possible to supply x and y evenly spaced 1d arrays for the grid? Is this a limitation of the algorithm? It would be very nice, if arbitrary 2d arrays for non evenly spaced grids could be supported like for quiver plots. (nice work btw.)

@tomflannaghan
tomflannaghan Feb 23, 2012

Limiting use to evenly spaced grids is required because interpolation
for trajectory integration is implemented using float to integer
conversion (this is the only way I could get acceptable performance in
pure python).

Interpolating the data onto an even high resolution grid before
trajectory integration might be a workaround, although care would be
needed to handle low resolution or very uneven data well (I am uneasy
about this). Obviously the ideal solution would be to write a fast
interpolation routine in C and retain the original grid.

On 23 February 2012 17:05, Manuel Jung
reply@reply.github.com
wrote:

+"""
+import numpy as np
+import matplotlib
+import matplotlib.patches as mpp
+
+
+all = ['streamplot']
+
+
+def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,

  •               cmap=None, arrowsize=1, arrowstyle='-|>', minlength=0.1):
  •    """Draws streamlines of a vector flow.
    +
  •    Parameters
  •    ----------
  •    x, y : 1d arrays

Why is it only possible to supply x and y evenly spaced 1d arrays for the grid? Is this a limitation of the algorithm? It would be very nice, if arbitrary 2d arrays for non evenly spaced grids could be supported like for quiver plots. (nice work btw.)


Reply to this email directly or view it on GitHub:
https://github.com/matplotlib/matplotlib/pull/664/files#r481081

@jdh2358 jdh2358 commented on an outdated diff Feb 26, 2012
lib/matplotlib/streamplot.py
@@ -0,0 +1,545 @@
+"""
+Streamline plotting for 2D vector fields.
+
+"""
+import numpy as np
+import matplotlib
+import matplotlib.patches as mpp
@jdh2358
jdh2358 Feb 26, 2012 Collaborator

I would like to be consistent in how we do these internal module names. Generally one of two styles is preferred:

import matplotlib.patches as patches

or:

import matplotlib.patches as mpatches

Use the latter if you are worried about local name clashes.

@jdh2358 jdh2358 commented on an outdated diff Feb 26, 2012
lib/matplotlib/streamplot.py
+ Each cell in the grid can have, at most, one traversing streamline.
+ For different densities in each direction, use [density_x, density_y].
+ linewidth : numeric or 2d array
+ vary linewidth when given a 2d array with the same shape as velocities.
+ color : matplotlib color code, or 2d array
+ Streamline color. When given an array with the same shape as
+ velocities, values are converted to color using cmap, norm, vmin and
+ vmax args.
+ cmap : Colormap
+ Colormap used to plot streamlines and arrows. Only necessary when using
+ an array input for `color`.
+ arrowsize : float
+ Factor scale arrow size.
+ arrowstyle : str
+ Arrow style specification. See `matplotlib.patches.FancyArrowPatch`.
+ minlength : float
@jdh2358
jdh2358 Feb 26, 2012 Collaborator

minor nit -- we usually use minlength when formatting rest docstring parameters.

@jdh2358
Collaborator
jdh2358 commented Feb 26, 2012

This looks like a well reviewed PR and a nice piece of code. @WeatherGod , do you want to do the honors and merge this?

@tonysyu tonysyu Address JDH's PR comments
* Surround parameter names with asterisks (e.g. *param*)
* Import `patches` instead of `mpp`
* Use `cbook.is_numlike` instead of `isinstance`.
* Fix outdated docstring.
6ca72da
@WeatherGod
Member

Gladly. Great work Tony!

@WeatherGod WeatherGod merged commit 6ca72da into matplotlib:master Feb 27, 2012
@tonysyu
Contributor
tonysyu commented Feb 27, 2012

Awesome! Thanks to every one for your suggestions, and especially @tomflannaghan for the original implementation.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment