Repository: pydata/xarray

Issue:
Plotting with 2d coordinates, error in transposed coordinate
This is similar to issue #1503, but specifically deals with requiring the correct order of dimensions for plotting 2d coordinates. An example of the issue is:

```
# create an example dataset
dates = pd.date_range('2000-01-01', '2001-12-31', name='time')
times = dates - dates[0]

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, :]

data = (np.sin(x) * np.cos(z)) * np.cos(np.asarray(times.days[:, np.newaxis, np.newaxis]))

# create an xarray dataset

ds = xr.Dataset({'data': (('time', 's', 'x'), data)}, 
                {'time':dates, 'x':x, 's':s, 'z':(('s', 'x'), z)})
ds.coords['zt'] = ds.z.transpose()
```
plotting with `z` and `zt` should give the same results, since the dimensions are clearly labeled and other than the order of the dimensions, they are exactly the same. 

This works:
```
ds.data[0].plot(x='x', y='z')
```

This does not:
```
ds.data[0].plot(x='x', y='zt')
```

Broadcasting, e.g., does not give a similar error. Specifically, this is true:
```
np.allclose((ds.data * ds.zt).values, (ds.data * ds.z).values)
```

Discussion:
max-sixty: I don't know this code well, but LGTM
dcherian: I'll merge tomorrow evening if nobody has any comments

Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 884c3cef..8930947a 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -28,10 +28,14 @@ New Features
 - :py:meth:`Dataset.quantile`, :py:meth:`DataArray.quantile` and ``GroupBy.quantile``
   now work with dask Variables.
   By `Deepak Cherian <https://github.com/dcherian>`_.
-
+- Added the :py:meth:`count` reduction method to both :py:class:`DatasetCoarsen`
+  and :py:class:`DataArrayCoarsen` objects. (:pull:`3500`)
+  By `Deepak Cherian <https://github.com/dcherian/>`_
 
 Bug fixes
 ~~~~~~~~~
+- Fix plotting with transposed 2D non-dimensional coordinates. (:issue:`3138`, :pull:`3441`)
+  By `Deepak Cherian <https://github.com/dcherian>`_.
 
 
 Documentation
diff --git a/xarray/core/nanops.py b/xarray/core/nanops.py
index 17240faf..f70e9621 100644
--- a/xarray/core/nanops.py
+++ b/xarray/core/nanops.py
@@ -25,7 +25,7 @@ def _maybe_null_out(result, axis, mask, min_count=1):
     """
     if hasattr(axis, "__len__"):  # if tuple or list
         raise ValueError(
-            "min_count is not available for reduction " "with more than one dimensions."
+            "min_count is not available for reduction with more than one dimensions."
         )
 
     if axis is not None and getattr(result, "ndim", False):
diff --git a/xarray/core/ops.py b/xarray/core/ops.py
index 78c4466f..b789f93b 100644
--- a/xarray/core/ops.py
+++ b/xarray/core/ops.py
@@ -347,13 +347,3 @@ def inject_all_ops_and_reduce_methods(cls, priority=50, array_only=True):
 
     inject_reduce_methods(cls)
     inject_cum_methods(cls)
-
-
-def inject_coarsen_methods(cls):
-    # standard numpy reduce methods
-    methods = [(name, getattr(duck_array_ops, name)) for name in NAN_REDUCE_METHODS]
-    for name, f in methods:
-        func = cls._reduce_method(f)
-        func.__name__ = name
-        func.__doc__ = _COARSEN_REDUCE_DOCSTRING_TEMPLATE.format(name=func.__name__)
-        setattr(cls, name, func)
diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py
index a1864332..ea6d72b2 100644
--- a/xarray/core/rolling.py
+++ b/xarray/core/rolling.py
@@ -1,12 +1,12 @@
 import functools
 import warnings
-from typing import Callable
+from typing import Any, Callable, Dict
 
 import numpy as np
 
 from . import dtypes, duck_array_ops, utils
 from .dask_array_ops import dask_rolling_wrapper
-from .ops import inject_coarsen_methods
+from .ops import inject_reduce_methods
 from .pycompat import dask_array_type
 
 try:
@@ -542,6 +542,11 @@ class Coarsen:
         self.side = side
         self.boundary = boundary
 
+        absent_dims = [dim for dim in windows.keys() if dim not in self.obj.dims]
+        if absent_dims:
+            raise ValueError(
+                f"Dimensions {absent_dims!r} not found in {self.obj.__class__.__name__}."
+            )
         if not utils.is_dict_like(coord_func):
             coord_func = {d: coord_func for d in self.obj.dims}
         for c in self.obj.coords:
@@ -565,18 +570,23 @@ class Coarsen:
 class DataArrayCoarsen(Coarsen):
     __slots__ = ()
 
+    _reduce_extra_args_docstring = """"""
+
     @classmethod
-    def _reduce_method(cls, func):
+    def _reduce_method(cls, func: Callable, include_skipna: bool, numeric_only: bool):
         """
