ENH SequencePlotter: specify columns/rows content

eelbrain/plot/_brain.py

@@ -1,9 +1,10 @@
 # Author: Christian Brodbeck <christianbrodbeck@nyu.edu>
 from dataclasses import dataclass
 from enum import Enum
-from functools import partial
+from functools import partial, reduce
 from itertools import product
 from numbers import Number
+import operator
 from typing import Any, Literal, Sequence, Union
 
 import matplotlib
@@ -652,7 +653,12 @@ def __init__(
         if column_labels:
             self.add_column_titles(column_labels)
 
-    def _add_ims(self, ims, cmap_params, cmap_data):
+    def _add_ims(
+            self,
+            ims: Sequence,  # shape = (rows, columns)
+            cmap_params,
+            cmap_data,
+    ):
         for row, column in product(range(self._n_rows), range(self._n_columns)):
             ax = self.axes[row * self._n_columns + column]
             ax.imshow(ims[row][column])
@@ -1442,10 +1448,12 @@ def plot_table(
             self,
             hemi: Literal['lh', 'rh', 'both'] = None,
             view: Union[str, Sequence] = ('lateral', 'medial'),
-            orientation: Literal['horizontal', 'vertical'] = 'horizontal',
+            orientation: Literal['horizontal', 'vertical'] = None,
             labels: Union[bool, Sequence[str]] = True,
             mode: Literal['rgb', 'rgba'] = 'rgb',
             antialiased: bool = False,
+            rows: Sequence[str] = ('view',),
+            columns: Sequence[str] = ('bin', 'hemi'),
             **kwargs,
     ) -> ImageTable:
         """Create a figure with the images
@@ -1459,14 +1467,19 @@ def plot_table(
             including parallel-view parameters ``(view, forward, up, scale)``,
             e.g., ``('lateral', 0, 10, 70)``.
         orientation
-            Direction of the time/case axis.
+            Direction of the time/case axis; overrides ``row`` and ``column``
+            settings.
         labels
             Headers for columns/rows of images (default is inferred from the data).
         mode
             Image mode (default ```'rgb'``, set to ``'rgba'`` to include alpha
             channel).
         antialiased
             Apply antialiasing to the images (default ``False``).
+        rows
+            What to alternate along the rows.
+        columns
+            What to alternate along the columns.
         ...
             Layout parameters for the figure.
 
@@ -1526,30 +1539,43 @@ def plot_table(
             bins = list(range(self._n_items()))
         else:
             bins = [None]
-        n_bins = len(bins)
-        if orientation == 'horizontal':
-            n_columns = n_bins
-            n_rows = n_views
-            transpose = False
-            column_labels = labels
-            row_labels = None
-        elif orientation == 'vertical':
-            n_columns = n_views
-            n_rows = n_bins
-            transpose = True
-            column_labels = None
-            row_labels = labels
+
+        # determine im-table layout
+        if orientation:
+            if orientation == 'vertical':
+                rows, columns = ('hemi', 'view'), ('bin',)
+            elif orientation == 'horizontal':
+                rows, columns = ('bin', 'hemi'), ('view',)
+            else:
+                raise ValueError(f"{orientation=}")
+        axis_sequences = {'bin': bins, 'hemi': hemis, 'view': views}
+        ns = {key: len(values) for key, values in axis_sequences.items()}
+        specified = set(columns).union(rows)
+        if specified != set(ns) or len(columns) + len(rows) != 3:
+            raise ValueError(f"{rows=}, {columns=}")
+        n_columns = reduce(operator.mul, (ns[key] for key in columns))
+        n_rows = reduce(operator.mul, (ns[key] for key in rows))
+        # labels
+        bin_in_column = 'bin' in columns
+        bin_axis = columns if bin_in_column else rows
+        if len(bin_axis) > 1:
+            components = [labels if key == 'bin' else ('',)*ns[key] for key in bin_axis]
+            index = bin_axis.index('bin')
+            labels = [items[index] for items in product(*components)]
+            for i in range(len(labels) - 1, 0, -1):
+                if labels[i] == labels[i - 1]:
+                    labels[i] = None
+        if bin_in_column:
+            column_labels, row_labels = labels, None
         else:
-            raise ValueError(f"orientation={orientation!r}")
+            column_labels, row_labels = None, labels
 
         figure = ImageTable(n_rows, n_columns, row_labels, column_labels, **kwargs)
 
-        im_rows = []
+        ims = {}
         cmap_params = None
         cmap_data = None
         for hemi in hemis:
-            hemi_rows = [[] for _ in views]
-
             # plot brain
             b = brain(self._source, hemi=hemi, views='lateral', w=figure._res_w, h=figure._res_h, time_label='', **self._brain_args)
 
@@ -1564,7 +1590,7 @@ def plot_table(
                 for i in bins:
                     if i is not None:
                         b.set_data_time_index(i)
-                        self._capture(b, hemi_rows, views, mode, antialiased)
+                        self._capture(b, views, mode, antialiased, ims, hemi, i)
             elif self._bin_kind in (SPLayer.ITEM, SPLayer.UNDEFINED):  # UNDEFINED: has overlays only
                 for i in bins:
                     for l in self._data:
@@ -1573,30 +1599,36 @@ def plot_table(
                             if cmap_params is None and l.plot_type == SPPlotType.DATA and l.kind == SPLayer.ITEM:
                                 cmap_params = b._get_cmap_params()
                                 cmap_data = l.data
-                    self._capture(b, hemi_rows, views, mode, antialiased)
+                    self._capture(b, views, mode, antialiased, ims, hemi, i)
                     b.remove_data()
                     b.remove_labels()
             else:
                 raise NotImplementedError(f"{self._bin_kind}")
-            im_rows += hemi_rows
             b.close()
 
-        if transpose:
-            im_rows = list(zip(*im_rows))
-        figure._add_ims(im_rows, cmap_params, cmap_data)
+        # reshape im-order
+        row_pattern = list(product(*[axis_sequences[key] for key in rows]))
+        column_pattern = list(product(*[axis_sequences[key] for key in columns]))
+        im_array = np.ndarray((n_rows, n_columns), object)
+        for i_row, key_row in enumerate(row_pattern):
+            key_dict = {key: value for key, value in zip(rows, key_row)}
+            for i_col, key_col in enumerate(column_pattern):
+                key_dict.update({key: value for key, value in zip(columns, key_col)})
+                key = tuple([key_dict[key] for key in ['hemi', 'bin', 'view']])
+                im_array[i_row, i_col] = ims[key]
+        figure._add_ims(im_array, cmap_params, cmap_data)
         return figure
 
-    def _capture(self, b, hemi_rows, views, mode, antialiased):
-        for row, view in zip(hemi_rows, views):
+    def _capture(self, b, views, mode, antialiased, ims, hemi, i):
+        for view in views:
             if isinstance(view, str):
                 b.show_view(view)
                 if self._parallel_view:
                     b.set_parallel_view(**self._parallel_view)
             else:
                 b.show_view(view[0])
                 b.set_parallel_view(*view[1:])
-            im = b.screenshot_single(mode, antialiased)
-            row.append(im)
+            ims[hemi, i, view] = b.screenshot_single(mode, antialiased)
 
 
 def connectivity(source):