From 4892c26c308f58ebf48a4e1f3969f8c2f213e42e Mon Sep 17 00:00:00 2001 From: caenrigen <31376402+caenrigen@users.noreply.github.com> Date: Fri, 8 Oct 2021 20:55:00 +0200 Subject: [PATCH 1/6] Display coords' units for slice plots --- xarray/core/dataarray.py | 7 ++++++- xarray/plot/utils.py | 30 ++++++++++++++++-------------- 2 files changed, 22 insertions(+), 15 deletions(-) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index ed8b393628d..b8958551f3b 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -22,6 +22,7 @@ import pandas as pd from ..plot.plot import _PlotMethods +from ..plot.utils import _get_units_from_attrs from . import ( computation, dtypes, @@ -3108,7 +3109,11 @@ def _title_for_slice(self, truncate: int = 50) -> str: for dim, coord in self.coords.items(): if coord.size == 1: one_dims.append( - "{dim} = {v}".format(dim=dim, v=format_item(coord.values)) + "{dim} = {v}{unit}".format( + dim=dim, + v=format_item(coord.values), + unit=_get_units_from_attrs(coord), + ) ) title = ", ".join(one_dims) diff --git a/xarray/plot/utils.py b/xarray/plot/utils.py index 594f2e5360e..9e7e78f4c44 100644 --- a/xarray/plot/utils.py +++ b/xarray/plot/utils.py @@ -468,6 +468,21 @@ def _maybe_gca(**kwargs): return plt.axes(**kwargs) +def _get_units_from_attrs(da): + """Extracts and formats the unit/units from a attributes.""" + pint_array_type = DuckArrayModule("pint").type + units = " [{}]" + if isinstance(da.data, pint_array_type): + units = units.format(str(da.data.units)) + elif da.attrs.get("units"): + units = units.format(da.attrs["units"]) + elif da.attrs.get("unit"): + units = units.format(da.attrs["unit"]) + else: + units = "" + return units + + def label_from_attrs(da, extra=""): """Makes informative labels if variable metadata (attrs) follows CF conventions.""" @@ -481,20 +496,7 @@ def label_from_attrs(da, extra=""): else: name = "" - def _get_units_from_attrs(da): - if da.attrs.get("units"): - units = " [{}]".format(da.attrs["units"]) - elif da.attrs.get("unit"): - units = " [{}]".format(da.attrs["unit"]) - else: - units = "" - return units - - pint_array_type = DuckArrayModule("pint").type - if isinstance(da.data, pint_array_type): - units = " [{}]".format(str(da.data.units)) - else: - units = _get_units_from_attrs(da) + units = _get_units_from_attrs(da) # Treat `name` differently if it's a latex sequence if name.startswith("$") and (name.count("$") % 2 == 0): From f9926fed82ec5bb0a8461e7aea28b4f24b758837 Mon Sep 17 00:00:00 2001 From: caenrigen <31376402+caenrigen@users.noreply.github.com> Date: Fri, 8 Oct 2021 21:20:00 +0200 Subject: [PATCH 2/6] Add what-new entry --- doc/whats-new.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 9c0b1bd04df..f2e19ed6638 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -36,6 +36,8 @@ New Features `Nathan Lis `_. - Histogram plots are set with a title displaying the scalar coords if any, similarly to the other plots (:issue:`5791`, :pull:`5792`). By `Maxime Liquet `_. +- Slice plots display the coords units in the same way as x/y/colorbar labels (:pull:`5847`). + By `Victor Negîrneac `_. Breaking changes ~~~~~~~~~~~~~~~~ From cd6f5ab5b7f9dfe06d1959dc4f23bcb82a8acaee Mon Sep 17 00:00:00 2001 From: caenrigen <31376402+caenrigen@users.noreply.github.com> Date: Fri, 15 Oct 2021 14:50:44 +0200 Subject: [PATCH 3/6] Fix unrelated typo --- xarray/core/utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 77d973f613f..ebf6d7e28ed 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -1,5 +1,4 @@ -"""Internal utilties; not for external use -""" +"""Internal utilities; not for external use""" import contextlib import functools import io From 4cb9448a49314e9fe1d3f628d620ec5071669f28 Mon Sep 17 00:00:00 2001 From: caenrigen <31376402+caenrigen@users.noreply.github.com> Date: Fri, 15 Oct 2021 14:51:46 +0200 Subject: [PATCH 4/6] Add test for slice --- xarray/tests/test_units.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py index 7bde6ce8b9f..3532e01e1f8 100644 --- a/xarray/tests/test_units.py +++ b/xarray/tests/test_units.py @@ -5614,6 +5614,21 @@ def test_units_in_line_plot_labels(self): assert ax.get_ylabel() == "pressure [pascal]" assert ax.get_xlabel() == "x [meters]" + def test_units_in_slice_line_plot_labels(self): + arr = xr.DataArray( + name="var_a", + data=np.array([[1, 2], [3, 4]]), + # TODO make coord a Quantity once unit-aware indexes supported + coords=dict( + a=("a", np.array([5, 6]), {"unit": "s"}), + b=("b", np.array([7, 8]), {"unit": "s"}), + ), + dims=("a", "b"), + ) + arr.sel(a=5).plot(marker="o") + ax = plt.gca() + assert ax.get_title() == "a = 5 [s]" + def test_units_in_2d_plot_labels(self): arr = np.ones((2, 3)) * unit_registry.Pa da = xr.DataArray(data=arr, dims=["x", "y"], name="pressure") From 4ddc1661323b679190980de458f0cb162c8d5c06 Mon Sep 17 00:00:00 2001 From: caenrigen <31376402+caenrigen@users.noreply.github.com> Date: Thu, 21 Oct 2021 20:25:06 +0100 Subject: [PATCH 5/6] Fix test --- xarray/tests/test_units.py | 60 ++++++++++++++++++++++++++++++-------- 1 file changed, 48 insertions(+), 12 deletions(-) diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py index 3532e01e1f8..74146ceb5ca 100644 --- a/xarray/tests/test_units.py +++ b/xarray/tests/test_units.py @@ -5600,34 +5600,70 @@ def test_duck_array_ops(self): @requires_matplotlib class TestPlots(PlotTestCase): - def test_units_in_line_plot_labels(self): + + xfail_coord_units = ( + "coord_unit, coord_attrs", + [ + (1, {"units": "meter"}), + pytest.param( + unit_registry.m, + {}, + marks=pytest.mark.xfail(reason="indexes don't support units"), + ), + ], + ) + + @pytest.mark.parametrize(*xfail_coord_units) + def test_units_in_line_plot_labels(self, coord_unit, coord_attrs): arr = np.linspace(1, 10, 3) * unit_registry.Pa - # TODO make coord a Quantity once unit-aware indexes supported - x_coord = xr.DataArray( - np.linspace(1, 3, 3), dims="x", attrs={"units": "meters"} - ) + coord_arr = np.linspace(1, 3, 3) * coord_unit + x_coord = xr.DataArray(coord_arr, dims="x", attrs=coord_attrs) da = xr.DataArray(data=arr, dims="x", coords={"x": x_coord}, name="pressure") da.plot.line() ax = plt.gca() assert ax.get_ylabel() == "pressure [pascal]" - assert ax.get_xlabel() == "x [meters]" + assert ax.get_xlabel() == "x [meter]" - def test_units_in_slice_line_plot_labels(self): + @pytest.mark.parametrize(*xfail_coord_units) + def test_units_in_slice_line_plot_labels_sel(self, coord_unit, coord_attrs): arr = xr.DataArray( name="var_a", data=np.array([[1, 2], [3, 4]]), - # TODO make coord a Quantity once unit-aware indexes supported coords=dict( - a=("a", np.array([5, 6]), {"unit": "s"}), - b=("b", np.array([7, 8]), {"unit": "s"}), + a=("a", np.array([5, 6]) * coord_unit, coord_attrs), + b=("b", np.array([7, 8]) * coord_unit, coord_attrs), ), dims=("a", "b"), ) arr.sel(a=5).plot(marker="o") - ax = plt.gca() - assert ax.get_title() == "a = 5 [s]" + + assert plt.gca().get_title() == "a = 5 [meter]" + + @pytest.mark.parametrize( + "coord_unit, coord_attrs", + [ + (1, {"units": "meter"}), + pytest.param( + unit_registry.m, + {}, + marks=pytest.mark.xfail(reason="pint.errors.UnitStrippedWarning"), + ), + ], + ) + def test_units_in_slice_line_plot_labels_isel(self, coord_unit, coord_attrs): + arr = xr.DataArray( + name="var_a", + data=np.array([[1, 2], [3, 4]]), + coords=dict( + a=("x", np.array([5, 6]) * coord_unit, coord_attrs), + b=("y", np.array([7, 8])), + ), + dims=("x", "y"), + ) + arr.isel(x=0).plot(marker="o") + assert plt.gca().get_title() == "a = 5 [meter]" def test_units_in_2d_plot_labels(self): arr = np.ones((2, 3)) * unit_registry.Pa From f043c643281cefcb4070dfc59469bb99b0522dfc Mon Sep 17 00:00:00 2001 From: caenrigen <31376402+caenrigen@users.noreply.github.com> Date: Fri, 22 Oct 2021 13:26:21 +0100 Subject: [PATCH 6/6] Keep tests parameterization explicit --- xarray/tests/test_units.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py index 74146ceb5ca..bce0e8e5ee6 100644 --- a/xarray/tests/test_units.py +++ b/xarray/tests/test_units.py @@ -5600,8 +5600,7 @@ def test_duck_array_ops(self): @requires_matplotlib class TestPlots(PlotTestCase): - - xfail_coord_units = ( + @pytest.mark.parametrize( "coord_unit, coord_attrs", [ (1, {"units": "meter"}), @@ -5612,8 +5611,6 @@ class TestPlots(PlotTestCase): ), ], ) - - @pytest.mark.parametrize(*xfail_coord_units) def test_units_in_line_plot_labels(self, coord_unit, coord_attrs): arr = np.linspace(1, 10, 3) * unit_registry.Pa coord_arr = np.linspace(1, 3, 3) * coord_unit @@ -5626,7 +5623,17 @@ def test_units_in_line_plot_labels(self, coord_unit, coord_attrs): assert ax.get_ylabel() == "pressure [pascal]" assert ax.get_xlabel() == "x [meter]" - @pytest.mark.parametrize(*xfail_coord_units) + @pytest.mark.parametrize( + "coord_unit, coord_attrs", + [ + (1, {"units": "meter"}), + pytest.param( + unit_registry.m, + {}, + marks=pytest.mark.xfail(reason="indexes don't support units"), + ), + ], + ) def test_units_in_slice_line_plot_labels_sel(self, coord_unit, coord_attrs): arr = xr.DataArray( name="var_a",