Repository: pydata/xarray

Issue:
Wrong hue assignment in scatter plot
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports: http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples: https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**What happened**:
When using the hue keyword in a scatter plot to color the points based on a string variable, the color assignment in the plot is wrong (whereas the legend is correct).

**What you expected to happen**:
In the example, data of category `"A"` ranges between 0 and 2 in u-direction and 0 and 0.5 in v-direction. Points in that square should be orange (the color for "A") but currently are blue.

**Minimal Complete Verifiable Example**:

```python
import xarray as xr
import numpy as np

u = np.random.rand(50, 2) * np.array([1, 2])
v = np.random.rand(50, 2) * np.array([1, 0.5])

ds = xr.Dataset(
    {
        "u": (("x", "category"), u),
        "v": (("x", "category"), v),
    },
    coords={"category": ["B", "A"],}
)

g = ds.plot.scatter(
    y="u",
    x="v",
    hue="category",
);
```

**Anything else we need to know?**:
I think that this might be related to sorting at some point. If the variable by which I color is sorted alphabetically (`["A", "B"]` instead of `["B", "A"]`), the color assignment is correct.

Not sure if this issue is related to https://github.com/pydata/xarray/issues/4126, bit it looks different to me (the problem is not the legend, but the colors in the plot itself).

**Environment**:

<details><summary>Output of <tt>xr.show_versions()</tt></summary>

<!-- Paste the output here xr.show_versions() here -->
INSTALLED VERSIONS
------------------
commit: None
python: 3.7.8 | packaged by conda-forge | (default, Jul 31 2020, 02:25:08) 
[GCC 7.5.0]
python-bits: 64
OS: Linux
OS-release: 4.15.0-122-generic
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: en_US.UTF-8
libhdf5: 1.10.6
libnetcdf: 4.7.4

xarray: 0.16.0
pandas: 1.1.2
numpy: 1.17.5
scipy: 1.5.2
netCDF4: 1.5.4
pydap: None
h5netcdf: None
h5py: 2.10.0
Nio: None
zarr: 2.4.0
cftime: 1.2.1
nc_time_axis: None
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: None
dask: 2.26.0
distributed: 2.26.0
matplotlib: 3.3.2
cartopy: None
seaborn: 0.11.0
numbagg: None
pint: None
setuptools: 49.6.0.post20200814
pip: 20.2.3
conda: 4.8.3
pytest: 6.0.1
IPython: 7.18.1
sphinx: None


</details>


Discussion:
keewis: the upstream-dev failure can be ignored, I guess (the next run of the scheduled CI should open a issue to track this).

Should we merge this PR?
keewis: thanks for the reviews, @dcherian, @shoyer.

Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 398c3324..db10ec65 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -64,6 +64,8 @@ Bug fixes
   By `Richard Kleijn <https://github.com/rhkleijn>`_ .
 - Remove dictionary unpacking when using ``.loc`` to avoid collision with ``.sel`` parameters (:pull:`4695`).
   By `Anderson Banihirwe <https://github.com/andersy005>`_
+- Fix the legend created by :py:meth:`Dataset.plot.scatter` (:issue:`4641`, :pull:`4723`).
+  By `Justus Magin <https://github.com/keewis>`_.
 - Fix a crash in orthogonal indexing on geographic coordinates with ``engine='cfgrib'`` (:issue:`4733` :pull:`4737`).
   By `Alessandro Amici <https://github.com/alexamici>`_
 - Coordinates with dtype ``str`` or ``bytes`` now retain their dtype on many operations,
diff --git a/xarray/plot/dataset_plot.py b/xarray/plot/dataset_plot.py
index 7ba0f93f..6d942e1b 100644
--- a/xarray/plot/dataset_plot.py
+++ b/xarray/plot/dataset_plot.py
@@ -291,7 +291,7 @@ def _dsplot(plotfunc):
             allargs = locals().copy()
             allargs["plotfunc"] = globals()[plotfunc.__name__]
             allargs["data"] = ds
-            # TODO dcherian: why do I need to remove kwargs?
+            # remove kwargs to avoid passing the information twice
             for arg in ["meta_data", "kwargs", "ds"]:
                 del allargs[arg]
 
@@ -422,7 +422,10 @@ def scatter(ds, x, y, ax, **kwargs):
 
     if hue_style == "discrete":
         primitive = []
-        for label in np.unique(data["hue"].values):
+        # use pd.unique instead of np.unique because that keeps the order of the labels,
+        # which is important to keep them in sync with the ones used in
+        # FacetGrid.add_legend
+        for label in pd.unique(data["hue"].values.ravel()):
             mask = data["hue"] == label
             if data["sizes"] is not None:
                 kwargs.update(s=data["sizes"].where(mask, drop=True).values.flatten())


Test patch:
diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py
index 471bbb70..47b15446 100644
--- a/xarray/tests/test_plot.py
+++ b/xarray/tests/test_plot.py
@@ -2290,6 +2290,17 @@ class TestDatasetScatterPlots(PlotTestCase):
         lines = ds2.plot.scatter(x="A", y="B", hue="hue")
         assert [t.get_text() for t in lines[0].axes.get_legend().texts] == ["a", "b"]
 
+    def test_legend_labels_facetgrid(self):
+        ds2 = self.ds.copy()
+        ds2["hue"] = ["d", "a", "c", "b"]
+        g = ds2.plot.scatter(x="A", y="B", hue="hue", col="col")
+        legend_labels = tuple(t.get_text() for t in g.figlegend.texts)
+        attached_labels = [
+            tuple(m.get_label() for m in mappables_per_ax)
+            for mappables_per_ax in g._mappables
+        ]
+        assert list(set(attached_labels)) == [legend_labels]
+
     def test_add_legend_by_default(self):
         sc = self.ds.plot.scatter(x="A", y="B", hue="hue")
         assert len(sc.figure.axes) == 2