diff --git a/doc/api/plotting.rst b/doc/api/plotting.rst index e2f6068351a..1378f9d0578 100644 --- a/doc/api/plotting.rst +++ b/doc/api/plotting.rst @@ -37,6 +37,14 @@ DataArray DataArray.plot.scatter DataArray.plot.surface +DataTree +--------- + +.. autosummary:: + :toctree: ../generated/ + :template: autosummary/accessor_method.rst + + DataTree.plot.scatter Faceting -------- diff --git a/doc/user-guide/plotting.md b/doc/user-guide/plotting.md index 020412d03be..8b4630aaff2 100644 --- a/doc/user-guide/plotting.md +++ b/doc/user-guide/plotting.md @@ -723,6 +723,27 @@ ds.plot.streamplot(x="x", y="y", u="A", v="B", col="w", row="z"); (plot-maps)= +## DataTree + +Xarray has support for plotting DataTree variables from different nodes against each other. +In order to plot, variables must have the same name along every node in the DataTree. +Currently the only supported plotting methods are scatter plots. + +Consider this DataTree + +```{code-cell} +dt = xr.tutorial.open_datatree('precipitation.nc4') +dt +``` + +### Scatter + +Let’s plot the "precipitation" variable in the "observed" and "reanalysis" groups. + +```{code-cell} +dt.plot.scatter('precipitation', add_legend=True) +``` + ## Maps To follow this section you'll need to have Cartopy installed and working. diff --git a/doc/whats-new.rst b/doc/whats-new.rst index d1505bfa081..b0028fcebe6 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -25,6 +25,9 @@ New Features silently being written uncompressed (:issue:`10657`, :pull:`11067`). By `Mark Harfouche `_. +- ``DataTree`` objects now have plotting functionality with the ``plot.scatter`` method + (:issue:`9348`, :pull:`11516`). + By `Eni Awowale `_. Breaking Changes ~~~~~~~~~~~~~~~~ diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index 98934f29b92..b186cc4285b 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -66,6 +66,7 @@ from xarray.core.variable import Variable from xarray.namedarray.parallelcompat import get_chunked_array_type from xarray.namedarray.pycompat import is_chunked_array +from xarray.plot.accessor import DataTreePlotAccessor from xarray.structure.alignment import align from xarray.structure.merge import dataset_update_method @@ -610,6 +611,8 @@ def _resolve_inherit( "Expected True, False, 'indexes', or 'all'." ) + plot = utils.UncachedAccessor(DataTreePlotAccessor) + @property def _dims(self) -> ChainMap[Hashable, int]: return ChainMap(self._node_dims, *(p._node_dims for p in self.parents)) diff --git a/xarray/plot/accessor.py b/xarray/plot/accessor.py index 2b4c28a9027..96eaf972015 100644 --- a/xarray/plot/accessor.py +++ b/xarray/plot/accessor.py @@ -7,7 +7,7 @@ import numpy as np # Accessor methods have the same name as plotting methods, so we need a different namespace -from xarray.plot import dataarray_plot, dataset_plot +from xarray.plot import dataarray_plot, dataset_plot, datatree_plot if TYPE_CHECKING: from matplotlib.axes import Axes @@ -23,6 +23,7 @@ from xarray.core.dataarray import DataArray from xarray.core.dataset import Dataset + from xarray.core.datatree import DataTree from xarray.core.types import AspectOptions, HueStyleOptions, ScaleOptions from xarray.plot.facetgrid import FacetGrid @@ -1270,3 +1271,131 @@ def streamplot( @functools.wraps(dataset_plot.streamplot, assigned=("__doc__",)) def streamplot(self, *args, **kwargs) -> LineCollection | FacetGrid[Dataset]: return dataset_plot.streamplot(self._ds, *args, **kwargs) + + +class DataTreePlotAccessor: + """ + Enables use of xarray.plot functions as attributes on a DataTree. + For example, DataTree.plot.scatter + """ + + _dt: DataTree + __slots__ = ("_dt",) + + def __init__(self, datatree: DataTree) -> None: + self._dt = datatree + + def __call__(self, *args, **kwargs) -> NoReturn: + raise ValueError( + "DataTree.plot cannot be called directly. Use " + "an explicit plot method, e.g. dt.plot.scatter(...)" + ) + + @overload + def scatter( + self, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + markersize: Hashable | None = None, + row: None = None, # no wrap -> primitive + col: None = None, # no wrap -> primitive + col_wrap: int | Literal["auto"] | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: tuple[float, float] | None = None, + ylim: tuple[float, float] | None = None, + cmap=None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend=None, + levels=None, + fig_kw: dict[str, Any] | None = None, + **kwargs: Any, + ) -> FacetGrid[DataArray]: ... + + @overload + def scatter( # type: ignore[misc,unused-ignore] # None is hashable :(s: DataTree, + self, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + markersize: Hashable | None = None, + row: Hashable | None = None, + col: Hashable, # wrap -> FacetGrid + col_wrap: int | Literal["auto"] | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: tuple[float, float] | None = None, + ylim: tuple[float, float] | None = None, + cmap=None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend=None, + levels=None, + fig_kw: dict[str, Any] | None = None, + **kwargs: Any, + ) -> PathCollection: ... + + @overload + def scatter( + self, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + markersize: Hashable | None = None, + row: Hashable, # wrap -> FacetGrid + col: Hashable | None = None, + col_wrap: int | Literal["auto"] | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: tuple[float, float] | None = None, + ylim: tuple[float, float] | None = None, + cmap=None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend=None, + levels=None, + fig_kw: dict[str, Any] | None = None, + **kwargs: Any, + ) -> FacetGrid[DataArray] | PathCollection: ... + + @functools.wraps(datatree_plot.scatter, assigned=("__doc__",)) + def scatter(self, *args, **kwargs) -> PathCollection | FacetGrid[DataArray]: + return datatree_plot.scatter(self._dt, *args, **kwargs) diff --git a/xarray/plot/datatree_plot.py b/xarray/plot/datatree_plot.py new file mode 100644 index 00000000000..5834cd1af70 --- /dev/null +++ b/xarray/plot/datatree_plot.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +import functools +from collections.abc import Callable, Hashable, Iterable +from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload + +from xarray.plot import dataarray_plot + +if TYPE_CHECKING: + from matplotlib.axes import Axes + from matplotlib.collections import PathCollection + from matplotlib.colors import Colormap, Normalize + from numpy.typing import ArrayLike + + from xarray.core.dataarray import DataArray + from xarray.core.datatree import DataTree + from xarray.core.types import ( + ExtendOptions, + ScaleOptions, + ) + from xarray.plot.facetgrid import FacetGrid + + +F = TypeVar("F", bound=Callable) + +# TODO Eni +# Added back _dtplot accessor for documentation specific to datatree + + +def _update_doc_to_datatree(datatree_plotfunc: Callable) -> Callable[[F], F]: + """ + Add a common docstring by reusing the DataArray one. + + Parameters + ---------- + datatree_plotfunc : Callable + Function that returns a finished plot primitive. + """ + + # Build on the original docstring + dt_doc = datatree_plotfunc.__doc__ + if dt_doc is None: + raise NotImplementedError("DataArray plot method requires a docstring") + + da_str = """ + Parameters + ---------- + darray : DataArray + """ + dt_str = """ + + The `y` DataArray will be used as base, any other variables are added as coordt. + + Parameters + ---------- + : DataTree + """ + # TODO: improve this? + # if dt_str in dt_doc: + # dt_doc = dt_doc.replace(dt_str, dt_str).replace("darray", "dt") + # else: + # dt_doc = dt_doc + + @functools.wraps(datatree_plotfunc) + def wrapper(datatree_plotfunc: F) -> F: + datatree_plotfunc.__doc__ = dt_doc + return datatree_plotfunc + + return wrapper # type: ignore[return-value] + + +@overload +def scatter( + dt: DataTree, + variable: str, + ax: Axes | None = None, + *, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + markersize: Hashable | None = None, + row: Hashable | None = None, + col: Hashable | None = None, + col_wrap: int | Literal["auto"] | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend: ExtendOptions = None, + levels: ArrayLike | None = None, + fig_kw: dict[str, Any] | None = None, + **kwargs: Any, +) -> FacetGrid[DataArray]: ... + + +@overload +def scatter( # type: ignore[misc,unused-ignore] # None is hashable :(s: DataTree, + dt: DataTree, + variable: str, + ax: Axes | None = None, + *, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + markersize: Hashable | None = None, + row: Hashable | None = None, + col: Hashable | None = None, + col_wrap: int | Literal["auto"] | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend: ExtendOptions = None, + levels: ArrayLike | None = None, + fig_kw: dict[str, Any] | None = None, + **kwargs: Any, +) -> PathCollection: ... + + +@_update_doc_to_datatree(dataarray_plot.scatter) +def scatter( # type: ignore[return] + dt: DataTree, + variable: str, + ax: Axes | None = None, + *, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + markersize: Hashable | None = None, + row: Hashable | None = None, + col: Hashable | None = None, + col_wrap: int | Literal["auto"] | None = None, + xincrease: bool | None = True, + yincrease: bool | None = True, + add_legend: bool | None = None, + add_colorbar: bool | None = None, + add_labels: bool | Iterable[bool] = True, + add_title: bool = True, + subplot_kws: dict[str, Any] | None = None, + xscale: ScaleOptions = None, + yscale: ScaleOptions = None, + xticks: ArrayLike | None = None, + yticks: ArrayLike | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + extend: ExtendOptions = None, + levels: ArrayLike | None = None, + fig_kw: dict[str, Any] | None = None, + **kwargs: Any, +) -> PathCollection | FacetGrid[DataArray]: + """Scatter DataTree data variables with the same node against each other.""" + + if fig_kw is None: + fig_kw = {} + + locals_ = locals() + del locals_["dt"] + locals_.update(locals_.pop("kwargs", {})) + locals_.pop("variable") + (locals_.pop("ax"),) + locals_.pop("fig_kw") + locals_.pop("add_legend") + import matplotlib.pyplot as plt + + fig, ax = plt.subplots(**fig_kw) + for node in dt.descendants: + try: + da = node[variable] + except KeyError as err: + raise KeyError(f"{variable} not found at node: {node.name}") from err + + da.plot.scatter( + ax=ax, + label=node.name, + **locals_, + ) + if add_legend: + ax.legend() diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py index d1531c16d8a..f7742f286e6 100644 --- a/xarray/tests/test_plot.py +++ b/xarray/tests/test_plot.py @@ -14,7 +14,7 @@ import xarray as xr import xarray.plot as xplt -from xarray import DataArray, Dataset +from xarray import DataArray, Dataset, DataTree from xarray.namedarray.utils import module_available from xarray.plot.dataarray_plot import _infer_interval_breaks from xarray.plot.dataset_plot import _infer_meta_data @@ -3649,3 +3649,43 @@ def test_facetgrid_figsize_rcparams() -> None: g = xplt.FacetGrid(da, col="z", figsize=explicit_size) actual_figsize = g.fig.get_size_inches() np.testing.assert_allclose(actual_figsize, explicit_size) + + +@requires_matplotlib +class TestDataTreeScatterPlots(PlotTestCase): + @pytest.fixture(autouse=True) + def setUp(self) -> None: + das = [ + DataArray( + np.random.randn(3, 3, 4, 4), + dims=["x", "row", "col", "hue"], + coords=[range(k) for k in [3, 3, 4, 4]], + ) + for _ in [1, 2, 3] + ] + dt = DataTree.from_dict( + { + "/": Dataset(das[0].coords), + "/group_1": Dataset({"A": das[0], "B": das[1]}), + "/group_2": Dataset({"A": das[2]}), + } + ) + dt.hue.name = "huename" + dt.hue.attrs["units"] = "hunits" + dt.x.attrs["units"] = "xunits" + dt.col.attrs["units"] = "colunits" + dt.row.attrs["units"] = "rowunits" + dt["group_1/A"].attrs["units"] = "Aunits" + dt["group_2/A"].attrs["units"] = "Aunits" + self.dt = dt + + def test_accessor(self) -> None: + from xarray.plot.accessor import DataTreePlotAccessor + + assert DataTree.plot is DataTreePlotAccessor + assert isinstance(self.dt.plot, DataTreePlotAccessor) + + def test_key_error(self): + """Assert that a KeyError is raised when accessing a variable that doesn't exist in a node.""" + with pytest.raises(KeyError): + self.dt.plot.scatter("B")