From c771eaef38ce3289d82af9be649ddc7e73bea9b3 Mon Sep 17 00:00:00 2001 From: dcherian Date: Wed, 13 Nov 2019 08:58:38 -0700 Subject: [PATCH] Squashed plot-transpose commit d430ae01014e2ae42df6180a556114ee4e0a4fdc Author: dcherian Date: Wed Nov 13 08:27:04 2019 -0700 proper fix. commit 7fd69be199a1fe832212769398635faf2243e555 Author: dcherian Date: Wed Nov 13 08:03:26 2019 -0700 fix whats-new merge. commit 448939412c572d04287f27d8aca929533c7c7713 Merge: 279ff1df b74f80ca Author: dcherian 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 (#3511) add missing pint integration tests (#3508) DOC: update bottleneck repo url (#3507) add drop_sel, drop_vars, map to api.rst (#3506) remove syntax warning (#3505) Dataset.map, GroupBy.map, Resample.map (#3459) tests for datasets with units (#3447) fix pandas-dev tests (#3491) unpin pseudonetcdf (#3496) whatsnew corrections (#3494) drop_vars; deprecate drop for variables (#3475) uamiv test using only raw uamiv variables (#3485) Optimize dask array equality checks. (#3453) Propagate indexes in DataArray binary operations. (#3481) python 3.8 tests (#3477) commit 279ff1dfc297d71431cfe9500fa5da6547f5e7ce Author: dcherian Date: Wed Nov 13 08:02:44 2019 -0700 Undo the transpose change and add test to make sure transposition is right. commit c9cc69832c657f23c5c9351f9ed1cda8609e978f Author: dcherian Date: Wed Nov 13 08:01:39 2019 -0700 Test to make sure transpose is right commit 9b35ecf3bb2878a5cedcd3943baa5bbbfa556aac Author: dcherian Date: Sat Nov 2 15:49:08 2019 -0600 Additional test. commit 7aed950777bf28eb7ad37a96b46f36c602eac5ec Author: dcherian Date: Sat Nov 2 15:20:07 2019 -0600 make plotting work with transposed nondim coords. --- doc/whats-new.rst | 2 ++ xarray/plot/plot.py | 25 ++++++++++++++++++------- xarray/tests/test_plot.py | 29 +++++++++++++++++++++++++++++ 3 files changed, 49 insertions(+), 7 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 9747fd46ad0..1cafb620c76 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -99,6 +99,8 @@ Bug fixes (:issue:`3402`). By `Deepak Cherian `_ - Allow appending datetime and bool data variables to zarr stores. (:issue:`3480`). By `Akihiro Matsukawa `_. +- Fix plotting with transposed 2D non-dimensional coordinates. (:issue:`3138`) + By `Deepak Cherian `_. Documentation ~~~~~~~~~~~~~ diff --git a/xarray/plot/plot.py b/xarray/plot/plot.py index 2acc8ddd452..33955038ba1 100644 --- a/xarray/plot/plot.py +++ b/xarray/plot/plot.py @@ -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 @@ -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) diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py index 730bf335dc8..e44ad2a5164 100644 --- a/xarray/tests/test_plot.py +++ b/xarray/tests/test_plot.py @@ -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() @@ -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