-        Return a wrapped function for injecting numpy methods.
-        see ops.inject_coarsen_methods
+        Return a wrapped function for injecting reduction methods.
+        see ops.inject_reduce_methods
         """
+        kwargs: Dict[str, Any] = {}
+        if include_skipna:
+            kwargs["skipna"] = None
 
         def wrapped_func(self, **kwargs):
             from .dataarray import DataArray
 
             reduced = self.obj.variable.coarsen(
-                self.windows, func, self.boundary, self.side
+                self.windows, func, self.boundary, self.side, **kwargs
             )
             coords = {}
             for c, v in self.obj.coords.items():
@@ -585,7 +595,11 @@ class DataArrayCoarsen(Coarsen):
                 else:
                     if any(d in self.windows for d in v.dims):
                         coords[c] = v.variable.coarsen(
-                            self.windows, self.coord_func[c], self.boundary, self.side
+                            self.windows,
+                            self.coord_func[c],
+                            self.boundary,
+                            self.side,
+                            **kwargs,
                         )
                     else:
                         coords[c] = v
@@ -597,12 +611,17 @@ class DataArrayCoarsen(Coarsen):
 class DatasetCoarsen(Coarsen):
     __slots__ = ()
 
+    _reduce_extra_args_docstring = """"""
+
     @classmethod
-    def _reduce_method(cls, func):
+    def _reduce_method(cls, func: Callable, include_skipna: bool, numeric_only: bool):
         """
-        Return a wrapped function for injecting numpy methods.
-        see ops.inject_coarsen_methods
+        Return a wrapped function for injecting reduction methods.
+        see ops.inject_reduce_methods
         """
+        kwargs: Dict[str, Any] = {}
+        if include_skipna:
+            kwargs["skipna"] = None
 
         def wrapped_func(self, **kwargs):
             from .dataset import Dataset
@@ -610,14 +629,18 @@ class DatasetCoarsen(Coarsen):
             reduced = {}
             for key, da in self.obj.data_vars.items():
                 reduced[key] = da.variable.coarsen(
-                    self.windows, func, self.boundary, self.side
+                    self.windows, func, self.boundary, self.side, **kwargs
                 )
 
             coords = {}
             for c, v in self.obj.coords.items():
                 if any(d in self.windows for d in v.dims):
                     coords[c] = v.variable.coarsen(
-                        self.windows, self.coord_func[c], self.boundary, self.side
+                        self.windows,
+                        self.coord_func[c],
+                        self.boundary,
+                        self.side,
+                        **kwargs,
                     )
                 else:
                     coords[c] = v.variable
@@ -626,5 +649,5 @@ class DatasetCoarsen(Coarsen):
         return wrapped_func
 
 
-inject_coarsen_methods(DataArrayCoarsen)
-inject_coarsen_methods(DatasetCoarsen)
+inject_reduce_methods(DataArrayCoarsen)
+inject_reduce_methods(DatasetCoarsen)
diff --git a/xarray/core/variable.py b/xarray/core/variable.py
index 041c303d..773dcef0 100644
--- a/xarray/core/variable.py
+++ b/xarray/core/variable.py
@@ -1864,9 +1864,9 @@ class Variable(
             ),
         )
 
-    def coarsen(self, windows, func, boundary="exact", side="left"):
+    def coarsen(self, windows, func, boundary="exact", side="left", **kwargs):
         """
-        Apply
+        Apply reduction function.
         """
         windows = {k: v for k, v in windows.items() if k in self.dims}
         if not windows:
@@ -1878,11 +1878,11 @@ class Variable(
             func = getattr(duck_array_ops, name, None)
             if func is None:
                 raise NameError(f"{name} is not a valid method.")
-        return type(self)(self.dims, func(reshaped, axis=axes), self._attrs)
+        return self._replace(data=func(reshaped, axis=axes, **kwargs))
 
     def _coarsen_reshape(self, windows, boundary, side):
         """
