Custom runoff grid and runoff to ocean mapping

ProConPy/config_var.py

@@ -334,6 +334,11 @@ def update_options_validities(self):
 
         return validities_changed
 
+    @property
+    def valid_options(self):
+        """Returns the list of valid options for this variable."""
+        return [opt for opt in self._options if self._options_validities.get(opt, False)]
+
     def _refresh_widget_options(self):
         """Refresh the widget options list based on information in the current self._options_validities."""
 

ProConPy/stage.py

@@ -352,7 +352,7 @@ def _proceed(self):
             return
 
         # Display the child stage and its siblings by appending them to the current stage's widget
-        if self.has_children():
+        if self.has_children() and next_stage.is_descendant_of(self):
             self._widget.add_child_stages(first_child=next_stage)
 
         # Proceed the csp solver before enabling the next stage
@@ -378,20 +378,24 @@ def get_next(self, full_dfs=False):
             The next stage to visit, if found. Otherwise, None.
         """
 
-        if self.has_children():
-            return self._get_child_to_enable(full_dfs)
-        elif self._right is not None:
+        # First try to get a child stage to enable
+        if (child_to_enable := self._get_child_to_enable(full_dfs)) is not None:
+            return child_to_enable
+
+        # No child stage to enable. Try to get the right sibling.
+        if self._right is not None:
             return self._right
-        else:  # Backtrack
-            ancestor = self._parent
-            while ancestor is not None:
-                if ancestor._right is not None and (
-                    full_dfs or not ancestor.has_condition()
-                ):
-                    return ancestor._right
-                else:
-                    ancestor = ancestor._parent
-            return None
+
+        # No child or right sibling. Backtrack to find the next stage.
+        ancestor = self._parent
+        while ancestor is not None:
+            if ancestor._right is not None and (
+                full_dfs or not ancestor.has_condition()
+            ):
+                return ancestor._right
+            else:
+                ancestor = ancestor._parent
+        return None
 
     def _get_child_to_enable(self, full_dfs):
         """Determine the child stage to activate.
