Merge pull request #794 from xylar/cull-restarts-in-files-for-e3sm

Add util for culling MPAS-Ocean and -Seaice restart files

compass/ocean/tests/utility/__init__.py

@@ -1,3 +1,4 @@
+from compass.ocean.tests.utility.cull_restarts import CullRestarts
 from compass.ocean.tests.utility.extrap_woa import ExtrapWoa
 from compass.testgroup import TestGroup
 
@@ -14,4 +15,5 @@ def __init__(self, mpas_core):
         """
         super().__init__(mpas_core=mpas_core, name='utility')
 
+        self.add_test_case(CullRestarts(test_group=self))
         self.add_test_case(ExtrapWoa(test_group=self))

compass/ocean/tests/utility/cull_restarts/__init__.py

@@ -0,0 +1,22 @@
+from compass.ocean.tests.utility.cull_restarts.cull import Cull
+from compass.testcase import TestCase
+
+
+class CullRestarts(TestCase):
+    """
+    A test case for culling MPAS-Ocean and -Seaice restart files to exclude
+    ice-shelf cavities
+    """
+
+    def __init__(self, test_group):
+        """
+        Create the test case
+
+        Parameters
+        ----------
+        test_group : compass.ocean.tests.utility.Utility
+            The test group that this test case belongs to
+        """
+        super().__init__(test_group=test_group, name='cull_restarts')
+
+        self.add_step(Cull(test_case=self))

compass/ocean/tests/utility/cull_restarts/cull.py

@@ -0,0 +1,106 @@
+import xarray as xr
+from mpas_tools.io import write_netcdf
+from mpas_tools.logging import check_call
+from mpas_tools.mesh.cull import write_culled_dataset, write_map_culled_to_base
+
+from compass.step import Step
+
+
+class Cull(Step):
+    """
+    A step for culling MPAS-Ocean and -Seaice restart files to exclude
+    ice-shelf cavities
+    """
+
+    def __init__(self, test_case):
+        """
+        Create a new step
+
+        Parameters
+        ----------
+        test_case : compass.ocean.tests.utility.cull_restarts.CullRestarts
+            The test case this step belongs to
+        """
+        super().__init__(test_case, name='cull', cpus_per_task=128,
+                         min_cpus_per_task=1)
+        self.add_output_file(filename='unculled_mesh.nc')
+        self.add_output_file(filename='mesh_no_isc.nc')
+        self.add_output_file(filename='culled_graph.info')
+        self.add_output_file(filename='no_isc_to_culled_map.nc')
+        self.add_output_file(filename='culled_ocean_restart.nc')
+        self.add_output_file(filename='culled_seaice_restart.nc')
+
+    def run(self):
+        """
+        Run this step of the test case
+        """
+        config = self.config
+        logger = self.logger
+
+        section = config['cull_restarts']
+        ocean_restart_filename = section['ocean_restart_filename']
+        if ocean_restart_filename == '<<<Missing>>>':
+            raise ValueError('It looks like you did not set the '
+                             'ocean_restart_filename config option.')
+
+        seaice_restart_filename = section['seaice_restart_filename']
+        if seaice_restart_filename == '<<<Missing>>>':
+            raise ValueError('It looks like you did not set the '
+                             'seaice_restart_filename config option.')
+
+        culled_mesh_filename = 'mesh_no_isc.nc'
+        map_filename = 'no_isc_to_culled_map.nc'
+
+        mesh_vars = [
+            'areaCell', 'cellsOnCell', 'edgesOnCell', 'indexToCellID',
+            'latCell', 'lonCell', 'meshDensity', 'nEdgesOnCell',
+            'verticesOnCell', 'xCell', 'yCell', 'zCell', 'angleEdge',
+            'cellsOnEdge', 'dcEdge', 'dvEdge', 'edgesOnEdge',
+            'indexToEdgeID', 'latEdge', 'lonEdge', 'nEdgesOnCell',
+            'nEdgesOnEdge', 'verticesOnEdge', 'weightsOnEdge', 'xEdge',
+            'yEdge', 'zEdge', 'areaTriangle', 'cellsOnVertex', 'edgesOnVertex',
+            'indexToVertexID', 'kiteAreasOnVertex', 'latVertex',
+            'lonVertex', 'xVertex', 'yVertex', 'zVertex']
+        ds_unculled_mesh = xr.open_dataset(ocean_restart_filename)
+        ds_unculled_mesh = ds_unculled_mesh[mesh_vars + ['landIceMask']]
+        # cull cells where landIceMask == 1
+        ds_unculled_mesh = ds_unculled_mesh.rename({'landIceMask': 'cullCell'})
+
+        write_netcdf(ds_unculled_mesh, 'unculled_mesh.nc')
+
+        args = ['MpasCellCuller.x', 'unculled_mesh.nc', culled_mesh_filename]
+        # also produces culled_graph.info
+        check_call(args=args, logger=logger)
+
+        write_map_culled_to_base(base_mesh_filename=ocean_restart_filename,
+                                 culled_mesh_filename=culled_mesh_filename,
+                                 out_filename=map_filename,
+                                 workers=self.cpus_per_task)
+
+        in_filename = ocean_restart_filename
+        out_filename = 'culled_ocean_restart.nc'
+
+        write_culled_dataset(in_filename=in_filename,
+                             out_filename=out_filename,
+                             base_mesh_filename=ocean_restart_filename,
+                             culled_mesh_filename=culled_mesh_filename,
+                             map_culled_to_base_filename=map_filename,
+                             logger=logger)
+
+        with xr.open_dataset('culled_ocean_restart.nc') as ds:
+            ds = ds.drop_vars('xtime')
+            write_netcdf(ds, 'culled_ocean_restart_no_xtime.nc')
+
+        in_filename = seaice_restart_filename
+        out_filename = 'culled_seaice_restart.nc'
+
+        write_culled_dataset(in_filename=in_filename,
+                             out_filename=out_filename,
+                             base_mesh_filename=ocean_restart_filename,
+                             culled_mesh_filename=culled_mesh_filename,
+                             map_culled_to_base_filename=map_filename,
+                             logger=logger)
+
+        with xr.open_dataset('culled_seaice_restart.nc') as ds:
+            ds = ds.drop_vars('xtime')
+            write_netcdf(ds, 'culled_seaice_restart_no_xtime.nc')

compass/ocean/tests/utility/cull_restarts/cull_restarts.cfg

@@ -0,0 +1,9 @@
+# config options related to culling MPAS-Ocean and -Seaice restart files
+# to exclude ice-shelf cavities
+[cull_restarts]
+
+# the ocean restart file
+ocean_restart_filename = <<<Missing>>>
+
+# the sea-ice restart file
+seaice_restart_filename = <<<Missing>>>

docs/developers_guide/ocean/api.rst

@@ -907,6 +907,10 @@ utility
 
    Utility
 
+   cull_restarts.CullRestarts
+   cull_restarts.Cull
+   cull_restarts.Cull.run
+
    extrap_woa.ExtrapWoa
    extrap_woa.Combine
    extrap_woa.Combine.setup

docs/developers_guide/ocean/test_groups/utility.rst

@@ -8,6 +8,21 @@ may create datasets for other test groups to use.  It is partly designed to
 provide provenance for data processing that may be more complex than a single,
 short script.
 
+cull_restarts
+-------------
+The class :py:class:`compass.ocean.tests.utility.cull_restarts.CullRestarts`
+defines a test case for culling ice-shelf cavities from MPAS-Ocean and -Seaice
+restart files.
+
+cull
+~~~~
+
+The class :py:class:`compass.ocean.tests.utility.cull_restarts.Cull` defines
+a step for culling ice-shelf cavities from MPAS-Ocean and -Seaice
+restart files.  The step produces a culled mesh file, graph file, maps between
+the unculled and culled meshes in addition to culled versions of the restart
+files.
+
 extrap_woa
 ----------
 The class :py:class:`compass.ocean.tests.utility.extrap_woa.ExtrapWoa`

docs/users_guide/ocean/test_groups/utility.rst

@@ -8,11 +8,18 @@ may create datasets for other test groups to use.  It is partly designed to
 provide provenance for data processing that may be more complex than a single,
 short script.
 
+cull_restarts
+-------------
+The ``ocean/utility/cull_restarts`` test case is used to cull ice-shelf
+cavities from MPAS-Ocean and -Seaice restart files.  It is intended for
+expert users wanting to create an E3SM branch run without ice-shelf cavities
+from a previous run that included cavities.
+
 extrap_woa
 ----------
 The ``ocean/utility/extrap_woa`` test case is used to extrapolate
 `WOA 2023 <https://www.ncei.noaa.gov/products/world-ocean-atlas>`_ data into
 ice-shelf cavities and coastal regions, then into land, grounded ice and below
 bathymetry.  It is provided mainly for developers to update for use on
 future datasets and is not intended for users, so we do not provide full
-documentation here.
+documentation here.