diff --git a/README.md b/README.md index b10ade0..02dfd59 100644 --- a/README.md +++ b/README.md @@ -349,3 +349,49 @@ canvas.show() (
, array([[]], dtype=object)) + +### xarray data + +Plot labelled [xarray](https://docs.xarray.dev) data directly +(`pip install maxplotlibx[xarray]`). Axes come from the coordinates, +labels from the `long_name` and `units` attributes, and titles from the +coordinates you selected. `import maxplotlib.xarray` adds a `.maxplot` +accessor that mirrors xarray’s own `.plot` API and returns an ordinary +`Canvas`, so the backend is still chosen when rendering: + +``` python +import xarray as xr + +import maxplotlib.xarray # registers da.maxplot and ds.maxplot + +t = np.linspace(0, 1.5, 6) +xs = np.linspace(0, 2 * np.pi, 80) +ys = np.linspace(-1, 1, 50) +wave = xr.DataArray( + np.sin(xs - 2 * t[:, None, None]) * np.exp(-3 * ys[None, :, None] ** 2), + dims=("t", "y", "x"), + coords={"t": ("t", t, {"units": "s"}), "y": ys, "x": ("x", xs, {"units": "m"})}, + name="phi", + attrs={"long_name": "Potential", "units": "V"}, +) + +wave.maxplot.pcolormesh(col="t", col_wrap=3, canvas_kwargs={"width": "16cm", "ratio": 0.6}).show() +``` + +![](README_files/figure-commonmark/cell-20-output-1.png) + + (
, + array([[, + , + ], + [, + , + ]], + dtype=object)) + +The same works through Canvas methods, +e.g. `canvas.plot(da, hue="species")`, +`ax.pcolormesh(da, xcoord="R", ycoord="Z")` for curvilinear grids, or +`Canvas.facet(da, col="t")`. `ds.maxplot.scatter(x=..., y=..., hue=...)` +plots one Dataset variable against another. See the [xarray +tutorial](tutorials/tutorial_17_xarray.ipynb) for more. diff --git a/README.qmd b/README.qmd index 531a7ce..cefdf37 100644 --- a/README.qmd +++ b/README.qmd @@ -276,3 +276,37 @@ Show all layers: ```{python} canvas.show() ``` + +### xarray data + +Plot labelled [xarray](https://docs.xarray.dev) data directly +(`pip install maxplotlibx[xarray]`). Axes come from the coordinates, labels +from the `long_name` and `units` attributes, and titles from the coordinates +you selected. `import maxplotlib.xarray` adds a `.maxplot` accessor that mirrors +xarray's own `.plot` API and returns an ordinary `Canvas`, so the backend is +still chosen when rendering: + +```{python} +import xarray as xr + +import maxplotlib.xarray # registers da.maxplot and ds.maxplot + +t = np.linspace(0, 1.5, 6) +xs = np.linspace(0, 2 * np.pi, 80) +ys = np.linspace(-1, 1, 50) +wave = xr.DataArray( + np.sin(xs - 2 * t[:, None, None]) * np.exp(-3 * ys[None, :, None] ** 2), + dims=("t", "y", "x"), + coords={"t": ("t", t, {"units": "s"}), "y": ys, "x": ("x", xs, {"units": "m"})}, + name="phi", + attrs={"long_name": "Potential", "units": "V"}, +) + +wave.maxplot.pcolormesh(col="t", col_wrap=3, canvas_kwargs={"width": "16cm", "ratio": 0.6}).show() +``` + +The same works through Canvas methods, e.g. `canvas.plot(da, hue="species")`, +`ax.pcolormesh(da, xcoord="R", ycoord="Z")` for curvilinear grids, or +`Canvas.facet(da, col="t")`. `ds.maxplot.scatter(x=..., y=..., hue=...)` plots +one Dataset variable against another. See the +[xarray tutorial](tutorials/tutorial_17_xarray.ipynb) for more. diff --git a/README_files/figure-commonmark/cell-20-output-1.png b/README_files/figure-commonmark/cell-20-output-1.png new file mode 100644 index 0000000..74215c7 Binary files /dev/null and b/README_files/figure-commonmark/cell-20-output-1.png differ diff --git a/docs/source/index.rst b/docs/source/index.rst index e1bd3ab..82a2734 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -31,3 +31,4 @@ documentation for details. tutorials/tutorial_14_axis_and_layout_controls tutorials/tutorial_15_tikzfigure_subplots tutorials/tutorial_16_plotext_advanced + tutorials/tutorial_17_xarray diff --git a/pyproject.toml b/pyproject.toml index e7c49cf..2ca1936 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,8 +25,13 @@ dependencies = [ test = [ "pytest", "coverage", + "xarray", +] +xarray = [ + "xarray", ] docs = [ + "xarray", "myst-parser", "sphinx", "sphinx-rtd-theme", diff --git a/src/maxplotlib/canvas/canvas.py b/src/maxplotlib/canvas/canvas.py index 43a882c..ee66087 100644 --- a/src/maxplotlib/canvas/canvas.py +++ b/src/maxplotlib/canvas/canvas.py @@ -25,6 +25,7 @@ _tikz_step_coordinates, _tikz_style_kwargs, ) +from maxplotlib.utils import xarray_support from maxplotlib.utils.options import Backends @@ -370,6 +371,10 @@ def __init__( self._supylabel_kwargs: dict = {} self._subplots_adjust_kwargs: dict = {} self._tight_layout_kwargs: dict | None = None + self._hide_empty_subplots = False + # Set by Canvas.facet: merge Plotly legend entries across subplots and + # leave room for the figure title above the subplot titles. + self._facet = False self._set_tight_layout = None self._align_labels = False self._align_titles = False @@ -486,6 +491,171 @@ def subplots( return canvas, [row[0] for row in axes] return canvas, axes + _FACET_KINDS = { + "plot": "plot", + "scatter": "scatter", + "pcolormesh": "pcolormesh", + "imshow": "add_imshow", + "contour": "contour", + "contourf": "contourf", + } + + @classmethod + def facet( + cls, + da, + col=None, + row=None, + col_wrap: int | None = None, + kind: str = "pcolormesh", + sharey: bool = True, + canvas_kwargs: dict | None = None, + **kwargs, + ): + """ + Create a Canvas with one subplot per value of a DataArray dimension. + + Parameters: + da (xarray.DataArray): The data. After removing ``col``/``row``, it + must have the dimensions ``kind`` needs: 1-D for ``"plot"`` and + ``"scatter"`` (2-D with ``hue=``), 2-D for the others. + col, row (str): Dimensions to lay out across columns and rows. + col_wrap (int): Wrap a ``col`` facet after this many columns. + kind (str): ``"plot"``, ``"scatter"``, ``"pcolormesh"``, ``"imshow"``, + ``"contour"`` or ``"contourf"``. + sharey (bool): For ``"plot"``/``"scatter"``, give every subplot the + value range of the whole array (on the x-axis with ``ycoord=``). + With ``False`` each subplot scales its own axis and keeps its y + tick labels. + canvas_kwargs (dict): Forwarded to the Canvas constructor. + **kwargs: Forwarded to each subplot's ``kind`` method. + + Color-mapped kinds share one color scale (``vmin``/``vmax`` default to + the data range, and contour levels are shared) and one colorbar for + the whole figure, which ``add_colorbar=False`` turns off. Axis and + tick labels are kept on the outer subplots only. Each subplot is + titled with the coordinate values that differ between subplots, and + the figure with the ones they share. + + Returns: + (canvas, axes): The Canvas and a 2-D list of LinePlots, with ``None`` + where a wrapped grid has no subplot. + + Examples: + >>> canvas, axes = Canvas.facet(da, col="t", col_wrap=3) + >>> canvas, axes = Canvas.facet(da, row="species", kind="plot") + """ + if not xarray_support.is_dataarray(da): + raise TypeError("facet() needs an xarray.DataArray") + if kind not in cls._FACET_KINDS: + raise ValueError( + f"kind must be one of {sorted(cls._FACET_KINDS)}, got {kind!r}" + ) + if col is None and row is None: + raise ValueError("facet() needs col= and/or row=") + if col == row: + raise ValueError("col and row must be different dimensions") + for dim in (col, row): + if dim is not None and dim not in da.dims: + raise ValueError(f"{dim!r} is not a dimension of {da.dims}") + if col_wrap is not None and (col is None or row is not None): + raise ValueError("col_wrap= needs col= and no row=") + + ncol_values = da.sizes[col] if col is not None else 1 + nrow_values = da.sizes[row] if row is not None else 1 + if col_wrap is not None: + ncols = max(1, min(col_wrap, ncol_values)) + nrows = -(-ncol_values // ncols) + panels = [(j // ncols, j % ncols, {col: j}) for j in range(ncol_values)] + else: + ncols, nrows = ncol_values, nrow_values + panels = [ + (i, j, {d: k for d, k in ((row, i), (col, j)) if d is not None}) + for i in range(nrows) + for j in range(ncols) + ] + + mapped = kind not in ("plot", "scatter") + sharey = sharey or mapped + add_colorbar = kwargs.pop("add_colorbar", kind != "contour") + values = xarray_support.magnitude(da) + if mapped: + # Shared limits for every panel, with xarray's color defaults. + xarray_support.color_limits(values, kwargs) + kwargs.setdefault("vmin", float(np.nanmin(values))) + kwargs.setdefault("vmax", float(np.nanmax(values))) + if kind in ("contour", "contourf") and "levels" not in kwargs: + from matplotlib.ticker import MaxNLocator + + kwargs["levels"] = MaxNLocator(10).tick_values( + kwargs["vmin"], kwargs["vmax"] + ) + # One colorbar for the figure, added below; "colorbar" hides the + # per-trace scale Plotly would otherwise show. + kwargs["add_colorbar"] = False + kwargs["colorbar"] = False + + canvas_kwargs = dict(canvas_kwargs or {}) + if not {"subplot_spacing", "gridspec_kw"} & canvas_kwargs.keys(): + canvas_kwargs["subplot_spacing"] = SubplotSpacing( + wspace=0.1 if sharey else 0.3, hspace=0.3 + ) + canvas = cls(nrows=nrows, ncols=ncols, **canvas_kwargs) + canvas._hide_empty_subplots = True + canvas._facet = True + # Coordinates shared by every panel go in the figure title, so the + # panel titles only show what differs between them. + common = [name for name, coord in da.coords.items() if coord.ndim == 0] + common_title = xarray_support.title(da) + if common_title: + canvas.suptitle(common_title) + axes = [[None] * ncols for _ in range(nrows)] + method = cls._FACET_KINDS[kind] + for r, c, selection in panels: + subplot = canvas.add_subplot(row=r, col=c) + panel = da.isel(selection) + subplot.set_title( + xarray_support.title(panel, exclude=common) + or ", ".join(f"{dim} = {index}" for dim, index in selection.items()) + ) + getattr(subplot, method)(panel, **kwargs) + axes[r][c] = subplot + + vertical = "ycoord" in kwargs + if not mapped and sharey: + low, high = float(np.nanmin(values)), float(np.nanmax(values)) + margin = 0.05 * (high - low) + for r, c, _ in panels: + subplot = axes[r][c] + tick_params = {} + if r + 1 < nrows and axes[r + 1][c] is not None: + subplot._xlabel = None + tick_params["labelbottom"] = False + if c > 0 and sharey: + subplot._ylabel = None + tick_params["labelleft"] = False + if tick_params: + subplot.tick_params(**tick_params) + if not mapped and sharey and vertical: + subplot.set_xlim(low - margin, high + margin) + elif not mapped and sharey: + subplot.set_ylim(low - margin, high + margin) + if (r, c) != panels[0][:2]: + subplot._legend = False + if mapped and add_colorbar: + first = axes[panels[0][0]][panels[0][1]] + first._add( + { + "label": xarray_support.value_label(da), + "layer": 0, + "plot_type": "colorbar", + "kwargs": {}, + "span_figure": True, + }, + 0, + ) + return canvas, axes + @property def _subplot_dict(self): return self._subplots @@ -609,7 +779,7 @@ def _get_or_create_subplot(self, row, col): def scatter( self, x, - y, + y=None, layer=0, row: int | None = None, col: int | None = None, @@ -624,6 +794,8 @@ def scatter( layer (int): Layer index (default 0). row, col (int): Subplot position (default top-left). **kwargs: Forwarded to the backend (e.g., color, marker, s, label). + + ``scatter(da)`` accepts an ``xarray.DataArray`` like ``plot(da)``. """ sp = self._get_or_create_subplot(row, col) sp.scatter(x, y, layer=layer, **kwargs) @@ -830,40 +1002,44 @@ def eventplot( def contour( self, x, - y, - z, + y=None, + z=None, layer=0, row: int | None = None, col: int | None = None, **kwargs, ): - """Add contour lines to a subplot.""" + """Add contour lines to a subplot; ``contour(da)`` takes a DataArray.""" self._get_or_create_subplot(row, col).contour(x, y, z, layer=layer, **kwargs) def contourf( self, x, - y, - z, + y=None, + z=None, layer=0, row: int | None = None, col: int | None = None, **kwargs, ): - """Add filled contours to a subplot.""" + """Add filled contours to a subplot; ``contourf(da)`` takes a DataArray.""" self._get_or_create_subplot(row, col).contourf(x, y, z, layer=layer, **kwargs) def pcolormesh( self, x, - y, - z, + y=None, + z=None, layer=0, row: int | None = None, col: int | None = None, **kwargs, ): - """Add a pseudocolor mesh to a subplot.""" + """Add a pseudocolor mesh to a subplot. + + ``pcolormesh(da)`` accepts a 2-D ``xarray.DataArray``; see + :meth:`LinePlot.pcolormesh`. + """ self._get_or_create_subplot(row, col).pcolormesh(x, y, z, layer=layer, **kwargs) def hexbin( @@ -1593,7 +1769,7 @@ def imshow( col: int | None = None, **kwargs, ): - """Add an image/matrix plot to a subplot.""" + """Add an image/matrix plot to a subplot; ``imshow(da)`` takes a DataArray.""" self._get_or_create_subplot(row, col).add_imshow(data, layer=layer, **kwargs) def add_image( @@ -2082,10 +2258,19 @@ def plot(self, *args, backend=None, **kwargs): """Add a line, or render when called with backend options. ``canvas.plot(x, y, **style)`` is the convenient direct plotting form. + ``canvas.plot(da)`` plots a 1-D ``xarray.DataArray`` against its + coordinate, labelling the axes from its attributes; ``hue=`` + draws a 2-D one as one line per value of ``dim``. Rendering is named explicitly by ``canvas.render(...)``; the legacy ``canvas.plot(backend=...)`` form remains supported. """ explicit_render = backend is not None or (args and isinstance(args[0], str)) + if len(args) == 1 and xarray_support.is_dataarray(args[0]): + layer = kwargs.pop("layer", 0) + row = kwargs.pop("row", None) + col = kwargs.pop("col", None) + self._get_or_create_subplot(row, col).plot(args[0], layer=layer, **kwargs) + return self if args and not isinstance(args[0], str): if len(args) < 2: raise TypeError("plot(x, y) requires both x and y data") @@ -2381,6 +2566,18 @@ def plot_matplotlib( else: subplot.plot_matplotlib(ax, layers=layers) + if self._hide_empty_subplots: + for row in range(self.nrows): + for col in range(self.ncols): + if (row, col) not in self._subplot_dict: + axes[row][col].set_visible(False) + figure_axes = [ax for ax in fig.axes if ax.get_visible()] + for subplot in self._subplot_dict.values(): + figure_colorbar = getattr(subplot, "_figure_colorbar", None) + if figure_colorbar is not None and figure_colorbar[0] is not None: + mappable, label = figure_colorbar + fig.colorbar(mappable, ax=figure_axes, label=label) + if verbose: print("Finished plotting subplots.") @@ -2393,6 +2590,10 @@ def plot_matplotlib( fig.supxlabel(self._supxlabel, **self._supxlabel_kwargs) if self._supylabel: fig.supylabel(self._supylabel, **self._supylabel_kwargs) + if self._facet and self._suptitle and "top" not in self._subplots_adjust_kwargs: + # About four font heights: the figure title plus subplot titles. + room = 4 * self.fontsize / 72 + fig.subplots_adjust(top=max(0.5, 1 - room / fig.get_figheight())) if self._subplots_adjust_kwargs: fig.subplots_adjust(**self._subplots_adjust_kwargs) if self._tight_layout_kwargs is not None: @@ -2580,8 +2781,8 @@ def plot_tikzfigure( ) if plot_type == "plot": # Extract and transform x, y data - x = (line_data["x"] + line_plot._xshift) * line_plot._xscale - y = (line_data["y"] + line_plot._yshift) * line_plot._yscale + x = line_plot._shift_x(line_data["x"]) + y = line_plot._shift_y(line_data["y"]) kwargs = line_data.get("kwargs", {}) if verbose: print(f"Line {kwargs = }") @@ -2592,8 +2793,8 @@ def plot_tikzfigure( **_tikz_style_kwargs(kwargs), ) elif plot_type == "scatter": - x = (line_data["x"] + line_plot._xshift) * line_plot._xscale - y = (line_data["y"] + line_plot._yshift) * line_plot._yscale + x = line_plot._shift_x(line_data["x"]) + y = line_plot._shift_y(line_data["y"]) kwargs = _tikz_style_kwargs(line_data.get("kwargs", {})) kwargs.setdefault("mark", "*") kwargs["line_width"] = 0 @@ -2926,12 +3127,30 @@ def plot_plotly( subplot_titles=subplot_titles, specs=specs, ) + if self._hide_empty_subplots: + for row in range(self.nrows): + for col in range(self.ncols): + if (row, col) not in self._subplot_dict: + fig.update_xaxes(visible=False, row=row + 1, col=col + 1) + fig.update_yaxes(visible=False, row=row + 1, col=col + 1) # Plot each subplot and propagate axis labels/scale + legend_names = set() for (row, col), line_plot in self._subplot_dict.items(): traces, shapes, annotations = line_plot.plot_plotly( layers=layers, allow_unsupported=allow_unsupported ) + if self._facet: + # One legend entry per label across all subplots; clicking it + # toggles the matching trace in every subplot. + for trace in traces: + name = getattr(trace, "name", None) + if not name or trace.type in ("pie", "table"): + continue + trace.legendgroup = name + if name in legend_names: + trace.showlegend = False + legend_names.add(name) for trace in traces: if trace.type in ("pie", "table"): fig.add_trace(trace) diff --git a/src/maxplotlib/subfigure/line_plot.py b/src/maxplotlib/subfigure/line_plot.py index 346db44..3107353 100644 --- a/src/maxplotlib/subfigure/line_plot.py +++ b/src/maxplotlib/subfigure/line_plot.py @@ -3,13 +3,24 @@ import matplotlib.pyplot as plt import numpy as np import plotly.graph_objects as go -from mpl_toolkits.axes_grid1 import make_axes_locatable from tikzfigure import TikzFigure +from maxplotlib.utils import xarray_support + # Keyword arguments that every drawing method accepts and that are handled by # maxplotlib itself rather than being forwarded to a backend drawing call. _NEUTRAL_KWARGS = ("hover", "meta") + +def _mpl_mappable_kwargs(line): + """Kwargs for a Matplotlib color-mapped artist. + + ``colorbar`` is a Plotly-only flag; Matplotlib adds colorbars through a + separate ``colorbar()`` entry, and would reject it as an artist property. + """ + return {k: v for k, v in line["kwargs"].items() if k != "colorbar"} + + _TIKZ_SUPPORTED_PLOT_TYPES = { "plot", "scatter", @@ -388,14 +399,96 @@ def add_line( } self._add(ld, layer) - def plot(self, x, y, layer=0, **kwargs): - """Matplotlib-style alias for :meth:`add_line`.""" + def plot(self, x, y=None, layer=0, **kwargs): + """Matplotlib-style alias for :meth:`add_line`. + + ``plot(da)`` with a 1-D ``xarray.DataArray`` plots it against its + coordinate; see :meth:`_plot_dataarray`. + """ + if y is None: + if not xarray_support.is_dataarray(x): + raise TypeError("plot(x, y) requires both x and y data") + self._plot_dataarray(self.add_line, x, layer, kwargs) + return self.add_line(x, y, layer=layer, **kwargs) - def scatter(self, x, y, layer=0, **kwargs): + def _set_default_labels(self, xlabel, ylabel, da=None): + """Set axis labels, and a title from ``da``, that the user has not set.""" + if not self._xlabel and xlabel: + self._xlabel = xlabel + if not self._ylabel and ylabel: + self._ylabel = ylabel + if da is not None and not self._title: + self._title = xarray_support.title(da) or self._title + + def _plot_dataarray(self, method, da, layer, kwargs): + """Draw a DataArray as lines or points with ``method(x, y, ...)``. + + Axes are labelled from the array's attributes and the title from its + single-value coordinates, unless already set. ``xcoord=`` names the + dimension or 1-D coordinate to plot against; ``ycoord=`` does the same + with the coordinate on the y-axis, e.g. for vertical profiles. + ``hue=`` draws a 2-D array as one labelled series per value of + ``dim``, colored ``C0``, ``C1``, ... unless ``color`` is given, and + shows the legend unless ``add_legend=False``. + """ + hue = kwargs.pop("hue", None) + add_legend = kwargs.pop("add_legend", True) + if hue is not None and "label" in kwargs: + raise TypeError("label= cannot be combined with hue=") + data = xarray_support.line_data( + da, hue=hue, x=kwargs.pop("xcoord", None), y=kwargs.pop("ycoord", None) + ) + if data.vertical: + self._set_default_labels(data.value_label, data.coord_label, da) + else: + self._set_default_labels(data.coord_label, data.value_label, da) + cycle_colors = hue is not None and not {"color", "c"} & kwargs.keys() + for i, (positions, values, label) in enumerate(data.lines): + line_kwargs = dict(kwargs) + if label is not None: + line_kwargs["label"] = label + if cycle_colors: + line_kwargs["color"] = f"C{i % 10}" + if data.vertical: + method(values, positions, layer=layer, **line_kwargs) + else: + method(positions, values, layer=layer, **line_kwargs) + if hue is not None and add_legend: + self._legend = True + + def _mesh_dataarray(self, method, da, layer, kwargs, add_colorbar=True, name=None): + """Draw a 2-D DataArray with ``method(x, y, z, ...)``. + + ``xcoord=``/``ycoord=`` name the dimension or coordinate for each + axis; 2-D coordinates give a curvilinear mesh. Colors follow xarray's + defaults (see :func:`xarray_support.color_limits`, including + ``robust=`` and ``center=``). A colorbar labelled from the array is + added unless ``add_colorbar=False``. + """ + add_colorbar = kwargs.pop("add_colorbar", add_colorbar) + mesh = self._mesh_of(da, kwargs, name or method.__name__) + self._set_default_labels(mesh.xlabel, mesh.ylabel, da) + xarray_support.color_limits(mesh.z, kwargs) + method(mesh.x, mesh.y, mesh.z, layer=layer, **kwargs) + if add_colorbar: + self.add_colorbar(label=mesh.zlabel, layer=layer) + + @staticmethod + def _mesh_of(da, kwargs, method): + return xarray_support.mesh_data( + da, + x=kwargs.pop("xcoord", None), + y=kwargs.pop("ycoord", None), + method=method, + ) + + def scatter(self, x, y=None, layer=0, **kwargs): """ Add a scatter plot to the subplot. + ``scatter(da)`` accepts an ``xarray.DataArray`` like :meth:`plot`. + Parameters: x (array-like): X-axis data. y (array-like): Y-axis data. @@ -403,6 +496,11 @@ def scatter(self, x, y, layer=0, **kwargs): **kwargs: Additional keyword arguments forwarded to the backend (e.g., color, marker, s, label). """ + if y is None: + if not xarray_support.is_dataarray(x): + raise TypeError("scatter(x, y) requires both x and y data") + self._plot_dataarray(self.scatter, x, layer, kwargs) + return ld = { "x": np.array(x), "y": np.array(y), @@ -580,8 +678,19 @@ def eventplot(self, positions, layer=0, **kwargs): layer, ) - def contour(self, x, y, z, layer=0, **kwargs): - """Add contour lines for a 2D scalar field.""" + def contour(self, x, y=None, z=None, layer=0, **kwargs): + """Add contour lines for a 2D scalar field. + + ``contour(da)`` accepts a 2-D ``xarray.DataArray`` like + :meth:`pcolormesh`, but adds no colorbar by default. + """ + if xarray_support.is_dataarray(x): + if y is not None or z is not None: + raise TypeError("contour(da) takes no y or z data") + self._mesh_dataarray(self.contour, x, layer, kwargs, add_colorbar=False) + return + if y is None or z is None: + raise TypeError("contour(x, y, z) requires x, y and z data") self._add( { "x": x, @@ -594,8 +703,19 @@ def contour(self, x, y, z, layer=0, **kwargs): layer, ) - def contourf(self, x, y, z, layer=0, **kwargs): - """Add filled contours for a 2D scalar field.""" + def contourf(self, x, y=None, z=None, layer=0, **kwargs): + """Add filled contours for a 2D scalar field. + + ``contourf(da)`` accepts a 2-D ``xarray.DataArray`` like + :meth:`pcolormesh`. + """ + if xarray_support.is_dataarray(x): + if y is not None or z is not None: + raise TypeError("contourf(da) takes no y or z data") + self._mesh_dataarray(self.contourf, x, layer, kwargs, add_colorbar=True) + return + if y is None or z is None: + raise TypeError("contourf(x, y, z) requires x, y and z data") self._add( { "x": x, @@ -608,8 +728,23 @@ def contourf(self, x, y, z, layer=0, **kwargs): layer, ) - def pcolormesh(self, x, y, z, layer=0, **kwargs): - """Add a pseudocolor mesh.""" + def pcolormesh(self, x, y=None, z=None, layer=0, **kwargs): + """Add a pseudocolor mesh. + + ``pcolormesh(da)`` with a 2-D ``xarray.DataArray`` uses its coordinates. + Axes are labelled from its attributes and the title from its + single-value coordinates, unless already set. ``xcoord=``/``ycoord=`` + pick the dimension or coordinate for each axis (2-D coordinates give + a curvilinear mesh), and a labelled colorbar is added unless + ``add_colorbar=False``. See :meth:`_mesh_dataarray` for colors. + """ + if xarray_support.is_dataarray(x): + if y is not None or z is not None: + raise TypeError("pcolormesh(da) takes no y or z data") + self._mesh_dataarray(self.pcolormesh, x, layer, kwargs) + return + if y is None or z is None: + raise TypeError("pcolormesh(x, y, z) requires x, y and z data") self._add( { "x": x, @@ -1381,6 +1516,25 @@ def text(self, x, y, s, layer=0, **kwargs): self._add(ld, layer) def add_imshow(self, data, layer=0, **kwargs): + """Add an image. + + ``add_imshow(da)`` with a 2-D ``xarray.DataArray`` places the image by + its evenly spaced coordinates (``origin="lower"``) and otherwise + behaves like :meth:`pcolormesh` with a DataArray. + """ + if xarray_support.is_dataarray(data): + mesh = self._mesh_of(data, dict(kwargs), "imshow") + if mesh.curvilinear: + raise ValueError( + "imshow() cannot draw on 2-D coordinates; use pcolormesh() instead" + ) + kwargs.setdefault("origin", "lower") + kwargs.setdefault( + "extent", + xarray_support.image_extent(mesh.x, mesh.y, mesh.xname, mesh.yname), + ) + self._mesh_dataarray(self._imshow_xyz, data, layer, kwargs, name="imshow") + return ld = { "data": np.asanyarray(data).copy(), "layer": layer, @@ -1389,6 +1543,13 @@ def add_imshow(self, data, layer=0, **kwargs): } self._add(ld, layer) + def _imshow_xyz(self, x, y, z, layer=0, **kwargs): + self.add_imshow(z, layer=layer, **kwargs) + + def imshow(self, data, layer=0, **kwargs): + """Matplotlib-style alias for :meth:`add_imshow`.""" + self.add_imshow(data, layer=layer, **kwargs) + def add_image(self, data, layer=0, **kwargs): """Matplotlib-style alias for ``imshow``.""" self.add_imshow(data, layer=layer, **kwargs) @@ -1431,6 +1592,7 @@ def plot_matplotlib( ax (matplotlib.axes.Axes): Axis on which to plot the lines. """ im = None + self._figure_colorbar = None self._import_rendered_artists = {} if hasattr(self, "_import_projection_artist_ids"): for name, identifiers in self._import_projection_artist_ids.items(): @@ -1454,8 +1616,8 @@ def plot_matplotlib( line["snapshot"].draw(ax, **line["kwargs"]) elif line["plot_type"] == "plot": ax.plot( - (line["x"] + self._xshift) * self._xscale, - (line["y"] + self._yshift) * self._yscale, + self._shift_x(line["x"]), + self._shift_y(line["y"]), **line["kwargs"], ) elif line["plot_type"] == "scatter": @@ -1465,14 +1627,14 @@ def plot_matplotlib( scatter_kwargs = { k: v for k, v in line["kwargs"].items() if k != "colorbar" } - ax.scatter( - (line["x"] + self._xshift) * self._xscale, - (line["y"] + self._yshift) * self._yscale, + im = ax.scatter( + self._shift_x(line["x"]), + self._shift_y(line["y"]), **scatter_kwargs, ) elif line["plot_type"] == "bar": ax.bar( - (line["x"] + self._xshift) * self._xscale, + self._shift_x(line["x"]), line["height"] * self._yscale, **line["kwargs"], ) @@ -1486,8 +1648,8 @@ def plot_matplotlib( ax.hist(line["x"], bins=line["bins"], **line["kwargs"]) elif line["plot_type"] == "step": ax.step( - (line["x"] + self._xshift) * self._xscale, - (line["y"] + self._yshift) * self._yscale, + self._shift_x(line["x"]), + self._shift_y(line["y"]), **line["kwargs"], ) elif line["plot_type"] == "stairs": @@ -1511,16 +1673,15 @@ def plot_matplotlib( ax.violinplot(line["dataset"], **line["kwargs"]) elif line["plot_type"] == "eventplot": ax.eventplot(line["positions"], **line["kwargs"]) - elif line["plot_type"] == "contour": - contour_sets.append( - ax.contour(line["x"], line["y"], line["z"], **line["kwargs"]) - ) - elif line["plot_type"] == "contourf": - contour_sets.append( - ax.contourf(line["x"], line["y"], line["z"], **line["kwargs"]) + elif line["plot_type"] in ("contour", "contourf"): + im = getattr(ax, line["plot_type"])( + line["x"], line["y"], line["z"], **_mpl_mappable_kwargs(line) ) + contour_sets.append(im) elif line["plot_type"] == "pcolormesh": - ax.pcolormesh(line["x"], line["y"], line["z"], **line["kwargs"]) + im = ax.pcolormesh( + line["x"], line["y"], line["z"], **_mpl_mappable_kwargs(line) + ) elif line["plot_type"] == "hexbin": ax.hexbin(line["x"], line["y"], **line["kwargs"]) elif line["plot_type"] == "matshow": @@ -1571,7 +1732,7 @@ def plot_matplotlib( ax.table(cellText=line["cellText"], **line["kwargs"]) elif line["plot_type"] == "gantt": tasks = line["tasks"] - start_times = (line["start_times"] + self._xshift) * self._xscale + start_times = self._shift_x(line["start_times"]) durations = line["durations"] * self._xscale y_positions = np.arange(len(tasks)) ax.barh(y_positions, durations, left=start_times, **line["kwargs"]) @@ -1589,7 +1750,7 @@ def plot_matplotlib( if start_times is None: start_times = np.zeros(n) else: - start_times = (start_times + self._xshift) * self._xscale + start_times = self._shift_x(start_times) # Calculate depths based on parent relationships for i in range(n): @@ -1641,21 +1802,21 @@ def plot_matplotlib( ax.set_ylabel("Stack Depth") elif line["plot_type"] == "fill_between": ax.fill_between( - (line["x"] + self._xshift) * self._xscale, + self._shift_x(line["x"]), ( line["y1"] if np.isscalar(line["y1"]) - else (line["y1"] + self._yshift) * self._yscale + else self._shift_y(line["y1"]) ), ( line["y2"] if np.isscalar(line["y2"]) - else (line["y2"] + self._yshift) * self._yscale + else self._shift_y(line["y2"]) ), **line["kwargs"], ) elif line["plot_type"] == "fill_betweenx": - y = (line["y"] + self._yshift) * self._yscale + y = self._shift_y(line["y"]) x1 = line["x1"] x2 = line["x2"] if np.isscalar(x1): @@ -1664,16 +1825,16 @@ def plot_matplotlib( x2 = np.full_like(y, x2, dtype=float) ax.fill_betweenx( y, - (np.asarray(x1) + self._xshift) * self._xscale, - (np.asarray(x2) + self._xshift) * self._xscale, + self._shift_x(np.asarray(x1)), + self._shift_x(np.asarray(x2)), **line["kwargs"], ) elif line["plot_type"] == "fill": ax.fill(*line["args"], **line["kwargs"]) elif line["plot_type"] == "errorbar": ax.errorbar( - (line["x"] + self._xshift) * self._xscale, - (line["y"] + self._yshift) * self._yscale, + self._shift_x(line["x"]), + self._shift_y(line["y"]), yerr=line["yerr"], xerr=line["xerr"], **line["kwargs"], @@ -1715,17 +1876,20 @@ def plot_matplotlib( elif line["plot_type"] == "imshow": im = ax.imshow( line["data"], - **line["kwargs"], + **_mpl_mappable_kwargs(line), ) elif line["plot_type"] == "patch": ax.add_patch( line["patch"], **line["kwargs"], ) + elif line["plot_type"] == "colorbar" and line.get("span_figure"): + # Drawn by the Canvas once every subplot exists. + self._figure_colorbar = (im, line["label"]) elif line["plot_type"] == "colorbar": - divider = make_axes_locatable(ax) - cax = divider.append_axes("right", size="5%", pad=0.05) - plt.colorbar(im, cax=cax, label="Potential (V)") + # Shrinks ``ax`` to make room, so the label stays inside + # the figure. + ax.figure.colorbar(im, ax=ax, label=line["label"]) if "source_artist_id" in line: created = [ @@ -1962,8 +2126,8 @@ def plot_tikzfigure(self, layers=None, verbose: bool = False) -> TikzFigure: f"{plot_type} is not supported by the tikzfigure backend" ) if plot_type == "plot": - x = (line["x"] + self._xshift) * self._xscale - y = (line["y"] + self._yshift) * self._yscale + x = self._shift_x(line["x"]) + y = self._shift_y(line["y"]) nodes = [[xi, yi] for xi, yi in zip(x, y)] tikz_figure.draw( @@ -1971,8 +2135,8 @@ def plot_tikzfigure(self, layers=None, verbose: bool = False) -> TikzFigure: **_tikz_style_kwargs(line["kwargs"]), ) elif plot_type == "scatter": - x = (line["x"] + self._xshift) * self._xscale - y = (line["y"] + self._yshift) * self._yscale + x = self._shift_x(line["x"]) + y = self._shift_y(line["y"]) style = _tikz_style_kwargs(line["kwargs"]) style.setdefault("mark", "*") style["line_width"] = 0 @@ -1989,7 +2153,7 @@ def plot_tikzfigure(self, layers=None, verbose: bool = False) -> TikzFigure: if plot_type == "bar": width = kwargs.get("width", 0.8) for x, height in zip(line["x"], line["height"]): - x = (x + self._xshift) * self._xscale + x = self._shift_x(x) height = height * self._yscale tikz_figure.draw( nodes=[ @@ -2004,7 +2168,7 @@ def plot_tikzfigure(self, layers=None, verbose: bool = False) -> TikzFigure: else: height = kwargs.get("height", 0.8) for y, width in zip(line["y"], line["width"]): - y = (y + self._yshift) * self._yscale + y = self._shift_y(y) width = width * self._xscale tikz_figure.draw( nodes=[ @@ -2017,7 +2181,7 @@ def plot_tikzfigure(self, layers=None, verbose: bool = False) -> TikzFigure: **style, ) elif plot_type == "fill_between": - x = (line["x"] + self._xshift) * self._xscale + x = self._shift_x(line["x"]) y1 = np.asarray(line["y1"]) y2 = np.broadcast_to(line["y2"], y1.shape) nodes = [[xi, yi] for xi, yi in zip(x, y1)] @@ -2028,8 +2192,8 @@ def plot_tikzfigure(self, layers=None, verbose: bool = False) -> TikzFigure: style["fill_opacity"] = kwargs.get("alpha", 0.25) tikz_figure.draw(nodes=nodes, cycle=True, **style) elif plot_type == "errorbar": - x = (line["x"] + self._xshift) * self._xscale - y = (line["y"] + self._yshift) * self._yscale + x = self._shift_x(line["x"]) + y = self._shift_y(line["y"]) style = _tikz_style_kwargs(line["kwargs"]) tikz_figure.draw(nodes=[[xi, yi] for xi, yi in zip(x, y)], **style) y_bounds = _tikz_error_bounds(line["yerr"], y) @@ -2057,15 +2221,15 @@ def plot_tikzfigure(self, layers=None, verbose: bool = False) -> TikzFigure: y = np.r_[values, values[-1]] where = "post" x, y = _tikz_step_coordinates(x, y, where=where) - x = (x + self._xshift) * self._xscale - y = (y + self._yshift) * self._yscale + x = self._shift_x(x) + y = self._shift_y(y) tikz_figure.draw( nodes=[[xi, yi] for xi, yi in zip(x, y)], **_tikz_style_kwargs(kwargs), ) elif plot_type == "stem": - x = (line["x"] + self._xshift) * self._xscale - y = (line["y"] + self._yshift) * self._yscale + x = self._shift_x(line["x"]) + y = self._shift_y(line["y"]) kwargs = line["kwargs"] style = _tikz_style_kwargs(kwargs) marker_style = dict(style) @@ -2128,7 +2292,7 @@ def plot_tikzfigure(self, layers=None, verbose: bool = False) -> TikzFigure: ) elif line["plot_type"] == "gantt": tasks = line["tasks"] - start_times = (line["start_times"] + self._xshift) * self._xscale + start_times = self._shift_x(line["start_times"]) durations = line["durations"] * self._xscale y_positions = np.arange(len(tasks)) @@ -2167,7 +2331,7 @@ def plot_tikzfigure(self, layers=None, verbose: bool = False) -> TikzFigure: if start_times is None: start_times = np.zeros(n) else: - start_times = (start_times + self._xshift) * self._xscale + start_times = self._shift_x(start_times) for i in range(n): if parents[i] is None: @@ -2345,6 +2509,16 @@ def bar_marker(kwargs): for line in self._iter_layer_lines(layers=layers): plot_type = line["plot_type"] + if plot_type in ("contour", "contourf", "pcolormesh") and ( + np.ndim(line["x"]) == 2 or np.ndim(line["y"]) == 2 + ): + if allow_unsupported: + continue + raise NotImplementedError( + f"Plotly cannot draw {plot_type} on 2-D (curvilinear) " + "coordinates; use the matplotlib backend, or pass " + "allow_unsupported=True to skip it for Plotly" + ) if plot_type in unsupported_plot_types: if allow_unsupported: continue @@ -2458,6 +2632,8 @@ def bar_marker(kwargs): marker=marker_dict, ) traces.append(trace) + if colorscale is not None: + last_heatmap_idx = len(traces) - 1 elif plot_type == "bar": kwargs = line["kwargs"] base = kwargs.get("bottom") @@ -2702,10 +2878,15 @@ def bar_marker(kwargs): y=line["y"], z=line["z"], contours=contours, - colorscale=kwargs.get("cmap", "Viridis"), + colorscale=_colormap_to_plotly_colorscale( + kwargs.get("cmap", "viridis") + ), showscale=kwargs.get("colorbar", True), + zmin=kwargs.get("vmin"), + zmax=kwargs.get("vmax"), ) ) + last_heatmap_idx = len(traces) - 1 elif plot_type == "contourf": kwargs = line["kwargs"] contours = {} @@ -2723,11 +2904,16 @@ def bar_marker(kwargs): x=line["x"], y=line["y"], z=line["z"], - colorscale=kwargs.get("cmap", "Viridis"), + colorscale=_colormap_to_plotly_colorscale( + kwargs.get("cmap", "viridis") + ), showscale=kwargs.get("colorbar", True), contours=contours, + zmin=kwargs.get("vmin"), + zmax=kwargs.get("vmax"), ) ) + last_heatmap_idx = len(traces) - 1 elif plot_type == "pcolormesh": kwargs = line["kwargs"] traces.append( @@ -2735,10 +2921,15 @@ def bar_marker(kwargs): x=line["x"], y=line["y"], z=line["z"], - colorscale=kwargs.get("cmap", "Viridis"), + colorscale=_colormap_to_plotly_colorscale( + kwargs.get("cmap", "viridis") + ), showscale=kwargs.get("colorbar", True), + zmin=kwargs.get("vmin"), + zmax=kwargs.get("vmax"), ) ) + last_heatmap_idx = len(traces) - 1 elif plot_type in ("pcolor", "pcolorfast"): # Plotly's heatmap is the closest equivalent to Matplotlib's # pseudocolor artists. The cell-centered rendering differs @@ -3526,11 +3717,13 @@ def error_spec(values, scale): ) elif plot_type == "colorbar": if last_heatmap_idx is not None: + trace = traces[last_heatmap_idx] + # Value-colored scatter markers carry their own scale. + target = trace.marker if trace.type == "scatter" else trace + target.update(showscale=True) label = line.get("label", "") or line["kwargs"].get("label", "") if label: - traces[last_heatmap_idx].update( - colorbar=dict(title=dict(text=label)) - ) + target.update(colorbar=dict(title=dict(text=label))) elif plot_type == "patch": kwargs = line["kwargs"] patch = line["patch"] @@ -3918,12 +4111,28 @@ def _symlog_inverse(self, values): def _plotext_axis_scale(self, axis: str): return self._xaxis_scale if axis == "x" else self._yaxis_scale + def _shift_x(self, values): + """Apply ``xshift``/``xscale``; values pass through unchanged by default. + + Skipping the no-op arithmetic keeps non-numeric data, such as + datetimes or category labels, plottable. + """ + if self._xshift == 0 and self._xscale == 1: + return values + return (values + self._xshift) * self._xscale + + def _shift_y(self, values): + """Apply ``yshift``/``yscale``; see :meth:`_shift_x`.""" + if self._yshift == 0 and self._yscale == 1: + return values + return (values + self._yshift) * self._yscale + def _plotext_axis_transform(self, values, axis: str): array = np.asarray(values) if axis == "x": - transformed = (array + self._xshift) * self._xscale + transformed = self._shift_x(array) else: - transformed = (array + self._yshift) * self._yscale + transformed = self._shift_y(array) if self._plotext_axis_scale(axis) == "symlog": return self._symlog_transform(transformed) return transformed diff --git a/src/maxplotlib/tests/test_xarray.py b/src/maxplotlib/tests/test_xarray.py new file mode 100644 index 0000000..0713d5e --- /dev/null +++ b/src/maxplotlib/tests/test_xarray.py @@ -0,0 +1,896 @@ +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +import numpy as np +import pytest + +xr = pytest.importorskip("xarray") + +from maxplotlib import Canvas # noqa: E402 +from maxplotlib.utils import xarray_support # noqa: E402 + + +def _line(): + t = np.linspace(0, 1, 5) + return xr.DataArray( + t**2, + dims="t", + coords={"t": ("t", t, {"long_name": "Time", "units": "s"})}, + name="energy", + attrs={"units": "J"}, + ) + + +def _mesh(): + x = np.linspace(0, 1, 4) + y = np.linspace(0, 2, 3) + return xr.DataArray( + np.arange(12.0).reshape(3, 4), + dims=("y", "x"), + coords={"x": ("x", x, {"units": "m"}), "y": y}, + name="phi", + attrs={"long_name": "Potential", "units": "V"}, + ) + + +def test_is_dataarray(): + assert xarray_support.is_dataarray(_line()) + assert not xarray_support.is_dataarray(np.zeros(3)) + assert not xarray_support.is_dataarray(_line().to_dataset()) + + +def test_labels_from_attrs(): + da = _mesh() + assert xarray_support.value_label(da) == "Potential [V]" + assert xarray_support.coord_label(da, "x") == "x [m]" + assert xarray_support.coord_label(da, "y") == "y" + assert xarray_support.value_label(xr.DataArray([1.0])) == "" + assert ( + xarray_support.value_label(xr.DataArray([1.0], attrs={"units": "K"})) == "[K]" + ) + + +def test_dimension_without_coordinate_uses_index(): + da = xr.DataArray([3.0, 4.0, 5.0], dims="i") + data = xarray_support.line_data(da) + np.testing.assert_array_equal(data.lines[0][0], [0, 1, 2]) + assert data.coord_label == "i" + + +def test_canvas_plot_dataarray_matplotlib(): + da = _line() + canvas = Canvas() + assert canvas.plot(da, color="red") is canvas + + fig, ax = canvas.get_matplotlib_figaxs() + ax = np.ravel(ax)[0] + line = ax.get_lines()[0] + np.testing.assert_allclose(line.get_xdata(), da.t.values) + np.testing.assert_allclose(line.get_ydata(), da.values) + assert ax.get_xlabel() == "Time [s]" + assert ax.get_ylabel() == "energy [J]" + plt.close(fig) + + +def test_subplot_plot_dataarray(): + canvas, ax = Canvas.subplots() + ax.plot(_line()) + fig = canvas.render(backend="plotly") + assert fig.layout.xaxis.title.text == "Time [s]" + np.testing.assert_allclose(fig.data[0].y, _line().values) + + +def test_explicit_labels_win_regardless_of_order(): + canvas, ax = Canvas.subplots() + ax.set_xlabel("before") + ax.plot(_line()) + ax.set_ylabel("after") + fig, axes = canvas.get_matplotlib_figaxs() + axes = np.ravel(axes)[0] + assert axes.get_xlabel() == "before" + assert axes.get_ylabel() == "after" + plt.close(fig) + + +def test_plot_rejects_wrong_ndim(): + canvas = Canvas() + with pytest.raises(ValueError, match="1-D DataArray.*sel"): + canvas.plot(_mesh()) + with pytest.raises(TypeError, match="both x and y"): + canvas.plot(np.arange(3)) + + +def test_plot_x_y_arrays_still_work_with_dataarrays(): + da = _line() + canvas = Canvas() + canvas.plot(da.t, da) + fig, ax = canvas.get_matplotlib_figaxs() + ax = np.ravel(ax)[0] + np.testing.assert_allclose(ax.get_lines()[0].get_ydata(), da.values) + assert ax.get_xlabel() == "" + plt.close(fig) + + +def test_pcolormesh_dataarray_matplotlib(): + da = _mesh() + canvas = Canvas() + canvas.pcolormesh(da, cmap="magma") + fig, axes = canvas.get_matplotlib_figaxs() + ax = np.ravel(axes)[0] + assert ax.get_xlabel() == "x [m]" + assert ax.get_ylabel() == "y" + mesh = ax.collections[0] + assert mesh.get_cmap().name == "magma" + np.testing.assert_allclose(mesh.get_array().reshape(3, 4), da.values) + colorbars = [a for a in fig.axes if a is not ax] + assert len(colorbars) == 1 + assert colorbars[0].get_ylabel() == "Potential [V]" + plt.close(fig) + + +def test_pcolormesh_dataarray_transpose_and_no_colorbar(): + da = _mesh() + canvas = Canvas() + canvas.pcolormesh(da, xcoord="y", add_colorbar=False) + fig, axes = canvas.get_matplotlib_figaxs() + ax = np.ravel(axes)[0] + assert ax.get_xlabel() == "y" + assert ax.get_ylabel() == "x [m]" + np.testing.assert_allclose(ax.collections[0].get_array().reshape(4, 3), da.values.T) + assert len(fig.axes) == 1 + plt.close(fig) + + +def test_pcolormesh_dataarray_plotly(): + da = _mesh() + canvas = Canvas() + canvas.pcolormesh(da) + fig = canvas.render(backend="plotly") + heatmap = fig.data[0] + np.testing.assert_allclose(heatmap.z, da.values) + np.testing.assert_allclose(heatmap.x, da.x.values) + assert heatmap.colorbar.title.text == "Potential [V]" + assert fig.layout.xaxis.title.text == "x [m]" + + +def test_pcolormesh_dataarray_errors(): + canvas = Canvas() + with pytest.raises(ValueError, match="2-D DataArray"): + canvas.pcolormesh(_line()) + with pytest.raises(ValueError, match="not a dimension or coordinate"): + canvas.pcolormesh(_mesh(), xcoord="t") + with pytest.raises(ValueError, match="different dimensions"): + canvas.pcolormesh(_mesh(), xcoord="x", ycoord="x") + with pytest.raises(TypeError, match="no y or z"): + canvas.pcolormesh(_mesh(), np.arange(3)) + with pytest.raises(TypeError, match="requires x, y and z"): + canvas.pcolormesh(np.arange(3), np.arange(3)) + + +def test_imshow_colorbar_uses_given_label(): + canvas = Canvas() + canvas.imshow(np.arange(4.0).reshape(2, 2)) + canvas.colorbar(label="density") + fig, axes = canvas.get_matplotlib_figaxs() + assert fig.axes[-1].get_ylabel() == "density" + plt.close(fig) + + +# --------------------------------------------------------------------------- +# Titles from single-value coordinates +# --------------------------------------------------------------------------- + + +def _cube(): + t = np.array([0.0, 0.5, 1.0]) + return xr.DataArray( + np.arange(36.0).reshape(3, 3, 4), + dims=("t", "y", "x"), + coords={ + "t": ("t", t, {"units": "s"}), + "x": np.linspace(0, 1, 4), + "y": np.linspace(0, 2, 3), + }, + name="phi", + attrs={"long_name": "Potential", "units": "V"}, + ) + + +def test_title_from_selected_coordinates(): + da = _cube().sel(t=0.5).isel(y=1) + assert xarray_support.title(da) == "t = 0.5 s, y = 1" + canvas = Canvas() + canvas.plot(da) + fig, axes = canvas.get_matplotlib_figaxs() + assert np.ravel(axes)[0].get_title() == "t = 0.5 s, y = 1" + plt.close(fig) + + +def test_title_formats_strings_and_datetimes(): + da = xr.DataArray( + np.zeros((2, 2)), + dims=("species", "time"), + coords={ + "species": ["ions", "electrons"], + "time": np.array(["2026-01-01", "2026-01-02"], dtype="datetime64[D]"), + }, + ) + assert xarray_support.title(da.isel(species=0, time=1)) == ( + "species = ions, time = 2026-01-02" + ) + + +def test_explicit_title_wins(): + canvas, ax = Canvas.subplots() + ax.set_title("mine") + ax.pcolormesh(_cube().isel(t=0)) + fig = canvas.render(backend="plotly") + assert "mine" in [a.text for a in fig.layout.annotations] + assert not any("t = " in (a.text or "") for a in fig.layout.annotations) + + +def test_no_title_without_scalar_coords(): + canvas, ax = Canvas.subplots() + ax.plot(_line()) + assert ax._title is None + + +# --------------------------------------------------------------------------- +# imshow, contour, contourf, scatter +# --------------------------------------------------------------------------- + + +def test_imshow_dataarray_extent_and_labels(): + da = _mesh() + canvas = Canvas() + canvas.imshow(da) + fig, axes = canvas.get_matplotlib_figaxs() + ax = np.ravel(axes)[0] + image = ax.get_images()[0] + # x: 0..1 in 4 steps of 1/3, y: 0..2 in 3 steps of 1. + np.testing.assert_allclose(image.get_extent(), [-1 / 6, 7 / 6, -0.5, 2.5]) + assert image.origin == "lower" + np.testing.assert_allclose(image.get_array(), da.values) + assert ax.get_xlabel() == "x [m]" + assert fig.axes[-1].get_ylabel() == "Potential [V]" + plt.close(fig) + + +def test_imshow_dataarray_plotly_shows_colorbar(): + canvas = Canvas() + canvas.imshow(_mesh()) + fig = canvas.render(backend="plotly") + heatmap = fig.data[0] + assert heatmap.showscale is True + assert heatmap.colorbar.title.text == "Potential [V]" + np.testing.assert_allclose(heatmap.x0, 0.0, atol=1e-12) + np.testing.assert_allclose(heatmap.dx, 1 / 3) + + +def test_imshow_dataarray_rejects_uneven_coordinates(): + da = _mesh().assign_coords(x=[0.0, 0.1, 0.5, 2.0]) + with pytest.raises(ValueError, match="'x' is not.*pcolormesh"): + Canvas().imshow(da) + + +def test_contour_and_contourf_dataarray(): + da = _mesh() + canvas, (left, right) = Canvas.subplots(ncols=2) + left.contour(da, levels=3) + right.contourf(da) + fig, axes = canvas.get_matplotlib_figaxs() + axes = np.ravel(axes) + assert axes[0].get_xlabel() == "x [m]" + # contour adds no colorbar by default, contourf does. + colorbars = [a for a in fig.axes if a not in axes] + assert [a.get_ylabel() for a in colorbars] == ["Potential [V]"] + plt.close(fig) + + +def test_contour_colorbar_flag_does_not_reach_matplotlib(): + canvas = Canvas() + canvas.contourf(_mesh(), colorbar=False) + fig, _ = canvas.get_matplotlib_figaxs() + plt.close(fig) + + +def test_scatter_dataarray(): + da = _line() + canvas = Canvas() + canvas.scatter(da, color="k") + fig, axes = canvas.get_matplotlib_figaxs() + ax = np.ravel(axes)[0] + np.testing.assert_allclose(ax.collections[0].get_offsets()[:, 1], da.values) + assert ax.get_ylabel() == "energy [J]" + plt.close(fig) + + +# --------------------------------------------------------------------------- +# hue +# --------------------------------------------------------------------------- + + +def _species(): + t = np.linspace(0, 1, 5) + return xr.DataArray( + np.stack([t, 2 * t]), + dims=("species", "t"), + coords={"species": ["ions", "electrons"], "t": ("t", t, {"units": "s"})}, + name="density", + ) + + +def test_plot_hue_draws_one_labelled_line_per_value(): + da = _species() + canvas = Canvas() + canvas.plot(da, hue="species") + fig, axes = canvas.get_matplotlib_figaxs() + ax = np.ravel(axes)[0] + lines = ax.get_lines() + assert [line.get_label() for line in lines] == [ + "species = ions", + "species = electrons", + ] + np.testing.assert_allclose(lines[1].get_ydata(), da.sel(species="electrons")) + assert ax.get_legend() is not None + assert ax.get_xlabel() == "t [s]" + plt.close(fig) + + +def test_hue_on_first_dimension_and_without_coordinate(): + da = _species().T.drop_vars("species") + data = xarray_support.line_data(da, hue="species") + assert [label for _, _, label in data.lines] == ["species = 0", "species = 1"] + positions, values, _ = data.lines[1] + np.testing.assert_allclose(values, 2 * positions) + + +def test_hue_errors_and_add_legend(): + canvas = Canvas() + with pytest.raises(ValueError, match="hue="): + canvas.plot(_species()) + with pytest.raises(ValueError, match="hue= needs a 2-D"): + canvas.plot(_line(), hue="t") + with pytest.raises(TypeError, match="label="): + canvas.plot(_species(), hue="species", label="x") + canvas, ax = Canvas.subplots() + ax.scatter(_species(), hue="species", add_legend=False) + assert not ax._legend + fig = canvas.render(backend="plotly") + assert [trace.name for trace in fig.data] == [ + "species = ions", + "species = electrons", + ] + + +# --------------------------------------------------------------------------- +# facets +# --------------------------------------------------------------------------- + + +def test_facet_col_wrap_shares_scale_and_colorbar(): + da = _cube() + canvas, axes = Canvas.facet(da, col="t", col_wrap=2) + assert len(axes) == 2 and len(axes[0]) == 2 + assert axes[1][1] is None + fig, mpl_axes = canvas.get_matplotlib_figaxs() + assert not mpl_axes[1][1].get_visible() + meshes = [mpl_axes[r][c].collections[0] for r, c in [(0, 0), (0, 1), (1, 0)]] + assert {mesh.norm.vmin for mesh in meshes} == {0.0} + assert {mesh.norm.vmax for mesh in meshes} == {35.0} + assert [ax.get_title() for ax in fig.axes[:4] if ax.get_visible()] == [ + "t = 0 s", + "t = 0.5 s", + "t = 1 s", + ] + # One figure-wide colorbar: 4 grid axes + 1 colorbar axis. + assert len(fig.axes) == 5 + assert fig.axes[-1].get_ylabel() == "Potential [V]" + # Outer labels only: (0, 1) has nothing below it, so keeps its x label. + assert mpl_axes[0][0].get_xlabel() == "" + assert mpl_axes[0][1].get_xlabel() == "x" + assert mpl_axes[1][0].get_ylabel() == "y" + assert mpl_axes[0][1].get_ylabel() == "" + plt.close(fig) + + +def test_facet_plotly_single_colorbar(): + canvas, _ = Canvas.facet(_cube(), col="t") + fig = canvas.render(backend="plotly") + assert [trace.showscale for trace in fig.data] == [True, False, False] + assert {trace.zmin for trace in fig.data} == {0.0} + assert {trace.zmax for trace in fig.data} == {35.0} + assert fig.data[0].colorbar.title.text == "Potential [V]" + + +def test_facet_row_and_col_lines(): + da = _cube() + canvas, axes = Canvas.facet(da.isel(y=0), row="t", kind="plot") + assert len(axes) == 3 and len(axes[0]) == 1 + # y = 0 is shared by every panel, so it goes in the figure title. + assert axes[2][0]._title == "t = 1 s" + assert canvas._suptitle == "y = 0" + canvas, axes = Canvas.facet(da, row="t", col="y", kind="plot", color="k") + assert len(axes) == 3 and len(axes[0]) == 3 + fig, _ = canvas.get_matplotlib_figaxs() + assert len(fig.axes) == 9 # no colorbar for lines + plt.close(fig) + + +def test_facet_lines_share_y_range_and_hide_inner_ticks(): + da = _cube().isel(y=0) # values 0..27 over t + canvas, axes = Canvas.facet(da, col="t", kind="plot") + fig, mpl_axes = canvas.get_matplotlib_figaxs() + limits = {tuple(ax.get_ylim()) for ax in np.ravel(mpl_axes)} + assert len(limits) == 1 + np.testing.assert_allclose(limits.pop(), (-1.35, 28.35)) # data range + 5 % + assert mpl_axes[0][0].yaxis.get_tick_params()["labelleft"] + assert not mpl_axes[0][1].yaxis.get_tick_params()["labelleft"] + plt.close(fig) + + canvas, axes = Canvas.facet(da, col="t", kind="plot", sharey=False) + fig, mpl_axes = canvas.get_matplotlib_figaxs() + assert len({tuple(ax.get_ylim()) for ax in np.ravel(mpl_axes)}) == 3 + assert mpl_axes[0][1].get_ylabel() == "Potential [V]" + plt.close(fig) + + +def test_facet_contourf_shares_levels_and_can_skip_colorbar(): + canvas, axes = Canvas.facet(_cube(), col="t", kind="contourf", add_colorbar=False) + fig, mpl_axes = canvas.get_matplotlib_figaxs() + levels = [ax.collections[0].levels for ax in np.ravel(mpl_axes)] + np.testing.assert_allclose(levels[0], levels[2]) + assert len(fig.axes) == 3 + plt.close(fig) + + +def test_facet_errors(): + da = _cube() + with pytest.raises(ValueError, match="col= and/or row="): + Canvas.facet(da) + with pytest.raises(ValueError, match="not a dimension"): + Canvas.facet(da, col="z") + with pytest.raises(ValueError, match="kind must be"): + Canvas.facet(da, col="t", kind="bar") + with pytest.raises(ValueError, match="col_wrap"): + Canvas.facet(da, row="t", col_wrap=2) + with pytest.raises(ValueError, match="2-D DataArray"): + Canvas.facet(da.isel(x=0), col="t") # leaves 1-D (y) panels + + +# --------------------------------------------------------------------------- +# pint units +# --------------------------------------------------------------------------- + + +def test_pint_units_label_and_values(): + pint = pytest.importorskip("pint") + ureg = pint.UnitRegistry() + t = np.linspace(0, 1, 4) + da = xr.DataArray( + ureg.Quantity(t * 3.0, "m/s"), + dims="t", + coords={"t": ("t", t, {"units": "s"})}, + name="speed", + ) + assert xarray_support.value_label(da) == "speed [m/s]" + canvas = Canvas() + canvas.plot(da) + fig, axes = canvas.get_matplotlib_figaxs() + ax = np.ravel(axes)[0] + np.testing.assert_allclose(ax.get_lines()[0].get_ydata(), t * 3.0) + assert ax.get_ylabel() == "speed [m/s]" + plt.close(fig) + + +def test_pint_units_override_attrs_and_mesh(): + pint = pytest.importorskip("pint") + ureg = pint.UnitRegistry() + da = _mesh().copy(data=ureg.Quantity(_mesh().values, "kV")) + # attrs still say "V"; the quantity is authoritative. + assert xarray_support.value_label(da) == "Potential [kV]" + x, y, z, *_ = xarray_support.mesh_data(da) + assert type(z) is np.ndarray + np.testing.assert_allclose(z, _mesh().values) + + +# --------------------------------------------------------------------------- +# da.maxplot accessor +# --------------------------------------------------------------------------- + + +def test_accessor_methods_return_canvases(): + import maxplotlib.xarray # noqa: F401 + + canvas = _line().maxplot.line(color="red") + assert isinstance(canvas, Canvas) + fig, axes = canvas.get_matplotlib_figaxs() + ax = np.ravel(axes)[0] + assert ax.get_lines()[0].get_color() == "red" + assert ax.get_xlabel() == "Time [s]" + plt.close(fig) + + for method in ("pcolormesh", "imshow", "contour", "contourf"): + canvas = getattr(_mesh().maxplot, method)() + fig = canvas.render(backend="plotly") + np.testing.assert_allclose(fig.data[0].z, _mesh().values) + + canvas = _species().maxplot.scatter(hue="species") + assert len(canvas.render(backend="plotly").data) == 2 + + +def test_accessor_canvas_kwargs_and_facets(): + import maxplotlib.xarray # noqa: F401 + + canvas = _mesh().maxplot.pcolormesh(canvas_kwargs={"fontsize": 14}) + assert canvas.fontsize == 14 + + canvas = _cube().maxplot.pcolormesh(col="t", col_wrap=2, cmap="magma") + assert (canvas.nrows, canvas.ncols) == (2, 2) + fig = canvas.render(backend="plotly") + assert [trace.showscale for trace in fig.data] == [True, False, False] + + canvas = _cube().isel(y=0).maxplot.line(col="t", sharey=False) + assert canvas.ncols == 3 + + +def test_accessor_call_picks_kind(): + import maxplotlib.xarray # noqa: F401 + + fig = _line().maxplot().render(backend="plotly") + assert fig.data[0].type == "scatter" + fig = _mesh().maxplot().render(backend="plotly") + assert fig.data[0].type == "heatmap" + fig = _species().maxplot(hue="species").render(backend="plotly") + assert [trace.type for trace in fig.data] == ["scatter", "scatter"] + canvas = _cube().maxplot(col="t") + assert canvas.ncols == 3 + with pytest.raises(ValueError, match="1-D or 2-D"): + _cube().maxplot() + + +def test_maxplotlib_does_not_import_xarray(): + import subprocess + import sys + + code = "import sys, maxplotlib; assert 'xarray' not in sys.modules" + subprocess.run([sys.executable, "-c", code], check=True) + + +# --------------------------------------------------------------------------- +# Non-numeric coordinates +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("kind", ["line", "scatter"]) +def test_datetime_and_string_coordinates(kind): + import maxplotlib.xarray # noqa: F401 + + dates = np.arange("2026-01-01", "2026-01-06", dtype="datetime64[D]") + for da in ( + xr.DataArray(np.arange(5.0), dims="time", coords={"time": dates}), + xr.DataArray([1.0, 2, 3], dims="species", coords={"species": ["a", "b", "c"]}), + ): + canvas = getattr(da.maxplot, kind)() + fig, _ = canvas.get_matplotlib_figaxs() + plt.close(fig) + np.testing.assert_array_equal( + canvas.render(backend="plotly").data[0].x, da[da.dims[0]].values + ) + + +def test_shift_and_scale_still_apply_to_numbers(): + canvas, ax = Canvas.subplots() + ax._xshift, ax._xscale = 1.0, 2.0 + ax.plot([0.0, 1.0], [0.0, 1.0]) + fig, axes = canvas.get_matplotlib_figaxs() + np.testing.assert_allclose(axes[0][0].get_lines()[0].get_xdata(), [2.0, 4.0]) + plt.close(fig) + + +# --------------------------------------------------------------------------- +# Choosing coordinates: xcoord=/ycoord= (x=/y= in the accessor) +# --------------------------------------------------------------------------- + + +def _profile(): + z = np.linspace(0, 10, 6) + return xr.DataArray( + np.linspace(300, 250, 6), + dims="i", + coords={ + "z": ("i", z, {"long_name": "Height", "units": "km"}), + "p": ("i", 1000 - 50 * z, {"units": "hPa"}), + }, + name="T", + attrs={"units": "K"}, + ) + + +def test_line_against_non_dimension_coordinate(): + da = _profile() + canvas = Canvas() + canvas.plot(da, xcoord="p") + fig, axes = canvas.get_matplotlib_figaxs() + ax = np.ravel(axes)[0] + np.testing.assert_allclose(ax.get_lines()[0].get_xdata(), da.p.values) + assert ax.get_xlabel() == "p [hPa]" + plt.close(fig) + + +def test_vertical_profile_with_ycoord(): + da = _profile() + canvas = Canvas() + canvas.plot(da, ycoord="z") + fig, axes = canvas.get_matplotlib_figaxs() + ax = np.ravel(axes)[0] + line = ax.get_lines()[0] + np.testing.assert_allclose(line.get_xdata(), da.values) + np.testing.assert_allclose(line.get_ydata(), da.z.values) + assert (ax.get_xlabel(), ax.get_ylabel()) == ("T [K]", "Height [km]") + plt.close(fig) + + +def test_line_coordinate_errors(): + canvas = Canvas() + with pytest.raises(ValueError, match="not both"): + canvas.plot(_profile(), xcoord="z", ycoord="p") + with pytest.raises(ValueError, match="not a dimension or coordinate"): + canvas.plot(_profile(), xcoord="q") + with pytest.raises(ValueError, match="other than hue"): + canvas.plot(_species(), hue="species", xcoord="species") + + +def _curvilinear(): + r = np.linspace(1, 2, 4)[:, None] + theta = np.linspace(0, np.pi / 2, 5)[None, :] + return xr.DataArray( + np.arange(20.0).reshape(4, 5), + dims=("e1", "e2"), + coords={ + "R": (("e1", "e2"), r * np.cos(theta), {"units": "m"}), + "Z": (("e1", "e2"), r * np.sin(theta), {"units": "m"}), + }, + name="n", + ) + + +def test_pcolormesh_on_2d_coordinates(): + da = _curvilinear() + canvas = Canvas() + canvas.pcolormesh(da, xcoord="R", ycoord="Z") + fig, axes = canvas.get_matplotlib_figaxs() + ax = np.ravel(axes)[0] + mesh = ax.collections[0] + # Matplotlib turns the cell centres into a (5, 6) grid of cell corners. + assert mesh.get_coordinates().shape == (5, 6, 2) + np.testing.assert_allclose(mesh.get_array().reshape(4, 5), da.values) + # The edges extend half a cell beyond the R values (0 to 2 m). + low, high = ax.dataLim.intervalx + assert low < da.R.min() and high > da.R.max() + assert (ax.get_xlabel(), ax.get_ylabel()) == ("R [m]", "Z [m]") + plt.close(fig) + # Transposed coordinates are broadcast to the data's dimension order. + t = da.assign_coords(R=da.R.T) + mesh_data = xarray_support.mesh_data(t, x="R", y="Z") + np.testing.assert_allclose(mesh_data.x, da.R.values) + + +def test_2d_coordinate_errors_and_plotly(): + da = _curvilinear() + with pytest.raises(ValueError, match="give both"): + Canvas().pcolormesh(da, xcoord="R") + with pytest.raises(ValueError, match="imshow.*pcolormesh"): + Canvas().imshow(da, xcoord="R", ycoord="Z") + canvas = Canvas() + canvas.contourf(da, xcoord="R", ycoord="Z") + with pytest.raises(NotImplementedError, match="curvilinear"): + canvas.render(backend="plotly") + canvas.render(backend="plotly", allow_unsupported=True) + + +def test_mesh_with_one_axis_coordinate(): + da = _mesh().assign_coords(xc=("x", np.arange(4) * 10.0)) + data = xarray_support.mesh_data(da, x="xc") + assert (data.xname, data.yname) == ("xc", "y") + np.testing.assert_allclose(data.x, [0, 10, 20, 30]) + + +def test_accessor_x_and_y(): + import maxplotlib.xarray # noqa: F401 + + fig, axes = _profile().maxplot.line(y="z").get_matplotlib_figaxs() + assert np.ravel(axes)[0].get_ylabel() == "Height [km]" + plt.close(fig) + fig, axes = _curvilinear().maxplot.pcolormesh(x="R", y="Z").get_matplotlib_figaxs() + assert np.ravel(axes)[0].get_xlabel() == "R [m]" + plt.close(fig) + + +# --------------------------------------------------------------------------- +# Color defaults: centering and robust limits +# --------------------------------------------------------------------------- + + +def test_color_limits(): + signed = np.array([-1.0, 0.5, 3.0]) + assert xarray_support.color_limits(signed, {}) == { + "cmap": "RdBu_r", + "vmin": -3.0, + "vmax": 3.0, + } + assert xarray_support.color_limits(signed, {"cmap": "magma"})["cmap"] == "magma" + assert "cmap" not in xarray_support.color_limits(signed, {"colors": "k"}) + assert xarray_support.color_limits(signed, {"center": False}) == {} + assert xarray_support.color_limits(np.arange(3.0), {}) == {} + assert xarray_support.color_limits(np.arange(3.0), {"center": 1.5}) == { + "cmap": "RdBu_r", + "vmin": 0.0, + "vmax": 3.0, + } + both = {"vmin": -1, "vmax": 5} + assert xarray_support.color_limits(signed, dict(both)) == both + robust = xarray_support.color_limits(np.r_[np.arange(100.0), 1e6], {"robust": True}) + assert robust["vmax"] < 1e3 + assert "norm" in xarray_support.color_limits(signed, {"norm": "x"}) + + +def test_signed_data_gets_diverging_colormap_in_both_backends(): + da = _mesh() - 5.5 + canvas = Canvas() + canvas.pcolormesh(da) + fig, axes = canvas.get_matplotlib_figaxs() + mesh = np.ravel(axes)[0].collections[0] + assert mesh.get_cmap().name == "RdBu_r" + assert (mesh.norm.vmin, mesh.norm.vmax) == (-5.5, 5.5) + plt.close(fig) + heatmap = canvas.render(backend="plotly").data[0] + assert (heatmap.zmin, heatmap.zmax) == (-5.5, 5.5) + # Matplotlib's RdBu_r runs blue -> red; Plotly's own "RdBu_r" is reversed, + # so the colorscale must be sampled from Matplotlib. + low, high = heatmap.colorscale[0][1], heatmap.colorscale[-1][1] + assert low.startswith("#05") and high.startswith("#67") + + +# --------------------------------------------------------------------------- +# Colorbar layout +# --------------------------------------------------------------------------- + + +def test_colorbar_label_stays_inside_figure(): + canvas = Canvas() + canvas.pcolormesh(_mesh()) + fig, _ = canvas.get_matplotlib_figaxs() + fig.canvas.draw() + colorbar_axes = fig.axes[-1] + label = colorbar_axes.yaxis.label.get_window_extent() + assert label.x1 <= fig.bbox.x1 + plt.close(fig) + + +# --------------------------------------------------------------------------- +# Facet titles, legends and colors +# --------------------------------------------------------------------------- + + +def test_hue_lines_get_consistent_colors(): + canvas, axes = Canvas.facet( + xr.concat([_species(), 2 * _species()], dim="run"), + col="run", + kind="plot", + hue="species", + ) + fig = canvas.render(backend="plotly") + colors = [trace.line.color for trace in fig.data] + assert colors[:2] == colors[2:] + # One legend entry per species, grouped across panels. + assert [trace.showlegend for trace in fig.data] == [True, True, False, False] + assert {trace.legendgroup for trace in fig.data} == { + "species = ions", + "species = electrons", + } + + +def test_facet_without_coordinate_titles_by_index(): + da = _cube().drop_vars("t") + canvas, axes = Canvas.facet(da, col="t") + assert [ax._title for ax in axes[0]] == ["t = 0", "t = 1", "t = 2"] + + +def test_facet_vertical_profiles_share_value_axis(): + da = xr.concat([_profile(), _profile() + 10], dim="run") + canvas, axes = Canvas.facet(da, col="run", kind="plot", ycoord="z") + assert axes[0][0]._xmin == axes[0][1]._xmin + assert axes[0][0]._ymin is None + + +def test_facet_robust_and_center(): + da = _cube() - 10 + canvas, axes = Canvas.facet(da, col="t", center=False) + heatmaps = canvas.render(backend="plotly").data + assert {trace.zmin for trace in heatmaps} == {-10.0} + canvas, axes = Canvas.facet(da, col="t") + assert {trace.zmin for trace in canvas.render(backend="plotly").data} == {-25.0} + + +# --------------------------------------------------------------------------- +# Dataset accessor +# --------------------------------------------------------------------------- + + +def _dataset(): + t = np.linspace(0, 1, 5) + return xr.Dataset( + { + "n": (("species", "t"), np.stack([t + 1, 2 * t + 1]), {"units": "m^-3"}), + "T": (("species", "t"), np.stack([10 * t, 20 * t]), {"units": "eV"}), + "q": (("species", "t"), np.stack([-t, t]), {"long_name": "Charge"}), + }, + coords={"species": ["ions", "electrons"], "t": t}, + ) + + +def test_dataset_line_and_scatter_with_dimension_hue(): + import maxplotlib.xarray # noqa: F401 + + ds = _dataset() + fig = ds.maxplot.line(x="n", y="T", hue="species").render(backend="plotly") + assert [trace.name for trace in fig.data] == [ + "species = ions", + "species = electrons", + ] + np.testing.assert_allclose(fig.data[1].x, ds.n.sel(species="electrons")) + assert fig.layout.xaxis.title.text == "n [m^-3]" + canvas = ds.isel(species=0).maxplot.scatter(x="n", y="T") + fig, axes = canvas.get_matplotlib_figaxs() + assert np.ravel(axes)[0].get_title() == "species = ions" + plt.close(fig) + + +def test_dataset_scatter_colored_by_variable(): + import maxplotlib.xarray # noqa: F401 + + ds = _dataset() + canvas = ds.maxplot.scatter(x="n", y="T", hue="q") + fig, axes = canvas.get_matplotlib_figaxs() + points = np.ravel(axes)[0].collections[0] + assert len(points.get_offsets()) == 10 + assert points.get_cmap().name == "RdBu_r" # q has both signs + assert fig.axes[-1].get_ylabel() == "Charge" + plt.close(fig) + marker = canvas.render(backend="plotly").data[0].marker + assert marker.showscale is True + assert marker.colorbar.title.text == "Charge" + + +def test_dataset_errors(): + import maxplotlib.xarray # noqa: F401 + + ds = _dataset() + with pytest.raises(ValueError, match="not a variable"): + ds.maxplot.scatter(x="nope", y="T") + with pytest.raises(ValueError, match="facets are not supported"): + ds.maxplot.scatter(x="n", y="T", hue="q", col="species") + with pytest.raises(ValueError, match="not all dims"): + ds.assign(u=("u", [1.0, 2.0])).maxplot.line(x="u", y="T", hue="species") + + +def test_subplot_imshow_alias(): + canvas, ax = Canvas.subplots() + ax.imshow(_mesh()) + fig, axes = canvas.get_matplotlib_figaxs() + assert len(np.ravel(axes)[0].get_images()) == 1 + plt.close(fig) + + +def test_facet_figure_title_does_not_overlap_panel_titles(): + canvas, _ = Canvas.facet(_cube().isel(y=0), col="t", kind="plot") + fig, axes = canvas.get_matplotlib_figaxs() + fig.canvas.draw() + renderer = fig.canvas.get_renderer() + suptitle = fig._suptitle.get_window_extent(renderer) + for ax in np.ravel(axes): + assert ax.title.get_window_extent(renderer).y1 <= suptitle.y0 + plt.close(fig) diff --git a/src/maxplotlib/utils/xarray_support.py b/src/maxplotlib/utils/xarray_support.py new file mode 100644 index 0000000..a36438d --- /dev/null +++ b/src/maxplotlib/utils/xarray_support.py @@ -0,0 +1,346 @@ +"""Optional support for plotting labelled ``xarray.DataArray`` objects. + +xarray is never imported here: an object can only be a ``DataArray`` if the +caller has already imported xarray, so detection looks it up in +``sys.modules``. Install it with ``pip install maxplotlibx[xarray]``. + +Units come from the ``units`` attribute, or from the data itself when it is a +pint ``Quantity`` (as with pint-xarray). +""" + +import sys +from typing import NamedTuple + +import numpy as np + + +def is_dataarray(obj) -> bool: + """Return True if ``obj`` is an ``xarray.DataArray``.""" + xr = sys.modules.get("xarray") + return xr is not None and isinstance(obj, xr.DataArray) + + +def _is_quantity(data) -> bool: + return hasattr(data, "magnitude") and hasattr(data, "units") + + +def _pint_units(data) -> str: + try: + return format(data.units, "~P") + except (TypeError, ValueError): + return str(data.units) + + +def magnitude(variable): + """Plain numpy values of a DataArray or coordinate, stripping pint units.""" + data = variable.data + return np.asarray(data.magnitude if _is_quantity(data) else variable.values) + + +def units(variable) -> str: + """Units of a DataArray or coordinate: pint units first, then ``attrs``.""" + if _is_quantity(variable.data): + return _pint_units(variable.data) + return str(variable.attrs.get("units", "")) + + +def _label(variable) -> str: + """Build ``"long_name [units]"`` from CF-style attributes.""" + attrs = variable.attrs + label = attrs.get("long_name") or attrs.get("standard_name") or variable.name + label = "" if label is None else str(label) + unit = units(variable) + if unit: + return f"{label} [{unit}]" if label else f"[{unit}]" + return label + + +def value_label(da) -> str: + """Label for the values of ``da``.""" + return _label(da) + + +def coord_label(da, dim) -> str: + """Label for the coordinate ``dim`` of ``da`` (the dimension name if none).""" + if dim in da.coords: + return _label(da.coords[dim]) + return str(dim) + + +def coord_values(da, dim): + """Values of the coordinate ``dim``, or ``0..n-1`` if it has none.""" + if dim in da.coords: + return magnitude(da.coords[dim]) + return np.arange(da.sizes[dim]) + + +def format_value(value) -> str: + """Short text for a single coordinate value.""" + value = np.asarray(value) + if np.issubdtype(value.dtype, np.floating): + return f"{float(value):.4g}" + if np.issubdtype(value.dtype, np.datetime64): + return np.datetime_as_string(value, unit="auto") + return str(value.item() if value.ndim == 0 else value) + + +def _coord_text(name, coord, value) -> str: + unit = units(coord) + text = f"{name} = {format_value(value)}" + return f"{text} {unit}" if unit else text + + +def title(da, exclude=()) -> str: + """Title from the single-value coordinates left by ``sel``/``isel``. + + For example ``da.sel(t=1.0)`` gives ``"t = 1 s"`` if ``t`` has units ``s``. + Coordinates in ``exclude`` are left out. + """ + return ", ".join( + _coord_text(name, coord, magnitude(coord)) + for name, coord in da.coords.items() + if coord.ndim == 0 and name not in exclude + ) + + +def _check_dims(da, *dims): + for dim in dims: + if dim is not None and dim not in da.dims: + raise ValueError(f"{dim!r} is not a dimension of {da.dims}") + + +def _resolve(da, name): + """Return ``(dims, variable)`` for a dimension or coordinate name. + + ``variable`` is the coordinate, or ``None`` for a dimension without one. + """ + if name in da.coords: + coord = da.coords[name] + if not set(coord.dims) <= set(da.dims): + raise ValueError(f"coordinate {name!r} has dims outside {da.dims}") + return coord.dims, coord + if name in da.dims: + return (name,), None + raise ValueError( + f"{name!r} is not a dimension or coordinate of the DataArray; " + f"dims are {da.dims}, coordinates {tuple(da.coords)}" + ) + + +def _axis(da, name): + """``(values, label)`` for a 1-D dimension or coordinate.""" + dims, coord = _resolve(da, name) + if coord is None: + return np.arange(da.sizes[name]), str(name) + return magnitude(coord), _label(coord) + + +class LineData(NamedTuple): + """Data for :func:`line_data`: one ``(positions, values, label)`` per line.""" + + lines: list + coord_label: str + value_label: str + vertical: bool + + +def line_data(da, hue=None, x=None, y=None): + """Return the positions, value arrays and labels for plotting lines. + + A 1-D array gives one unlabelled line. A 2-D array needs ``hue``, the + dimension to draw one line per value of. ``x`` names the dimension or + coordinate to plot against; ``y`` does the same but puts the coordinate + on the y-axis (``vertical`` is then True), e.g. for profiles. With + ``hue``, the coordinate may also vary along ``hue`` (2-D), giving each + line its own positions. + """ + if x is not None and y is not None: + raise ValueError("give x= or y= for a line plot, not both") + _check_dims(da, hue) + if hue is None and da.ndim != 1: + raise ValueError( + f"a 1-D DataArray is needed, got {da.ndim}-D with dims {da.dims}. " + "Select a slice first, e.g. da.sel(...) or da.isel(...), or pass " + "hue= to draw one line per value of a dimension." + ) + if hue is not None and da.ndim != 2: + raise ValueError( + f"hue= needs a 2-D DataArray, got {da.ndim}-D with dims {da.dims}" + ) + name = x if x is not None else y + if name is None: + (dim,) = [d for d in da.dims if d != hue] + name = dim + coord = None + else: + dims, coord = _resolve(da, name) + along = [d for d in dims if d != hue] + if len(along) != 1 or len(dims) > 1 + (hue is not None): + raise ValueError( + f"{name!r} must be a dimension, or a coordinate along one " + f"dimension other than hue (optionally also along hue); it " + f"has dims {dims}" + ) + (dim,) = along + if coord is not None and coord.ndim == 2: + positions = magnitude(coord.transpose(hue, dim)) + position_label = _label(coord) + else: + positions, position_label = _axis(da, name) + if hue is None: + lines = [(positions, magnitude(da), None)] + else: + hue_coord = da.coords[hue] if hue in da.coords else None + values = magnitude(da.transpose(hue, dim)) + lines = [] + for i in range(da.sizes[hue]): + if hue_coord is None: + label = f"{hue} = {i}" + else: + label = _coord_text(hue, hue_coord, magnitude(hue_coord)[i]) + line_positions = positions[i] if np.ndim(positions) == 2 else positions + lines.append((line_positions, values[i], label)) + return LineData(lines, position_label, value_label(da), y is not None) + + +class MeshData(NamedTuple): + """Data for :func:`mesh_data`; ``x``/``y`` are 2-D on curvilinear grids.""" + + x: np.ndarray + y: np.ndarray + z: np.ndarray + xlabel: str + ylabel: str + zlabel: str + xname: str + yname: str + + @property + def curvilinear(self) -> bool: + return np.ndim(self.x) == 2 + + +def mesh_data(da, x=None, y=None, method="pcolormesh"): + """Return the coordinates, values and labels for a 2-D DataArray. + + ``x``/``y`` name a dimension or a coordinate, which may be 2-D, e.g. + physical ``R(e1, e2)`` over logical dimensions. Like + ``xarray.DataArray.plot.pcolormesh``, the first dimension goes on the + y-axis and the second on the x-axis by default. + """ + if da.ndim != 2: + raise ValueError( + f"{method}() needs a 2-D DataArray, got {da.ndim}-D with dims " + f"{da.dims}. Select a slice first, e.g. da.sel(...) or da.isel(...)." + ) + if x is None and y is None: + y, x = da.dims + resolved = {name: _resolve(da, name) for name in (x, y) if name is not None} + if any(len(dims) == 2 for dims, _ in resolved.values()): + if x is None or y is None: + raise ValueError( + "give both x= and y= to plot on a 2-D coordinate " + f"({', '.join(resolved)})" + ) + return _curvilinear_mesh(da, x, y) + if x is None: + (ydim,) = resolved[y][0] + (x,) = [d for d in da.dims if d != ydim] + elif y is None: + (xdim,) = resolved[x][0] + (y,) = [d for d in da.dims if d != xdim] + (xdim,), _ = _resolve(da, x) + (ydim,), _ = _resolve(da, y) + if xdim == ydim: + raise ValueError(f"x and y must lie along different dimensions ({x!r}, {y!r})") + xvalues, xlabel = _axis(da, x) + yvalues, ylabel = _axis(da, y) + z = magnitude(da.transpose(ydim, xdim)) + return MeshData(xvalues, yvalues, z, xlabel, ylabel, value_label(da), x, y) + + +def _curvilinear_mesh(da, x, y): + def grid(name): + dims, coord = _resolve(da, name) + if coord is None: + coord = type(da)(np.arange(da.sizes[name]), dims=dims) + return magnitude(coord.broadcast_like(da).transpose(*da.dims)), coord + + xvalues, xcoord = grid(x) + yvalues, ycoord = grid(y) + return MeshData( + xvalues, + yvalues, + magnitude(da), + _label(xcoord) if xcoord.name is not None else str(x), + _label(ycoord) if ycoord.name is not None else str(y), + value_label(da), + x, + y, + ) + + +def _edges(values, name): + """Outer edges of evenly spaced cell centres.""" + if values.size == 1: + return values[0] - 0.5, values[0] + 0.5 + if not np.issubdtype(values.dtype, np.number): + raise ValueError( + f"imshow() needs a numeric coordinate {name!r}; use pcolormesh() instead" + ) + step = np.diff(values.astype(float)) + if not np.allclose(step, step[0], rtol=1e-3, atol=0): + raise ValueError( + f"imshow() needs an evenly spaced coordinate, but {name!r} is not; " + "use pcolormesh() instead" + ) + return values[0] - step[0] / 2, values[-1] + step[0] / 2 + + +def image_extent(x, y, xname="x", yname="y"): + """``(left, right, bottom, top)`` for ``imshow(..., origin="lower")``.""" + return (*_edges(x, xname), *_edges(y, yname)) + + +def color_limits(values, kwargs): + """Apply xarray's color defaults to ``kwargs`` for color-mapped ``values``. + + Pops ``robust`` and ``center``. ``robust=True`` takes ``vmin``/``vmax`` + from the 2nd and 98th percentiles. Data that crosses ``center`` (default: + 0, when it has both signs and ``vmin``/``vmax`` are not both given) gets + limits symmetric about it and, unless ``cmap`` or ``colors`` is given, + the diverging ``"RdBu_r"`` colormap. ``center=False`` turns that off. + A ``norm`` is left alone. + """ + robust = kwargs.pop("robust", False) + center = kwargs.pop("center", None) + if kwargs.get("norm") is not None: + return kwargs + values = np.asarray(values) + if not np.issubdtype(values.dtype, np.number): + return kwargs + finite = values[np.isfinite(values)] + if finite.size == 0: + return kwargs + if robust: + low, high = np.percentile(finite, [2, 98]) + else: + low, high = finite.min(), finite.max() + vmin, vmax = kwargs.get("vmin"), kwargs.get("vmax") + low = low if vmin is None else vmin + high = high if vmax is None else vmax + if center is None: + both_given = vmin is not None and vmax is not None + divergent = not both_given and low < 0 < high + center = 0.0 + else: + divergent = center is not False + if divergent: + half_range = max(abs(low - center), abs(high - center)) + low = center - half_range if vmin is None else vmin + high = center + half_range if vmax is None else vmax + if kwargs.get("colors") is None: # contour(colors=...) excludes cmap + kwargs.setdefault("cmap", "RdBu_r") + if divergent or robust: + kwargs["vmin"], kwargs["vmax"] = float(low), float(high) + return kwargs diff --git a/src/maxplotlib/xarray.py b/src/maxplotlib/xarray.py new file mode 100644 index 0000000..e95f4a0 --- /dev/null +++ b/src/maxplotlib/xarray.py @@ -0,0 +1,170 @@ +"""The ``DataArray.maxplot`` and ``Dataset.maxplot`` accessors. + +Importing this module registers them:: + + import maxplotlib.xarray # noqa: F401 + + canvas = da.maxplot.line() + canvas.show(backend="plotly") + +Every method returns a new :class:`~maxplotlib.Canvas`; choose the backend +when rendering it. ``x=``/``y=`` name the dimension or coordinate for each +axis, as in xarray. ``col=``, ``row=``, ``col_wrap=`` and ``sharey=`` lay the +data out over several subplots with :meth:`Canvas.facet`, and +``canvas_kwargs=`` is forwarded to the Canvas constructor. Everything else is +forwarded to the Canvas method of the same name (``line`` uses ``plot``). +""" + +try: + import xarray as xr +except ImportError as error: # pragma: no cover - depends on the environment + raise ImportError( + "maxplotlib.xarray needs xarray; install it with " + "`pip install maxplotlibx[xarray]`" + ) from error + +from maxplotlib.canvas.canvas import Canvas +from maxplotlib.utils import xarray_support + +__all__ = ["MaxplotAccessor", "MaxplotDatasetAccessor"] + +_FACET_KWARGS = ("col", "row", "col_wrap", "sharey") + + +@xr.register_dataarray_accessor("maxplot") +class MaxplotAccessor: + """Plot a DataArray with maxplotlib: ``da.maxplot.line()`` etc.""" + + def __init__(self, da): + self._da = da + + def __call__(self, **kwargs): + """Plot with a kind chosen from the dimensions, like ``da.plot()``. + + After removing any ``col``/``row`` dimensions, 1-D data (or 2-D data + with ``hue=``) is drawn with :meth:`line` and 2-D data with + :meth:`pcolormesh`. + """ + facet_dims = {kwargs.get("col"), kwargs.get("row")} - {None} + ndim = self._da.ndim - len(facet_dims & set(self._da.dims)) + if ndim == 1 or (ndim == 2 and "hue" in kwargs): + return self.line(**kwargs) + if ndim == 2: + return self.pcolormesh(**kwargs) + raise ValueError( + f"da.maxplot() plots 1-D or 2-D data, got dims {self._da.dims}. " + "Select a slice first, e.g. da.sel(...), or facet with col=/row=." + ) + + def line(self, **kwargs): + """Lines against the coordinate; ``hue=`` for one per value.""" + return self._draw("plot", kwargs) + + def scatter(self, **kwargs): + """Points against the coordinate; ``hue=`` for one series per value.""" + return self._draw("scatter", kwargs) + + def pcolormesh(self, **kwargs): + """A pseudocolor mesh of 2-D data, with a labelled colorbar.""" + return self._draw("pcolormesh", kwargs) + + def imshow(self, **kwargs): + """An image of 2-D data on evenly spaced coordinates.""" + return self._draw("imshow", kwargs) + + def contour(self, **kwargs): + """Contour lines of 2-D data.""" + return self._draw("contour", kwargs) + + def contourf(self, **kwargs): + """Filled contours of 2-D data, with a labelled colorbar.""" + return self._draw("contourf", kwargs) + + def _draw(self, kind, kwargs): + canvas_kwargs = kwargs.pop("canvas_kwargs", None) + # xarray's x=/y= are xcoord=/ycoord= on Canvas methods, where x and y + # are the positional data arguments. + for axis in ("x", "y"): + if axis in kwargs: + kwargs[f"{axis}coord"] = kwargs.pop(axis) + if any(name in kwargs for name in _FACET_KWARGS): + canvas, _ = Canvas.facet( + self._da, kind=kind, canvas_kwargs=canvas_kwargs, **kwargs + ) + return canvas + canvas = Canvas(**(canvas_kwargs or {})) + getattr(canvas, kind)(self._da, **kwargs) + return canvas + + +@xr.register_dataset_accessor("maxplot") +class MaxplotDatasetAccessor: + """Plot one Dataset variable against another: ``ds.maxplot.scatter(x=, y=)``.""" + + def __init__(self, ds): + self._ds = ds + + def line(self, x, y, hue=None, **kwargs): + """Lines of variable ``y`` against variable or coordinate ``x``. + + ``hue`` names a dimension to draw one line per value of. Other + arguments are as for ``DataArray.maxplot.line``, including facets. + """ + return self._as_dataarray(x, y).maxplot.line(x=x, hue=hue, **kwargs) + + def scatter(self, x, y, hue=None, **kwargs): + """Points of variable ``y`` against variable or coordinate ``x``. + + ``hue`` names a dimension, for one series per value, or a variable, + which colors the points by value with a labelled colorbar (turned off + by ``add_colorbar=False``; ``cmap``, ``robust`` and ``center`` apply). + """ + if hue is None or hue in self._ds.dims: + return self._as_dataarray(x, y).maxplot.scatter(x=x, hue=hue, **kwargs) + if any(name in kwargs for name in _FACET_KWARGS): + raise ValueError("facets are not supported with a variable as hue=") + return self._scatter_colored(x, y, hue, kwargs) + + def _variable(self, name): + if name not in self._ds.variables: + raise ValueError( + f"{name!r} is not a variable or coordinate of the Dataset; " + f"variables are {tuple(self._ds.data_vars)}" + ) + return self._ds[name] + + def _as_dataarray(self, x, y): + """``y`` with ``x`` attached as a coordinate, to plot against it.""" + xvar, yvar = self._variable(x), self._variable(y) + if x in yvar.coords: + return yvar + if not set(xvar.dims) <= set(yvar.dims): + raise ValueError( + f"{x!r} has dims {xvar.dims}, which are not all dims of " + f"{y!r} {yvar.dims}" + ) + return yvar.assign_coords({x: xvar}) + + def _scatter_colored(self, x, y, hue, kwargs): + canvas_kwargs = kwargs.pop("canvas_kwargs", None) + add_colorbar = kwargs.pop("add_colorbar", True) + xvar, yvar, cvar = xr.broadcast( + self._variable(x), self._variable(y), self._variable(hue) + ) + colors = xarray_support.magnitude(cvar).ravel() + xarray_support.color_limits(colors, kwargs) + canvas = Canvas(**(canvas_kwargs or {})) + canvas.scatter( + xarray_support.magnitude(xvar).ravel(), + xarray_support.magnitude(yvar).ravel(), + c=colors, + **kwargs, + ) + canvas.set_xlabel(xarray_support.value_label(xvar)) + canvas.set_ylabel(xarray_support.value_label(yvar)) + title = xarray_support.title(yvar) + if title: + canvas.set_title(title) + if add_colorbar: + canvas.colorbar(label=xarray_support.value_label(cvar)) + return canvas diff --git a/tutorials/tutorial_17_xarray.ipynb b/tutorials/tutorial_17_xarray.ipynb new file mode 100644 index 0000000..b2d37c0 --- /dev/null +++ b/tutorials/tutorial_17_xarray.ipynb @@ -0,0 +1,600 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Tutorial 17 – Plotting xarray data\n", + "\n", + "maxplotlib understands labelled [xarray](https://docs.xarray.dev) data. Hand it a\n", + "`DataArray` and it takes the coordinates for the axes, builds axis labels from the\n", + "`long_name` and `units` attributes, and titles the plot with the coordinates you\n", + "selected. Because the result is an ordinary `Canvas`, the same plot renders with\n", + "every backend: Matplotlib, Plotly, plotext or TikZ.\n", + "\n", + "There are two ways in:\n", + "\n", + "- **Canvas methods** such as `canvas.plot(da)`, `canvas.pcolormesh(da)` and\n", + " `Canvas.facet(da, ...)`, which fit in with the rest of maxplotlib.\n", + "- **The `.maxplot` accessor**, `da.maxplot.line()` and friends, which mirrors\n", + " xarray's own `da.plot.line()` API. Enable it with `import maxplotlib.xarray`.\n", + "\n", + "xarray is optional: install it with `pip install maxplotlibx[xarray]`.\n", + "maxplotlib never imports it for you." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import xarray as xr\n", + "\n", + "import maxplotlib.xarray # registers da.maxplot and ds.maxplot\n", + "from maxplotlib import Canvas" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## A labelled line\n", + "\n", + "A 1-D `DataArray` is plotted against its coordinate. The axis labels come from the\n", + "attributes: `long_name` (falling back to the variable name) plus `[units]`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "t = np.linspace(0, 10, 400)\n", + "energy = xr.DataArray(\n", + " np.exp(-0.2 * t) * np.cos(3 * t) ** 2,\n", + " dims=\"t\",\n", + " coords={\"t\": (\"t\", t, {\"long_name\": \"Time\", \"units\": \"s\"})},\n", + " name=\"energy\",\n", + " attrs={\"long_name\": \"Field energy\", \"units\": \"J\"},\n", + ")\n", + "\n", + "canvas = Canvas(width=\"12cm\", ratio=0.5)\n", + "canvas.plot(energy, color=\"steelblue\")\n", + "canvas.show()" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": {}, + "source": [ + "Labels you set yourself always win, whether you set them before or after\n", + "plotting the array." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "canvas = Canvas(width=\"12cm\", ratio=0.5)\n", + "canvas.set_ylabel(\"$W$ [J]\")\n", + "canvas.plot(energy, color=\"steelblue\")\n", + "canvas.show()" + ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "## Choose the backend when rendering\n", + "\n", + "`da.maxplot.line()` returns a `Canvas`; it doesn't draw anything yet. Choose\n", + "the backend when you show or save it, so the same object gives you an interactive\n", + "Plotly figure, a terminal plot or TikZ code." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "canvas = energy.maxplot.line(canvas_kwargs={\"width\": \"12cm\", \"ratio\": 0.5})\n", + "fig = canvas.show(backend=\"plotly\")\n", + "fig" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "canvas.show(backend=\"plotext\")" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "## Titles from selections\n", + "\n", + "After `sel` or `isel`, the selected values stay on the array as single-value\n", + "coordinates. maxplotlib turns them into the title, with units, unless you set one." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "x = np.linspace(0, 2 * np.pi, 80)\n", + "y = np.linspace(-1, 1, 50)\n", + "ts = np.linspace(0, 1.5, 6)\n", + "wave = xr.DataArray(\n", + " np.sin(x[None, None, :] - 2 * ts[:, None, None])\n", + " * np.exp(-3 * y[None, :, None] ** 2),\n", + " dims=(\"t\", \"y\", \"x\"),\n", + " coords={\n", + " \"t\": (\"t\", ts, {\"units\": \"s\"}),\n", + " \"y\": (\"y\", y, {\"long_name\": \"Height\", \"units\": \"m\"}),\n", + " \"x\": (\"x\", x, {\"units\": \"m\"}),\n", + " },\n", + " name=\"phi\",\n", + " attrs={\"long_name\": \"Potential\", \"units\": \"V\"},\n", + ")\n", + "\n", + "wave.sel(t=0.6).isel(y=25).maxplot.line().show()" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "## Time series and categories\n", + "\n", + "Datetime and text coordinates work as they are." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "days = np.arange(\"2026-01-01\", \"2026-03-01\", dtype=\"datetime64[D]\")\n", + "temperature = xr.DataArray(\n", + " 5\n", + " + 3 * np.sin(np.arange(days.size) / 9)\n", + " + np.random.default_rng(0).normal(0, 0.6, days.size),\n", + " dims=\"time\",\n", + " coords={\"time\": days},\n", + " name=\"temperature\",\n", + " attrs={\"units\": \"°C\"},\n", + ")\n", + "canvas = temperature.maxplot.line(canvas_kwargs={\"width\": \"12cm\", \"ratio\": 0.45})\n", + "canvas.autofmt_xdate()\n", + "canvas.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "yields = xr.DataArray(\n", + " [4.2, 3.1, 5.6, 2.4],\n", + " dims=\"crop\",\n", + " coords={\"crop\": [\"wheat\", \"barley\", \"maize\", \"oats\"]},\n", + " name=\"yield\",\n", + " attrs={\"units\": \"t/ha\"},\n", + ")\n", + "yields.maxplot.scatter(s=60).show()" + ] + }, + { + "cell_type": "markdown", + "id": "14", + "metadata": {}, + "source": [ + "## One line per value: `hue=`\n", + "\n", + "A 2-D array draws as one line per value of the `hue` dimension. The lines are\n", + "labelled `dim = value`, coloured `C0`, `C1`, … and the legend is switched on.\n", + "Pass `add_legend=False` to hide it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "density = xr.DataArray(\n", + " np.stack([np.exp(-t / 3), 0.8 * np.exp(-t / 5), 0.5 * np.exp(-t)]),\n", + " dims=(\"species\", \"t\"),\n", + " coords={\"species\": [\"ions\", \"electrons\", \"fast\"], \"t\": energy.t},\n", + " name=\"density\",\n", + " attrs={\"units\": \"m^-3\"},\n", + ")\n", + "density.maxplot.line(hue=\"species\").show()" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "## Choosing the coordinate for each axis\n", + "\n", + "With the accessor, `x=` and `y=` name a dimension or coordinate, exactly as in\n", + "xarray. On Canvas methods, where `x` and `y` are already the positional data\n", + "arguments, the same options are called `xcoord=` and `ycoord=`.\n", + "\n", + "`x=` can pick a non-dimension coordinate, and `y=` puts the coordinate on the\n", + "y-axis for a vertical profile:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "z = np.linspace(0, 12, 40)\n", + "temperature_profile = xr.DataArray(\n", + " 288 - 6.5 * np.minimum(z, 11) + 0.4 * np.sin(z),\n", + " dims=\"level\",\n", + " coords={\n", + " \"z\": (\"level\", z, {\"long_name\": \"Altitude\", \"units\": \"km\"}),\n", + " \"p\": (\n", + " \"level\",\n", + " 1013 * np.exp(-z / 8),\n", + " {\"long_name\": \"Pressure\", \"units\": \"hPa\"},\n", + " ),\n", + " },\n", + " name=\"T\",\n", + " attrs={\"long_name\": \"Temperature\", \"units\": \"K\"},\n", + ")\n", + "\n", + "canvas, (left, right) = Canvas.subplots(ncols=2, width=\"14cm\", ratio=0.45, wspace=0.35)\n", + "left.plot(temperature_profile, xcoord=\"p\")\n", + "right.plot(temperature_profile, ycoord=\"z\")\n", + "canvas.show()" + ] + }, + { + "cell_type": "markdown", + "id": "18", + "metadata": {}, + "source": [ + "## 2-D fields\n", + "\n", + "`pcolormesh`, `imshow`, `contourf` and `contour` take a 2-D array. Like xarray,\n", + "the first dimension goes on the y-axis. A labelled colorbar is added, except for\n", + "`contour`; `add_colorbar=` changes that. Other keyword arguments go to the\n", + "backend as usual, such as `aspect=\"auto\"` to stretch the image to its axes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", + "metadata": {}, + "outputs": [], + "source": [ + "field = wave.sel(t=0)\n", + "canvas, axes = Canvas.subplots(\n", + " nrows=2, ncols=2, width=\"16cm\", ratio=0.75, wspace=0.7, hspace=0.5\n", + ")\n", + "axes[0][0].pcolormesh(field)\n", + "axes[0][1].imshow(field, aspect=\"auto\")\n", + "axes[1][0].contourf(field)\n", + "axes[1][1].contour(field, add_colorbar=True)\n", + "canvas.show()" + ] + }, + { + "cell_type": "markdown", + "id": "20", + "metadata": {}, + "source": [ + "Swap the axes by naming the coordinates. Plots also combine: here are contour\n", + "lines drawn over a mesh." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21", + "metadata": {}, + "outputs": [], + "source": [ + "canvas = Canvas(width=\"10cm\", ratio=0.8)\n", + "canvas.pcolormesh(field, xcoord=\"y\", ycoord=\"x\", cmap=\"magma\", center=False)\n", + "canvas.contour(field, xcoord=\"y\", ycoord=\"x\", colors=\"white\", levels=6)\n", + "canvas.show()" + ] + }, + { + "cell_type": "markdown", + "id": "22", + "metadata": {}, + "source": [ + "## Colour defaults\n", + "\n", + "Data with both signs gets a diverging colour map (`RdBu_r`) centred on zero, as in\n", + "xarray. `center=` moves the centre and `center=False` turns this off. `robust=True`\n", + "takes the colour limits from the 2nd and 98th percentiles, so a few outliers\n", + "don't wash out the plot." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "noisy = field.copy()\n", + "noisy[10, 10] = 25 # a single outlier\n", + "\n", + "canvas, (left, right) = Canvas.subplots(ncols=2, width=\"16cm\", ratio=0.4, wspace=0.7)\n", + "left.pcolormesh(noisy)\n", + "left.set_title(\"default: the outlier sets the scale\")\n", + "right.pcolormesh(noisy, robust=True)\n", + "right.set_title(\"robust=True\")\n", + "canvas.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "24", + "metadata": {}, + "outputs": [], + "source": [ + "canvas = Canvas(width=\"8cm\", ratio=0.8)\n", + "canvas.pcolormesh(field, center=False)\n", + "canvas.set_title(\"center=False: a sequential colour map\")\n", + "canvas.show()" + ] + }, + { + "cell_type": "markdown", + "id": "25", + "metadata": {}, + "source": [ + "## Curvilinear grids\n", + "\n", + "Simulation output often lives on logical dimensions with 2-D physical coordinates,\n", + "for example `R(e1, e2)` and `Z(e1, e2)`. Name both coordinates to draw the field\n", + "where it really is. This needs the Matplotlib backend; Plotly's heatmaps and\n", + "contours only support rectilinear grids." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "26", + "metadata": {}, + "outputs": [], + "source": [ + "e1 = np.linspace(0.2, 1, 30)[:, None]\n", + "e2 = np.linspace(0, 2 * np.pi, 90)[None, :]\n", + "R = 3 + e1 * np.cos(e2)\n", + "Z = 1.6 * e1 * np.sin(e2)\n", + "pressure = xr.DataArray(\n", + " (1 - e1**2) * (1 + 0.15 * np.cos(3 * e2)),\n", + " dims=(\"e1\", \"e2\"),\n", + " coords={\n", + " \"R\": ((\"e1\", \"e2\"), R, {\"units\": \"m\"}),\n", + " \"Z\": ((\"e1\", \"e2\"), Z, {\"units\": \"m\"}),\n", + " },\n", + " name=\"p\",\n", + " attrs={\"long_name\": \"Pressure\", \"units\": \"kPa\"},\n", + ")\n", + "\n", + "canvas = pressure.maxplot.pcolormesh(\n", + " x=\"R\", y=\"Z\", canvas_kwargs={\"width\": \"8cm\", \"ratio\": 1.1}\n", + ")\n", + "canvas.set_aspect(\"equal\")\n", + "canvas.show()" + ] + }, + { + "cell_type": "markdown", + "id": "27", + "metadata": {}, + "source": [ + "## Small multiples: facets\n", + "\n", + "`col=`, `row=` and `col_wrap=` lay the data out over a grid, one subplot per value.\n", + "Colour-mapped plots share one colour scale and one colorbar. Axis and tick labels\n", + "appear on the outer subplots only, and each subplot is titled with its value." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "28", + "metadata": {}, + "outputs": [], + "source": [ + "wave.maxplot.pcolormesh(\n", + " col=\"t\", col_wrap=3, canvas_kwargs={\"width\": \"16cm\", \"ratio\": 0.6}\n", + ").show()" + ] + }, + { + "cell_type": "markdown", + "id": "29", + "metadata": {}, + "source": [ + "Line facets share their value range too, which makes panels easy to compare.\n", + "`sharey=False` gives each panel its own scale. Coordinates that all panels share\n", + "(here the selected height) move into the figure title." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "30", + "metadata": {}, + "outputs": [], + "source": [ + "wave.sel(y=0, method=\"nearest\").maxplot.line(\n", + " col=\"t\", col_wrap=3, canvas_kwargs={\"width\": \"16cm\", \"ratio\": 0.5}\n", + ").show()" + ] + }, + { + "cell_type": "markdown", + "id": "31", + "metadata": {}, + "source": [ + "The accessor methods are shortcuts for `Canvas.facet`, which you can also call\n", + "directly. Like `Canvas.subplots`, it returns the canvas and a grid of subplots,\n", + "which you can keep editing. In Plotly, a label shared across subplots gets a\n", + "single legend entry that toggles all of them." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "32", + "metadata": {}, + "outputs": [], + "source": [ + "runs = xr.concat(\n", + " [density, 1.5 * density],\n", + " dim=xr.DataArray([\"baseline\", \"heated\"], dims=\"run\", name=\"run\"),\n", + ")\n", + "canvas, axes = Canvas.facet(\n", + " runs.isel(t=slice(0, None, 4)),\n", + " col=\"run\",\n", + " kind=\"plot\",\n", + " hue=\"species\",\n", + " canvas_kwargs={\"width\": \"15cm\", \"ratio\": 0.4},\n", + ")\n", + "axes[0][0].set_yscale(\"log\")\n", + "axes[0][1].set_yscale(\"log\")\n", + "fig = canvas.show(backend=\"plotly\")\n", + "fig" + ] + }, + { + "cell_type": "markdown", + "id": "33", + "metadata": {}, + "source": [ + "## Datasets\n", + "\n", + "`ds.maxplot.line()` and `ds.maxplot.scatter()` plot one variable against another.\n", + "`hue=` may be a dimension, for one series per value, or, for scatter, another\n", + "variable that colours the points." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "34", + "metadata": {}, + "outputs": [], + "source": [ + "rng = np.random.default_rng(1)\n", + "n = rng.uniform(1, 5, (3, 60))\n", + "plasma = xr.Dataset(\n", + " {\n", + " \"n\": ((\"species\", \"sample\"), n, {\"long_name\": \"Density\", \"units\": \"1e19 m^-3\"}),\n", + " \"T\": (\n", + " (\"species\", \"sample\"),\n", + " 2 * n + rng.normal(0, 0.5, n.shape),\n", + " {\"long_name\": \"Temperature\", \"units\": \"keV\"},\n", + " ),\n", + " \"v\": (\n", + " (\"species\", \"sample\"),\n", + " rng.normal(0, 1, n.shape),\n", + " {\"long_name\": \"Flow\", \"units\": \"km/s\"},\n", + " ),\n", + " },\n", + " coords={\"species\": [\"D\", \"T\", \"e\"]},\n", + ")\n", + "\n", + "plasma.maxplot.scatter(x=\"n\", y=\"T\", hue=\"species\").show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "35", + "metadata": {}, + "outputs": [], + "source": [ + "plasma.maxplot.scatter(x=\"n\", y=\"T\", hue=\"v\").show()" + ] + }, + { + "cell_type": "markdown", + "id": "36", + "metadata": {}, + "source": [ + "## Units with pint\n", + "\n", + "If the data is a pint `Quantity`, as with\n", + "[pint-xarray](https://pint-xarray.readthedocs.io), maxplotlib reads the units from\n", + "it and plots the plain magnitudes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "37", + "metadata": {}, + "outputs": [], + "source": [ + "import pint\n", + "\n", + "ureg = pint.UnitRegistry()\n", + "speed = xr.DataArray(\n", + " ureg.Quantity(3 * np.sqrt(t), \"m/s\"),\n", + " dims=\"t\",\n", + " coords={\"t\": energy.t},\n", + " name=\"speed\",\n", + ")\n", + "speed.maxplot.line().show()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}