Merge pull request #22707 from aragilar/no_broken_streamlines

oscargus · web-flow · commit c33557d120ee · 2022-04-12T18:07:58.000+02:00
Proposed ENH: Allow user to turn off breaking of streamlines in streamplot (rebased)
diff --git a/doc/users/next_whats_new/no_broken_streamlines.rst b/doc/users/next_whats_new/no_broken_streamlines.rst
@@ -0,0 +1,25 @@
+Add option to plt.streamplot to not break streamlines
+-----------------------------------------------------
+
+It is now possible to specify that streamplots have continuous, unbroken
+streamlines. Previously streamlines would end to limit the number of lines
+within a single grid cell. See the difference between the plots below:
+
+.. plot::
+
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    w = 3
+    Y, X = np.mgrid[-w:w:100j, -w:w:100j]
+    U = -1 - X**2 + Y
+    V = 1 + X - Y**2
+    speed = np.sqrt(U**2 + V**2)
+
+    fig, (ax0, ax1) = plt.subplots(1, 2, sharex=True)
+
+    ax0.streamplot(X, Y, U, V, broken_streamlines=True)
+    ax0.set_title('broken_streamlines=True')
+
+    ax1.streamplot(X, Y, U, V, broken_streamlines=False)
+    ax1.set_title('broken_streamlines=False')
diff --git a/examples/images_contours_and_fields/plot_streamplot.py b/examples/images_contours_and_fields/plot_streamplot.py
@@ -11,6 +11,8 @@
 * Varying the line width along a streamline.
 * Controlling the starting points of streamlines.
 * Streamlines skipping masked regions and NaN values.
+* Unbroken streamlines even when exceeding the limit of lines within a single
+  grid cell.
 """
 import numpy as np
 import matplotlib.pyplot as plt
@@ -61,13 +63,17 @@
 U[:20, :20] = np.nan
 U = np.ma.array(U, mask=mask)
 
-ax4 = fig.add_subplot(gs[2:, :])
+ax4 = fig.add_subplot(gs[2, 0])
 ax4.streamplot(X, Y, U, V, color='r')
 ax4.set_title('Streamplot with Masking')
 
 ax4.imshow(~mask, extent=(-w, w, -w, w), alpha=0.5, cmap='gray', aspect='auto')
 ax4.set_aspect('equal')
 
+ax5 = fig.add_subplot(gs[2, 1])
+ax5.streamplot(X, Y, U, V, broken_streamlines=False)
+ax5.set_title('Streamplot with with unbroken streamlines')
+
 plt.tight_layout()
 plt.show()
 #############################################################################
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -2914,14 +2914,16 @@ def streamplot(
         x, y, u, v, density=1, linewidth=None, color=None, cmap=None,
         norm=None, arrowsize=1, arrowstyle='-|>', minlength=0.1,
         transform=None, zorder=None, start_points=None, maxlength=4.0,
-        integration_direction='both', *, data=None):
+        integration_direction='both', broken_streamlines=True, *,
+        data=None):
     __ret = gca().streamplot(
         x, y, u, v, density=density, linewidth=linewidth, color=color,
         cmap=cmap, norm=norm, arrowsize=arrowsize,
         arrowstyle=arrowstyle, minlength=minlength,
         transform=transform, zorder=zorder, start_points=start_points,
         maxlength=maxlength,
         integration_direction=integration_direction,
+        broken_streamlines=broken_streamlines,
         **({"data": data} if data is not None else {}))
     sci(__ret.lines)
     return __ret
diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py
@@ -18,7 +18,8 @@
 def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
                cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',
                minlength=0.1, transform=None, zorder=None, start_points=None,
-               maxlength=4.0, integration_direction='both'):
+               maxlength=4.0, integration_direction='both',
+               broken_streamlines=True):
     """
     Draw streamlines of a vector flow.
 
@@ -70,6 +71,10 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
         Integrate the streamline in forward, backward or both directions.
     data : indexable object, optional
         DATA_PARAMETER_PLACEHOLDER
+    broken_streamlines : boolean, default: True
+        If False, forces streamlines to continue until they
+        leave the plot domain.  If True, they may be terminated if they
+        come too close to another streamline.
 
     Returns
     -------
@@ -149,7 +154,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
         for xm, ym in _gen_starting_points(mask.shape):
             if mask[ym, xm] == 0:
                 xg, yg = dmap.mask2grid(xm, ym)
-                t = integrate(xg, yg)
+                t = integrate(xg, yg, broken_streamlines)
                 if t is not None:
                     trajectories.append(t)
     else:
@@ -177,7 +182,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
             xg = np.clip(xg, 0, grid.nx - 1)
             yg = np.clip(yg, 0, grid.ny - 1)
 
-            t = integrate(xg, yg)
+            t = integrate(xg, yg, broken_streamlines)
             if t is not None:
                 trajectories.append(t)
 
@@ -294,19 +299,19 @@ def data2grid(self, xd, yd):
     def grid2data(self, xg, yg):
         return xg / self.x_data2grid, yg / self.y_data2grid
 
-    def start_trajectory(self, xg, yg):
+    def start_trajectory(self, xg, yg, broken_streamlines=True):
         xm, ym = self.grid2mask(xg, yg)
