Merge pull request #616 from matthewhoffman/landice/ensemble_projection

Create ensemble of branch runs for projections

This PR introduces a branch_ensemble test within the ensemble_generator test group that can be used to set up an ensemble of branch runs from the spinup ensemble for the purpose of running projections

compass/landice/ais_observations.py

@@ -43,7 +43,14 @@
         'input': [73, 10],
         'outflow': [77, 4],
         'net': [-4, 11],
-        'shelf_melt': [35.5, 23.0]},
+        'shelf_melt': [35.5, 23.0],
+        'filter criteria': {
+            'total area change rate': {
+                'values': [-99.7361194, 75.85552239]},  # km2
+            'grd area change rate': {
+                'values': [-29.1723, 31.1477]},  # km2/yr
+            'grd vol change rate': {
+                'values': [-24.54, 21.228]}}},  # Gt/yr
     'ISMIP6BasinCCp': {
         'name': 'Phillipi, Denman',
         'input': [81, 13],

compass/landice/tests/ensemble_generator/__init__.py

@@ -1,4 +1,9 @@
-from compass.landice.tests.ensemble_generator.ensemble import Ensemble
+from compass.landice.tests.ensemble_generator.branch_ensemble import (
+    BranchEnsemble,
+)
+from compass.landice.tests.ensemble_generator.spinup_ensemble import (
+    SpinupEnsemble,
+)
 from compass.testgroup import TestGroup
 
 
@@ -15,4 +20,5 @@ def __init__(self, mpas_core):
         super().__init__(mpas_core=mpas_core,
                          name='ensemble_generator')
 
-        self.add_test_case(Ensemble(test_group=self))
+        self.add_test_case(SpinupEnsemble(test_group=self))
+        self.add_test_case(BranchEnsemble(test_group=self))

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

@@ -0,0 +1,103 @@
+import os
+import pickle
+import sys
+
+import numpy as np
+
+from compass.landice.tests.ensemble_generator.branch_ensemble.branch_run import (  # noqa
+    BranchRun,
+)
+from compass.landice.tests.ensemble_generator.ensemble_manager import (
+    EnsembleManager,
+)
+from compass.testcase import TestCase
+
+
+class BranchEnsemble(TestCase):
+    """
+    A test case for performing an ensemble of
+    simulations for uncertainty quantification studies.
+    """
+
+    def __init__(self, test_group):
+        """
+        Create the test case
+
+        Parameters
+        ----------
+        test_group : compass.landice.tests.ensemble_generator.EnsembleGenerator
+            The test group that this test case belongs to
+
+        """
+        name = 'branch_ensemble'
+        super().__init__(test_group=test_group, name=name)
+
+        # We don't want to initialize all the individual runs
+        # So during init, we only add the run manager
+        self.add_step(EnsembleManager(test_case=self))
+
+    def configure(self):
+        """
+        Configure a parameter ensemble of MALI simulations.
+
+        Start by identifying the start and end run numbers to set up
+        from the config.
+
+        Next, read a pre-defined unit parameter vector that can be used
+        for assigning parameter values to each ensemble member.
+
+        The main work is using the unit parameter vector to set parameter
+        values for each parameter to be varied, over prescribed ranges.
+
+        Then create the ensemble member as a step in the test case by calling
+        the EnsembleMember constructor.
+
+        Finally, add this step to the test case's step_to_run.  This normally
+        happens automatically if steps are added to the test case in the test
+        case constructor, but because we waited to add these steps until this
+        configure phase, we must explicitly add the steps to steps_to_run.
+        """
+
+        config = self.config
+        section = config['branch_ensemble']
+
+        spinup_test_dir = section.get('spinup_test_dir')
+        branch_year = section.getint('branch_year')
+
+        # Determine start and end run numbers being requested
+        self.start_run = section.getint('start_run')
+        self.end_run = section.getint('end_run')
+
+        # Determine whether to only set up filtered runs
+        self.set_up_filtered_only = section.getboolean('set_up_filtered_only')
+        self.ensemble_pickle_file = section.get('ensemble_pickle_file')
+        if self.set_up_filtered_only:
+            with open(self.ensemble_pickle_file, 'rb') as f:
+                [param_info, qoi_info] = pickle.load(f)
+            filtered_runs = np.isfinite(qoi_info['VAF change']['values'])
+        else:
+            filtered_runs = np.ones((self.end_run + 1,))
+
+        for run_num in range(self.start_run, self.end_run + 1):
+            run_name = f'run{run_num:03}'
+            if (filtered_runs[run_num] and
+                os.path.isfile(os.path.join(spinup_test_dir, run_name,
+                                            f'rst.{branch_year}-01-01.nc'))):
+                if os.path.exists(os.path.join(self.work_dir, run_name)):
+                    print(f"WARNING: {run_name} path already exists; "
+                          "skipping.  ")
+                else:
+                    print(f"Adding {run_name}")
+                    # use this run
+                    self.add_step(BranchRun(test_case=self, run_num=run_num))
+                    # Note: do not add to steps_to_run; ensemble_manager
+                    # will handle submitting and running the runs
+
+        # Have compass run only run the run_manager but not any actual runs.
+        # This is because the individual runs will be submitted as jobs
+        # by the ensemble manager.
+        self.steps_to_run = ['ensemble_manager',]
+
+    # no run() method is needed
+
+    # no validate() method is needed

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

@@ -0,0 +1,27 @@
+# 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

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

@@ -0,0 +1,184 @@
+import os
+import shutil
+
+import netCDF4
+
+import compass.namelist
+from compass.io import symlink
+from compass.job import write_job_script
+from compass.model import run_model
+from compass.step import Step
+
+
+class BranchRun(Step):
+    """
+    A step for setting up a single ensemble member
+
+    Attributes
+    ----------
+    run_num : integer
+        the run number for this ensemble member
+
+    name : str
+        the name of the run being set up in this step
+
+    ntasks : integer
+        the number of parallel (MPI) tasks the step would ideally use
+
+    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):
+        """
+        Creates a new run within an ensemble
+
+        Parameters
+        ----------
+        test_case : compass.TestCase
+            The test case this step belongs to
+
+        run_num : integer
+            the run number for this ensemble member
+        """
+        self.run_num = run_num
+
+        # define step (run) name
+        self.name = f'run{run_num:03}'
+
+        super().__init__(test_case=test_case, name=self.name)
+
+    def setup(self):
+        """
+        Set up this run by setting up a baseline run configuration
+        and then modifying parameters for this ensemble member based on
+        an externally provided unit parameter vector
+        """
+
+        print(f'Setting up run number {self.run_num}')
+
+        config = self.config
+        section = config['branch_ensemble']
+
+        spinup_test_dir = section.get('spinup_test_dir')
+        branch_year = section.getint('branch_year')
+
+        spinup_dir = os.path.join(os.path.join(spinup_test_dir, self.name))
+
+        # copy over the following:
+        # restart file - but change year
+        rst_file = os.path.join(spinup_dir, f'rst.{branch_year:04}-01-01.nc')
+        shutil.copy(rst_file, os.path.join(self.work_dir,
+                                           'rst.2015-01-01.nc'))
+        f = netCDF4.Dataset(os.path.join(self.work_dir,
+                                         'rst.2015-01-01.nc'), 'r+')
+        xtime = f.variables['xtime']
+        xtime[0, :] = list('2015-01-01_00:00:00'.ljust(64))
+        f.close()
+
+        # create restart_timestamp
+        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)
+
+        # set up namelist
+        # start with the namelist from the spinup
+        # Note: this differs from most compass tests, which would start with
+        # the default namelist from the mpas build dir
+        namelist = compass.namelist.ingest(os.path.join(spinup_dir,
+                                                        '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')
+        namelist = compass.namelist.replace(namelist, options)
+        compass.namelist.write(namelist, os.path.join(self.work_dir,
+                                                      'namelist.landice'))
+
+        # set up streams
+        stream_replacements = {}
+        TF_file_path = section.get('TF_file_path')
+        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'
+        self.add_streams_file(strm_src,
+                              'streams.landice',
+                              out_name='streams.landice',
+                              template_replacements=stream_replacements)
+
+        # copy run_info file
+        shutil.copy(os.path.join(spinup_dir, 'run_info.cfg'), self.work_dir)
+
+        # copy graph files
+        shutil.copy(os.path.join(spinup_dir, 'graph.info'), self.work_dir)
+
+        # save number of tasks to use
+        # eventually compass could determine this, but for now we want
+        # explicit control
+        self.ntasks = self.config.getint('ensemble', 'ntasks')
+        self.min_tasks = self.ntasks
+
+        self.add_model_as_input()
+
+        # set job name to run number so it will get set in batch script
+        # Note: currently, for this to work right, one has to delete/comment
+        # the call to write_job_script at line 316-7 in compass/setup.py
+        self.config.set('job', 'job_name', f'uq_{self.name}')
+        machine = self.config.get('deploy', 'machine')
+        pre_run_cmd = ('LOGDIR=previous_logs_`date +"%Y-%m-%d_%H-%M-%S"`;'
+                       'mkdir $LOGDIR; cp log* $LOGDIR; date')
+        post_run_cmd = "date"
+        write_job_script(self.config, machine,
+                         target_cores=self.ntasks, min_cores=self.min_tasks,
+                         work_dir=self.work_dir,
+                         pre_run_commands=pre_run_cmd,
+                         post_run_commands=post_run_cmd)
+
+        # COMPASS does not create symlinks for the load script in step dirs,
+        # so use the standard approach for creating that symlink in each
+        # step dir.
+        if 'LOAD_COMPASS_ENV' in os.environ:
+            script_filename = os.environ['LOAD_COMPASS_ENV']
+            # make a symlink to the script for loading the compass conda env.
+            symlink(script_filename, os.path.join(self.work_dir,
+                                                  'load_compass_env.sh'))
+
+    def run(self):
+        """
+        Run this member of the ensemble.
+        Eventually we want this function to handle restarts.
+        """
+
+        run_model(self)

compass/landice/tests/ensemble_generator/branch_ensemble/copy_existing_runs.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+# This script can be used to copy over all of the contents of an existing ensemble
+# to a fresh generation of the ensemble directory structure.
+# This is useful if your compass environment gets corrupted or similar
+# situation requiring you to start over with creating the environment.
+
+SRC_PATH=/pscratch/sd/h/hoffman2/AMERY_corrected_forcing_6param_ensemble_2023-03-18_branch_runs_yr2050_2023-07-28_noGroundedCalving_noCalvingError_calvingMetrics_200runsFiltered/landice/ensemble_generator/branch_ensemble
+DEST_PATH=/pscratch/sd/h/hoffman2/AMERY_corrected_forcing_6param_ensemble_2023-03-18_branch_runs_yr2050_2023-08-31/landice/ensemble_generator/branch_ensemble
+
+for dir in ${SRC_PATH}/run* ; do
+    echo $dir
+    run=`basename $dir`
+    cp ${dir}/log.landice* ${DEST_PATH}/${run}/
+    cp ${dir}/restart_timestamp ${DEST_PATH}/${run}/
+    cp ${dir}/rst*nc ${DEST_PATH}/${run}/
+    cp ${dir}/uq_run*.o* ${DEST_PATH}/${run}/
+    cp -R ${dir}/output ${DEST_PATH}/${run}/
+done

compass/landice/tests/ensemble_generator/branch_ensemble/namelist.landice

@@ -0,0 +1,8 @@
+config_do_restart = .true.
+config_start_time = 'file'
+config_stop_time = '2300-01-01_00:00:00'
+config_grounded_von_Mises_threshold_stress = 1.0e9
+config_min_adaptive_timestep = 21600
+config_calving_error_threshold = 1.0e9
+config_front_mass_bal_grounded = 'ismip6'
+config_use_3d_thermal_forcing_for_face_melt = .true.