Fixed plotter storing correct tab when switching plot types

src/ert/gui/tools/plot/plot_window.py

@@ -1,6 +1,6 @@
 import logging
 import time
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
@@ -35,6 +35,10 @@
 STATISTICS = "Statistics"
 STD_DEV = "Std Dev"
 
+RESPONSE_DEFAULT = 0
+GEN_KW_DEFAULT = 2
+STD_DEV_DEFAULT = 6
+
 logger = logging.getLogger(__name__)
 
 from qtpy.QtWidgets import (
@@ -137,7 +141,15 @@ def __init__(self, config_file: str, parent: QWidget | None):
         self.addPlotWidget(CROSS_ENSEMBLE_STATISTICS, CrossEnsembleStatisticsPlot())
         self.addPlotWidget(STD_DEV, StdDevPlot())
         self._central_tab.currentChanged.connect(self.currentTabChanged)
-        self._prev_tab_widget: QWidget | None = None
+
+        self._prev_tab_widget_index = -1
+        self._current_tab_index = -1
+        self._prev_key_dimensionality = -1
+        self._prev_tab_widget_index_map: dict[int, int] = {
+            2: RESPONSE_DEFAULT,
+            1: GEN_KW_DEFAULT,
+            3: STD_DEV_DEFAULT,
+        }
 
         QApplication.setOverrideCursor(Qt.CursorShape.WaitCursor)
         try:
@@ -165,7 +177,8 @@ def __init__(self, config_file: str, parent: QWidget | None):
         logger.info(f"PlotWindow __init__ done. time={time.perf_counter() -t}")
 
     @Slot(int)
-    def currentTabChanged(self, index: Any) -> None:
+    def currentTabChanged(self, index: int) -> None:
+        self._current_tab_index = index
         self.updatePlot()
 
     @Slot(int)
@@ -334,7 +347,6 @@ def keySelected(self) -> None:
             for widget in self._plot_widgets
             if widget._plotter.dimensionality == key_def.dimensionality
         ]
-
         current_widget = self._central_tab.currentWidget()
 
         # Enabling/disabling tab triggers the
@@ -350,17 +362,20 @@ def keySelected(self) -> None:
             )
         self._central_tab.currentChanged.connect(self.currentTabChanged)
 
-        # Remember which tab widget was selected when switching between
-        # both same and different data-types.
-        if current_widget in available_widgets:
-            self._central_tab.setCurrentWidget(current_widget)
-        else:
-            if self._prev_tab_widget is None:
-                self._central_tab.setCurrentWidget(available_widgets[0])
-            else:
-                self._central_tab.setCurrentWidget(self._prev_tab_widget)
-            self._prev_tab_widget = current_widget
+        if 0 < self._prev_key_dimensionality != key_def.dimensionality:
+            if self._current_tab_index == -1:
+                self._current_tab_index = self._prev_tab_widget_index
+            self._prev_tab_widget_index_map[self._prev_key_dimensionality] = (
+                self._current_tab_index
+            )
+            current_widget = self._central_tab.widget(
+                self._prev_tab_widget_index_map[key_def.dimensionality]
+            )
+            self._current_tab_index = -1
 
+        self._central_tab.setCurrentWidget(current_widget)
+        self._prev_tab_widget_index = self._central_tab.currentIndex()
+        self._prev_key_dimensionality = key_def.dimensionality
         self.updatePlot()
 
     def toggleCustomizeDialog(self) -> None:

tests/ert/ui_tests/gui/test_main_window.py

@@ -46,7 +46,12 @@
 from ert.gui.tools.plot.plot_ensemble_selection_widget import (
     EnsembleSelectListWidget,
 )
-from ert.gui.tools.plot.plot_window import PlotApi, PlotWindow
+from ert.gui.tools.plot.plot_window import (
+    GEN_KW_DEFAULT,
+    RESPONSE_DEFAULT,
+    PlotApi,
+    PlotWindow,
+)
 from ert.plugins import ErtPluginManager
 from ert.run_models import (
     EnsembleExperiment,
@@ -341,21 +346,53 @@ def test_that_the_plot_window_contains_the_expected_elements(
         "Std Dev",
     }
 
-    # Cycle through showing all the tabs and plot each data key
-
     model = data_keys.model()
     assert model is not None
+
+    def click_plotter_item(pos: int) -> None:
+        center = data_keys.visualRect(model.index(pos, 0)).center()
+        viewport = data_keys.viewport()
+        center = viewport.mapToGlobal(center)
+        local_pos = viewport.mapFromGlobal(center)
+        qtbot.mouseClick(data_keys.viewport(), Qt.LeftButton, pos=local_pos)
+
+    def click_tab_index(pos: int) -> None:
+        tab_bar = plot_window._central_tab.tabBar()
+        tab_center = tab_bar.tabRect(pos).center()
+        qtbot.mouseClick(tab_bar, Qt.LeftButton, pos=tab_center)
+
+    # make sure plotter remembers plot types selected previously
+    response_index = 0
+    gen_kw_index = 1
+    response_alternate_index = 1
+    gen_kw_alternate_index = 3
+
+    # check default selections
+    click_plotter_item(response_index)
+    assert plot_window._central_tab.currentIndex() == RESPONSE_DEFAULT
+    click_plotter_item(gen_kw_index)
+    assert plot_window._central_tab.currentIndex() == GEN_KW_DEFAULT
+
+    # alter selections
+    click_plotter_item(response_index)
+    click_tab_index(response_alternate_index)
+    click_plotter_item(gen_kw_index)
+    click_tab_index(gen_kw_alternate_index)
+
+    # verify previous selections still valid
+    click_plotter_item(response_index)
+    assert plot_window._central_tab.currentIndex() == response_alternate_index
+    click_plotter_item(gen_kw_index)
+    assert plot_window._central_tab.currentIndex() == gen_kw_alternate_index
+
+    # finally click all items
     for i in range(model.rowCount()):
-        index = model.index(i, 0)
-        qtbot.mouseClick(
-            data_types.data_type_keys_widget,
-            Qt.LeftButton,
-            pos=data_types.data_type_keys_widget.visualRect(index).center(),
-        )
+        click_plotter_item(i)
         for tab_index in range(plot_window._central_tab.count()):
             if not plot_window._central_tab.isTabEnabled(tab_index):
                 continue
-            plot_window._central_tab.setCurrentIndex(tab_index)
+            click_tab_index(tab_index)
+
     plot_window.close()