Display Stats in the Topo Editor on zoom

mom6_forge/_supergrid.py

@@ -220,6 +220,14 @@ def from_ds(cls, ds: xr.Dataset) -> "SupergridBase":
 
         Does not dispatch to subclasses
         """
+        radius = ds.attrs.get("radius")
+        if radius is None:
+            print(
+                "Warning: 'radius' attribute not found in dataset; "
+                "defaulting to Earth's radius (6.378e6 m). "
+                "dx, dy, and area cannot be recalculated exactly without the original radius."
+            )
+            radius = 6.378e6
         return cls(
             ds.x.data,
             ds.y.data,
@@ -229,7 +237,7 @@ def from_ds(cls, ds: xr.Dataset) -> "SupergridBase":
             ds.angle_dx.data,
             ds.x.attrs.get("units", "degrees"),
             grid_type=ds.attrs.get("grid_type"),
-            radius=ds.attrs["radius"],
+            radius=radius,
         )
 
     @staticmethod

mom6_forge/topo.py

@@ -1369,7 +1369,7 @@ def generate_mask_from_landfrac_file(
         new_values = binary_mask.values.ravel().tolist()
 
         # Get old values (current mask or 0 if no mask set)
-        old_mask = self.mask
+        old_mask = self.tmask
         old_values = (
             old_mask.values.ravel().tolist()
             if isinstance(old_mask, xr.DataArray)

mom6_forge/topo_editor.py

@@ -6,6 +6,7 @@
 from matplotlib.ticker import MaxNLocator
 from mom6_forge.edit_command import *
 from mom6_forge.git_utils import *
+from mom6_forge.utils import normalize_deg
 
 
 class TopoEditor(widgets.HBox):
@@ -171,6 +172,15 @@ def construct_control_panel(self):
             layout={"width": "80%"},
             style={"description_width": "auto"},
         )
+        self._set_to_mean_button = widgets.Button(
+            description="Mean", disabled=True, layout={"width": "30%"}
+        )
+        self._set_to_max_button = widgets.Button(
+            description="Max", disabled=True, layout={"width": "30%"}
+        )
+        self._set_to_min_button = widgets.Button(
+            description="Min", disabled=True, layout={"width": "30%"}
+        )
 
         # --- Basin editing widgets ---
         self._basin_specifier_toggle = widgets.Button(
@@ -244,13 +254,32 @@ def construct_control_panel(self):
                 self._min_depth_specifier,
             ]
         )
-        cell_editing_section = widgets.VBox(
-            [
-                widgets.HTML("<h3>Cell Editing</h3>"),
-                self._selected_cell_label,
-                self._depth_specifier,
-            ]
-        )
+        cell_editing_section_children = [
+            widgets.HTML("<h3>Cell Editing</h3>"),
+            self._selected_cell_label,
+            self._depth_specifier,
+        ]
+
+        # Only add stats section if statistics are available
+        has_stats = self.topo.stats is not None
+        if has_stats:
+            cell_editing_section_children.extend(
+                [
+                    widgets.HTML(
+                        "<p style='margin: 5px 0; font-size: 12px;'>Set to statistic:</p>"
+                    ),
+                    widgets.HBox(
+                        [
+                            self._set_to_mean_button,
+                            self._set_to_max_button,
+                            self._set_to_min_button,
+                        ],
+                        layout={"justify_content": "space-between"},
+                    ),
+                ]
+            )
+
+        cell_editing_section = widgets.VBox(cell_editing_section_children)
         basin_section = widgets.VBox(
             [
                 widgets.HTML("<h3>Basin Selector</h3>"),
@@ -343,6 +372,118 @@ def trigger_refresh(self):
         self._min_depth_specifier.value = self.topo.min_depth
         self.update_undo_redo_buttons()
 
+    def _draw_cell_stats(self, visible):
+        """Draw per-cell stat annotations directly on each visible cell."""
+        # Clear previous annotations
+        if hasattr(self, "_cell_stat_texts"):
+            for t in self._cell_stat_texts:
+                t.remove()
+        self._cell_stat_texts = []
+
+        if self.topo.stats is None:
+            return
+
+        ds = self.topo.stats
+        js, is_ = np.where(visible)
+
+        qlon = self.topo._grid.qlon.data
+        qlat = self.topo._grid.qlat.data
+
+        for idx in range(len(js)):
+            j, i = js[idx], is_[idx]
+
+            # Cell centre
+            cx = 0.25 * (
+                qlon[j, i] + qlon[j, i + 1] + qlon[j + 1, i] + qlon[j + 1, i + 1]
+            )
+            cy = 0.25 * (
+                qlat[j, i] + qlat[j, i + 1] + qlat[j + 1, i] + qlat[j + 1, i + 1]
+            )
+
+            lines = []
+            for var in ds.data_vars:
+                if var == "D2_mean":
+                    continue  # Skip this variable for now since it's not super informative
+                label = var
+                if var.startswith("D_"):
+                    label = var[2:]  # Get rid of the "D_" prefix for display purposes
+                val = ds[var].data[j, i]
+                units = ds[var].attrs.get("units", "")
+                if units == "1":  # Unitless
+                    units = ""
+                lines.append(f"{label}: {val:.2f} {units}")
+
+            t = self.ax.text(
+                cx,
+                cy,
+                "\n".join(lines),
+                fontsize=7,
+                ha="center",
+                va="center",
+                transform=ccrs.PlateCarree(),
+                bbox=dict(
+                    boxstyle="round,pad=0.2",
+                    facecolor="none",
+                    alpha=0.7,
+                    edgecolor="none",
+                ),
+                zorder=11,
+            )
+            self._cell_stat_texts.append(t)
+
+        self.fig.canvas.draw_idle()
+
+    def _on_zoom_change(self, ax):
+        xlim = ax.get_xlim()
+        ylim = ax.get_ylim()
+
+        qlon = normalize_deg(self.topo._grid.qlon.data)
+        qlat = self.topo._grid.qlat.data
+        xlo = normalize_deg(xlim[0])
+        xhi = normalize_deg(xlim[1])
+
+        # Check if all four corners of each cell are within the view limits
+        if xlo <= xhi:
+            lon_visible = (
+                (qlon[:-1, :-1] >= xlo)
+                & (qlon[:-1, :-1] <= xhi)
+                & (qlon[:-1, 1:] >= xlo)
+                & (qlon[:-1, 1:] <= xhi)
+                & (qlon[1:, :-1] >= xlo)
+                & (qlon[1:, :-1] <= xhi)
+                & (qlon[1:, 1:] >= xlo)
+                & (qlon[1:, 1:] <= xhi)
+            )
+        else:  # view crosses the 0/360 boundary after normalization
+            lon_visible = (
+                ((qlon[:-1, :-1] >= xlo) | (qlon[:-1, :-1] <= xhi))
+                & ((qlon[:-1, 1:] >= xlo) | (qlon[:-1, 1:] <= xhi))
+                & ((qlon[1:, :-1] >= xlo) | (qlon[1:, :-1] <= xhi))
+                & ((qlon[1:, 1:] >= xlo) | (qlon[1:, 1:] <= xhi))
+            )
+        visible = (
+            lon_visible
+            & (qlat[:-1, :-1] >= ylim[0])
+            & (qlat[:-1, :-1] <= ylim[1])
+            & (qlat[:-1, 1:] >= ylim[0])
+            & (qlat[:-1, 1:] <= ylim[1])
+            & (qlat[1:, :-1] >= ylim[0])
+            & (qlat[1:, :-1] <= ylim[1])
+            & (qlat[1:, 1:] >= ylim[0])
+            & (qlat[1:, 1:] <= ylim[1])
+        )
+
+        # Always clear old annotations first
+        if hasattr(self, "_cell_stat_texts"):
+            for t in self._cell_stat_texts:
+                t.remove()
+        self._cell_stat_texts = []
+
+        if visible.sum() <= 40:
+            self._draw_cell_stats(visible)
+        else:
+            self.fig.canvas.draw_idle()
+
     def _select_cell(self, i, j):
         """Select a cell in the topography grid and update the UI accordingly."""
         # Remove old patch if it exists
@@ -393,6 +534,16 @@ def _select_cell(self, i, j):
         if hasattr(self, "_depth_specifier"):
             self._depth_specifier.disabled = False
             self._depth_specifier.value = self.topo.masked_depth.data[j, i]
+
+        # Enable statistic buttons if statistics are available
+        has_stats = self.topo.stats is not None
+        for btn in [
+            self._set_to_mean_button,
+            self._set_to_max_button,
+            self._set_to_min_button,
+        ]:
+            btn.disabled = not has_stats
+
         if hasattr(self, "_basin_specifier"):
             label = self.topo.basintmask.data[j, i]
             self._basin_specifier.value = f"Basin Label Number: {str(label)}"
@@ -415,7 +566,9 @@ def construct_observances(self):
 
         # Double click event for cell selection on the plot
         self.fig.canvas.mpl_connect("button_press_event", self.on_double_click)
-
+        # Zoom-dependent stats overlay
+        self.ax.callbacks.connect("xlim_changed", self._on_zoom_change)
+        self.ax.callbacks.connect("ylim_changed", self._on_zoom_change)
         # Min depth change observer
         self._min_depth_specifier.observe(
             self.on_min_depth_change, names="value", type="change"
@@ -430,6 +583,11 @@ def construct_observances(self):
             self.on_depth_change, names="value", type="change"
         )
 
+        # Statistic buttons
+        self._set_to_mean_button.on_click(self.set_depth_to_mean)
+        self._set_to_max_button.on_click(self.set_depth_to_max)
+        self._set_to_min_button.on_click(self.set_depth_to_min)
+
         if self.has_version_control:
             # Undo/Redo/Reset buttons
             self._undo_button.on_click(self.undo_last_edit)
@@ -507,6 +665,37 @@ def on_depth_change(self, change):
         self.apply_edit(cmd)
         self.update_undo_redo_buttons()
 
+    def _get_statistic_value(self, stat_name):
+        """Get a statistic value for the selected cell."""
+        if self._selected_cell is None or self.topo.stats is None:
+            return None
+
+        i, j, _ = self._selected_cell
+        ds = self.topo.stats
+
+        if ds is None or stat_name not in ds.data_vars:
+            return None
+
+        return float(ds[stat_name].data[j, i])
+
+    def set_depth_to_mean(self, b):
+        """Set the selected cell's depth to the mean value."""
+        val = self._get_statistic_value("D_mean")
+        if val is not None:
+            self._depth_specifier.value = val
+
+    def set_depth_to_max(self, b):
+        """Set the selected cell's depth to the max value."""
+        val = self._get_statistic_value("D_max")
+        if val is not None:
+            self._depth_specifier.value = val
+
+    def set_depth_to_min(self, b):
+        """Set the selected cell's depth to the min value."""
+        val = self._get_statistic_value("D_min")
+        if val is not None:
+            self._depth_specifier.value = val
+
     def on_git_create_branch(self, b):
         """Create a new git branch"""
         name = self._git_branch_name.value.strip()