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

compass/landice/tests/ensemble_generator/ensemble_manager.py

@@ -1,5 +1,6 @@
 import glob
 import os
+import re
 from importlib.resources import path
 
 from mpas_tools.logging import check_call
@@ -9,6 +10,49 @@
 from compass.step import Step
 
 
+def _write_restart_job_script(spinup_run_dir, step_work_dir):
+    """
+    Rewrite ``job_script.sh`` in *spinup_run_dir* so that the SLURM job
+    changes into the compass step work directory before calling ``compass run``.
+
+    The existing script's SBATCH headers (and any ``source`` / environment
+    setup lines) are preserved verbatim.  Only the ``compass run`` invocation
+    at the end is replaced with::
+
+        cd <step_work_dir>
+        compass run
+
+    Parameters
+    ----------
+    spinup_run_dir : str
+        Path to the original spinup run directory that contains the
+        ``job_script.sh`` to be rewritten.
+
+    step_work_dir : str
+        Absolute path to the compass step work directory where
+        ``step.pickle`` lives (i.e. ``InPlaceRestartMember.work_dir``).
+    """
+    job_script_path = os.path.join(spinup_run_dir, 'job_script.sh')
+
+    with open(job_script_path, 'r') as f:
+        lines = f.readlines()
+
+    new_lines = []
+    compass_run_pattern = re.compile(r'^(\s*)compass\s+run(.*)')
+    for line in lines:
+        m = compass_run_pattern.match(line)
+        if m:
+            indent = m.group(1)
+            trailing = m.group(2)
+            new_lines.append(f'{indent}cd {step_work_dir}\n')
+            new_lines.append(f'{indent}compass run{trailing}\n')
+        else:
+            new_lines.append(line)
+
+    with open(job_script_path, 'w') as f:
+        f.writelines(new_lines)
+
+
 class EnsembleManager(Step):
     """
     A step for an ensemble manager.
@@ -64,7 +108,8 @@ def run(self):
         for run in runs:
             # Get step object from 'steps' dictionary
             runStep = self.test_case.steps[run]
-            os.chdir(runStep.work_dir)
+            run_dir = getattr(runStep, 'spinup_run_dir', runStep.work_dir)
+            os.chdir(run_dir)
 
             # Skip run if an error is detected
             err_files = glob.glob('log.landice.*.err')
@@ -83,6 +128,14 @@ def run(self):
                     print(f'{run} has reached stop time.  Skipping.')
                     continue
             # If we didn't skip this run, submit it
+            # For InPlaceRestartMember steps the step.pickle lives in
+            # runStep.work_dir, not in spinup_run_dir.  Rewrite the
+            # job_script.sh to cd into the compass step directory before
+            # calling 'compass run' so SLURM picks up the correct pickle.
+            if hasattr(runStep, 'spinup_run_dir'):
+                _write_restart_job_script(
+                    spinup_run_dir=run_dir,
+                    step_work_dir=runStep.work_dir)
             check_call(['sbatch', 'job_script.sh'], logger)
             logger.info(f'{run} submitted.')
         logger.info('All runs submitted.')