-        self.mask._start_trajectory(xm, ym)
+        self.mask._start_trajectory(xm, ym, broken_streamlines)
 
     def reset_start_point(self, xg, yg):
         xm, ym = self.grid2mask(xg, yg)
         self.mask._current_xy = (xm, ym)
 
-    def update_trajectory(self, xg, yg):
+    def update_trajectory(self, xg, yg, broken_streamlines=True):
         if not self.grid.within_grid(xg, yg):
             raise InvalidIndexError
         xm, ym = self.grid2mask(xg, yg)
-        self.mask._update_trajectory(xm, ym)
+        self.mask._update_trajectory(xm, ym, broken_streamlines)
 
     def undo_trajectory(self):
         self.mask._undo_trajectory()
@@ -396,17 +401,17 @@ def __init__(self, density):
     def __getitem__(self, args):
         return self._mask[args]
 
-    def _start_trajectory(self, xm, ym):
+    def _start_trajectory(self, xm, ym, broken_streamlines=True):
         """Start recording streamline trajectory"""
         self._traj = []
-        self._update_trajectory(xm, ym)
+        self._update_trajectory(xm, ym, broken_streamlines)
 
     def _undo_trajectory(self):
         """Remove current trajectory from mask"""
         for t in self._traj:
             self._mask[t] = 0
 
-    def _update_trajectory(self, xm, ym):
+    def _update_trajectory(self, xm, ym, broken_streamlines=True):
         """
         Update current trajectory position in mask.
 
@@ -418,7 +423,10 @@ def _update_trajectory(self, xm, ym):
                 self._mask[ym, xm] = 1
                 self._current_xy = (xm, ym)
             else:
-                raise InvalidIndexError
+                if broken_streamlines:
+                    raise InvalidIndexError
+                else:
+                    pass
 
 
 class InvalidIndexError(Exception):
@@ -457,9 +465,9 @@ def backward_time(xi, yi):
         dxi, dyi = forward_time(xi, yi)
         return -dxi, -dyi
 
-    def integrate(x0, y0):
+    def integrate(x0, y0, broken_streamlines=True):
         """
-        Return (N, 2) grid-coordinates of trajectory based on starting point.
+        Return x, y grid-coordinates of trajectory based on starting point.
 
         Integrate both forward and backward in time from starting point in
         grid coordinates.
@@ -472,17 +480,19 @@ def integrate(x0, y0):
         stotal, xy_traj = 0., []
 
         try:
-            dmap.start_trajectory(x0, y0)
+            dmap.start_trajectory(x0, y0, broken_streamlines)
         except InvalidIndexError:
             return None
         if integration_direction in ['both', 'backward']:
-            s, xyt = _integrate_rk12(x0, y0, dmap, backward_time, maxlength)
+            s, xyt = _integrate_rk12(x0, y0, dmap, backward_time, maxlength,
+                                     broken_streamlines)
             stotal += s
             xy_traj += xyt[::-1]
 
         if integration_direction in ['both', 'forward']:
             dmap.reset_start_point(x0, y0)
-            s, xyt = _integrate_rk12(x0, y0, dmap, forward_time, maxlength)
+            s, xyt = _integrate_rk12(x0, y0, dmap, forward_time, maxlength,
+                                     broken_streamlines)
             stotal += s
             xy_traj += xyt[1:]
 
@@ -508,7 +518,7 @@ class OutOfBounds(IndexError):
     pass
 
 
-def _integrate_rk12(x0, y0, dmap, f, maxlength):
+def _integrate_rk12(x0, y0, dmap, f, maxlength, broken_streamlines=True):
     """
     2nd-order Runge-Kutta algorithm with adaptive step size.
 
@@ -588,7 +598,7 @@ def _integrate_rk12(x0, y0, dmap, f, maxlength):
             xi += dx2
             yi += dy2
             try:
-                dmap.update_trajectory(xi, yi)
+                dmap.update_trajectory(xi, yi, broken_streamlines)
             except InvalidIndexError:
                 break
             if stotal + ds > maxlength:
diff --git a/lib/matplotlib/tests/baseline_images/test_streamplot/streamplot_maxlength_no_broken.png b/lib/matplotlib/tests/baseline_images/test_streamplot/streamplot_maxlength_no_broken.png
diff --git a/lib/matplotlib/tests/test_streamplot.py b/lib/matplotlib/tests/test_streamplot.py
@@ -77,6 +77,18 @@ def test_maxlength():
     ax.set(xlim=(None, 3.2555988021882305), ylim=(None, 3.078326760195413))
 
 
+@image_comparison(['streamplot_maxlength_no_broken.png'],
+                  remove_text=True, style='mpl20', tol=0.302)
+def test_maxlength_no_broken():
+    x, y, U, V = swirl_velocity_field()
+    ax = plt.figure().subplots()
+    ax.streamplot(x, y, U, V, maxlength=10., start_points=[[0., 1.5]],
+                  linewidth=2, density=2, broken_streamlines=False)
+    assert ax.get_xlim()[-1] == ax.get_ylim()[-1] == 3
+    # Compatibility for old test image
+    ax.set(xlim=(None, 3.2555988021882305), ylim=(None, 3.078326760195413))
+
+
 @image_comparison(['streamplot_direction.png'],
                   remove_text=True, style='mpl20', tol=0.073)
 def test_direction():
