Skip to content
forked from pydata/xarray

Commit

Permalink
Squashed plot-transpose
Browse files Browse the repository at this point in the history
commit d430ae0
Author: dcherian <deepak@cherian.net>
Date:   Wed Nov 13 08:27:04 2019 -0700

    proper fix.

commit 7fd69be
Author: dcherian <deepak@cherian.net>
Date:   Wed Nov 13 08:03:26 2019 -0700

    fix whats-new merge.

commit 4489394
Merge: 279ff1d b74f80c
Author: dcherian <deepak@cherian.net>
Date:   Wed Nov 13 08:03:06 2019 -0700

    Merge remote-tracking branch 'upstream/master' into fix/plot-broadcast

    * upstream/master:
      format indexing.rst code with black (pydata#3511)
      add missing pint integration tests (pydata#3508)
      DOC: update bottleneck repo url (pydata#3507)
      add drop_sel, drop_vars, map to api.rst (pydata#3506)
      remove syntax warning (pydata#3505)
      Dataset.map, GroupBy.map, Resample.map (pydata#3459)
      tests for datasets with units (pydata#3447)
      fix pandas-dev tests (pydata#3491)
      unpin pseudonetcdf (pydata#3496)
      whatsnew corrections (pydata#3494)
      drop_vars; deprecate drop for variables (pydata#3475)
      uamiv test using only raw uamiv variables (pydata#3485)
      Optimize dask array equality checks. (pydata#3453)
      Propagate indexes in DataArray binary operations. (pydata#3481)
      python 3.8 tests (pydata#3477)

commit 279ff1d
Author: dcherian <deepak@cherian.net>
Date:   Wed Nov 13 08:02:44 2019 -0700

    Undo the transpose change and add test to make sure transposition is right.

commit c9cc698
Author: dcherian <deepak@cherian.net>
Date:   Wed Nov 13 08:01:39 2019 -0700

    Test to make sure transpose is right

commit 9b35ecf
Author: dcherian <deepak@cherian.net>
Date:   Sat Nov 2 15:49:08 2019 -0600

    Additional test.

commit 7aed950
Author: dcherian <deepak@cherian.net>
Date:   Sat Nov 2 15:20:07 2019 -0600

    make plotting work with transposed nondim coords.
  • Loading branch information
dcherian committed Nov 13, 2019
1 parent ecc928c commit c771eae
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 7 deletions.
2 changes: 2 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,8 @@ Bug fixes
(:issue:`3402`). By `Deepak Cherian <https://github.com/dcherian/>`_
- Allow appending datetime and bool data variables to zarr stores.
(:issue:`3480`). By `Akihiro Matsukawa <https://github.com/amatsukawa/>`_.
- Fix plotting with transposed 2D non-dimensional coordinates. (:issue:`3138`)
By `Deepak Cherian <https://github.com/dcherian>`_.

Documentation
~~~~~~~~~~~~~
Expand Down
25 changes: 18 additions & 7 deletions xarray/plot/plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -698,10 +698,22 @@ def newplotfunc(

# check if we need to broadcast one dimension
if xval.ndim < yval.ndim:
xval = np.broadcast_to(xval, yval.shape)
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)

if yval.ndim < xval.ndim:
yval = np.broadcast_to(yval, xval.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
else:
dims = (darray[ylab].dims[0], darray[xlab].dims[0])

# 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 All @@ -711,10 +723,9 @@ def newplotfunc(
# we transpose to (y, x, color) to make this work.
yx_dims = (ylab, xlab)
dims = yx_dims + tuple(d for d in darray.dims if d not in yx_dims)
if dims != darray.dims:
darray = darray.transpose(*dims, transpose_coords=True)
elif darray[xlab].dims[-1] == darray.dims[0]:
darray = darray.transpose(transpose_coords=True)

if dims != darray.dims:
darray = darray.transpose(*dims, transpose_coords=True)

# Pass the data as a masked ndarray too
zval = darray.to_masked_array(copy=False)
Expand Down
29 changes: 29 additions & 0 deletions xarray/tests/test_plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,7 @@ def test2d_1d_2d_coordinates_contourf(self):
)

a.plot.contourf(x="time", y="depth")
a.plot.contourf(x="depth", y="time")

def test3d(self):
self.darray.plot()
Expand Down Expand Up @@ -2156,6 +2157,34 @@ def test_yticks_kwarg(self, da):
assert np.all(plt.gca().get_yticks() == expected)


@requires_matplotlib
@pytest.mark.parametrize("plotfunc", ["pcolormesh", "contourf", "contour"])
def test_plot_transposed_nondim_coord(plotfunc):
x = np.linspace(0, 10, 101)
h = np.linspace(3, 7, 101)
s = np.linspace(0, 1, 51)
z = s[:, np.newaxis] * h[np.newaxis, :]
da = xr.DataArray(
np.sin(x) * np.cos(z),
dims=["s", "x"],
coords={"x": x, "s": s, "z": (("s", "x"), z), "zt": (("x", "s"), z.T)},
)
getattr(da.plot, plotfunc)(x="x", y="zt")
getattr(da.plot, plotfunc)(x="zt", y="x")


@requires_matplotlib
@pytest.mark.parametrize("plotfunc", ["pcolormesh", "imshow"])
def test_plot_transposes_properly(plotfunc):
# test that we aren't mistakenly transposing when the 2 dimensions have equal sizes.
da = xr.DataArray([np.sin(2 * np.pi / 10 * np.arange(10))] * 10, dims=("y", "x"))
hdl = getattr(da.plot, plotfunc)(x="x", y="y")
# get_array doesn't work for contour, contourf. It returns the colormap intervals.
# pcolormesh returns 1D array but imshow returns a 2D array so it is necessary
# to ravel() on the LHS
assert np.all(hdl.get_array().ravel() == da.to_masked_array().ravel())


@requires_matplotlib
class TestDataArrayGroupByPlot:
@requires_cftime
Expand Down

0 comments on commit c771eae

Please sign in to comment.