Adding Scatter Plot functionality to DataTree

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
 --------

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.

doc/whats-new.rst

@@ -25,6 +25,9 @@ New Features
   silently being written uncompressed (:issue:`10657`, :pull:`11067`).
   By `Mark Harfouche <https://github.com/hmaarrfk>`_.
 
+- ``DataTree`` objects now have plotting functionality with the ``plot.scatter`` method
+  (:issue:`9348`, :pull:`11516`).
+  By `Eni Awowale <https://github.com/eni-awowale>`_.
 
 Breaking Changes
 ~~~~~~~~~~~~~~~~

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))

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)