diff --git a/compass/ocean/iceshelf.py b/compass/ocean/iceshelf.py index d9d9000e35..19bf6d90ea 100644 --- a/compass/ocean/iceshelf.py +++ b/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,6 +58,14 @@ 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 @@ -62,17 +73,17 @@ def adjust_ssh(variable, iteration_count, step): 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") diff --git a/compass/ocean/tests/global_ocean/forward.py b/compass/ocean/tests/global_ocean/forward.py index 69abcdc988..15776e5422 100644 --- a/compass/ocean/tests/global_ocean/forward.py +++ b/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') diff --git a/compass/ocean/tests/global_ocean/global_ocean.cfg b/compass/ocean/tests/global_ocean/global_ocean.cfg index 3526590435..3e97acfc6d 100644 --- a/compass/ocean/tests/global_ocean/global_ocean.cfg +++ b/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 diff --git a/compass/ocean/tests/global_ocean/init/initial_state.py b/compass/ocean/tests/global_ocean/init/initial_state.py index d824c6279a..d9ea96eedf 100644 --- a/compass/ocean/tests/global_ocean/init/initial_state.py +++ b/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 diff --git a/compass/ocean/tests/global_ocean/init/ssh_adjustment.py b/compass/ocean/tests/global_ocean/init/ssh_adjustment.py index 8625708196..7e8e4a445b 100644 --- a/compass/ocean/tests/global_ocean/init/ssh_adjustment.py +++ b/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) diff --git a/docs/developers_guide/machines/index.rst b/docs/developers_guide/machines/index.rst index cee7ac556a..c84ec310b7 100644 --- a/docs/developers_guide/machines/index.rst +++ b/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`_ +`this page `_ and correct it later if necessary. You may need to load a system module to get the compilers and potentially other diff --git a/docs/users_guide/ocean/test_groups/global_ocean.rst b/docs/users_guide/ocean/test_groups/global_ocean.rst index 67e779ee0b..936be87916 100644 --- a/docs/users_guide/ocean/test_groups/global_ocean.rst +++ b/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