Set Depth from Computed Statistics

mom6_forge/topo.py

@@ -873,6 +873,30 @@ def _compute_stats(self, nx_sub, ny_sub, mask_hmin):
         )
         return self.stats
 
+    def set_depth_from_stats(self, statistic):
+        """
+        Set the topo depth to a statistic computed by _compute_stats.
+
+        Parameters
+        ----------
+        statistic : str
+            Which depth statistic to use. Must be one of the "D_*" keys
+            in self.src.stats (e.g. "mean", "min", "max").
+        """
+
+        assert (
+            self.src.stats is not None
+        ), "Source bathymetry must have topo stats computed, please call _compute_stats first if you have not already"
+        approved_list = []
+        for key in self.src.stats.data_vars:
+            if key.startswith("D_"):
+                approved_list.append(key[2:])
+        assert (
+            statistic in approved_list
+        ), f"Invalid statistic {statistic}, must be one of {approved_list}"
+
+        self.send_entire_depth_change_to_tcm(self.src.stats[f"D_{statistic}"])
+
     def set_from_dataset(
         self,
         bathymetry_path,

tests/test_topo_bathymetry_workflows.py

@@ -54,3 +54,25 @@ def test_compute_topo_stats(get_rect_topo, synthetic_bathy_file):
         stats2 = topo._compute_stats(nx_sub=nx_sub, ny_sub=ny_sub, mask_hmin=0.0)
         # Should be the exact same object (cached)
         assert stats2 is stats
+
+
+def test_set_depth_from_stats(get_rect_topo, synthetic_bathy_file):
+    """Test set_depth_from_stats sets topo depth to the chosen statistic from _compute_stats."""
+    topo = get_rect_topo
+
+    # Load source bathymetry and slice to topo domain
+    src = SourceBathy(
+        topo,
+        synthetic_bathy_file,
+        depth_name="elevation",
+        is_input_positive_below_msl=False,
+    )
+    topo.src = src
+    topo._compute_stats(nx_sub=2, ny_sub=2, mask_hmin=0.0)
+
+    topo.set_depth_from_stats("mean")
+
+    mask = ~np.isnan(topo.depth.values)
+    assert np.isclose(
+        topo.depth.values[mask], topo.src.stats["D_mean"].values[mask]
+    ).all()