Merge pull request matplotlib#18947 from aitikgupta/streamplot-grid

MNT: Strictly increasing check with test coverage for streamplot grid

Author: tacaswell

lib/matplotlib/streamplot.py

@@ -25,8 +25,8 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
 
     Parameters
     ----------
-    x, y : 1D arrays
-        An evenly spaced grid.
+    x, y : 1D/2D arrays
+        Evenly spaced strictly increasing arrays to make a grid.
     u, v : 2D arrays
         *x* and *y*-velocities. The number of rows and columns must match
         the length of *y* and *x*, respectively.
@@ -333,6 +333,11 @@ def __init__(self, x, y):
         else:
             raise ValueError("'y' can have at maximum 2 dimensions")
 
+        if not (np.diff(x) > 0).all():
+            raise ValueError("'x' must be strictly increasing")
+        if not (np.diff(y) > 0).all():
+            raise ValueError("'y' must be strictly increasing")
+
         self.nx = len(x)
         self.ny = len(y)
 

lib/matplotlib/tests/test_streamplot.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 from numpy.testing import assert_array_almost_equal
+import pytest
 import matplotlib.pyplot as plt
 from matplotlib.testing.decorators import image_comparison
 import matplotlib.transforms as mtransforms
@@ -114,3 +115,49 @@ def test_streamplot_limits():
     # datalim.
     assert_array_almost_equal(ax.dataLim.bounds, (20, 30, 15, 6),
                               decimal=1)
+
+
+def test_streamplot_grid():
+    u = np.ones((2, 2))
+    v = np.zeros((2, 2))
+
+    # Test for same rows and columns
+    x = np.array([[10, 20], [10, 30]])
+    y = np.array([[10, 10], [20, 20]])
+
+    with pytest.raises(ValueError, match="The rows of 'x' must be equal"):
+        plt.streamplot(x, y, u, v)
+
+    x = np.array([[10, 20], [10, 20]])
+    y = np.array([[10, 10], [20, 30]])
+
+    with pytest.raises(ValueError, match="The columns of 'y' must be equal"):
+        plt.streamplot(x, y, u, v)
+
+    x = np.array([[10, 20], [10, 20]])
+    y = np.array([[10, 10], [20, 20]])
+    plt.streamplot(x, y, u, v)
+
+    # Test for maximum dimensions
+    x = np.array([0, 10])
+    y = np.array([[[0, 10]]])
+
+    with pytest.raises(ValueError, match="'y' can have at maximum "
+                                         "2 dimensions"):
+        plt.streamplot(x, y, u, v)
+
+    # Test for equal spacing
+    u = np.ones((3, 3))
+    v = np.zeros((3, 3))
+    x = np.array([0, 10, 20])
+    y = np.array([0, 10, 30])
+
+    with pytest.raises(ValueError, match="'y' values must be equally spaced"):
+        plt.streamplot(x, y, u, v)
+
+    # Test for strictly increasing
+    x = np.array([0, 20, 40])
+    y = np.array([0, 20, 10])
+
+    with pytest.raises(ValueError, match="'y' must be strictly increasing"):
+        plt.streamplot(x, y, u, v)