Clean up global ocean test group for large meshes

compass/ocean/iceshelf.py

@@ -1,7 +1,9 @@
-import numpy
-import xarray
+import os
 import shutil
+import subprocess
 
+import numpy
+import xarray
 from mpas_tools.cime.constants import constants
 from mpas_tools.io import write_netcdf
 
@@ -34,12 +36,13 @@ def compute_land_ice_pressure_and_draft(ssh, modify_mask, ref_density):
     """
     gravity = constants['SHR_CONST_G']
     landIcePressure = \
-        modify_mask*numpy.maximum(-ref_density * gravity * ssh, 0.)
+        modify_mask * numpy.maximum(-ref_density * gravity * ssh, 0.)
     landIceDraft = ssh
     return landIcePressure, landIceDraft
 
 
-def adjust_ssh(variable, iteration_count, step):
+def adjust_ssh(variable, iteration_count, step, update_pio=True,
+               convert_to_cdf5=False):
     """
     Adjust the sea surface height or land-ice pressure to be dynamically
     consistent with one another.  A series of short model runs are performed,
@@ -55,24 +58,32 @@ def adjust_ssh(variable, iteration_count, step):
 
     step : compass.Step
         the step for performing SSH or land-ice pressure adjustment
+
+    update_pio : bool, optional
+        Whether to update PIO tasks and stride
+
+    convert_to_cdf5 : bool, optional
+        Whether to convert files to CDF5 format with ncks after writing them
+        out.  This is intended to improve MPAS-Ocean performance, since reading
+        in NETCDF4 format files can be very slow.
     """
     ntasks = step.ntasks
     config = step.config
     logger = step.logger
     out_filename = None
 
     if variable not in ['ssh', 'landIcePressure']:
-        raise ValueError("Unknown variable to modify: {}".format(variable))
+        raise ValueError(f"Unknown variable to modify: {variable}")
 
-    step.update_namelist_pio('namelist.ocean')
+    if update_pio:
+        step.update_namelist_pio('namelist.ocean')
     partition(ntasks, config, logger)
 
     for iterIndex in range(iteration_count):
-        logger.info(" * Iteration {}/{}".format(iterIndex + 1,
-                                                iteration_count))
+        logger.info(f" * Iteration {iterIndex + 1}/{iteration_count}")
 
-        in_filename = 'adjusting_init{}.nc'.format(iterIndex)
-        out_filename = 'adjusting_init{}.nc'.format(iterIndex+1)
+        in_filename = f'adjusting_init{iterIndex}.nc'
+        out_filename = f'adjusting_init{iterIndex + 1}.nc'
         symlink(in_filename, 'adjusting_init.nc')
 
         logger.info("   * Running forward model")
@@ -92,7 +103,7 @@ def adjust_ssh(variable, iteration_count, step):
 
             initSSH = ds.ssh
             if 'minLevelCell' in ds:
-                minLevelCell = ds.minLevelCell-1
+                minLevelCell = ds.minLevelCell - 1
             else:
                 minLevelCell = xarray.zeros_like(ds.maxLevelCell)
 
@@ -139,10 +150,18 @@ def adjust_ssh(variable, iteration_count, step):
 
                 finalSSH = initSSH
 
-            write_netcdf(ds_out, out_filename)
+            if convert_to_cdf5:
+                name, ext = os.path.splitext(out_filename)
+                write_filename = f'{name}_before_cdf5{ext}'
+            else:
+                write_filename = out_filename
+            write_netcdf(ds_out, write_filename)
+            if convert_to_cdf5:
+                args = ['ncks', '-5', write_filename, out_filename]
+                subprocess.check_call(args)
 
             # Write the largest change in SSH and its lon/lat to a file
-            with open('maxDeltaSSH_{:03d}.log'.format(iterIndex), 'w') as \
+            with open(f'maxDeltaSSH_{iterIndex:03d}.log', 'w') as \
                     log_file:
 
                 mask = landIcePressure > 0.
@@ -151,22 +170,21 @@ def adjust_ssh(variable, iteration_count, step):
                 ds_cell = ds.isel(nCells=iCell)
 
                 if on_a_sphere:
-                    coords = 'lon/lat: {:f} {:f}'.format(
-                        numpy.rad2deg(ds_cell.lonCell.values),
-                        numpy.rad2deg(ds_cell.latCell.values))
+                    coords = (f'lon/lat: '
+                              f'{numpy.rad2deg(ds_cell.lonCell.values):f} '
+                              f'{numpy.rad2deg(ds_cell.latCell.values):f}')
                 else:
-                    coords = 'x/y: {:f} {:f}'.format(
-                        1e-3 * ds_cell.xCell.values,
-                        1e-3 * ds_cell.yCell.values)
-                string = 'deltaSSHMax: {:g}, {}'.format(
-                    deltaSSH.isel(nCells=iCell).values, coords)
-                logger.info('     {}'.format(string))
-                log_file.write('{}\n'.format(string))
-                string = 'ssh: {:g}, landIcePressure: {:g}'.format(
-                    finalSSH.isel(nCells=iCell).values,
-                    landIcePressure.isel(nCells=iCell).values)
-                logger.info('     {}'.format(string))
-                log_file.write('{}\n'.format(string))
+                    coords = (f'x/y: {1e-3 * ds_cell.xCell.values:f} '
+                              f'{1e-3 * ds_cell.yCell.values:f}')
+                string = (f'deltaSSHMax: '
+                          f'{deltaSSH.isel(nCells=iCell).values:g}, {coords}')
+                logger.info(f'     {string}')
+                log_file.write(f'{string}\n')
+                string = (f'ssh: {finalSSH.isel(nCells=iCell).values:g}, '
+                          f'landIcePressure: '
+                          f'{landIcePressure.isel(nCells=iCell).values:g}')
+                logger.info(f'     {string}')
+                log_file.write(f'{string}\n')
 
         logger.info("   - Complete\n")
 

compass/ocean/tests/global_ocean/forward.py

@@ -198,7 +198,9 @@ def run(self):
         """
         Run this step of the testcase
         """
-        run_model(self)
+        update_pio = self.config.getboolean('global_ocean',
+                                            'forward_update_pio')
+        run_model(self, update_pio=update_pio)
         add_mesh_and_init_metadata(self.outputs, self.config,
                                    init_filename='init.nc')
 

compass/ocean/tests/global_ocean/global_ocean.cfg

@@ -10,6 +10,18 @@ cull_mesh_min_cpus_per_task = 1
 cull_mesh_max_memory = 1000
 
 
+# Options relate to adjusting the sea-surface height or land-ice pressure
+# below ice shelves to they are dynamically consistent with one another
+[ssh_adjustment]
+
+# the number of iterations of ssh adjustment to perform
+iterations = 10
+
+# Whether to convert adjusted initial condition files to CDF5 format during
+# ssh adjustment under ice shelves
+convert_to_cdf5 = False
+
+
 # options for global ocean testcases
 [global_ocean]
 
@@ -40,10 +52,17 @@ init_ntasks = 36
 init_min_tasks = 8
 # number of threads
 init_threads = 1
+# The number of cores per task in init mode -- used to avoid running out of
+# memory where needed
+init_cpus_per_task = 1
+# whether to update PIO tasks and stride
+init_update_pio = True
 
 ## config options related to the forward steps
 # number of threads
 forward_threads = 1
+# whether to update PIO tasks and stride
+forward_update_pio = True
 
 ## metadata related to the mesh
 # whether to add metadata to output files

compass/ocean/tests/global_ocean/init/initial_state.py

@@ -168,7 +168,8 @@ def run(self):
         plot_vertical_grid(grid_filename='vertical_grid.nc', config=config,
                            out_filename='vertical_grid.png')
 
-        run_model(self)
+        update_pio = config.getboolean('global_ocean', 'init_update_pio')
+        run_model(self, update_pio=update_pio)
 
         add_mesh_and_init_metadata(self.outputs, config,
                                    init_filename='initial_state.nc')
@@ -182,3 +183,6 @@ def _get_resources(self):
         self.ntasks = config.getint('global_ocean', 'init_ntasks')
         self.min_tasks = config.getint('global_ocean', 'init_min_tasks')
         self.openmp_threads = config.getint('global_ocean', 'init_threads')
+        self.cpus_per_task = config.getint('global_ocean',
+                                           'init_cpus_per_task')
+        self.min_cpus_per_task = self.cpus_per_task

compass/ocean/tests/global_ocean/init/ssh_adjustment.py

@@ -1,3 +1,5 @@
+from importlib.resources import contents
+
 from compass.ocean.iceshelf import adjust_ssh
 from compass.ocean.tests.global_ocean.forward import ForwardStep
 
@@ -28,6 +30,16 @@ def __init__(self, test_case):
         self.add_streams_file('compass.ocean.tests.global_ocean.init',
                               'streams.ssh_adjust')
 
+        mesh_package = test_case.mesh.package
+        mesh_package_contents = list(contents(mesh_package))
+        mesh_namelist = 'namelist.ssh_adjust'
+        if mesh_namelist in mesh_package_contents:
+            self.add_namelist_file(mesh_package, mesh_namelist)
+
+        mesh_streams = 'streams.ssh_adjust'
+        if mesh_streams in mesh_package_contents:
+            self.add_streams_file(mesh_package, mesh_streams)
+
         mesh_path = test_case.mesh.get_cull_mesh_path()
         init_path = test_case.steps['initial_state'].path
 
@@ -49,5 +61,9 @@ def run(self):
         """
         config = self.config
         iteration_count = config.getint('ssh_adjustment', 'iterations')
