Repository: hyperspy/exspy (language: python)

Issue:
Background Window Plotting not working with 2D signal
#### Describe the bug
Running the following line from the example note book:
```python
cs.isig[5.:12.5].plot(True, integration_windows=iw, background_windows=bw)
```
Results in the following error:

```python
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[15], line 1
----> 1 cs.isig[5.:12.5].plot(True, integration_windows=iw, background_windows=bw)

File ~\AppData\Local\anaconda3\lib\site-packages\exspy\signals\eds.py:933, in EDSSpectrum.plot(self, xray_lines, only_lines, only_one, background_windows, integration_windows, navigator, plot_markers, autoscale, norm, axes_manager, navigator_kwds, **kwargs)
    857 """Plot the EDS spectrum. The following markers can be added
    858 
    859 - The position of the X-ray lines and their names.
   (...)
    922 get_lines_intensity, estimate_background_windows
    923 """
    924 super().plot(
    925     navigator=navigator,
    926     plot_markers=plot_markers,
   (...)
    931     **kwargs,
    932 )
--> 933 self._plot_xray_lines(
    934     xray_lines,
    935     only_lines,
    936     only_one,
    937     background_windows,
    938     integration_windows,
    939     render_figure=False,
    940 )
    941 self._render_figure(plot=["signal_plot"])

File ~\AppData\Local\anaconda3\lib\site-packages\exspy\signals\eds.py:989, in EDSSpectrum._plot_xray_lines(self, xray_lines, only_lines, only_one, background_windows, integration_windows, render_figure)
    987 self.add_xray_lines_markers(xray_lines, render_figure=False)
    988 if background_windows is not None:
--> 989     self._add_background_windows_markers(
    990         background_windows, render_figure=False
    991     )
    992 if integration_windows is not None:
    993     if integration_windows == "auto":

File ~\AppData\Local\anaconda3\lib\site-packages\exspy\signals\eds.py:1142, in EDSSpectrum._add_background_windows_markers(self, windows_position, render_figure)
   1140 segments = np.array(segments)
   1141 lines = hs.plot.markers.Lines(segments=segments, color="black")
-> 1142 self.add_marker(lines, render_figure=False)
   1143 if render_figure:
   1144     self._render_figure(plot=["signal_plot"])

File ~\AppData\Local\anaconda3\lib\site-packages\hyperspy\signal.py:6505, in BaseSignal.add_marker(self, marker, plot_on_signal, plot_marker, permanent, plot_signal, render_figure)
   6502     raise ValueError("Markers can not be added to several signals")
   6504 m._plot_on_signal = plot_on_signal
-> 6505 m._signal = self
   6507 if plot_marker:
   6508     if self._plot is None or not self._plot.is_active:

File ~\AppData\Local\anaconda3\lib\site-packages\hyperspy\drawing\markers.py:283, in Markers._signal(self, signal)
    278         nav_shape = value.shape if is_iterating(value) else ()
    279         if (
    280             len(nav_shape) != 0
    281             and nav_shape != signal.axes_manager.navigation_shape
    282         ):
--> 283             raise ValueError(
    284                 "The shape of the variable length argument must match "
    285                 "the navigation shape of the signal."
    286             )
    287 self.__signal = signal

ValueError: The shape of the variable length argument must match the navigation shape of the signal.
```

#### Expected behavior
Should run without error and allow the user to navigate over pixels and see how the background fitting should be performing

#### Python environement:
 - eXSpy version: 0.2
 - HyperSpy version: 2.1.1
 - Python version: 3.10.9




Discussion / hints:

