fixup! Add RFT QC tool

Author: HakonSohoel

src/ert/gui/tools/manage_experiments/rft_qc_widget.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import html
+import math
 from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass, field
 from enum import StrEnum
@@ -708,8 +709,66 @@ def _facet_counts(self, col: str, dataframes: list[pl.DataFrame]) -> dict[str, i
         return counts
 
 
+class DisplayPoints:
+    def __init__(self) -> None:
+        self._use_utm: bool = False
+        self._df: pl.DataFrame = pl.DataFrame(schema=self.schema())
+
+    @staticmethod
+    def schema() -> dict[str, Any]:
+        return {
+            "well_connection_cell": pl.Array(pl.Int64, 3),
+            "east": pl.Float32,
+            "north": pl.Float32,
+            "tvd": pl.Float32,
+            "status": pl.String,
+        }
+
+    def update_display_points(
+        self, display_points: pl.DataFrame, use_utm: bool
+    ) -> None:
+        self._use_utm = use_utm
+        self._df = display_points.select(self.schema().keys())
+
+    def get_point(self, index: int | None) -> dict[str, Any] | None:
+        if index is not None and 0 <= index < len(self._df):
+            return self._df.row(index, named=True)
+        return None
+
+    def get_point_coordinates(
+        self, index: int | None
+    ) -> tuple[float, float, float] | None:
+        point = self.get_point(index)
+        if point is None:
+            return None
+        if self._use_utm:
+            return (
+                point["east"],
+                point["north"],
+                point["tvd"],
+            )
+        return tuple(point["well_connection_cell"])
+
+    def get_points_to_plot(self) -> tuple[list[float], list[float], list[float]]:
+        if self._use_utm:
+            return (
+                self._df["east"].to_list(),
+                self._df["north"].to_list(),
+                self._df["tvd"].to_list(),
+            )
+        return (
+            self._df["well_connection_cell"].arr.get(0).to_list(),
+            self._df["well_connection_cell"].arr.get(1).to_list(),
+            self._df["well_connection_cell"].arr.get(2).to_list(),
+        )
+
+
 class RftPlot:
-    def __init__(self, show_details: Callable[[dict[str, Any]], None]) -> None:
+    def __init__(
+        self,
+        show_details: Callable[[dict[str, Any]], None],
+        clear_details: Callable[[], None],
+    ) -> None:
         figure = Figure()
         self._canvas: FigureCanvas = FigureCanvas(figure)
         self._canvas.setSizePolicy(
@@ -723,16 +782,14 @@ def __init__(self, show_details: Callable[[dict[str, Any]], None]) -> None:
         self._autoscaled_limits: tuple[Any, Any, Any] | None = None
 
         self._point_artist: PathCollection | None = None
-        self._point_coords: list[tuple[float, float, float]] = []
-        self._display_points: pl.DataFrame = pl.DataFrame(
-            schema=self._display_point_schema()
-        )
+        self._display_points: DisplayPoints = DisplayPoints()
         self._selection_artist: Any = None
         self._hover_artist: Any = None
         self._selected_index: int | None = None
         self._hover_index: int | None = None
 
         self._show_details = show_details
+        self._clear_details = clear_details
 
     @property
     def canvas(self) -> FigureCanvas:
@@ -759,20 +816,21 @@ def redraw(
                 self._ax.get_ylim(),
                 self._ax.get_zlim(),
             )
+        previous_selected_point = self._display_points.get_point(self._selected_index)
         self._ax.cla()
         self._ax.invert_zaxis()
         self._point_artist = None
         self._selection_artist = None
         self._hover_artist = None
-        self._point_coords = []
         self._selected_index = None
+        self._clear_details()
         self._hover_index = None
 
         if obs_df.is_empty() and response_df.is_empty() and file_rft_df.is_empty():
             self._canvas.draw()
             return
 
-        point_columns = self._display_point_schema().keys()
+        point_columns = self._display_points.schema().keys()
 
         points = pl.concat(
             [
@@ -876,11 +934,6 @@ def _get_observation_cell_center_overlay(
                 )
                 self._ax.add_collection3d(lc)
 
-            xs, ys, zs = (
-                points["east"].to_list(),
-                points["north"].to_list(),
-                points["tvd"].to_list(),
-            )
         else:
             self._ax.set_xlabel("i", labelpad=6)
             self._ax.set_ylabel("j", labelpad=6)
@@ -897,11 +950,8 @@ def _get_observation_cell_center_overlay(
                 _point_style(statuses), statuses, coords
             )
 
-            xs, ys, zs = (
-                points["well_connection_cell"].arr.get(0).to_list(),
-                points["well_connection_cell"].arr.get(1).to_list(),
-                points["well_connection_cell"].arr.get(2).to_list(),
-            )
+        self._display_points.update_display_points(points, use_utm)
+        xs, ys, zs = self._display_points.get_points_to_plot()
         self._point_artist = self._ax.scatter(
             xs,
             ys,
@@ -910,10 +960,8 @@ def _get_observation_cell_center_overlay(
             picker=5,
             depthshade=False,
         )
-        self._point_coords = list(zip(xs, ys, zs, strict=True))
-        self._display_points = points
 
-        displayed_statuses = self._display_points["status"].unique().to_list()
+        displayed_statuses = points["status"].unique().to_list()
         for status, style in _POINT_STYLE.items():
             if status in displayed_statuses:
                 self._ax.scatter(
@@ -926,6 +974,39 @@ def _get_observation_cell_center_overlay(
 
         self._ax.legend(loc="upper left", fontsize="x-small")
         self._create_overlay_artists()
+
+        if previous_selected_point is not None:
+            # Try to restore the previously selected point
+
+            fallback_index = None
+            fallback_point = None
+            for i, point in enumerate(points.rows(named=True)):
+                if (
+                    point["well_connection_cell"]
+                    == previous_selected_point["well_connection_cell"]
+                ):
+                    fallback_index = i
+                    fallback_point = point
+                    if (
+                        math.isclose(
+                            point["east"], previous_selected_point["east"], rel_tol=1e-5
+                        )
+                        and math.isclose(
+                            point["north"],
+                            previous_selected_point["north"],
+                            rel_tol=1e-5,
+                        )
+                        and math.isclose(
+                            point["tvd"], previous_selected_point["tvd"], rel_tol=1e-5
+                        )
+                    ):
+                        self._update_selected_point(i, point)
+                        break
+            # If we didn't find an exact match, but did find a point with the same cell
+            # center, use that as a fallback.
+            if fallback_index is not None and fallback_point is not None:
+                self._update_selected_point(fallback_index, fallback_point)
+
         self._canvas.draw()
 
         # Store the autoscale limits after the redraw, so we can restore them later if
@@ -943,16 +1024,6 @@ def _get_observation_cell_center_overlay(
             self._ax.set_zlim(*zlim)
             self._canvas.draw_idle()
 
-    @staticmethod
-    def _display_point_schema() -> dict[str, Any]:
-        return {
-            "well_connection_cell": pl.Array(pl.Int64, 3),
-            "east": pl.Float32,
-            "north": pl.Float32,
-            "tvd": pl.Float32,
-            "status": pl.String,
-        }
-
     def _fit_view_to_displayed_points(self) -> None:
         if self._autoscaled_limits is None:
             return
@@ -963,11 +1034,12 @@ def _fit_view_to_displayed_points(self) -> None:
         self._canvas.draw_idle()
 
     def _center_on_selected(self) -> None:
-        if self._selected_index is None or not (
-            0 <= self._selected_index < len(self._point_coords)
-        ):
+        if self._selected_index is None:
+            return
+        coordinates = self._display_points.get_point_coordinates(self._selected_index)
+        if coordinates is None:
             return
-        cx, cy, cz = self._point_coords[self._selected_index]
+        cx, cy, cz = coordinates
         for getter, setter, center in (
             (self._ax.get_xlim, self._ax.set_xlim, cx),
             (self._ax.get_ylim, self._ax.set_ylim, cy),
@@ -1011,9 +1083,10 @@ def _refresh_overlays(self) -> None:
         def _coords_for(
             idx: int | None,
         ) -> tuple[list[float], list[float], list[float]]:
-            if idx is None or not (0 <= idx < len(self._point_coords)):
+            coordinates = self._display_points.get_point_coordinates(idx)
+            if coordinates is None:
                 return ([], [], [])
-            x, y, z = self._point_coords[idx]
+            x, y, z = coordinates
             return ([x], [y], [z])
 
         hover_idx = (
@@ -1033,11 +1106,15 @@ def _on_pick(self, event: PickEvent) -> None:
         if not hasattr(event, "ind") or len(event.ind) == 0:
             return
         idx = int(event.ind[0])
-        if not (0 <= idx < len(self._point_coords)):
-            return
-        self._selected_index = idx
-        if 0 <= idx < len(self._display_points):
-            self._show_details(self._display_points.row(idx, named=True))
+        selected_point = self._display_points.get_point(idx)
+        if selected_point is not None:
+            self._update_selected_point(idx, selected_point)
+
+    def _update_selected_point(
+        self, index: int, selected_point: dict[str, Any]
+    ) -> None:
+        self._selected_index = index
+        self._show_details(selected_point)
         self._refresh_overlays()
 
     def _on_hover(self, event: MouseEvent) -> None:
@@ -1077,12 +1154,9 @@ def _on_scroll(self, event: MouseEvent) -> None:
         if event.inaxes is not self._ax:
             return
         scale = 0.8 if event.button == "up" else 1.25
-        if self._selected_index is not None and 0 <= self._selected_index < len(
-            self._point_coords
-        ):
-            cx, cy, cz = self._point_coords[self._selected_index]
-        else:
-            cx, cy, cz = None, None, None
+        coordinates = self._display_points.get_point_coordinates(self._selected_index)
+        (cx, cy, cz) = (None, None, None) if coordinates is None else coordinates
+
         # Assert to reassure mypy that self._ax is indeed an Axes3D instance:
         assert isinstance(self._ax, Axes3D)
         for getter, setter, center in (
@@ -1141,7 +1215,7 @@ def __init__(self, ert_config: ErtConfig | None = None) -> None:
         self._load_rft_file: bool = False
 
         self._use_utm = False
-        self._plot = RftPlot(self._show_details)
+        self._plot = RftPlot(self._show_details, self._clear_details)
         self._filter_panel = FilterPanel(
             self._apply_filter_and_redraw,
             self._plot._fit_view_to_displayed_points,
@@ -1697,3 +1771,6 @@ def _obs_coords(row: dict[str, Any]) -> str | None:
                     {file_point_details_html}
                 </table>
         """)
+
+    def _clear_details(self) -> None:
+        self._details.clear()