+        update_pio = config.getboolean('global_ocean', 'forward_update_pio')
+        convert_to_cdf5 = config.getboolean('ssh_adjustment',
+                                            'convert_to_cdf5')
         adjust_ssh(variable='landIcePressure', iteration_count=iteration_count,
-                   step=self)
+                   step=self, update_pio=update_pio,
+                   convert_to_cdf5=convert_to_cdf5)

docs/developers_guide/machines/index.rst

@@ -311,7 +311,7 @@ Finally, you might need to update the ``target`` and ``operating_system``.
 This is a bit of a "catch 22" in that you can use Spack to find this out but
 compass is designed to clone and set up Spack for you so we assume you don't
 have it yet.  For now, make your best guess using the info on
-`this page<https://spack.readthedocs.io/en/latest/basic_usage.html#architecture-specifiers>`_
+`this page <https://spack.readthedocs.io/en/latest/basic_usage.html#architecture-specifiers>`_
 and correct it later if necessary.
 
 You may need to load a system module to get the compilers and potentially other

docs/users_guide/ocean/test_groups/global_ocean.rst

@@ -45,6 +45,18 @@ Note that meshes and test cases may modify these options, as noted below.
     cull_mesh_max_memory = 1000
 
 
+    # Options relate to adjusting the sea-surface height or land-ice pressure
+    # below ice shelves to they are dynamically consistent with one another
+    [ssh_adjustment]
+
+    # the number of iterations of ssh adjustment to perform
+    iterations = 10
+
+    # Whether to convert adjusted initial condition files to CDF5 format during
+    # ssh adjustment under ice shelves
+    convert_to_cdf5 = False
+
+
     # options for global ocean testcases
     [global_ocean]
 
@@ -65,16 +77,27 @@ Note that meshes and test cases may modify these options, as noted below.
     btr_dt_per_km = 1.5
 
     ## config options related to the initial_state step
+
+    # Maximum allowed Haney number for configurations with ice-shelf cavities
+    rx1_max = 20
+
     # number of cores to use
     init_ntasks = 36
     # minimum of cores, below which the step fails
     init_min_tasks = 8
     # number of threads
     init_threads = 1
+    # The number of cores per task in init mode -- used to avoid running out of
+    # memory where needed
+    init_cpus_per_task = 1
+    # whether to update PIO tasks and stride
+    init_update_pio = True
 
     ## config options related to the forward steps
     # number of threads
     forward_threads = 1
+    # whether to update PIO tasks and stride
+    forward_update_pio = True
 
     ## metadata related to the mesh
     # whether to add metadata to output files