Fixed color.shape != grid.shape in streamplot when nx != ny (#10411)

scaomath · web-flow · commit 5c396407ef01 · 2025-06-14T10:20:51.000Z
diff --git a/xarray/plot/dataset_plot.py b/xarray/plot/dataset_plot.py
@@ -630,6 +630,8 @@ def streamplot(
     cmap_params = kwargs.pop("cmap_params")
 
     if hue:
+        if xdim is not None and ydim is not None:
+            ds[hue] = ds[hue].transpose(ydim, xdim)
         kwargs["color"] = ds[hue].values
 
         # TODO: Fix this by always returning a norm with vmin, vmax in cmap_params
diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py
@@ -2698,9 +2698,9 @@ class TestDatasetStreamplotPlots(PlotTestCase):
     def setUp(self) -> None:
         das = [
             DataArray(
-                np.random.randn(3, 3, 2, 2),
+                np.random.randn(3, 4, 2, 2),
                 dims=["x", "y", "row", "col"],
-                coords=[range(k) for k in [3, 3, 2, 2]],
+                coords=[range(k) for k in [3, 4, 2, 2]],
             )
             for _ in [1, 2]
         ]

Original file line number	Diff line number	Diff line change
`@@ -2698,9 +2698,9 @@ class TestDatasetStreamplotPlots(PlotTestCase):`
`2698`	`2698`	`def setUp(self) -> None:`
`2699`	`2699`	`das = [`
`2700`	`2700`	`DataArray(`
`2701`		`- np.random.randn(3, 3, 2, 2),`
	`2701`	`+ np.random.randn(3, 4, 2, 2),`
`2702`	`2702`	`dims=["x", "y", "row", "col"],`
`2703`		`- coords=[range(k) for k in [3, 3, 2, 2]],`
	`2703`	`+ coords=[range(k) for k in [3, 4, 2, 2]],`
`2704`	`2704`	`)`
`2705`	`2705`	`for _ in [1, 2]`
`2706`	`2706`	`]`