@@ -401,6 +405,9 @@ def _get_child_to_enable(self, full_dfs):
         full_dfs : bool
             If True, visit all the stages in the stage tree. Otherwise, skip stages whose guards
             are not satisfied."""
+
+        if self.has_children() is False:
+            return None
 
         child_to_activate = None
 
@@ -412,18 +419,19 @@ def _get_child_to_enable(self, full_dfs):
                         child_to_activate is None
                     ), "Only one child stage can be activated at a time."
                     child_to_activate = child
+
+            if child_to_activate is None:
+                # No child guard's condition is satisfied.
+                # Let the caller handle this case (by backtracking).
+                return None
         else:
             # If children are not guards, the first child is activated.
             # Note the remaining children will be activated in sequence by their siblings.
             child_to_activate = self._children[0]
 
         # If the child to activate is a Guard, return it's first child
         if child_to_activate.has_condition():
-            return child_to_activate._children[0]
-
-        assert (
-            child_to_activate is not None
-        ), "At least one child stage must be activated."
+            child_to_activate = child_to_activate._children[0]
 
         return child_to_activate
 

environment.yml

@@ -7,6 +7,7 @@ channels:
 dependencies:
   - python>=3.11.10,<3.12
   - libxml2>=2.13,<2.14
+  - xesmf>=0.8.10,<0.9
   - pip
   - pip:
     - -e ./external/mom6_bathy/

tests/3_system/test_custom_mom6_grid.py

@@ -141,7 +141,10 @@ def test_custom_mom6_grid():
         assert Stage.active().title.startswith("Simple Initial Conditions")
         cvars["T_REF"].value = 10.0
 
-        # Since land grid gets set automatically, we should be in the Launch stage:
+        # Since land grid and runoff grid get set automatically, we should be in the runoff to ocn mapping:
+        assert Stage.active().title.startswith("Runoff to Ocean Mapping")
+        cvars["ROF_OCN_MAPPING_STATUS"].value = "skip"
+
         assert Stage.active().title.startswith("3. Launch")
         launch_stage = Stage.active()
 

tests/3_system/test_f2000_custom_grid.py

@@ -159,6 +159,9 @@ def construct_custom_res_from_std_grids(cime):
         assert Stage.active().title.startswith("Land Grid")
         cvars["CUSTOM_LND_GRID"].value = "0.9x1.25"
 
+        assert Stage.active().title.startswith("Runoff Grid")
+        cvars["CUSTOM_ROF_GRID"].value = "r05"
+
         assert Stage.active().title.startswith("3. Launch")
         launch_stage = Stage.active()
 
@@ -248,6 +251,9 @@ def construct_custom_res_from_modified_clm_grid(cime):
         # click the "Run Surface Data Modifier" button
         fsurdat_modifier_launcher._on_launch_clicked(b=None)
 
+        assert Stage.active().title.startswith("Runoff Grid")
+        cvars["CUSTOM_ROF_GRID"].value = "r05"
+
         assert Stage.active().title.startswith("3. Launch")
         launch_stage = Stage.active()
 
@@ -364,6 +370,9 @@ def construct_custom_res_from_new_mom6_grid_modified_clm_grid(cime):
         # click the "Run Surface Data Modifier" button
         fsurdat_modifier_launcher._on_launch_clicked(b=None)
 
+        assert Stage.active().title.startswith("Runoff Grid")
+        cvars["CUSTOM_ROF_GRID"].value = "r05"
+
         assert Stage.active().title.startswith("3. Launch")
         launch_stage = Stage.active()
 

tests/3_system/test_fhist_custom_grid.py

@@ -129,6 +129,9 @@ def construct_custom_res_from_std_grids(cime):
         assert Stage.active().title.startswith("Land Grid")
         cvars["CUSTOM_LND_GRID"].value = "0.9x1.25"
 
+        assert Stage.active().title.startswith("Runoff Grid")
+        cvars["CUSTOM_ROF_GRID"].value = "r05"
+
         assert Stage.active().title.startswith("3. Launch")
         launch_stage = Stage.active()
 
@@ -218,6 +221,9 @@ def construct_custom_res_from_modified_clm_grid(cime):
         # click the "Run Surface Data Modifier" button
         fsurdat_modifier_launcher._on_launch_clicked(b=None)
 
+        assert Stage.active().title.startswith("Runoff Grid")
+        cvars["CUSTOM_ROF_GRID"].value = "r05"
+
         assert Stage.active().title.startswith("3. Launch")
         launch_stage = Stage.active()
 

tests/3_system/test_rof_ocn_map.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python3
+
+import pytest
+import os
+import tempfile
+
+from ProConPy.config_var import ConfigVar, cvars
+from ProConPy.stage import Stage
+from ProConPy.csp_solver import csp
+from visualCaseGen.cime_interface import CIME_interface
+from visualCaseGen.initialize_configvars import initialize_configvars
+from visualCaseGen.initialize_widgets import initialize_widgets
+from visualCaseGen.initialize_stages import initialize_stages
+from visualCaseGen.specs.options import set_options
+from visualCaseGen.specs.relational_constraints import get_relational_constraints
+
+# do not show logger output
+import logging
+
+logger = logging.getLogger()
+logger.setLevel(logging.CRITICAL)
+
+base_temp_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "temp"))
+
+def configure_for_rof_map(compset_alias):
+    """This function configures a standard compset with custom grids suitable for testing
+    the runoff to ocean mapping generator."""
+
+    ConfigVar.reboot()
+    Stage.reboot()
+    cime = CIME_interface()
+    initialize_configvars(cime)
+    initialize_widgets(cime)
+    initialize_stages(cime)
+    set_options(cime)
+    csp.initialize(cvars, get_relational_constraints(cvars), Stage.first())
+
+    assert os.path.exists(base_temp_dir), "temp testing directory does not exist"
+
+    # At initialization, the first stage should be enabled
+    assert Stage.first().enabled
+    cvars["COMPSET_MODE"].value = "Standard"
+
+    # COMPSET_MODE is the only variable in the first stage, so assigning a value to it should disable the first stage
+    assert not Stage.first().enabled
+
+    # The next stage is Custom Component Set, whose first child is Model Time Period
+    assert Stage.active().title.startswith("Support Level")
+    cvars["SUPPORT_LEVEL"].value = "All"
+
+    # Apply filters
+    for comp_class in cime.comp_classes:
+        cvars[f"COMP_{comp_class}_FILTER"].value = "any"
+
+    ## Pick a standard compset
+    cvars["COMPSET_ALIAS"].value = compset_alias
+
+    # Create a custom grid
+    assert Stage.active().title.startswith("2. Grid")
+    cvars["GRID_MODE"].value = "Custom"
+
+    # Set the custom grid path
+    assert Stage.active().title.startswith("Custom Grid")
+
+def test_standard_rof_to_ocn_mapping():
+    """This test configures a case with a standard runoff to ocean mapping for a custom resolution"""
+
+    configure_for_rof_map("C_JRA")
+
+    assert Stage.active().title.startswith("Custom Grid")
+
+    with tempfile.TemporaryDirectory(dir=base_temp_dir) as temp_grid_path:
+
+        cvars["CUSTOM_GRID_PATH"].value = temp_grid_path
+        # since this is a JRA run, the atmosphere grid must automatically be set to TL319
+        assert cvars["CUSTOM_ATM_GRID"].value == "TL319"
+
+        # Set the custom ocean grid mode
+        assert Stage.active().title.startswith("Ocean")
+        cvars["OCN_GRID_MODE"].value = "Standard"
+
+        assert Stage.active().title.startswith("Ocean Grid")
+        cvars["CUSTOM_OCN_GRID"].value = "tx2_3v2"
+
+        # Land Grid and Runoff grid should be set automatically
+        assert Stage.active().title.startswith("Runoff to Ocean Mapping")
+
+        # Currently, smoothing parameters should be unset
+        assert cvars["ROF_OCN_MAPPING_RMAX"].value is None, "ROF_OCN_MAPPING_RMAX should be None"
+        assert cvars["ROF_OCN_MAPPING_FOLD"].value is None, "ROF_OCN_MAPPING_FOLD should be None"
+
+        # *Click* "Use Standard Map" button
+        runoffMappingGenerator = Stage.active()._widget._supplementary_widgets[0]
+        runoffMappingGenerator._btn_use_standard.click()
+
+        map_status = cvars["ROF_OCN_MAPPING_STATUS"].value
+        assert map_status == "Standard", f"ROF_OCN_MAPPING_STATUS should be 'Standard', but got: {map_status}"
+
+        # revert stage
+        assert Stage.active().title.startswith("3. Launch")
+        Stage.active().revert()
+        assert Stage.active().title.startswith("Runoff to Ocean Mapping")
+
+# WARNING: for this test to run successfully, MPI must be available. As such, this test
+# cannot be run through the derecho login nodes.
+@pytest.mark.slow
+def test_custom_rof_to_ocn_mapping():
+    """This test configures a case with a custom runoff to ocean mapping for a custom resolution"""
+
+    configure_for_rof_map('C_IAF')
+
+    assert Stage.active().title.startswith("Custom Grid")
+
+    with tempfile.TemporaryDirectory(dir=base_temp_dir) as temp_grid_path:
+
+        cvars["CUSTOM_GRID_PATH"].value = temp_grid_path
+        # since this is a JRA run, the atmosphere grid must automatically be set to TL319
+        assert cvars["CUSTOM_ATM_GRID"].value == "T62"
+
+        # Set the custom ocean grid mode
+        assert Stage.active().title.startswith("Ocean")
+        cvars["OCN_GRID_MODE"].value = "Standard"
+
+        assert Stage.active().title.startswith("Ocean Grid")
+        cvars["CUSTOM_OCN_GRID"].value = "tx2_3v2"
+
+        # Land Grid and Runoff grid should be set automatically
+        assert Stage.active().title.startswith("Runoff to Ocean Mapping")
+
+        # Currently, smoothing parameters should be unset
+        assert cvars["ROF_OCN_MAPPING_RMAX"].value is None, "ROF_OCN_MAPPING_RMAX should be None"
+        assert cvars["ROF_OCN_MAPPING_FOLD"].value is None, "ROF_OCN_MAPPING_FOLD should be None"
+
+        # *Click* the Generate New Map button
+        runoffMappingGenerator = Stage.active()._widget._supplementary_widgets[0]
+        runoffMappingGenerator._btn_generate_new.click()
+
+        # After clicking the button, the smoothing parameters should be set to suggested values
+        assert cvars["ROF_OCN_MAPPING_RMAX"].value is not None, "ROF_OCN_MAPPING_RMAX should have been set to a suggested value"
+        assert cvars["ROF_OCN_MAPPING_FOLD"].value is not None, "ROF_OCN_MAPPING_FOLD should have been set to a suggested value"
+
+        # *Click* the Run mapping generator button
+        runoffMappingGenerator._btn_run_generate.click()
+
+        map_status = cvars["ROF_OCN_MAPPING_STATUS"].value
+        assert map_status.startswith("CUSTOM:"), f"ROF_OCN_MAPPING_STATUS should indicate a custom mapping, but got: {map_status}"
+        map_paths = map_status.split("CUSTOM:")[1]
+        nn_map_path, nnsm_map_path = map_paths.split(",")
+
+        # check if the mapping file was actually created
+        assert os.path.isfile(nn_map_path), f"Nearest neighbor map file was not created at {nn_map_path}"
+        assert os.path.isfile(nnsm_map_path), f"Smoothed map file was not created at {nnsm_map_path}"
+
+        assert Stage.active().title.startswith("3. Launch")
+
+if __name__ == "__main__":
+    test_standard_rof_to_ocn_mapping()
+    test_custom_rof_to_ocn_mapping()
+    print("All tests passed!")