Skip to content

Commit

Permalink
Use broadcast_like for 2d plot coordinates
Browse files Browse the repository at this point in the history
Use broadcast_like if either `x` or `y` inputs are 2d to ensure that
both have dimensions in the same order as the DataArray being plotted.
Convert to numpy arrays after possibly using broadcast_like. Simplifies
code, and fixes pydata#5097 (bug when dimensions have the same size).
  • Loading branch information
johnomotani committed Mar 31, 2021
1 parent 57a4479 commit 3ad10d9
Showing 1 changed file with 14 additions and 21 deletions.
35 changes: 14 additions & 21 deletions xarray/plot/plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -671,28 +671,21 @@ def newplotfunc(
darray=darray, x=x, y=y, imshow=imshow_rgb, rgb=rgb
)

# better to pass the ndarrays directly to plotting functions
xval = darray[xlab].values
yval = darray[ylab].values

# check if we need to broadcast one dimension
if xval.ndim < yval.ndim:
dims = darray[ylab].dims
if xval.shape[0] == yval.shape[0]:
xval = np.broadcast_to(xval[:, np.newaxis], yval.shape)
else:
xval = np.broadcast_to(xval[np.newaxis, :], yval.shape)

elif yval.ndim < xval.ndim:
dims = darray[xlab].dims
if yval.shape[0] == xval.shape[0]:
yval = np.broadcast_to(yval[:, np.newaxis], xval.shape)
else:
yval = np.broadcast_to(yval[np.newaxis, :], xval.shape)
elif xval.ndim == 2:
dims = darray[xlab].dims
xval = darray[xlab]
yval = darray[ylab]

if xval.ndim > 1 or yval.ndim > 1:
# Passing 2d coordinate values, need to ensure they are transposed the same
# way as darray
xval = xval.broadcast_like(darray)
yval = yval.broadcast_like(darray)
dims = darray.dims
else:
dims = (darray[ylab].dims[0], darray[xlab].dims[0])
dims = (yval.dims[0], xval.dims[0])

# better to pass the ndarrays directly to plotting functions
xval = xval.values
yval = yval.values

# May need to transpose for correct x, y labels
# xlab may be the name of a coord, we have to check for dim names
Expand Down

0 comments on commit 3ad10d9

Please sign in to comment.