index_mappings is working I think, lightly tested on p dim

kushalkolar · kushalkolar · commit 7d4e42024796 · 2026-02-01T19:12:43.000-05:00
diff --git a/fastplotlib/widgets/nd_widget/_nd_positions.py b/fastplotlib/widgets/nd_widget/_nd_positions.py
@@ -184,6 +184,9 @@ def get(self, indices: tuple[Any, ...]):
         Note that we do not use __getitem__ here since the index is a tuple specifying a single integer
         index for each dimension. Slices are not allowed, therefore __getitem__ is not suitable here.
         """
+        # apply any slider index mappings
+        indices = tuple([m(i) for m, i in zip(self.index_mappings, indices)])
+
         if len(indices) > 1:
             # there are dims in addition to the n_datapoints dim
             # apply window funcs
@@ -195,7 +198,8 @@ def get(self, indices: tuple[Any, ...]):
         # TODO: window function on the `p` n_datapoints dimension
 
         if self.display_window is not None:
-            dw = self.display_window
+            # display window is interpreted using the index mapping for the `p` dim
+            dw = self.index_mappings[-1](self.display_window)
 
             if dw == 1:
                 slices = [slice(indices[-1], indices[-1] + 1)]
@@ -244,7 +248,7 @@ def get(self, indices: tuple[Any, ...]):
             apply_dims = self.datapoints_window_func[1]
             ws = self.datapoints_window_size
 
-            # apply user's window func and return
+            # apply user's window func
             # result will be of shape [n, p, 2 | 3]
             if apply_dims == "all":
                 # windows will be of shape [n, p, 1 | 2 | 3, ws]
@@ -260,11 +264,11 @@ def get(self, indices: tuple[Any, ...]):
             ).squeeze()
 
             # this reshape is required to reshape wf outputs of shape [n, p] -> [n, p, 1] only when necessary
-            graphic_data[..., : self.display_window, dims] = wf(
+            graphic_data[..., : dw, dims] = wf(
                 windows, axis=-1
-            ).reshape(graphic_data.shape[0], self.display_window, len(dims))
+            ).reshape(graphic_data.shape[0], dw, len(dims))
 
-            return graphic_data[..., : self.display_window, :]
+            return graphic_data[..., : dw, :]
 
         return graphic_data
 
@@ -285,6 +289,7 @@ def __init__(
         display_window: int = 10,
         window_funcs: tuple[WindowFuncCallable | None] | None = None,
         window_sizes: tuple[int | None] | None = None,
+        index_mappings: tuple[Callable[[Any], int] | None] | None = None,
     ):
         if issubclass(graphic, LineCollection):
             multi = True
@@ -295,6 +300,7 @@ def __init__(
             display_window=display_window,
             window_funcs=window_funcs,
             window_sizes=window_sizes,
+            index_mappings=index_mappings,
         )
 
         self._indices = tuple([0] * self._processor.n_slider_dims)
diff --git a/fastplotlib/widgets/nd_widget/_processor_base.py b/fastplotlib/widgets/nd_widget/_processor_base.py
@@ -11,6 +11,10 @@
 WindowFuncCallable = Callable[[ArrayLike, int, bool], ArrayLike]
 
 
+def identity(index: int) -> int:
+    return index
+
+
 class NDProcessor:
     def __init__(
         self,
@@ -23,7 +27,7 @@ def __init__(
         spatial_func: Callable[[ArrayProtocol], ArrayProtocol] | None = None,
     ):
         self._data = self._validate_data(data)
-        self._index_mappings = self._validate_index_mappings(index_mappings)
+        self._index_mappings = tuple(self._validate_index_mappings(index_mappings))
 
         self.window_funcs = window_funcs
         self.window_sizes = window_sizes
@@ -205,19 +209,30 @@ def spatial_func(self) -> Callable[[ArrayProtocol], ArrayProtocol] | None:
     #     pass
 
     @property
-    def index_mappings(self) -> tuple[Callable[[Any], int] | None, ...]:
+    def index_mappings(self) -> tuple[Callable[[Any], int]]:
         return self._index_mappings
 
     @index_mappings.setter
-    def index_mappings(self, maps):
-        self._index_mappings = self._validate_index_mappings(maps)
+    def index_mappings(self, maps: tuple[Callable[[Any], int] | None] | None):
+        self._index_mappings = tuple(self._validate_index_mappings(maps))
 
     def _validate_index_mappings(self, maps):
-        if maps is not None:
-            if not all([callable(m) or m is None for m in maps]):
+        if maps is None:
+            return tuple([identity] * self.n_slider_dims)
+
+        if len(maps) != self.n_slider_dims:
+            raise IndexError
+
+        _maps = list()
+        for m in maps:
+            if m is None:
+                _maps.append(identity)
+            elif callable(m):
+                _maps.append(identity)
+            else:
                 raise TypeError
 
-        return maps
+        return tuple(maps)
 
     def __getitem__(self, item: tuple[Any, ...]) -> ArrayProtocol:
         pass
