From 366ad663b0da17b5becb1017ad4f59f0922f44ca Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Mon, 20 Jul 2026 18:06:30 -0400 Subject: [PATCH 01/16] precommit --- xarray/core/datatree.py | 3 + xarray/plot/accessor.py | 146 +++++++++- xarray/plot/datatree_plot.py | 544 +++++++++++++++++++++++++++++++++++ 3 files changed, 692 insertions(+), 1 deletion(-) create mode 100644 xarray/plot/datatree_plot.py 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..e4b36f69ce5 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,146 @@ 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 Dataset. + 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( # type: ignore[misc,unused-ignore] # None is hashable :( + self, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | 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, + **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, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | 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, + **kwargs: Any, + ) -> FacetGrid[DataTree]: ... + + @overload + def scatter( + self, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | 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, + **kwargs: Any, + ) -> FacetGrid[DataTree]: ... + + @functools.wraps(datatree_plot.scatter, assigned=("__doc__",)) + def scatter(self, *args, **kwargs) -> PathCollection | FacetGrid[DataTree]: + 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..6dfce2807e5 --- /dev/null +++ b/xarray/plot/datatree_plot.py @@ -0,0 +1,544 @@ +from __future__ import annotations + +import functools +import warnings +from collections.abc import Callable, Hashable, Iterable +from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload + +from xarray.plot import dataarray_plot +from xarray.plot.facetgrid import _easy_facetgrid +from xarray.plot.utils import ( + _add_colorbar, + _get_nice_quiver_magnitude, + _infer_meta_data, + _process_cmap_cbar_kwargs, + get_axis, +) + +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 ( + AspectOptions, + ExtendOptions, + HueStyleOptions, + ScaleOptions, + ) + from xarray.plot.facetgrid import FacetGrid + + +def _dtplot(plotfunc): + commondoc = """ +Parameters +---------- +dt : DataTree +x : Hashable or None, optional + Variable name for x-axis. +y : Hashable or None, optional + Variable name for y-axis. +u : Hashable or None, optional + Variable name for the *u* velocity (in *x* direction). + quiver/streamplot plots only. +v : Hashable or None, optional + Variable name for the *v* velocity (in *y* direction). + quiver/streamplot plots only. +hue: Hashable or None, optional + Variable by which to color scatter points or arrows. +hue_style: {'continuous', 'discrete'} or None, optional + How to use the ``hue`` variable: + + - ``'continuous'`` -- continuous color scale + (default for numeric ``hue`` variables) + - ``'discrete'`` -- a color for each unique value, using the default color cycle + (default for non-numeric ``hue`` variables) + +row : Hashable or None, optional + If passed, make row faceted plots on this dimension name. +col : Hashable or None, optional + If passed, make column faceted plots on this dimension name. +col_wrap : int, None or "auto", optional + "Wrap" the grid for the column variable after this number of columns, + adding rows if ``col_wrap`` is less than the number of facets. + If "auto" align the grid to the figsize or keep it as square as possible. +ax : matplotlib axes object or None, optional + If ``None``, use the current axes. Not applicable when using facets. +figsize : Iterable[float] or None, optional + A tuple (width, height) of the figure in inches. + Mutually exclusive with ``size`` and ``ax``. +size : scalar, optional + If provided, create a new figure for the plot with the given size. + Height (in inches) of each plot. See also: ``aspect``. +aspect : "auto", "equal", scalar or None, optional + Aspect ratio of plot, so that ``aspect * size`` gives the width in + inches. Only used if a ``size`` is provided. +sharex : bool or None, optional + If True all subplots share the same x-axis. +sharey : bool or None, optional + If True all subplots share the same y-axis. +add_guide: bool or None, optional + Add a guide that dependt on ``hue_style``: + + - ``'continuous'`` -- build a colorbar + - ``'discrete'`` -- build a legend + +subplot_kws : dict or None, optional + Dictionary of keyword arguments for Matplotlib subplots + (see :py:meth:`matplotlib:matplotlib.figure.Figure.add_subplot`). + Only applies to FacetGrid plotting. +cbar_kwargs : dict, optional + Dictionary of keyword arguments to pass to the colorbar + (see :meth:`matplotlib:matplotlib.figure.Figure.colorbar`). +cbar_ax : matplotlib axes object, optional + Axes in which to draw the colorbar. +cmap : matplotlib colormap name or colormap, optional + The mapping from data values to color space. Either a + Matplotlib colormap name or object. If not provided, this will + be either ``'viridis'`` (if the function infers a sequential + dataset) or ``'RdBu_r'`` (if the function infers a diverging + dataset). + See :doc:`Choosing Colormaps in Matplotlib ` + for more information. + + If *seaborn* is installed, ``cmap`` may also be a + `seaborn color palette `_. + Note: if ``cmap`` is a seaborn color palette, + ``levels`` must also be specified. +vmin : float or None, optional + Lower value to anchor the colormap, otherwise it is inferred from the + data and other keyword arguments. When a diverging dataset is inferred, + setting `vmin` or `vmax` will fix the other by symmetry around + ``center``. Setting both values prevents use of a diverging colormap. + If discrete levels are provided as an explicit list, both of these + values are ignored. +vmax : float or None, optional + Upper value to anchor the colormap, otherwise it is inferred from the + data and other keyword arguments. When a diverging dataset is inferred, + setting `vmin` or `vmax` will fix the other by symmetry around + ``center``. Setting both values prevents use of a diverging colormap. + If discrete levels are provided as an explicit list, both of these + values are ignored. +norm : matplotlib.colors.Normalize, optional + If ``norm`` has ``vmin`` or ``vmax`` specified, the corresponding + kwarg must be ``None``. +infer_intervals: bool | None + If True the intervals are inferred. +center : float, optional + The value at which to center the colormap. Passing this value implies + use of a diverging colormap. Setting it to ``False`` prevents use of a + diverging colormap. +robust : bool, optional + If ``True`` and ``vmin`` or ``vmax`` are absent, the colormap range is + computed with 2nd and 98th percentiles instead of the extreme values. +colors : str or array-like of color-like, optional + A single color or a list of colors. The ``levels`` argument + is required. +extend : {'neither', 'both', 'min', 'max'}, optional + How to draw arrows extending the colorbar beyond its limits. If not + provided, ``extend`` is inferred from ``vmin``, ``vmax`` and the data limits. +levels : int or array-like, optional + Split the colormap (``cmap``) into discrete color intervals. If an integer + is provided, "nice" levels are chosen based on the data range: this can + imply that the final number of levels is not exactly the expected one. + Setting ``vmin`` and/or ``vmax`` with ``levels=N`` is equivalent to + setting ``levels=np.linspace(vmin, vmax, N)``. +**kwargs : optional + Additional keyword arguments to wrapped Matplotlib function. + """ + + # Build on the original docstring + plotfunc.__doc__ = f"{plotfunc.__doc__}\n{commondoc}" + + @functools.wraps( + plotfunc, assigned=("__module__", "__name__", "__qualname__", "__doc__") + ) + def newplotfunc( + dt: DataTree, + *args: Any, + x: Hashable | None = None, + y: Hashable | None = None, + u: Hashable | None = None, + v: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + row: Hashable | None = None, + col: Hashable | None = None, + col_wrap: int | Literal["auto"] | None = None, + ax: Axes | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: AspectOptions = None, + sharex: bool = True, + sharey: bool = True, + add_guide: bool | None = None, + subplot_kws: dict[str, Any] | None = None, + cbar_kwargs: dict[str, Any] | None = None, + cbar_ax: Axes | None = None, + cmap: str | Colormap | None = None, + vmin: float | None = None, + vmax: float | None = None, + norm: Normalize | None = None, + infer_intervals: bool | None = None, + center: float | None = None, + robust: bool | None = None, + colors: str | ArrayLike | None = None, + extend: ExtendOptions = None, + levels: ArrayLike | None = None, + **kwargs: Any, + ) -> Any: + if args: + # TODO: Deprecated since 2022.10: + msg = "Using positional arguments is deprecated for plot methods, use keyword arguments instead." + assert x is None + x = args[0] + if len(args) > 1: + assert y is None + y = args[1] + if len(args) > 2: + assert u is None + u = args[2] + if len(args) > 3: + assert v is None + v = args[3] + if len(args) > 4: + assert hue is None + hue = args[4] + if len(args) > 5: + raise ValueError(msg) + else: + warnings.warn(msg, FutureWarning, stacklevel=2) + del msg + del args + + _is_facetgrid = kwargs.pop("_is_facetgrid", False) + if _is_facetgrid: # facetgrid call + meta_data = kwargs.pop("meta_data") + else: + meta_data = _infer_meta_data( + dt, x, y, hue, hue_style, add_guide, funcname=plotfunc.__name__ + ) + + hue_style = meta_data["hue_style"] + + # handle facetgridt first + if col or row: + allargs = locals().copy() + allargs["plotfunc"] = globals()[plotfunc.__name__] + allargs["data"] = dt + # remove kwargs to avoid passing the information twice + for arg in ["meta_data", "kwargs", "dt"]: + del allargs[arg] + + return _easy_facetgrid(kind="datatree", **allargs, **kwargs) + + figsize = kwargs.pop("figsize", None) + ax = get_axis(figsize, size, aspect, ax) + + if hue_style == "continuous" and hue is not None: + if _is_facetgrid: + cbar_kwargs = meta_data["cbar_kwargs"] + cmap_params = meta_data["cmap_params"] + else: + cmap_params, cbar_kwargs = _process_cmap_cbar_kwargs( + plotfunc, dt[hue].values, **locals() + ) + + # subset that can be passed to scatter, hist2d + cmap_params_subset = { + vv: cmap_params[vv] for vv in ["vmin", "vmax", "norm", "cmap"] + } + + else: + cmap_params_subset = {} + + if (u is not None or v is not None) and plotfunc.__name__ not in ( + "quiver", + "streamplot", + ): + raise ValueError("u, v are only allowed for quiver or streamplot plots.") + + primitive = plotfunc( + dt=dt, + x=x, + y=y, + ax=ax, + u=u, + v=v, + hue=hue, + hue_style=hue_style, + cmap_params=cmap_params_subset, + **kwargs, + ) + + if _is_facetgrid: # if this was called from Facetgrid.map_datatree, + return primitive # finish here. Else, make labels + + if meta_data.get("xlabel", None): + ax.set_xlabel(meta_data.get("xlabel")) + if meta_data.get("ylabel", None): + ax.set_ylabel(meta_data.get("ylabel")) + + if meta_data["add_legend"]: + ax.legend(handles=primitive, title=meta_data.get("hue_label", None)) + if meta_data["add_colorbar"]: + cbar_kwargs = {} if cbar_kwargs is None else cbar_kwargs + if "label" not in cbar_kwargs: + cbar_kwargs["label"] = meta_data.get("hue_label", None) + _add_colorbar(primitive, ax, cbar_ax, cbar_kwargs, cmap_params) + + if meta_data["add_quiverkey"]: + magnitude = _get_nice_quiver_magnitude(dt[u], dt[v]) + units = dt[u].attrs.get("units", "") + ax.quiverkey( + primitive, + X=0.85, + Y=0.9, + U=magnitude, + label=f"{magnitude}\n{units}", + labelpos="E", + coordinates="figure", + ) + + if plotfunc.__name__ in ("quiver", "streamplot"): + title = dt[u]._title_for_slice() + else: + title = dt[x]._title_for_slice() + ax.set_title(title) + + return primitive + + # we want to actually expose the signature of newplotfunc + # and not the copied **kwargs from the plotfunc which + # functools.wraps addt, so delete the wrapped attr + del newplotfunc.__wrapped__ + + return newplotfunc + + +F = TypeVar("F", bound=Callable) + + +def _update_doc_to_datatree(dataarray_plotfunc: Callable) -> Callable[[F], F]: + """ + Add a common docstring by reusing the DataArray one. + + TODO: Reduce code duplication. + + * The goal is to reduce code duplication by mov all DataTree + specific plots to the DataArray side and use this thin wrapper to + handle the converts between DataTree and DataArray. + * Improve docstring handling, maybe reword the DataArray versions to explain DataTrees better. + + Parameters + ---------- + dataarray_plotfunc : Callable + Function that returns a finished plot primitive. + """ + + # Build on the original docstring + da_doc = dataarray_plotfunc.__doc__ + if da_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 da_str in da_doc: + dt_doc = da_doc.replace(da_str, dt_str).replace("darray", "dt") + else: + dt_doc = da_doc + + @functools.wraps(dataarray_plotfunc) + def wrapper(datatree_plotfunc: F) -> F: + datatree_plotfunc.__doc__ = dt_doc + return datatree_plotfunc + + return wrapper # type: ignore[return-value] + + +@overload +def scatter( # type: ignore[misc,unused-ignore] # None is hashable :(s: DataTree, + *, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | 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: 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, + **kwargs: Any, +) -> PathCollection: ... + + +@overload +def scatter( + s: DataTree, + *, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | 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: 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, + **kwargs: Any, +) -> FacetGrid[DataArray]: ... + + +@overload +def scatter( + s: DataTree, + *, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | 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: 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, + **kwargs: Any, +) -> FacetGrid[DataArray]: ... + + +@_update_doc_to_datatree(dataarray_plot.scatter) +def scatter( + dt: DataTree, + variable, + *, + x: Hashable | None = None, + y: Hashable | None = None, + z: Hashable | None = None, + hue: Hashable | None = None, + hue_style: HueStyleOptions = None, + markersize: Hashable | None = None, + linewidth: Hashable | None = None, + figsize: Iterable[float] | None = None, + size: float | None = None, + aspect: float | None = None, + ax: Axes | 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, + **kwargs: Any, +) -> PathCollection | FacetGrid[DataArray]: + """Scat plot DataTree data variables against each other.""" + locals_ = locals() + del locals_["dt"] + locals_.update(locals_.pop("kwargs", {})) + locals_.pop("variable") + for node in dt.descendants: + da = node[variable] + print(da) + return da.plot.scatter(**locals_) From 6cab3f5c8c208ebebbe2dc495763e84a87eb063e Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Tue, 21 Jul 2026 16:22:05 -0400 Subject: [PATCH 02/16] test: raise keyerror --- xarray/plot/datatree_plot.py | 17 ++++++++---- xarray/tests/test_plot.py | 54 +++++++++++++++++++++++++++++++++++- 2 files changed, 65 insertions(+), 6 deletions(-) diff --git a/xarray/plot/datatree_plot.py b/xarray/plot/datatree_plot.py index 6dfce2807e5..8d7fa3a68b8 100644 --- a/xarray/plot/datatree_plot.py +++ b/xarray/plot/datatree_plot.py @@ -496,7 +496,8 @@ def scatter( @_update_doc_to_datatree(dataarray_plot.scatter) def scatter( dt: DataTree, - variable, + variable: str, + ax: Axes | None = None, *, x: Hashable | None = None, y: Hashable | None = None, @@ -508,7 +509,6 @@ def scatter( figsize: Iterable[float] | None = None, size: float | None = None, aspect: float | None = None, - ax: Axes | None = None, row: Hashable | None = None, col: Hashable | None = None, col_wrap: int | Literal["auto"] | None = None, @@ -538,7 +538,14 @@ def scatter( del locals_["dt"] locals_.update(locals_.pop("kwargs", {})) locals_.pop("variable") + locals_.pop("ax") + for node in dt.descendants: - da = node[variable] - print(da) - return da.plot.scatter(**locals_) + import matplotlib.pyplot as plt + + fig, ax = plt.subplots() + try: + da = node[variable] + da.plot.scatter(*locals_.pop("args", ()), **locals_, ax=ax) + except KeyError as err: + raise KeyError(f"{variable} not found at node: {node.name}") from err diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py index 21c4efe7f0c..87016a3f4a2 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 @@ -3639,3 +3639,55 @@ 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" + 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) + + # @pytest.mark.parametrize( + # "add_guide, hue_style, legend, colorbar", + # [ + # (None, None, False, True), + # (False, None, False, False), + # (True, None, False, True), + # (True, "continuous", False, True), + # (False, "discrete", False, False), + # (True, "discrete", True, False), + # ], + # ) + + def test_key_error(self): + """Assert that a KeyError is raised when accessing a variable that doesn't exist in node.""" + with pytest.raises(KeyError) as exc_info: + self.dt.plot.scatter("B") + print(exc_info.value) From b8a6b0bbc93ceadeeb924d930918b17df443e29c Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Tue, 21 Jul 2026 17:26:59 -0400 Subject: [PATCH 03/16] plot all --- xarray/plot/datatree_plot.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/xarray/plot/datatree_plot.py b/xarray/plot/datatree_plot.py index 8d7fa3a68b8..9e995e502d2 100644 --- a/xarray/plot/datatree_plot.py +++ b/xarray/plot/datatree_plot.py @@ -506,8 +506,8 @@ def scatter( hue_style: HueStyleOptions = None, markersize: Hashable | None = None, linewidth: Hashable | None = None, - figsize: Iterable[float] | None = None, - size: float | None = None, + # figsize: Iterable[float] | None = None, + # size: float | None = None, aspect: float | None = None, row: Hashable | None = None, col: Hashable | None = None, @@ -539,13 +539,18 @@ def scatter( locals_.update(locals_.pop("kwargs", {})) locals_.pop("variable") locals_.pop("ax") + print(locals_) - for node in dt.descendants: - import matplotlib.pyplot as plt + import matplotlib.pyplot as plt - fig, ax = plt.subplots() + fig, ax = plt.subplots() + for node in dt.descendants: try: da = node[variable] - da.plot.scatter(*locals_.pop("args", ()), **locals_, ax=ax) + da.plot.scatter( + ax=ax, + label=node.name, + **locals_, + ) except KeyError as err: raise KeyError(f"{variable} not found at node: {node.name}") from err From 616640e2ddb05fa8420efdc05921d6b4e383f5a7 Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Tue, 28 Jul 2026 21:10:27 -0400 Subject: [PATCH 04/16] feat: fig_kw --- xarray/plot/datatree_plot.py | 16 +++++++++------- xarray/tests/test_plot.py | 18 +++--------------- 2 files changed, 12 insertions(+), 22 deletions(-) diff --git a/xarray/plot/datatree_plot.py b/xarray/plot/datatree_plot.py index 9e995e502d2..abf7723313a 100644 --- a/xarray/plot/datatree_plot.py +++ b/xarray/plot/datatree_plot.py @@ -505,9 +505,7 @@ def scatter( hue: Hashable | None = None, hue_style: HueStyleOptions = None, markersize: Hashable | None = None, - linewidth: Hashable | None = None, - # figsize: Iterable[float] | None = None, - # size: float | None = None, + size: float | None = None, aspect: float | None = None, row: Hashable | None = None, col: Hashable | None = None, @@ -531,19 +529,23 @@ def scatter( norm: Normalize | None = None, extend: ExtendOptions = None, levels: ArrayLike | None = None, + fig_kw: Hashable | None = None, **kwargs: Any, ) -> PathCollection | FacetGrid[DataArray]: """Scat plot DataTree data variables 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") - print(locals_) - + (locals_.pop("ax"),) + locals_.pop("fig_kw") import matplotlib.pyplot as plt - fig, ax = plt.subplots() + fig, ax = plt.subplots(**fig_kw) for node in dt.descendants: try: da = node[variable] diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py index 87016a3f4a2..b957d0b1018 100644 --- a/xarray/tests/test_plot.py +++ b/xarray/tests/test_plot.py @@ -3666,6 +3666,7 @@ def setUp(self) -> None: 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: @@ -3674,20 +3675,7 @@ def test_accessor(self) -> None: assert DataTree.plot is DataTreePlotAccessor assert isinstance(self.dt.plot, DataTreePlotAccessor) - # @pytest.mark.parametrize( - # "add_guide, hue_style, legend, colorbar", - # [ - # (None, None, False, True), - # (False, None, False, False), - # (True, None, False, True), - # (True, "continuous", False, True), - # (False, "discrete", False, False), - # (True, "discrete", True, False), - # ], - # ) - def test_key_error(self): - """Assert that a KeyError is raised when accessing a variable that doesn't exist in node.""" - with pytest.raises(KeyError) as exc_info: + """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") - print(exc_info.value) From ff472c323e162fe1e0d860bd8f3b9eacc746beae Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Fri, 7 Aug 2026 16:36:39 -0400 Subject: [PATCH 05/16] docs: clean up autodocs --- xarray/plot/datatree_plot.py | 180 +++++++---------------------------- 1 file changed, 35 insertions(+), 145 deletions(-) diff --git a/xarray/plot/datatree_plot.py b/xarray/plot/datatree_plot.py index abf7723313a..413ada89e08 100644 --- a/xarray/plot/datatree_plot.py +++ b/xarray/plot/datatree_plot.py @@ -1,7 +1,6 @@ from __future__ import annotations import functools -import warnings from collections.abc import Callable, Hashable, Iterable from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload @@ -9,7 +8,6 @@ from xarray.plot.facetgrid import _easy_facetgrid from xarray.plot.utils import ( _add_colorbar, - _get_nice_quiver_magnitude, _infer_meta_data, _process_cmap_cbar_kwargs, get_axis, @@ -24,7 +22,6 @@ from xarray.core.dataarray import DataArray from xarray.core.datatree import DataTree from xarray.core.types import ( - AspectOptions, ExtendOptions, HueStyleOptions, ScaleOptions, @@ -37,16 +34,14 @@ def _dtplot(plotfunc): Parameters ---------- dt : DataTree +variable : str + name of the variable in multiple nodes, x : Hashable or None, optional Variable name for x-axis. y : Hashable or None, optional Variable name for y-axis. -u : Hashable or None, optional - Variable name for the *u* velocity (in *x* direction). - quiver/streamplot plots only. -v : Hashable or None, optional - Variable name for the *v* velocity (in *y* direction). - quiver/streamplot plots only. +z : Hashable or None, optional + if specified plot 3D and use this coordinate for z axis. hue: Hashable or None, optional Variable by which to color scatter points or arrows. hue_style: {'continuous', 'discrete'} or None, optional @@ -65,36 +60,10 @@ def _dtplot(plotfunc): "Wrap" the grid for the column variable after this number of columns, adding rows if ``col_wrap`` is less than the number of facets. If "auto" align the grid to the figsize or keep it as square as possible. -ax : matplotlib axes object or None, optional - If ``None``, use the current axes. Not applicable when using facets. -figsize : Iterable[float] or None, optional - A tuple (width, height) of the figure in inches. - Mutually exclusive with ``size`` and ``ax``. -size : scalar, optional - If provided, create a new figure for the plot with the given size. - Height (in inches) of each plot. See also: ``aspect``. -aspect : "auto", "equal", scalar or None, optional - Aspect ratio of plot, so that ``aspect * size`` gives the width in - inches. Only used if a ``size`` is provided. -sharex : bool or None, optional - If True all subplots share the same x-axis. -sharey : bool or None, optional - If True all subplots share the same y-axis. -add_guide: bool or None, optional - Add a guide that dependt on ``hue_style``: - - - ``'continuous'`` -- build a colorbar - - ``'discrete'`` -- build a legend - subplot_kws : dict or None, optional Dictionary of keyword arguments for Matplotlib subplots (see :py:meth:`matplotlib:matplotlib.figure.Figure.add_subplot`). Only applies to FacetGrid plotting. -cbar_kwargs : dict, optional - Dictionary of keyword arguments to pass to the colorbar - (see :meth:`matplotlib:matplotlib.figure.Figure.colorbar`). -cbar_ax : matplotlib axes object, optional - Axes in which to draw the colorbar. cmap : matplotlib colormap name or colormap, optional The mapping from data values to color space. Either a Matplotlib colormap name or object. If not provided, this will @@ -125,18 +94,6 @@ def _dtplot(plotfunc): norm : matplotlib.colors.Normalize, optional If ``norm`` has ``vmin`` or ``vmax`` specified, the corresponding kwarg must be ``None``. -infer_intervals: bool | None - If True the intervals are inferred. -center : float, optional - The value at which to center the colormap. Passing this value implies - use of a diverging colormap. Setting it to ``False`` prevents use of a - diverging colormap. -robust : bool, optional - If ``True`` and ``vmin`` or ``vmax`` are absent, the colormap range is - computed with 2nd and 98th percentiles instead of the extreme values. -colors : str or array-like of color-like, optional - A single color or a list of colors. The ``levels`` argument - is required. extend : {'neither', 'both', 'min', 'max'}, optional How to draw arrows extending the colorbar beyond its limits. If not provided, ``extend`` is inferred from ``vmin``, ``vmax`` and the data limits. @@ -146,6 +103,8 @@ def _dtplot(plotfunc): imply that the final number of levels is not exactly the expected one. Setting ``vmin`` and/or ``vmax`` with ``levels=N`` is equivalent to setting ``levels=np.linspace(vmin, vmax, N)``. +fig_kw : Hashable or None, optional + Matplotlib kwargs that get passed to pyplot.figure **kwargs : optional Additional keyword arguments to wrapped Matplotlib function. """ @@ -161,65 +120,28 @@ def newplotfunc( *args: Any, x: Hashable | None = None, y: Hashable | None = None, - u: Hashable | None = None, - v: Hashable | None = None, + z: Hashable | None = None, hue: Hashable | None = None, hue_style: HueStyleOptions = None, row: Hashable | None = None, col: Hashable | None = None, col_wrap: int | Literal["auto"] | None = None, - ax: Axes | None = None, - figsize: Iterable[float] | None = None, - size: float | None = None, - aspect: AspectOptions = None, - sharex: bool = True, - sharey: bool = True, - add_guide: bool | None = None, subplot_kws: dict[str, Any] | None = None, - cbar_kwargs: dict[str, Any] | None = None, - cbar_ax: Axes | None = None, cmap: str | Colormap | None = None, vmin: float | None = None, vmax: float | None = None, norm: Normalize | None = None, - infer_intervals: bool | None = None, - center: float | None = None, - robust: bool | None = None, - colors: str | ArrayLike | None = None, extend: ExtendOptions = None, levels: ArrayLike | None = None, **kwargs: Any, ) -> Any: - if args: - # TODO: Deprecated since 2022.10: - msg = "Using positional arguments is deprecated for plot methods, use keyword arguments instead." - assert x is None - x = args[0] - if len(args) > 1: - assert y is None - y = args[1] - if len(args) > 2: - assert u is None - u = args[2] - if len(args) > 3: - assert v is None - v = args[3] - if len(args) > 4: - assert hue is None - hue = args[4] - if len(args) > 5: - raise ValueError(msg) - else: - warnings.warn(msg, FutureWarning, stacklevel=2) - del msg - del args _is_facetgrid = kwargs.pop("_is_facetgrid", False) if _is_facetgrid: # facetgrid call meta_data = kwargs.pop("meta_data") else: meta_data = _infer_meta_data( - dt, x, y, hue, hue_style, add_guide, funcname=plotfunc.__name__ + dt, x, y, z, hue, hue_style, funcname=plotfunc.__name__ ) hue_style = meta_data["hue_style"] @@ -236,7 +158,7 @@ def newplotfunc( return _easy_facetgrid(kind="datatree", **allargs, **kwargs) figsize = kwargs.pop("figsize", None) - ax = get_axis(figsize, size, aspect, ax) + ax = get_axis(figsize) if hue_style == "continuous" and hue is not None: if _is_facetgrid: @@ -255,19 +177,11 @@ def newplotfunc( else: cmap_params_subset = {} - if (u is not None or v is not None) and plotfunc.__name__ not in ( - "quiver", - "streamplot", - ): - raise ValueError("u, v are only allowed for quiver or streamplot plots.") - primitive = plotfunc( dt=dt, x=x, y=y, ax=ax, - u=u, - v=v, hue=hue, hue_style=hue_style, cmap_params=cmap_params_subset, @@ -288,25 +202,9 @@ def newplotfunc( cbar_kwargs = {} if cbar_kwargs is None else cbar_kwargs if "label" not in cbar_kwargs: cbar_kwargs["label"] = meta_data.get("hue_label", None) - _add_colorbar(primitive, ax, cbar_ax, cbar_kwargs, cmap_params) - - if meta_data["add_quiverkey"]: - magnitude = _get_nice_quiver_magnitude(dt[u], dt[v]) - units = dt[u].attrs.get("units", "") - ax.quiverkey( - primitive, - X=0.85, - Y=0.9, - U=magnitude, - label=f"{magnitude}\n{units}", - labelpos="E", - coordinates="figure", - ) + _add_colorbar(primitive, ax, cbar_kwargs, cmap_params) - if plotfunc.__name__ in ("quiver", "streamplot"): - title = dt[u]._title_for_slice() - else: - title = dt[x]._title_for_slice() + title = dt[x]._title_for_slice() ax.set_title(title) return primitive @@ -373,20 +271,17 @@ def wrapper(datatree_plotfunc: F) -> F: @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, - hue_style: HueStyleOptions = None, markersize: Hashable | None = None, - linewidth: Hashable | None = None, - figsize: Iterable[float] | None = None, - size: float | None = None, - aspect: float | None = None, - ax: Axes | None = None, - row: None = None, # no wrap -> primitive - col: None = None, # no wrap -> primitive + row: Hashable | None = None, + col: Hashable | None = None, col_wrap: int | Literal["auto"] | None = None, xincrease: bool | None = True, yincrease: bool | None = True, @@ -407,27 +302,24 @@ def scatter( # type: ignore[misc,unused-ignore] # None is hashable :(s: DataTr norm: Normalize | None = None, extend: ExtendOptions = None, levels: ArrayLike | None = None, + fig_kw: Hashable | None = None, **kwargs: Any, ) -> PathCollection: ... @overload def scatter( - 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, - hue_style: HueStyleOptions = None, markersize: Hashable | None = None, - linewidth: Hashable | None = None, - figsize: Iterable[float] | None = None, - size: float | None = None, - aspect: float | None = None, - ax: Axes | None = None, row: Hashable | None = None, - col: Hashable, # wrap -> FacetGrid + col: Hashable | None = None, col_wrap: int | Literal["auto"] | None = None, xincrease: bool | None = True, yincrease: bool | None = True, @@ -448,26 +340,23 @@ def scatter( norm: Normalize | None = None, extend: ExtendOptions = None, levels: ArrayLike | None = None, + fig_kw: Hashable | None = None, **kwargs: Any, ) -> FacetGrid[DataArray]: ... @overload def scatter( - 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, - hue_style: HueStyleOptions = None, markersize: Hashable | None = None, - linewidth: Hashable | None = None, - figsize: Iterable[float] | None = None, - size: float | None = None, - aspect: float | None = None, - ax: Axes | None = None, - row: Hashable, # wrap -> FacetGrid + row: Hashable | None = None, col: Hashable | None = None, col_wrap: int | Literal["auto"] | None = None, xincrease: bool | None = True, @@ -489,6 +378,7 @@ def scatter( norm: Normalize | None = None, extend: ExtendOptions = None, levels: ArrayLike | None = None, + fig_kw: Hashable | None = None, **kwargs: Any, ) -> FacetGrid[DataArray]: ... @@ -503,10 +393,7 @@ def scatter( y: Hashable | None = None, z: Hashable | None = None, hue: Hashable | None = None, - hue_style: HueStyleOptions = None, markersize: Hashable | None = None, - size: float | None = None, - aspect: float | None = None, row: Hashable | None = None, col: Hashable | None = None, col_wrap: int | Literal["auto"] | None = None, @@ -549,10 +436,13 @@ def scatter( for node in dt.descendants: try: da = node[variable] - da.plot.scatter( - ax=ax, - label=node.name, - **locals_, - ) 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() From 18de0b1a23e2e4b3cc80ccb0d8fe309ec1ce0ef0 Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Fri, 7 Aug 2026 17:03:02 -0400 Subject: [PATCH 06/16] docs: plotting docs --- doc/api/plotting.rst | 8 ++++++++ doc/user-guide/plotting.md | 21 +++++++++++++++++++++ 2 files changed, 29 insertions(+) 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..d9c266fc6f6 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') +``` + ## Maps To follow this section you'll need to have Cartopy installed and working. From d83d05226488f3c2a21d3d5c1e27b355bb490d7a Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Wed, 12 Aug 2026 12:33:08 -0400 Subject: [PATCH 07/16] docs: added legend to scatter plot example --- doc/user-guide/plotting.md | 2 +- xarray/plot/datatree_plot.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/doc/user-guide/plotting.md b/doc/user-guide/plotting.md index d9c266fc6f6..8b4630aaff2 100644 --- a/doc/user-guide/plotting.md +++ b/doc/user-guide/plotting.md @@ -741,7 +741,7 @@ dt Let’s plot the "precipitation" variable in the "observed" and "reanalysis" groups. ```{code-cell} -dt.plot.scatter('precipitation') +dt.plot.scatter('precipitation', add_legend=True) ``` ## Maps diff --git a/xarray/plot/datatree_plot.py b/xarray/plot/datatree_plot.py index 413ada89e08..c419cd4cee0 100644 --- a/xarray/plot/datatree_plot.py +++ b/xarray/plot/datatree_plot.py @@ -430,6 +430,7 @@ def scatter( 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) @@ -444,5 +445,5 @@ def scatter( label=node.name, **locals_, ) - if add_legend: - ax.legend() + if add_legend: + ax.legend() From 43e711e0f06bd8872745f79a32723db93242b074 Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Wed, 12 Aug 2026 12:54:28 -0400 Subject: [PATCH 08/16] docs: what is new --- doc/whats-new.rst | 3 +++ 1 file changed, 3 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index c4cc7c9898a..c62bc07c88a 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -21,6 +21,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 ~~~~~~~~~~~~~~~~ From 8d9d01088e97a26ef47c3722f4620113194b2192 Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Wed, 12 Aug 2026 13:07:33 -0400 Subject: [PATCH 09/16] feat: DataTree accesor methods --- xarray/plot/accessor.py | 32 +++++++++++--------------------- 1 file changed, 11 insertions(+), 21 deletions(-) diff --git a/xarray/plot/accessor.py b/xarray/plot/accessor.py index e4b36f69ce5..09b86aed44f 100644 --- a/xarray/plot/accessor.py +++ b/xarray/plot/accessor.py @@ -1275,7 +1275,7 @@ def streamplot(self, *args, **kwargs) -> LineCollection | FacetGrid[Dataset]: class DataTreePlotAccessor: """ - Enables use of xarray.plot functions as attributes on a Dataset. + Enables use of xarray.plot functions as attributes on a DataTree. For example, DataTree.plot.scatter """ @@ -1339,15 +1339,9 @@ def scatter( y: Hashable | None = None, z: Hashable | None = None, hue: Hashable | None = None, - hue_style: HueStyleOptions = None, markersize: Hashable | None = None, - linewidth: Hashable | None = None, - figsize: Iterable[float] | None = None, - size: float | None = None, - aspect: float | None = None, - ax: Axes | None = None, row: Hashable | None = None, - col: Hashable, # wrap -> FacetGrid + col: Hashable | None = None, col_wrap: int | Literal["auto"] | None = None, xincrease: bool | None = True, yincrease: bool | None = True, @@ -1360,14 +1354,15 @@ def scatter( yscale: ScaleOptions = None, xticks: ArrayLike | None = None, yticks: ArrayLike | None = None, - xlim: tuple[float, float] | None = None, - ylim: tuple[float, float] | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, cmap=None, vmin: float | None = None, vmax: float | None = None, norm: Normalize | None = None, extend=None, - levels=None, + levels: ArrayLike | None = None, + fig_kw: Hashable | None = None, **kwargs: Any, ) -> FacetGrid[DataTree]: ... @@ -1379,14 +1374,8 @@ def scatter( y: Hashable | None = None, z: Hashable | None = None, hue: Hashable | None = None, - hue_style: HueStyleOptions = None, markersize: Hashable | None = None, - linewidth: Hashable | None = None, - figsize: Iterable[float] | None = None, - size: float | None = None, - aspect: float | None = None, - ax: Axes | None = None, - row: Hashable, # wrap -> FacetGrid + row: Hashable | None = None, col: Hashable | None = None, col_wrap: int | Literal["auto"] | None = None, xincrease: bool | None = True, @@ -1400,14 +1389,15 @@ def scatter( yscale: ScaleOptions = None, xticks: ArrayLike | None = None, yticks: ArrayLike | None = None, - xlim: tuple[float, float] | None = None, - ylim: tuple[float, float] | None = None, + xlim: ArrayLike | None = None, + ylim: ArrayLike | None = None, cmap=None, vmin: float | None = None, vmax: float | None = None, norm: Normalize | None = None, extend=None, - levels=None, + levels: ArrayLike | None = None, + fig_kw: Hashable | None = None, **kwargs: Any, ) -> FacetGrid[DataTree]: ... From 43a0f30fc229ae4df1d45b065b7c18bda12334d3 Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Wed, 12 Aug 2026 13:42:34 -0400 Subject: [PATCH 10/16] chore: mypy --- xarray/plot/accessor.py | 6 +++--- xarray/plot/datatree_plot.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/xarray/plot/accessor.py b/xarray/plot/accessor.py index 09b86aed44f..3d73188b669 100644 --- a/xarray/plot/accessor.py +++ b/xarray/plot/accessor.py @@ -1364,7 +1364,7 @@ def scatter( levels: ArrayLike | None = None, fig_kw: Hashable | None = None, **kwargs: Any, - ) -> FacetGrid[DataTree]: ... + ) -> FacetGrid[DataArray]: ... @overload def scatter( @@ -1399,8 +1399,8 @@ def scatter( levels: ArrayLike | None = None, fig_kw: Hashable | None = None, **kwargs: Any, - ) -> FacetGrid[DataTree]: ... + ) -> FacetGrid[DataArray]: ... @functools.wraps(datatree_plot.scatter, assigned=("__doc__",)) - def scatter(self, *args, **kwargs) -> PathCollection | FacetGrid[DataTree]: + 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 index c419cd4cee0..99a54098092 100644 --- a/xarray/plot/datatree_plot.py +++ b/xarray/plot/datatree_plot.py @@ -4,7 +4,6 @@ from collections.abc import Callable, Hashable, Iterable from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload -from xarray.plot import dataarray_plot from xarray.plot.facetgrid import _easy_facetgrid from xarray.plot.utils import ( _add_colorbar, @@ -155,7 +154,7 @@ def newplotfunc( for arg in ["meta_data", "kwargs", "dt"]: del allargs[arg] - return _easy_facetgrid(kind="datatree", **allargs, **kwargs) + return _easy_facetgrid(kind="dataarray", **allargs, **kwargs) figsize = kwargs.pop("figsize", None) ax = get_axis(figsize) @@ -383,7 +382,8 @@ def scatter( ) -> FacetGrid[DataArray]: ... -@_update_doc_to_datatree(dataarray_plot.scatter) +# @_update_doc_to_datatree(dataarray_plot.scatter) +@_dtplot def scatter( dt: DataTree, variable: str, From 487ecafa8fb3630f313d54b3ce389666a89b0411 Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Wed, 12 Aug 2026 14:06:04 -0400 Subject: [PATCH 11/16] remove _dt_plot function not necessary --- xarray/plot/datatree_plot.py | 207 +---------------------------------- 1 file changed, 2 insertions(+), 205 deletions(-) diff --git a/xarray/plot/datatree_plot.py b/xarray/plot/datatree_plot.py index 99a54098092..f2998a00156 100644 --- a/xarray/plot/datatree_plot.py +++ b/xarray/plot/datatree_plot.py @@ -4,13 +4,7 @@ from collections.abc import Callable, Hashable, Iterable from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload -from xarray.plot.facetgrid import _easy_facetgrid -from xarray.plot.utils import ( - _add_colorbar, - _infer_meta_data, - _process_cmap_cbar_kwargs, - get_axis, -) +from xarray.plot import dataarray_plot if TYPE_CHECKING: from matplotlib.axes import Axes @@ -22,200 +16,11 @@ from xarray.core.datatree import DataTree from xarray.core.types import ( ExtendOptions, - HueStyleOptions, ScaleOptions, ) from xarray.plot.facetgrid import FacetGrid -def _dtplot(plotfunc): - commondoc = """ -Parameters ----------- -dt : DataTree -variable : str - name of the variable in multiple nodes, -x : Hashable or None, optional - Variable name for x-axis. -y : Hashable or None, optional - Variable name for y-axis. -z : Hashable or None, optional - if specified plot 3D and use this coordinate for z axis. -hue: Hashable or None, optional - Variable by which to color scatter points or arrows. -hue_style: {'continuous', 'discrete'} or None, optional - How to use the ``hue`` variable: - - - ``'continuous'`` -- continuous color scale - (default for numeric ``hue`` variables) - - ``'discrete'`` -- a color for each unique value, using the default color cycle - (default for non-numeric ``hue`` variables) - -row : Hashable or None, optional - If passed, make row faceted plots on this dimension name. -col : Hashable or None, optional - If passed, make column faceted plots on this dimension name. -col_wrap : int, None or "auto", optional - "Wrap" the grid for the column variable after this number of columns, - adding rows if ``col_wrap`` is less than the number of facets. - If "auto" align the grid to the figsize or keep it as square as possible. -subplot_kws : dict or None, optional - Dictionary of keyword arguments for Matplotlib subplots - (see :py:meth:`matplotlib:matplotlib.figure.Figure.add_subplot`). - Only applies to FacetGrid plotting. -cmap : matplotlib colormap name or colormap, optional - The mapping from data values to color space. Either a - Matplotlib colormap name or object. If not provided, this will - be either ``'viridis'`` (if the function infers a sequential - dataset) or ``'RdBu_r'`` (if the function infers a diverging - dataset). - See :doc:`Choosing Colormaps in Matplotlib ` - for more information. - - If *seaborn* is installed, ``cmap`` may also be a - `seaborn color palette `_. - Note: if ``cmap`` is a seaborn color palette, - ``levels`` must also be specified. -vmin : float or None, optional - Lower value to anchor the colormap, otherwise it is inferred from the - data and other keyword arguments. When a diverging dataset is inferred, - setting `vmin` or `vmax` will fix the other by symmetry around - ``center``. Setting both values prevents use of a diverging colormap. - If discrete levels are provided as an explicit list, both of these - values are ignored. -vmax : float or None, optional - Upper value to anchor the colormap, otherwise it is inferred from the - data and other keyword arguments. When a diverging dataset is inferred, - setting `vmin` or `vmax` will fix the other by symmetry around - ``center``. Setting both values prevents use of a diverging colormap. - If discrete levels are provided as an explicit list, both of these - values are ignored. -norm : matplotlib.colors.Normalize, optional - If ``norm`` has ``vmin`` or ``vmax`` specified, the corresponding - kwarg must be ``None``. -extend : {'neither', 'both', 'min', 'max'}, optional - How to draw arrows extending the colorbar beyond its limits. If not - provided, ``extend`` is inferred from ``vmin``, ``vmax`` and the data limits. -levels : int or array-like, optional - Split the colormap (``cmap``) into discrete color intervals. If an integer - is provided, "nice" levels are chosen based on the data range: this can - imply that the final number of levels is not exactly the expected one. - Setting ``vmin`` and/or ``vmax`` with ``levels=N`` is equivalent to - setting ``levels=np.linspace(vmin, vmax, N)``. -fig_kw : Hashable or None, optional - Matplotlib kwargs that get passed to pyplot.figure -**kwargs : optional - Additional keyword arguments to wrapped Matplotlib function. - """ - - # Build on the original docstring - plotfunc.__doc__ = f"{plotfunc.__doc__}\n{commondoc}" - - @functools.wraps( - plotfunc, assigned=("__module__", "__name__", "__qualname__", "__doc__") - ) - def newplotfunc( - dt: DataTree, - *args: Any, - x: Hashable | None = None, - y: Hashable | None = None, - z: Hashable | None = None, - hue: Hashable | None = None, - hue_style: HueStyleOptions = None, - row: Hashable | None = None, - col: Hashable | None = None, - col_wrap: int | Literal["auto"] | None = None, - subplot_kws: dict[str, Any] | 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, - **kwargs: Any, - ) -> Any: - - _is_facetgrid = kwargs.pop("_is_facetgrid", False) - if _is_facetgrid: # facetgrid call - meta_data = kwargs.pop("meta_data") - else: - meta_data = _infer_meta_data( - dt, x, y, z, hue, hue_style, funcname=plotfunc.__name__ - ) - - hue_style = meta_data["hue_style"] - - # handle facetgridt first - if col or row: - allargs = locals().copy() - allargs["plotfunc"] = globals()[plotfunc.__name__] - allargs["data"] = dt - # remove kwargs to avoid passing the information twice - for arg in ["meta_data", "kwargs", "dt"]: - del allargs[arg] - - return _easy_facetgrid(kind="dataarray", **allargs, **kwargs) - - figsize = kwargs.pop("figsize", None) - ax = get_axis(figsize) - - if hue_style == "continuous" and hue is not None: - if _is_facetgrid: - cbar_kwargs = meta_data["cbar_kwargs"] - cmap_params = meta_data["cmap_params"] - else: - cmap_params, cbar_kwargs = _process_cmap_cbar_kwargs( - plotfunc, dt[hue].values, **locals() - ) - - # subset that can be passed to scatter, hist2d - cmap_params_subset = { - vv: cmap_params[vv] for vv in ["vmin", "vmax", "norm", "cmap"] - } - - else: - cmap_params_subset = {} - - primitive = plotfunc( - dt=dt, - x=x, - y=y, - ax=ax, - hue=hue, - hue_style=hue_style, - cmap_params=cmap_params_subset, - **kwargs, - ) - - if _is_facetgrid: # if this was called from Facetgrid.map_datatree, - return primitive # finish here. Else, make labels - - if meta_data.get("xlabel", None): - ax.set_xlabel(meta_data.get("xlabel")) - if meta_data.get("ylabel", None): - ax.set_ylabel(meta_data.get("ylabel")) - - if meta_data["add_legend"]: - ax.legend(handles=primitive, title=meta_data.get("hue_label", None)) - if meta_data["add_colorbar"]: - cbar_kwargs = {} if cbar_kwargs is None else cbar_kwargs - if "label" not in cbar_kwargs: - cbar_kwargs["label"] = meta_data.get("hue_label", None) - _add_colorbar(primitive, ax, cbar_kwargs, cmap_params) - - title = dt[x]._title_for_slice() - ax.set_title(title) - - return primitive - - # we want to actually expose the signature of newplotfunc - # and not the copied **kwargs from the plotfunc which - # functools.wraps addt, so delete the wrapped attr - del newplotfunc.__wrapped__ - - return newplotfunc - - F = TypeVar("F", bound=Callable) @@ -223,13 +28,6 @@ def _update_doc_to_datatree(dataarray_plotfunc: Callable) -> Callable[[F], F]: """ Add a common docstring by reusing the DataArray one. - TODO: Reduce code duplication. - - * The goal is to reduce code duplication by mov all DataTree - specific plots to the DataArray side and use this thin wrapper to - handle the converts between DataTree and DataArray. - * Improve docstring handling, maybe reword the DataArray versions to explain DataTrees better. - Parameters ---------- dataarray_plotfunc : Callable @@ -382,8 +180,7 @@ def scatter( ) -> FacetGrid[DataArray]: ... -# @_update_doc_to_datatree(dataarray_plot.scatter) -@_dtplot +@_update_doc_to_datatree(dataarray_plot.scatter) def scatter( dt: DataTree, variable: str, From 27e66164fdc25a513e4acc90d368d4fcf820d918 Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Wed, 12 Aug 2026 14:50:06 -0400 Subject: [PATCH 12/16] mypy --- xarray/plot/datatree_plot.py | 52 +++++------------------------------- 1 file changed, 7 insertions(+), 45 deletions(-) diff --git a/xarray/plot/datatree_plot.py b/xarray/plot/datatree_plot.py index f2998a00156..7a02d3ef9d1 100644 --- a/xarray/plot/datatree_plot.py +++ b/xarray/plot/datatree_plot.py @@ -66,44 +66,6 @@ def wrapper(datatree_plotfunc: F) -> F: return wrapper # type: ignore[return-value] -@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: Hashable | None = None, - **kwargs: Any, -) -> PathCollection: ... - - @overload def scatter( dt: DataTree, @@ -137,13 +99,13 @@ def scatter( norm: Normalize | None = None, extend: ExtendOptions = None, levels: ArrayLike | None = None, - fig_kw: Hashable | None = None, + fig_kw: dict[str, Any] | None = None, **kwargs: Any, ) -> FacetGrid[DataArray]: ... @overload -def scatter( +def scatter( # type: ignore[misc,unused-ignore] # None is hashable :(s: DataTree, dt: DataTree, variable: str, ax: Axes | None = None, @@ -175,13 +137,13 @@ def scatter( norm: Normalize | None = None, extend: ExtendOptions = None, levels: ArrayLike | None = None, - fig_kw: Hashable | None = None, + fig_kw: dict[str, Any] | None = None, **kwargs: Any, -) -> FacetGrid[DataArray]: ... +) -> PathCollection: ... @_update_doc_to_datatree(dataarray_plot.scatter) -def scatter( +def scatter( # type: ignore[return] dt: DataTree, variable: str, ax: Axes | None = None, @@ -213,10 +175,10 @@ def scatter( norm: Normalize | None = None, extend: ExtendOptions = None, levels: ArrayLike | None = None, - fig_kw: Hashable | None = None, + fig_kw: dict[str, Any] | None = None, **kwargs: Any, ) -> PathCollection | FacetGrid[DataArray]: - """Scat plot DataTree data variables against each other.""" + """Scatter DataTree data variables with the same node against each other.""" if fig_kw is None: fig_kw = {} From 3cb2420777ba36ebdccecdcd05bdfe4fdb233954 Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Wed, 12 Aug 2026 14:53:17 -0400 Subject: [PATCH 13/16] chore: mypy --- xarray/plot/accessor.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/xarray/plot/accessor.py b/xarray/plot/accessor.py index 3d73188b669..c92931365bf 100644 --- a/xarray/plot/accessor.py +++ b/xarray/plot/accessor.py @@ -1328,6 +1328,7 @@ def scatter( # type: ignore[misc,unused-ignore] # None is hashable :( norm: Normalize | None = None, extend=None, levels=None, + fig_kw: dict[str, Any] | None = None, **kwargs: Any, ) -> PathCollection: ... @@ -1362,7 +1363,7 @@ def scatter( norm: Normalize | None = None, extend=None, levels: ArrayLike | None = None, - fig_kw: Hashable | None = None, + fig_kw: dict[str, Any] | None = None, **kwargs: Any, ) -> FacetGrid[DataArray]: ... @@ -1397,7 +1398,7 @@ def scatter( norm: Normalize | None = None, extend=None, levels: ArrayLike | None = None, - fig_kw: Hashable | None = None, + fig_kw: dict[str, Any] | None = None, **kwargs: Any, ) -> FacetGrid[DataArray]: ... From eb1485d8cc77d83982259db37f3cdb95852dbe4d Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Wed, 12 Aug 2026 16:01:45 -0400 Subject: [PATCH 14/16] mypy plot accessor --- xarray/plot/accessor.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/xarray/plot/accessor.py b/xarray/plot/accessor.py index c92931365bf..ae67e6990c8 100644 --- a/xarray/plot/accessor.py +++ b/xarray/plot/accessor.py @@ -1292,14 +1292,13 @@ def __call__(self, *args, **kwargs) -> NoReturn: ) @overload - def scatter( # type: ignore[misc,unused-ignore] # None is hashable :( + 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, - hue_style: HueStyleOptions = None, markersize: Hashable | None = None, linewidth: Hashable | None = None, figsize: Iterable[float] | None = None, From 072777f34fa268771417117269a34179fb17dce2 Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Wed, 12 Aug 2026 18:12:37 -0400 Subject: [PATCH 15/16] fix: mypy --- xarray/plot/accessor.py | 31 +++++++++++++------------------ 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/xarray/plot/accessor.py b/xarray/plot/accessor.py index ae67e6990c8..96eaf972015 100644 --- a/xarray/plot/accessor.py +++ b/xarray/plot/accessor.py @@ -1292,7 +1292,7 @@ def __call__(self, *args, **kwargs) -> NoReturn: ) @overload - def scatter( # type: ignore[misc,unused-ignore] # None is hashable :(s: DataTree, + def scatter( self, *args: Any, x: Hashable | None = None, @@ -1300,11 +1300,6 @@ def scatter( # type: ignore[misc,unused-ignore] # None is hashable :(s: DataTr z: Hashable | None = None, hue: Hashable | None = None, markersize: Hashable | None = None, - linewidth: Hashable | None = None, - figsize: Iterable[float] | None = None, - size: float | None = None, - aspect: float | None = None, - ax: Axes | None = None, row: None = None, # no wrap -> primitive col: None = None, # no wrap -> primitive col_wrap: int | Literal["auto"] | None = None, @@ -1329,10 +1324,10 @@ def scatter( # type: ignore[misc,unused-ignore] # None is hashable :(s: DataTr levels=None, fig_kw: dict[str, Any] | None = None, **kwargs: Any, - ) -> PathCollection: ... + ) -> FacetGrid[DataArray]: ... @overload - def scatter( + def scatter( # type: ignore[misc,unused-ignore] # None is hashable :(s: DataTree, self, *args: Any, x: Hashable | None = None, @@ -1341,7 +1336,7 @@ def scatter( hue: Hashable | None = None, markersize: Hashable | None = None, row: Hashable | None = None, - col: Hashable | None = None, + col: Hashable, # wrap -> FacetGrid col_wrap: int | Literal["auto"] | None = None, xincrease: bool | None = True, yincrease: bool | None = True, @@ -1354,17 +1349,17 @@ def scatter( yscale: ScaleOptions = None, xticks: ArrayLike | None = None, yticks: ArrayLike | None = None, - xlim: ArrayLike | None = None, - ylim: 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: ArrayLike | None = None, + levels=None, fig_kw: dict[str, Any] | None = None, **kwargs: Any, - ) -> FacetGrid[DataArray]: ... + ) -> PathCollection: ... @overload def scatter( @@ -1375,7 +1370,7 @@ def scatter( z: Hashable | None = None, hue: Hashable | None = None, markersize: Hashable | None = None, - row: Hashable | None = None, + row: Hashable, # wrap -> FacetGrid col: Hashable | None = None, col_wrap: int | Literal["auto"] | None = None, xincrease: bool | None = True, @@ -1389,17 +1384,17 @@ def scatter( yscale: ScaleOptions = None, xticks: ArrayLike | None = None, yticks: ArrayLike | None = None, - xlim: ArrayLike | None = None, - ylim: 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: ArrayLike | None = None, + levels=None, fig_kw: dict[str, Any] | None = None, **kwargs: Any, - ) -> FacetGrid[DataArray]: ... + ) -> FacetGrid[DataArray] | PathCollection: ... @functools.wraps(datatree_plot.scatter, assigned=("__doc__",)) def scatter(self, *args, **kwargs) -> PathCollection | FacetGrid[DataArray]: From 1a5d5f7a918bd6457ebb75f5a2d86ca97b17e409 Mon Sep 17 00:00:00 2001 From: Olufunke Awowale Date: Wed, 19 Aug 2026 17:43:14 -0400 Subject: [PATCH 16/16] docs: for datatree plotting --- xarray/plot/datatree_plot.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/xarray/plot/datatree_plot.py b/xarray/plot/datatree_plot.py index 7a02d3ef9d1..5834cd1af70 100644 --- a/xarray/plot/datatree_plot.py +++ b/xarray/plot/datatree_plot.py @@ -23,20 +23,23 @@ F = TypeVar("F", bound=Callable) +# TODO Eni +# Added back _dtplot accessor for documentation specific to datatree -def _update_doc_to_datatree(dataarray_plotfunc: Callable) -> Callable[[F], F]: + +def _update_doc_to_datatree(datatree_plotfunc: Callable) -> Callable[[F], F]: """ Add a common docstring by reusing the DataArray one. Parameters ---------- - dataarray_plotfunc : Callable + datatree_plotfunc : Callable Function that returns a finished plot primitive. """ # Build on the original docstring - da_doc = dataarray_plotfunc.__doc__ - if da_doc is None: + dt_doc = datatree_plotfunc.__doc__ + if dt_doc is None: raise NotImplementedError("DataArray plot method requires a docstring") da_str = """ @@ -53,12 +56,12 @@ def _update_doc_to_datatree(dataarray_plotfunc: Callable) -> Callable[[F], F]: : DataTree """ # TODO: improve this? - if da_str in da_doc: - dt_doc = da_doc.replace(da_str, dt_str).replace("darray", "dt") - else: - dt_doc = da_doc + # 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(dataarray_plotfunc) + @functools.wraps(datatree_plotfunc) def wrapper(datatree_plotfunc: F) -> F: datatree_plotfunc.__doc__ = dt_doc return datatree_plotfunc