Sgh ensemble generator template

compass/landice/tests/ensemble_generator/__init__.py

@@ -1,6 +1,18 @@
+"""
+Ensemble generator test group for MALI simulations.
+
+Provides test cases for generating ensembles for UQ and sensitivity studies.
+"""
+
 from compass.landice.tests.ensemble_generator.branch_ensemble import (
     BranchEnsemble,
 )
+from compass.landice.tests.ensemble_generator.sgh_ensemble_analysis import (
+    AnalysisEnsemble,
+)
+from compass.landice.tests.ensemble_generator.sgh_restart_ensemble import (
+    RestartEnsemble,
+)
 from compass.landice.tests.ensemble_generator.spinup_ensemble import (
     SpinupEnsemble,
 )
@@ -14,6 +26,8 @@ class EnsembleGenerator(TestGroup):
     """
     def __init__(self, mpas_core):
         """
+        Parameters
+        ----------
         mpas_core : compass.landice.Landice
             the MPAS core that this test group belongs to
         """
@@ -22,3 +36,5 @@ def __init__(self, mpas_core):
 
         self.add_test_case(SpinupEnsemble(test_group=self))
         self.add_test_case(BranchEnsemble(test_group=self))
+        self.add_test_case(AnalysisEnsemble(test_group=self))
+        self.add_test_case(RestartEnsemble(test_group=self))

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',