Refactor ensemble generator for flexibility

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

@@ -10,6 +10,10 @@
 from compass.landice.tests.ensemble_generator.ensemble_manager import (
     EnsembleManager,
 )
+from compass.landice.tests.ensemble_generator.ensemble_template import (
+    add_template_file,
+    get_branch_template_package,
+)
 from compass.testcase import TestCase
 
 
@@ -59,6 +63,9 @@ def configure(self):
         """
 
         config = self.config
+        resource_module = get_branch_template_package(config)
+        add_template_file(config, resource_module, 'branch_ensemble.cfg')
+
         section = config['branch_ensemble']
 
         spinup_test_dir = section.get('spinup_test_dir')
@@ -89,7 +96,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',

compass/landice/tests/ensemble_generator/ensemble_member.py

@@ -40,12 +40,6 @@ class EnsembleMember(Step):
     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.
 
@@ -54,11 +48,12 @@ class EnsembleMember(Step):
     """
 
     def __init__(self, test_case, run_num,
+                 resource_module,
+                 namelist_option_values=None,
+                 namelist_parameter_values=None,
                  basal_fric_exp=None,
                  mu_scale=None,
                  stiff_scale=None,
-                 von_mises_threshold=None,
-                 calv_spd_lim=None,
                  gamma0=None,
                  meltflux=None,
                  deltaT=None):
