Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 4 commits into from
Jun 14, 2025
Merged

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

merged 4 commits into from
Jun 14, 2025

Conversation

scaomath
Copy link
Contributor

@scaomath scaomath commented Jun 11, 2025 •
edited
Loading

In

xarray/xarray/plot/dataset_plot.py

Lines 622 to 633 in f1a2c08

if xdim is not None and ydim is not None:
# Need to ensure the arrays are transposed correctly
dx = dx.transpose(ydim, xdim)
dy = dy.transpose(ydim, xdim)
du = du.transpose(ydim, xdim)
dv = dv.transpose(ydim, xdim)
hue = kwargs.pop("hue")
cmap_params = kwargs.pop("cmap_params")
if hue:
kwargs["color"] = ds[hue].values

x, y, u, v are all transposed, if the dataset only gives a 1D array for both x and y, which is a pretty nice addition to mpl's behavior. However, the color was left out. If one has a spatial grid of shape (nx, ny) and uses hue, and nx is different from ny, then matplotlib.streamplot will raises a value error caused by color.shape != grid.shape.

This PR simply adds the same condition to transpose hue if x, y, u, v are transposed.

I could not find a unittest corresponding to streamplot. But here is a minimal working example that does not work under current code but works after the simple change.

import numpy as np
import xarray

nx = 64
ny = 128
nt = 5

t = np.linspace(0, 1, nt, dtype=np.float64)
x = np.linspace(0, np.pi, nx, dtype=np.float64)
y = np.linspace(0, np.pi, ny, dtype=np.float64)

coords = {
    "t": t,
    "x": x,
    "y": y,
}

T, X, Y = np.meshgrid(t, x, y, indexing="ij")
u = np.sin(X) * np.cos(T*Y)
v = -np.cos(T*X) * np.sin(Y)

u_dataset = xarray.DataArray(
    u, dims=["t", "x", "y"], coords=coords
).to_dataset(name="u")
v_dataset = xarray.DataArray(
    v, dims=["t", "x", "y"], coords=coords
).to_dataset(name="v")

velocity_mag = np.sqrt(u**2 + v**2)
vmag_dataset = xarray.DataArray(
    velocity_mag, dims=["t", "x", "y"], coords=coords
).to_dataset(name="magnitude")

data = xarray.merge([u_dataset, v_dataset, vmag_dataset])

_ = data.plot.streamplot(
    x="x",
    y="y",
    u="u",
    v="v",
    hue="magnitude",
    col="t",
)

@welcome Welcome
Copy link

welcome bot commented Jun 11, 2025

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

@dcherian
Copy link
Contributor

Nice! Can you add a test here please:

xarray/xarray/tests/test_plot.py

Line 2696 in 60bc816

class TestDatasetStreamplotPlots(PlotTestCase):
?

@scaomath
Copy link
Contributor Author

Nice! Can you add a test here please:

xarray/xarray/tests/test_plot.py

Line 2696 in 60bc816

class TestDatasetStreamplotPlots(PlotTestCase):

?

@dcherian Sounds good. Do you want me to change only minimally the current self.ds to have a different DataArray by simply changing

xarray/xarray/tests/test_plot.py

Lines 2699 to 2705 in 60bc816

das = [
DataArray(
np.random.randn(3, 3, 2, 2),
dims=["x", "y", "row", "col"],
coords=[range(k) for k in [3, 3, 2, 2]],
)
for _ in [1, 2]

to be asymmetric in x and y? or adding a new self.ds_new and add accordingly its streamplot in test_facetgrid and test_streamline?

@dcherian
Copy link
Contributor

Yeah I guess if you change it so that x and y have different sizes, that should be good enough? Cna you make sure the modified test fails without your change please

@scaomath
Copy link
Contributor Author

Yeah I guess if you change it so that x and y have different sizes, that should be good enough? Cna you make sure the modified test fails without your change please

Done. I used the dev version==0.1.dev6002+g8172ee2. I asked Github Copilot to review this locally, and it suggested a much more complicated test which does not quite conform with the style in test_plot.py. So I kept this simple minimal change.

Here is what happens when no transpose hue:

Screenshot from 2025-06-12 10-20-40

Here is hue transposed when x, y, u, v are transposed.

Screenshot from 2025-06-12 10-21-35

@dcherian
Copy link
Contributor

Nice, thanks for taking the time to contribute! Welcome to Xarray.

@dcherian dcherian added the plan to merge Final call for comments label Jun 12, 2025
@dcherian dcherian enabled auto-merge (squash) June 13, 2025 20:33
@dcherian dcherian merged commit 5c39640 into pydata:main Jun 14, 2025
59 of 62 checks passed
@welcome Welcome
Copy link

welcome bot commented Jun 14, 2025

Congratulations on completing your first pull request! Welcome to Xarray! We are proud of you, and hope to see you again! celebration gif

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

@dcherian dcherian dcherian approved these changes

Assignees
No one assigned
Labels
plan to merge Final call for comments topic-documentation topic-plotting
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@scaomath @dcherian
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy