Merge pull request #139 from scipp/voxelviewer-kwargs

Support figure kwargs in FlatVoxelViewer

Author: nvaytet

src/ess/dream/diagnostics.py

@@ -5,6 +5,7 @@
 import math
 from collections.abc import Callable, Iterable, Mapping
 from functools import reduce
+from typing import Any
 
 import ipywidgets as ipw
 import numpy as np
@@ -24,7 +25,11 @@ class FlatVoxelViewer(ipw.VBox):
     """
 
     def __init__(
-        self, data: Mapping[str, sc.DataArray], *, rasterized: bool = True
+        self,
+        data: Mapping[str, sc.DataArray],
+        *,
+        rasterized: bool = True,
+        **kwargs: Any,
     ) -> None:
         """Create a new viewer.
 
@@ -35,14 +40,16 @@ def __init__(
         rasterized:
             If ``True``, the figure is rasterized which improves rendering
             speed but reduces resolution.
+        **kwargs:
+            Additional arguments passed to the plotting function.
         """
         self._data = self._prepare_data(data)
         self._bank_selector = _make_bank_selector(data.keys())
         self._bank = self._data[self._bank_selector.value]
 
         self._dim_selector = _DimensionSelector(self._bank.dims, self._update_view)
 
-        self._fig_kwargs = {'rasterized': rasterized}
+        self._fig_kwargs = {'rasterized': rasterized} | kwargs
         self._figure_box = ipw.HBox([self._make_figure()])
         self._bank_selector.observe(self._select_bank, names='value')
 
@@ -54,12 +61,12 @@ def __init__(
             ]
         )
 
-    def _select_bank(self, *_args: object, **_kwargs: object) -> None:
+    def _select_bank(self, *_args: Any, **_kwargs: Any) -> None:
         self._bank = self._data[self._bank_selector.value]
         self._dim_selector.set_dims(self._bank.dims)
         self._update_view()
 
-    def _update_view(self, *_args: object, **_kwargs: object) -> None:
+    def _update_view(self, *_args: Any, **_kwargs: Any) -> None:
         self._figure_box.children = [self._make_figure()]
 
     def _make_figure(self) -> FigureLike:
@@ -149,8 +156,7 @@ def _flat_voxel_figure(
     data: sc.DataArray,
     horizontal_dim: str,
     vertical_dim: str,
-    *,
-    rasterized: bool = True,
+    **kwargs: Any,
 ) -> FigureLike:
     kept_dims = {horizontal_dim, vertical_dim}
 
@@ -182,7 +188,7 @@ def _flat_voxel_figure(
     h_labels = [str(value) for value in h_coord.values]
     v_labels = [str(value) for value in v_coord.values]
 
-    fig = flat.plot(rasterized=rasterized, cbar=True)
+    fig = flat.plot(**kwargs)
 
     fig.ax.xaxis.set_ticks(ticks=h_ticks, labels=h_labels)
     fig.ax.yaxis.set_ticks(ticks=v_ticks, labels=v_labels)