-        Construct a reshaped-array for corsen
+        Construct a reshaped-array for coarsen
         """
         if not utils.is_dict_like(boundary):
             boundary = {d: boundary for d in windows.keys()}
diff --git a/xarray/plot/plot.py b/xarray/plot/plot.py
index 16a49436..d38c9765 100644
--- a/xarray/plot/plot.py
+++ b/xarray/plot/plot.py
@@ -672,10 +672,22 @@ def _plot2d(plotfunc):
 
         # 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
@@ -685,10 +697,9 @@ def _plot2d(plotfunc):
             # 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)


Test patch:
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index d8282f58..7db19116 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -5497,6 +5497,11 @@ def ds(request):
         )
 
 
+def test_coarsen_absent_dims_error(ds):
+    with raises_regex(ValueError, "not found in Dataset."):
+        ds.coarsen(foo=2)
+
+
 @pytest.mark.parametrize("dask", [True, False])
 @pytest.mark.parametrize(("boundary", "side"), [("trim", "left"), ("pad", "right")])
 def test_coarsen(ds, dask, boundary, side):
@@ -5505,12 +5510,11 @@ def test_coarsen(ds, dask, boundary, side):
 
     actual = ds.coarsen(time=2, x=3, boundary=boundary, side=side).max()
     assert_equal(
-        actual["z1"], ds["z1"].coarsen(time=2, x=3, boundary=boundary, side=side).max()
+        actual["z1"], ds["z1"].coarsen(x=3, boundary=boundary, side=side).max()
     )
     # coordinate should be mean by default
     assert_equal(
-        actual["time"],
-        ds["time"].coarsen(time=2, x=3, boundary=boundary, side=side).mean(),
+        actual["time"], ds["time"].coarsen(time=2, boundary=boundary, side=side).mean()
     )
 
 
@@ -5521,8 +5525,8 @@ def test_coarsen_coords(ds, dask):
 
     # check if coord_func works
     actual = ds.coarsen(time=2, x=3, boundary="trim", coord_func={"time": "max"}).max()
-    assert_equal(actual["z1"], ds["z1"].coarsen(time=2, x=3, boundary="trim").max())
-    assert_equal(actual["time"], ds["time"].coarsen(time=2, x=3, boundary="trim").max())
+    assert_equal(actual["z1"], ds["z1"].coarsen(x=3, boundary="trim").max())
+    assert_equal(actual["time"], ds["time"].coarsen(time=2, boundary="trim").max())
 
     # raise if exact
     with pytest.raises(ValueError):
diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py
index 6e283ea0..a10f0d9a 100644
--- a/xarray/tests/test_plot.py
+++ b/xarray/tests/test_plot.py
@@ -265,6 +265,7 @@ class TestPlot(PlotTestCase):
         )
 
         a.plot.contourf(x="time", y="depth")
+        a.plot.contourf(x="depth", y="time")
 
     def test3d(self):
         self.darray.plot()
@@ -2149,3 +2150,31 @@ class TestAxesKwargs:
         da.plot(yticks=np.arange(5))
         expected = np.arange(5)
         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())
diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py
index 245dc1ac..5b5aa1a5 100644
--- a/xarray/tests/test_variable.py
+++ b/xarray/tests/test_variable.py
@@ -1833,6 +1833,26 @@ class TestVariable(VariableSubclassobjects):
         expected[1, 1] *= 12 / 11
         assert_allclose(actual, expected)
 
+        v = self.cls(("x", "y"), np.arange(4 * 4, dtype=np.float32).reshape(4, 4))
+        actual = v.coarsen(dict(x=2, y=2), func="count", boundary="exact")
+        expected = self.cls(("x", "y"), 4 * np.ones((2, 2)))
+        assert_equal(actual, expected)
+
+        v[0, 0] = np.nan
+        v[-1, -1] = np.nan
+        expected[0, 0] = 3
+        expected[-1, -1] = 3
+        actual = v.coarsen(dict(x=2, y=2), func="count", boundary="exact")
+        assert_equal(actual, expected)
+
+        actual = v.coarsen(dict(x=2, y=2), func="sum", boundary="exact", skipna=False)
+        expected = self.cls(("x", "y"), [[np.nan, 18], [42, np.nan]])
+        assert_equal(actual, expected)
+
+        actual = v.coarsen(dict(x=2, y=2), func="sum", boundary="exact", skipna=True)
+        expected = self.cls(("x", "y"), [[10, 18], [42, 35]])
+        assert_equal(actual, expected)
+
 
 @requires_dask
 class TestVariableWithDask(VariableSubclassobjects):