Refactor ensemble generator for flexibility

compass/landice/extrapolate.py

@@ -4,7 +4,9 @@
 from netCDF4 import Dataset
 
 
-def extrapolate_variable(nc_file, var_name, extrap_method, set_value=None):
+def extrapolate_variable(nc_file, var_name, extrap_method='auto',  # noqa: C901
+                         valid_region_method='auto',
+                         set_value=None):
     """
     Function to extrapolate variable values into undefined regions
 
@@ -19,6 +21,9 @@ def extrapolate_variable(nc_file, var_name, extrap_method, set_value=None):
     extrap_method : str
         idw, min, or value method of extrapolation
 
+    valid_region_method : str
+        choice of how to define region of valid data
+
     set_value : float
         value to set variable to outside keepCellMask
         when using -v value
@@ -36,13 +41,33 @@ def extrapolate_variable(nc_file, var_name, extrap_method, set_value=None):
     xCell = dataset.variables["yCell"][:]
     yCell = dataset.variables["xCell"][:]
 
+    # Define extrap method
+    if extrap_method == 'auto':
+        if var_name in ["effectivePressure", "beta", "muFriction"]:
+            extrap_method = 'min'
+        elif var_name in ["floatingBasalMassBal"]:
+            extrap_method = 'idw'
+        else:
+            extrap_method = 'idw'
+
     # Define region of good data to extrapolate from.
-    # Different methods for different variables
-    if var_name in ["effectivePressure", "beta", "muFriction"]:
+    if valid_region_method == 'auto':
+        if var_name in ["effectivePressure", "beta", "muFriction"]:
+            valid_region_method = 'grounded_ice'
+        elif var_name in ["floatingBasalMassBal"]:
+            valid_region_method = 'floating_ice'
+        else:
+            valid_region_method = 'ice'
+
+    print(f"Start {var_name} extrapolation using {extrap_method} method, "
+          f"defining good data using {valid_region_method} method")
+
+    # Calculate seed mask baed on method for defining valid data
+    if valid_region_method == 'positive_val':
+        keepCellMask = (varValue[:] > 0.0)
+    elif valid_region_method == 'grounded_ice':
         groundedMask = (thickness > (-1028.0 / 910.0 * bed))
         keepCellMask = np.copy(groundedMask)
-        extrap_method == "min"
-
         # grow mask by one cell oceanward of GL
         for iCell in range(nCells):
             for n in range(nEdgesOnCell[iCell]):
@@ -51,13 +76,19 @@ def extrapolate_variable(nc_file, var_name, extrap_method, set_value=None):
                     keepCellMask[iCell] = 1
                     continue
         # ensure zero muFriction does not get extrapolated
-        keepCellMask *= (varValue > 0)
-    elif var_name in ["floatingBasalMassBal"]:
+        keepCellMask *= (varValue > 0.0)
+    elif valid_region_method == 'floating_ice':
         floatingMask = (thickness <= (-1028.0 / 910.0 * bed))
         keepCellMask = floatingMask * (varValue != 0.0)
-        extrap_method == "idw"
-    else:
+    elif valid_region_method == 'ice':
         keepCellMask = (thickness > 0.0)
+    else:
+        sys.exit("Unexpected value of valid_region_method encountered "
+                 "in landice/extrapolate.py")
+
+    if keepCellMask.sum() == 0:
+        sys.exit("In landice/extrapolate.py, initial keepCellMask has "
+                 "no valid cells")
 
     # make a copy to edit that will be used later
     keepCellMaskNew = np.copy(keepCellMask)
@@ -72,8 +103,6 @@ def extrapolate_variable(nc_file, var_name, extrap_method, set_value=None):
     # 5) Update mask
     # 6) go to step 1)
 
-    print("Start {} extrapolation using {} method".format(var_name,
-                                                          extrap_method))
     if extrap_method == 'value':
         varValue[np.where(np.logical_not(keepCellMask))] = float(set_value)
     else:

compass/landice/tests/ensemble_generator/branch_ensemble/__init__.py

@@ -10,6 +10,9 @@
 from compass.landice.tests.ensemble_generator.ensemble_manager import (
     EnsembleManager,
 )
+from compass.landice.tests.ensemble_generator.ensemble_template import (
+    get_branch_template_package,
+)
 from compass.testcase import TestCase
 
 
@@ -59,6 +62,8 @@ def configure(self):
         """
 
         config = self.config
+        resource_module = get_branch_template_package(config)
+
         section = config['branch_ensemble']
 
         spinup_test_dir = section.get('spinup_test_dir')
@@ -89,7 +94,8 @@ def configure(self):
                 else:
                     print(f"Adding {run_name}")
                     # use this run
-                    self.add_step(BranchRun(test_case=self, run_num=run_num))
+                    self.add_step(BranchRun(test_case=self, run_num=run_num,
+                                            resource_module=resource_module))
                     # Note: do not add to steps_to_run; ensemble_manager
                     # will handle submitting and running the runs
 