Fix patch:
diff --git a/exspy/signals/eds.py b/exspy/signals/eds.py
index c8e537bf7..db399d9b9 100644
--- a/exspy/signals/eds.py
+++ b/exspy/signals/eds.py
@@ -987,7 +987,7 @@ def _plot_xray_lines(
             self.add_xray_lines_markers(xray_lines, render_figure=False)
             if background_windows is not None:
                 self._add_background_windows_markers(
-                    background_windows, render_figure=False
+                    background_windows, linestyles="--", render_figure=False
                 )
             if integration_windows is not None:
                 if integration_windows == "auto":
@@ -997,7 +997,7 @@ def _plot_xray_lines(
                         windows_width=integration_windows, xray_lines=xray_lines
                     )
                 self._add_vertical_lines_groups(
-                    integration_windows, linestyle="--", render_figure=False
+                    integration_windows, render_figure=False
                 )
         # Render figure only at the end
         if render_figure:
@@ -1012,16 +1012,22 @@ def _add_vertical_lines_groups(self, position, render_figure=True, **kwargs):
         position: 2D array of float
             The position on the signal axis. Each row corresponds to a
             group.
-        kwargs
+        **kwargs : dict
             keywords argument for :class:`hyperspy.api.plot.markers.VerticalLine`
         """
-        colors = itertools.cycle(
+        position = np.array(position)
+        length = position.shape[1]
+        colors_cycle = itertools.cycle(
             np.sort(plt.rcParams["axes.prop_cycle"].by_key()["color"])
         )
+        colors = np.array(
+            [[c] * length for c, w in zip(colors_cycle, position)]
+        ).flatten()
 
-        for x, color in zip(position, colors):
-            line = hs.plot.markers.VerticalLines(offsets=x, color=color, **kwargs)
-            self.add_marker(line, render_figure=False)
+        line = hs.plot.markers.VerticalLines(
+            offsets=position.flatten(), color=colors, **kwargs
+        )
+        self.add_marker(line, render_figure=False)
         if render_figure:
             self._render_figure(plot=["signal_plot"])
 
@@ -1102,7 +1108,9 @@ def remove_xray_lines_markers(self, xray_lines, render_figure=True):
         if render_figure:
             self._render_figure(plot=["signal_plot"])
 
-    def _add_background_windows_markers(self, windows_position, render_figure=True):
+    def _add_background_windows_markers(
+        self, windows_position, render_figure=True, **kwargs
+    ):
         """
         Plot the background windows associated with each X-ray lines.
 
@@ -1121,25 +1129,50 @@ def _add_background_windows_markers(self, windows_position, render_figure=True):
         --------
         estimate_background_windows, get_lines_intensity
         """
-        self._add_vertical_lines_groups(windows_position)
-        ax = self.axes_manager.signal_axes[0]
-        segments = []
-        for bw in windows_position:
-            # TODO: test to prevent slicing bug. To be removed when fixed
-            if ax.value2index(bw[0]) == ax.value2index(bw[1]):
-                y1 = self.isig[bw[0]].data
-            else:
-                y1 = self.isig[bw[0] : bw[1]].mean(-1).data
-            if ax.value2index(bw[2]) == ax.value2index(bw[3]):
-                y2 = self.isig[bw[2]].data
-            else:
-                y2 = self.isig[bw[2] : bw[3]].mean(-1).data
-            x1 = (bw[0] + bw[1]) / 2.0
-            x2 = (bw[2] + bw[3]) / 2.0
-            segments.append([[x1, y1[0]], [x2, y2[0]]])
-        segments = np.array(segments)
-        lines = hs.plot.markers.Lines(segments=segments, color="black")
+        self._add_vertical_lines_groups(windows_position, **kwargs)
+
+        # Calculate the start and end of segments for each window
+        # nav_dim + (number of x-ray lines, vertices positions)
+        # vertices positions are (x0, y0), (x1, x2)
+        segments_ = np.zeros(
+            self.axes_manager.navigation_shape + (len(windows_position), 2, 2)
+        )
+
+        for i, bw in enumerate(windows_position):
+            # Check that all background windows are within the energy range
+            if any(v < self.axes_manager[-1].low_value for v in bw) or any(
+                v > self.axes_manager[-1].high_value for v in bw
+            ):
+                raise ValueError("Background windows is outside of the signal range.")
+
+            # calculate the position of the segments
+            y0 = self.isig[bw[0] : bw[1]].mean(-1).data
+            y1 = self.isig[bw[2] : bw[3]].mean(-1).data
+            x0 = (bw[0] + bw[1]) / 2.0
+            x1 = (bw[2] + bw[3]) / 2.0
+
+            segments_[..., i, 0, 0] = x0
+            segments_[..., i, 0, 1] = y0.T
+            segments_[..., i, 1, 0] = x1
+            segments_[..., i, 1, 1] = y1.T
+
+        # convert to ragged array to comply with requirement for
+        # navigation position dependent markers
+        # 2000 x 2000 navigation shape takes ~2s
+        # 1000 x 1000 navigation shape takes ~0.5s
+        # 500 x 500 navigation shape takes ~0.01s
+        segments = np.empty(self.axes_manager.navigation_shape, dtype=object)
+        for i in np.ndindex(self.axes_manager.navigation_shape):
+            segments[i] = segments_[i]
+
+        colors_cycle = itertools.cycle(
+            np.sort(plt.rcParams["axes.prop_cycle"].by_key()["color"])
+        )
+        colors = np.array([c for c, w in zip(colors_cycle, windows_position)]).flatten()
+
+        lines = hs.plot.markers.Lines(segments=segments, color=colors, **kwargs)
         self.add_marker(lines, render_figure=False)
+
         if render_figure:
             self._render_figure(plot=["signal_plot"])
 
diff --git a/upcoming_changes/86.bugfix.rst b/upcoming_changes/86.bugfix.rst
new file mode 100644
index 000000000..958ebee60
--- /dev/null
+++ b/upcoming_changes/86.bugfix.rst
@@ -0,0 +1,4 @@
+Plot background and integration windows fixes and improvements:
+
+- fix plotting windows with navigation dimension >=2,
+- raise improved error message when windows are out of the range of the signal,


Test patch:
diff --git a/exspy/tests/signals/test_eds_tem.py b/exspy/tests/signals/test_eds_tem.py
index 2771008b0..aa98e4de9 100644
--- a/exspy/tests/signals/test_eds_tem.py
+++ b/exspy/tests/signals/test_eds_tem.py
@@ -21,6 +21,7 @@
 import numpy as np
 import pytest
 
+import hyperspy.api as hs
 from hyperspy.components1d import Gaussian
 from hyperspy.decorators import lazifyTestClass
 
@@ -600,6 +601,29 @@ def test_intensity_dtype_uint(self):
             rtol=0.03,
         )
 
+    def test_outside_range_background_windows(self):
+        s = self.signal
+        bw = s.estimate_background_windows()
+        with pytest.raises(ValueError):
+            s.isig[2.0:].plot(True, background_windows=bw)
+
+
+def test_plot_windows():
+    s = exspy.data.EDS_TEM_FePt_nanoparticles()
+
+    rng = np.random.default_rng()
+
+    [s * v * 10 for v in rng.random((10))]
+
+    s = hs.stack([s * v * 10 for v in rng.random((10))])
+    s = hs.stack([s * v * 10 for v in rng.random((5))])
+
+    bw = s.estimate_background_windows(line_width=[5.0, 2.0])
+    iw = s.estimate_integration_windows(windows_width=3)
+
+    s.plot(True, background_windows=bw, integration_windows=iw)
+    s.axes_manager.indices = (1, 2)
+
 
 def test_with_signals_examples():
     sig = exspy.data.EDS_TEM_FePt_nanoparticles()