Fix misc issues

Author: oerc0122

MDANSE_GUI/Src/MDANSE_GUI/Tabs/Models/PlottingContext.py

@@ -17,12 +17,11 @@
 
 import copy
 import functools
-from collections.abc import Generator, Iterable, Sequence
 from contextlib import suppress
 from itertools import islice
 from math import prod
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal, NamedTuple, overload
+from typing import TYPE_CHECKING, NamedTuple
 
 import h5py
 import matplotlib.pyplot as mpl
@@ -32,15 +31,15 @@
 from matplotlib.colors import to_hex as mpl_to_hex
 from matplotlib.lines import lineStyles
 from matplotlib.markers import MarkerStyle
-from more_itertools import first, locate, nth, nth_product, sort_together, unzip
+from more_itertools import first, locate, nth, nth_product
 from qtpy.QtCore import QModelIndex, Qt, Signal, Slot
 from qtpy.QtGui import QColor, QStandardItem, QStandardItemModel
 
 from MDANSE.IO.IOUtils import summarise_array
 from MDANSE.MLogging import LOG
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
+    from collections.abc import Generator, Iterable, Sequence
 
 NUMBERS_FOR_SLICE = 3
 NUMBERS_FOR_RANGE = 2
@@ -619,7 +618,7 @@ def curves_vs_axis(
         x_axis = self.x_axis(axis_label)
 
         if self._data.ndim == 1:
-            yield axis_label, (x_axis, self.data)
+            yield "", (x_axis, self.data)
             return
 
         data_shape = self._data.shape
@@ -686,67 +685,94 @@ def curves_vs_axis(
 
     def planes_vs_axis(
         self,
-        axis_number: int,
-        max_limit: int = 1,
-    ) -> Generator[tuple[str, npt.NDArray[np.floating]]]:
+        main_axis: str,
+        max_limit: int = 9,
+    ) -> Generator[tuple[str, npt.NDArray[np.floating], tuple[str, str]]]:
         """Prepare for plotting 2D subsets of an ND array.
 
         Parameters
         ----------
-        axis_number : int
-            index of the axis perpendicular to the plotted array
+        main_axis : str
+            Label of the axis perpendicular to the plotted array.
 
         Yields
         ------
-        str
+        main_label : str
             Grid label.
-        npt.NDArray[np.floating]
+        image_array : npt.NDArray[np.floating]
             2D array.
-
+        axis_labels : tuple[str, ...]
+            Labels for each axis.
         """
+        main_axis_index = self.main_axis_index(main_axis)
+        other_labels = self._axis_labels(main_axis)
+
         match self._data.ndim:
             case 1:
                 pass
-            case 2 if axis_number == 1:
-                yield self._labels["medium"], self.data.T
+            case 2 if main_axis_index == 1:
+                yield self._labels["medium"], self.data.T, (main_axis, other_labels[0])
             case 2:
-                yield self._labels["medium"], self.data
+                yield self._labels["medium"], self.data, (main_axis, other_labels[0])
             case 3:
                 perpendicular_axis_name, perpendicular_axis = nth(
-                    self._axes.items(), axis_number, default=(None, None)
+                    self._axes.items(), main_axis_index, default=(None, None)
                 )
 
                 if perpendicular_axis is None:
                     return
 
-                reordered_view = np.moveaxis(self.data, axis_number, 0)
+                reordered_view = np.moveaxis(self.data, main_axis_index, 0)
 
                 for plane_number in self.curve_ind(max_limit):
+                    if plane_number > len(reordered_view):
+                        continue
+
                     yield (
                         f"{self._labels['minimal']}:{perpendicular_axis_name}={perpendicular_axis[plane_number]}",
                         reordered_view[plane_number],
+                        other_labels,
                     )
             case _:
                 raise NotImplementedError(
                     f"Cannot handle {self._data.ndim}-dimensional data."
                 )
 
-    def main_axis_index(self, main_axis: str | None, *, default: int) -> int:
+    def main_axis_index(
+        self, main_axis: str | None, *, default: int | None = None
+    ) -> int:
         """Find index of main axis.
 
         Parameters
         ----------
         main_axis : str
             Main axis name to search for.
-        default : int
+        default : int, optional
             Index if ``main_axis`` not found.
 
         Returns
         -------
         int
             Index of main axis.
+
+        Raises
+        ------
+        ValueError
+            Axis not found and no default.
         """
-        return first(locate(self._axes, pred=lambda x: x == main_axis), default)
+        ind = first(locate(self._axes, pred=main_axis.__eq__), default)
+        if ind is None:
+            raise ValueError(
+                f"Cannot find axis {main_axis} in {','.join(self._axes.keys())}"
+            )
+        return ind
+
+    def _axis_labels(self, main_axis: str) -> tuple[str] | tuple[str, str]:
+        main_axis_index = self.main_axis_index(main_axis)
+
+        return tuple(
+            label for i, label in enumerate(self._axes) if i != main_axis_index
+        )
 
 
 plotting_column_labels = [
@@ -1057,18 +1083,16 @@ def delete_dataset(self, index: QModelIndex):
         self._datasets.pop(dkey, None)
 
     def planes(
-        self, default_axis: int = 0, planes_per_dataset: int | None = None
-    ) -> Generator[tuple[PlotArgs, str, npt.NDArray[np.floating]]]:
+        self, default_axis: int | None = None, planes_per_dataset: int | None = None
+    ) -> Generator[tuple[PlotArgs, str, npt.NDArray[np.floating], tuple[str, str]]]:
         for databundle in self.datasets().values():
             ds = databundle.dataset
 
-            for label, plane in islice(
-                ds.planes_vs_axis(
-                    ds.main_axis_index(databundle.main_axis, default=default_axis)
-                ),
+            for label, plane, axis_labels in islice(
+                ds.planes_vs_axis(databundle.main_axis),
                 planes_per_dataset,
             ):
-                yield databundle, label, plane
+                yield databundle, label, plane, axis_labels
 
     def curves(
         self, curves_per_dataset: int | None = None

MDANSE_GUI/Src/MDANSE_GUI/Tabs/Plotters/Grid.py

@@ -146,27 +146,26 @@ def plot(
         self._normalisation_errors = []
         self.apply_settings(plotting_context)
 
-        nplots = min(
+        self._n_curves = min(
             sum(db.dataset.n_curves for db in plotting_context.datasets().values()),
             self._plot_limit,
         )
-        grid_size = self.grid_size(nplots)
+        grid_size = self.grid_size(self.n_curves)
         gs = self._figure.add_gridspec(*grid_size)
 
         for ind, (databundle, label, curve) in enumerate(
             islice(flatten(plotting_context.curves()), self._plot_limit)
         ):
             axes = target.add_subplot(gs[ind])
-            self._axes_titles.append(databundle.dataset._name)
+            self._axes_titles.append(f"{databundle.dataset._name} {label}")
 
             self._plot_single(
                 axes,
                 curve,
                 databundle,
-                label=label,
+                label="",
                 colour=databundle.colour,
             )
-            axes.legend()
             self._axes.append(axes)
 
         self.toggle_legend(plotting_context.use_legend)

MDANSE_GUI/Src/MDANSE_GUI/Tabs/Plotters/Grouped.py

@@ -201,14 +201,17 @@ def plot(
 
         self.height_max, self.length_max = 0.0, 0.0
         self._figure = target
+
         self._axes = []
         self._axes_titles = []
         self._backup_curves = []
         self._active_curves = []
         self._normalisation_errors = []
+
         self.apply_settings(plotting_context)
 
-        nplots = min(ilen(plotting_context.curves()), self._plot_limit)
+        self._n_curves = sum(ilen(curves) for curves in plotting_context.curves())
+        nplots = min(len(plotting_context.datasets()), self._plot_limit)
         grid_size = self.grid_size(nplots)
         gs = self._figure.add_gridspec(*grid_size)
         limits = [(0.0, 0.0, 0.0, 0.0) for _ in range(nplots)]
@@ -229,16 +232,18 @@ def plot(
 
             colours = self.colours(db.colour, ds.n_curves)
 
-            for (databundle, label, curve), colour in zip(
-                dataclump, colours, strict=True
+            for curve_ind, ((databundle, label, curve), colour) in enumerate(
+                zip(dataclump, colours, strict=True)
             ):
                 self._plot_single(
                     axes,
                     curve,
                     databundle,
+                    ind=curve_ind,
                     label=label,
                     colour=colour,
                 )
+
             axes.legend()
             self._axes.append(axes)
             limits[ind] = (*axes.get_xlim(), *axes.get_ylim())
@@ -283,6 +288,7 @@ def _plot_single(
         curve: tuple[np.ndarray, np.ndarray] | tuple[np.ndarray],
         databundle: PlotArgs,
         *,
+        ind: int,
         label: str,
         colour: tuple[float, float, float] | str,
     ):
@@ -304,7 +310,7 @@ def _plot_single(
         lines: list[Line2D] = axes.plot(
             *curve,
             linestyle=databundle.line_style,
-            label=label,
+            label=self.label(label, ind, n_curves=databundle.dataset.n_curves),
             color=colour,
         )