From 8bdd78f9e6f91b85e187053e1d76289bd87e777a Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 29 Jan 2025 13:27:34 +0100 Subject: [PATCH 1/4] move the colorizing code to a separate function --- xdggs/plotting.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/xdggs/plotting.py b/xdggs/plotting.py index 37989fb3..174f39ce 100644 --- a/xdggs/plotting.py +++ b/xdggs/plotting.py @@ -39,6 +39,14 @@ def normalize(var, center=None): return normalizer(var.data) +def colorize(var, *, center, colormap, alpha): + from lonboard.colormap import apply_continuous_cmap + + normalized_data = normalize(var, center=center) + + return apply_continuous_cmap(normalized_data, colormap, alpha=alpha) + + def explore( arr, cell_dim="cells", @@ -48,7 +56,6 @@ def explore( ): import lonboard from lonboard import SolidPolygonLayer - from lonboard.colormap import apply_continuous_cmap from matplotlib import colormaps if len(arr.dims) != 1 or cell_dim not in arr.dims: @@ -61,10 +68,8 @@ def explore( polygons = grid_info.cell_boundaries(cell_ids, backend="geoarrow") - normalized_data = normalize(arr.variable, center=center) - colormap = colormaps[cmap] if isinstance(cmap, str) else cmap - colors = apply_continuous_cmap(normalized_data, colormap, alpha=alpha) + colors = colorize(arr.variable, center=center, alpha=alpha, colormap=colormap) table = create_arrow_table(polygons, arr) layer = SolidPolygonLayer(table=table, filled=True, get_fill_color=colors) From e649815600697a162af51ddd508d2b291aba3cd6 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 29 Jan 2025 13:28:53 +0100 Subject: [PATCH 2/4] remove the limitation on 1D data --- xdggs/plotting.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xdggs/plotting.py b/xdggs/plotting.py index 174f39ce..4b1d5024 100644 --- a/xdggs/plotting.py +++ b/xdggs/plotting.py @@ -58,9 +58,9 @@ def explore( from lonboard import SolidPolygonLayer from matplotlib import colormaps - if len(arr.dims) != 1 or cell_dim not in arr.dims: + if cell_dim not in arr.dims: raise ValueError( - f"exploration only works with a single dimension ('{cell_dim}')" + f"exploration plotting only works with a spatial dimension ('{cell_dim}')" ) cell_ids = arr.dggs.coord.data From bf5c06145f789a913bac4f29980b8cea266678ef Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 29 Jan 2025 13:45:57 +0100 Subject: [PATCH 3/4] remove the explicitly passed cell dim --- xdggs/accessor.py | 2 -- xdggs/plotting.py | 7 +++++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/xdggs/accessor.py b/xdggs/accessor.py index f7fcf839..12f2accf 100644 --- a/xdggs/accessor.py +++ b/xdggs/accessor.py @@ -200,10 +200,8 @@ def explore(self, *, cmap="viridis", center=None, alpha=None): if isinstance(self._obj, xr.Dataset): raise ValueError("does not work with Dataset objects, yet") - cell_dim = self._obj[self._name].dims[0] return explore( self._obj, - cell_dim=cell_dim, cmap=cmap, center=center, alpha=alpha, diff --git a/xdggs/plotting.py b/xdggs/plotting.py index 4b1d5024..1c1fe456 100644 --- a/xdggs/plotting.py +++ b/xdggs/plotting.py @@ -49,7 +49,6 @@ def colorize(var, *, center, colormap, alpha): def explore( arr, - cell_dim="cells", cmap="viridis", center=None, alpha=None, @@ -58,12 +57,16 @@ def explore( from lonboard import SolidPolygonLayer from matplotlib import colormaps + # guaranteed to be 1D + cell_id_coord = arr.dggs.coord + [cell_dim] = cell_id_coord.dims + if cell_dim not in arr.dims: raise ValueError( f"exploration plotting only works with a spatial dimension ('{cell_dim}')" ) - cell_ids = arr.dggs.coord.data + cell_ids = cell_id_coord.data grid_info = arr.dggs.grid_info polygons = grid_info.cell_boundaries(cell_ids, backend="geoarrow") From 32082738dbacc04b1ad358f281ccf531a6a97391 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 29 Jan 2025 14:17:53 +0100 Subject: [PATCH 4/4] support n-dimensional data by subsetting using sliders --- xdggs/plotting.py | 75 +++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 72 insertions(+), 3 deletions(-) diff --git a/xdggs/plotting.py b/xdggs/plotting.py index 1c1fe456..ef438935 100644 --- a/xdggs/plotting.py +++ b/xdggs/plotting.py @@ -1,4 +1,45 @@ +from dataclasses import dataclass +from functools import partial +from typing import Any + +import ipywidgets import numpy as np +import xarray as xr +from lonboard import Map + + +def on_slider_change(change, container): + owner = change["owner"] + dim = owner.description + + indexers = { + slider.description: slider.value + for slider in container.dimension_sliders.children + if slider.description != dim + } | {dim: change["new"]} + new_slice = container.obj.isel(indexers) + + colors = colorize(new_slice.variable, **container.colorize_kwargs) + + layer = container.map.layers[0] + layer.get_fill_color = colors + + +@dataclass +class MapContainer: + """container for the map, any control widgets and the data object""" + + dimension_sliders: ipywidgets.VBox + map: Map + obj: xr.DataArray + + colorize_kwargs: dict[str, Any] + + def render(self): + # add any additional control widgets here + control_box = ipywidgets.HBox([self.dimension_sliders]) + + return ipywidgets.VBox([self.map, control_box]) def create_arrow_table(polygons, arr, coords=None): @@ -71,10 +112,38 @@ def explore( polygons = grid_info.cell_boundaries(cell_ids, backend="geoarrow") + initial_indexers = {dim: 0 for dim in arr.dims if dim != cell_dim} + initial_arr = arr.isel(initial_indexers) + colormap = colormaps[cmap] if isinstance(cmap, str) else cmap - colors = colorize(arr.variable, center=center, alpha=alpha, colormap=colormap) + colors = colorize(initial_arr, center=center, alpha=alpha, colormap=colormap) - table = create_arrow_table(polygons, arr) + table = create_arrow_table(polygons, initial_arr) layer = SolidPolygonLayer(table=table, filled=True, get_fill_color=colors) - return lonboard.Map(layer) + map_ = lonboard.Map(layer) + + if not initial_indexers: + # 1D data + return map_ + + sliders = ipywidgets.VBox( + [ + ipywidgets.IntSlider(min=0, max=arr.sizes[dim] - 1, description=dim) + for dim in arr.dims + if dim != cell_dim + ] + ) + + container = MapContainer( + sliders, + map_, + arr, + colorize_kwargs={"alpha": alpha, "center": center, "colormap": colormap}, + ) + + # connect slider with map + for slider in sliders.children: + slider.observe(partial(on_slider_change, container=container), names="value") + + return container.render()