@@ -73,6 +68,18 @@ def __init__(self, test_case, run_num,
         run_num : integer
             the run number for this ensemble member
 
+        resource_module : str
+            Package containing configuration-specific namelist, streams,
+            and albany input files
+
+        namelist_option_values : dict, optional
+            A dictionary of namelist option names and values to be
+            overridden for this ensemble member
+
+        namelist_parameter_values : dict, optional
+            A dictionary of run-info parameter names and values that
+            correspond to entries in ``namelist_option_values``
+
         basal_fric_exp : float
             value of basal friction exponent to use
 
@@ -82,27 +89,25 @@ def __init__(self, test_case, run_num,
         stiff_scale : float
             value to scale stiffnessFactor by
 
-        von_mises_threshold : float
-            value of von Mises stress threshold to use
-            assumes same value for grounded and floating ice
-
-        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.
         """
         self.run_num = run_num
+        self.resource_module = resource_module
+        if namelist_option_values is None:
+            namelist_option_values = {}
+        if namelist_parameter_values is None:
+            namelist_parameter_values = {}
+        self.namelist_option_values = dict(namelist_option_values)
+        self.namelist_parameter_values = dict(namelist_parameter_values)
 
         # store assigned param values for this run
         self.basal_fric_exp = basal_fric_exp
         self.mu_scale = mu_scale
         self.stiff_scale = stiff_scale
-        self.von_mises_threshold = von_mises_threshold
-        self.calv_spd_lim = calv_spd_lim
         self.gamma0 = gamma0
         self.meltflux = meltflux
         self.deltaT = deltaT
@@ -127,11 +132,12 @@ def setup(self):
                   "'compass setup' again to set this experiment up.")
             return
 
-        resource_module = 'compass.landice.tests.ensemble_generator'
+        resource_module = self.resource_module
 
         # Get config for info needed for setting up simulation
         config = self.config
-        section = config['ensemble']
+        section = config['ensemble_generator']
+        spinup_section = config['spinup_ensemble']
 
         # Create a python config (not compass config) file
         # for run-specific info useful for analysis/viz
@@ -151,14 +157,17 @@ def setup(self):
         # Set up base run configuration
         self.add_namelist_file(resource_module, 'namelist.landice')
 
-        # copy over albany yaml file
-        # cannot use add_input functionality because need to modify the file
-        # in this function, and inputs don't get processed until after this
-        # function
-        with resources.path(resource_module,
-                            'albany_input.yaml') as package_path:
-            target = str(package_path)
-            shutil.copy(target, self.work_dir)
+        # albany_input.yaml is optional unless fric_exp perturbations are used.
+        albany_input_name = 'albany_input.yaml'
+        albany_input_path = os.path.join(self.work_dir, albany_input_name)
+        albany_source = resources.files(resource_module).joinpath(
+            albany_input_name)
+        # Materialize a real filesystem path in case the package is not
+        # directly on the filesystem (e.g., zip/loader-backed).
+        with resources.as_file(albany_source) as albany_source_path:
+            has_albany_input = albany_source_path.is_file()
+            if has_albany_input:
+                shutil.copy(str(albany_source_path), self.work_dir)
 
         self.add_model_as_input()
 
@@ -171,25 +180,15 @@ def setup(self):
         options['config_adaptive_timestep_CFL_fraction'] = \
             f'{self.cfl_fraction}'
 
-        # von Mises stress threshold
-        if self.von_mises_threshold is not None:
-            options['config_grounded_von_Mises_threshold_stress'] = \
-                f'{self.von_mises_threshold}'
-            options['config_floating_von_Mises_threshold_stress'] = \
-                f'{self.von_mises_threshold}'
-            run_info_cfg.set('run_info', 'von_mises_threshold',
-                             f'{self.von_mises_threshold}')
-
-        # calving speed limit
-        if self.calv_spd_lim is not None:
-            options['config_calving_speed_limit'] = \
-                f'{self.calv_spd_lim}'
-            run_info_cfg.set('run_info', 'calv_spd_limit',
-                             f'{self.calv_spd_lim}')
+        # apply generic namelist float parameter perturbations
+        for option_name, value in self.namelist_option_values.items():
+            options[option_name] = f'{value}'
+        for parameter_name, value in self.namelist_parameter_values.items():
+            run_info_cfg.set('run_info', parameter_name, f'{value}')
 
         # adjust basal friction exponent
         # rename and copy base file
-        input_file_path = section.get('input_file_path')
+        input_file_path = spinup_section.get('input_file_path')
         input_file_name = input_file_path.split('/')[-1]
         base_fname = input_file_name.split('.')[:-1][0]
         new_input_fname = f'{base_fname}_MODIFIED.nc'
@@ -199,13 +198,16 @@ def setup(self):
         # set input filename in streams and create streams file
         stream_replacements = {'input_file_init_cond': new_input_fname}
         if self.basal_fric_exp is not None:
+            if not has_albany_input:
+                raise ValueError(
+                    "Parameter 'fric_exp' requires 'albany_input.yaml' "
+                    f"in template package '{resource_module}'.")
             # adjust mu and exponent
-            orig_fric_exp = section.getfloat('orig_fric_exp')
+            orig_fric_exp = spinup_section.getfloat('orig_fric_exp')
             _adjust_friction_exponent(orig_fric_exp, self.basal_fric_exp,
                                       os.path.join(self.work_dir,
                                                    new_input_fname),
-                                      os.path.join(self.work_dir,
-                                                   'albany_input.yaml'))
+                                      albany_input_path)
             run_info_cfg.set('run_info', 'basal_fric_exp',
                              f'{self.basal_fric_exp}')
 
@@ -227,7 +229,8 @@ def setup(self):
 
         # adjust gamma0 and deltaT
         # (only need to check one of these params)
-        basal_melt_param_file_path = section.get('basal_melt_param_file_path')
+        basal_melt_param_file_path = spinup_section.get(
+            'basal_melt_param_file_path')
         basal_melt_param_file_name = basal_melt_param_file_path.split('/')[-1]
         base_fname = basal_melt_param_file_name.split('.')[:-1][0]
         new_fname = f'{base_fname}_MODIFIED.nc'
@@ -243,9 +246,9 @@ def setup(self):
             run_info_cfg.set('run_info', 'deltaT', f'{self.deltaT}')
 
         # set up forcing files (unmodified)
-        TF_file_path = section.get('TF_file_path')
+        TF_file_path = spinup_section.get('TF_file_path')
         stream_replacements['TF_file_path'] = TF_file_path
-        SMB_file_path = section.get('SMB_file_path')
+        SMB_file_path = spinup_section.get('SMB_file_path')
         stream_replacements['SMB_file_path'] = SMB_file_path
 
         # store accumulated namelist and streams options