compass/landice/tests/ensemble_generator/branch_ensemble/branch_ensemble.cfg

@@ -1,27 +1,3 @@
-# config options for branching an ensemble
-[branch_ensemble]
-
-# start and end numbers for runs to set up and run
-# branch runs.
-# It is assumed that spinup runs have already been
-# conducted for these runs.
-start_run = 0
-end_run = 3
-
-# Path to thermal forcing file for the mesh to be used in the branch run
-TF_file_path = /global/cfs/cdirs/fanssie/MALI_projects/Amery_UQ/Amery_4to20km_from_whole_AIS/forcing/ocean_thermal_forcing/UKESM1-0-LL_SSP585/1995-2300/Amery_4to20km_TF_UKESM1-0-LL_SSP585_2300.nc
-
-# Path to SMB forcing file for the mesh to be used in the branch run
-SMB_file_path = /global/cfs/cdirs/fanssie/MALI_projects/Amery_UQ/Amery_4to20km_from_whole_AIS/forcing/atmosphere_forcing/UKESM1-0-LL_SSP585/1995-2300/Amery_4to20km_SMB_UKESM1-0-LL_SSP585_2300_noBareLandAdvance.nc
-
-# location of spinup ensemble to branch from
-spinup_test_dir = /pscratch/sd/h/hoffman2/AMERY_corrected_forcing_6param_ensemble_2023-03-18/landice/ensemble_generator/ensemble
-
-# year of spinup simulation from which to branch runs
-branch_year = 2050
-
-# whether to only set up branch runs for filtered runs or all runs
-set_up_filtered_only = True
-
-# path to pickle file containing filtering information generated by plot_ensemble.py
-ensemble_pickle_file = None
+# branch_ensemble options are loaded from the selected model configuration
+# package under:
+# compass.landice.tests.ensemble_generator.ensemble_templates.<name>.branch

compass/landice/tests/ensemble_generator/branch_ensemble/branch_run.py

@@ -28,36 +28,9 @@ class BranchRun(Step):
     input_file_name : str
         name of the input file that was read from the config
 
-    basal_fric_exp : float
-        value of basal friction exponent to use
-
-    mu_scale : float
-        value to scale muFriction by
-
-    stiff_scale : float
-        value to scale stiffnessFactor by
-
-    von_mises_threshold : float
-        value of von Mises stress threshold to use
-
-    calv_spd_lim : float
-        value of calving speed limit to use
-
-    gamma0 : float
-        value of gamma0 to use in ISMIP6 ice-shelf basal melt param.
-
-    deltaT : float
-        value of deltaT to use in ISMIP6 ice-shelf basal melt param.
     """
 
-    def __init__(self, test_case, run_num,
-                 basal_fric_exp=None,
-                 mu_scale=None,
-                 stiff_scale=None,
-                 von_mises_threshold=None,
-                 calv_spd_lim=None,
-                 gamma0=None,
-                 deltaT=None):
+    def __init__(self, test_case, run_num, resource_module):
         """
         Creates a new run within an ensemble
 
@@ -68,8 +41,13 @@ def __init__(self, test_case, run_num,
 
         run_num : integer
             the run number for this ensemble member
+
+        resource_module : str
+            Package containing configuration-specific branch namelist and
+            streams templates
         """
         self.run_num = run_num
+        self.resource_module = resource_module
 
         # define step (run) name
         self.name = f'run{run_num:03}'
@@ -108,9 +86,10 @@ def setup(self):
         with open(os.path.join(self.work_dir, 'restart_timestamp'), 'w') as f:
             f.write('2015-01-01_00:00:00')
 
-        # yaml file
-        shutil.copy(os.path.join(spinup_dir, 'albany_input.yaml'),
-                    self.work_dir)
+        # albany_input.yaml may be absent in templates that do not use Albany.
+        albany_input = os.path.join(spinup_dir, 'albany_input.yaml')
+        if os.path.isfile(albany_input):
+            shutil.copy(albany_input, self.work_dir)
 
         # set up namelist
         # start with the namelist from the spinup
@@ -120,8 +99,7 @@ def setup(self):
                                                         'namelist.landice'))
         # use the namelist in this module to update the spinup namelist
         options = compass.namelist.parse_replacements(
-            'compass.landice.tests.ensemble_generator.branch_ensemble',
-            'namelist.landice')
+            self.resource_module, 'namelist.landice')
         namelist = compass.namelist.replace(namelist, options)
         compass.namelist.write(namelist, os.path.join(self.work_dir,
                                                       'namelist.landice'))
@@ -132,7 +110,7 @@ def setup(self):
         stream_replacements['TF_file_path'] = TF_file_path
         SMB_file_path = section.get('SMB_file_path')
         stream_replacements['SMB_file_path'] = SMB_file_path
-        strm_src = 'compass.landice.tests.ensemble_generator.branch_ensemble'
+        strm_src = self.resource_module
         self.add_streams_file(strm_src,
                               'streams.landice',
                               out_name='streams.landice',