diff --git a/compass/ocean/tests/hurricane/__init__.py b/compass/ocean/tests/hurricane/__init__.py index 4601e44914..0a2eb79245 100644 --- a/compass/ocean/tests/hurricane/__init__.py +++ b/compass/ocean/tests/hurricane/__init__.py @@ -1,3 +1,4 @@ +from compass.ocean.tests.hurricane.files_for_e3sm import FilesForE3SM from compass.ocean.tests.hurricane.forward import Forward from compass.ocean.tests.hurricane.init import Init from compass.ocean.tests.hurricane.mesh import Mesh @@ -43,3 +44,8 @@ def __init__(self, mpas_core): use_lts=use_lts, wetdry=wetdry) self.add_test_case(forward) + + if not use_lts: + files_for_e3sm = FilesForE3SM(test_group=self, + mesh=mesh) + self.add_test_case(files_for_e3sm) diff --git a/compass/ocean/tests/hurricane/files_for_e3sm/__init__.py b/compass/ocean/tests/hurricane/files_for_e3sm/__init__.py new file mode 100644 index 0000000000..ca00690e5e --- /dev/null +++ b/compass/ocean/tests/hurricane/files_for_e3sm/__init__.py @@ -0,0 +1,48 @@ +import os +from datetime import datetime + +from compass.ocean.tests.hurricane.configure import configure_hurricane +from compass.ocean.tests.hurricane.files_for_e3sm.domain_files import ( + DomainFiles, +) +from compass.ocean.tests.hurricane.files_for_e3sm.forcing_maps import ( + ForcingMaps, +) +from compass.testcase import TestCase + + +class FilesForE3SM(TestCase): + """ + A test case for assembling files needed for MPAS-Ocean forcing in E3SM + + Attributes + ---------- + mesh : compass.ocean.tests.hurricane.mesh.Mesh + The test case that produces the mesh for this run + """ + def __init__(self, test_group, mesh=None): + """ + Create test case for creating a global MPAS-Ocean mesh + + Parameters + ---------- + test_group : compass.ocean.tests.hurricane.Hurricane + The global ocean test group that this test case belongs to + + mesh : compass.ocean.tests.hurricane.mesh.Mesh, optional + The test case that produces the mesh for this run + """ + + name = 'files_for_e3sm' + subdir = os.path.join(mesh.mesh_name, name) + super().__init__(test_group=test_group, name=name, subdir=subdir) + self.mesh = mesh + self.creation_date = datetime.now().strftime('%Y%m%d') + self.add_step(ForcingMaps(test_case=self)) + self.add_step(DomainFiles(test_case=self)) + + def configure(self): + """ + Modify the configuration options for this test case + """ + configure_hurricane(test_case=self, mesh=self.mesh) diff --git a/compass/ocean/tests/hurricane/files_for_e3sm/domain_files.py b/compass/ocean/tests/hurricane/files_for_e3sm/domain_files.py new file mode 100644 index 0000000000..eaa1737ae5 --- /dev/null +++ b/compass/ocean/tests/hurricane/files_for_e3sm/domain_files.py @@ -0,0 +1,87 @@ +import os + +from mpas_tools.logging import check_call + +from compass.step import Step + + +class DomainFiles(Step): + """ + A step for building domain files for meshes tailored to hurricane-induced + coastal flooding + + Attributes + ---------- + mesh : compass.mesh.spherical.SphericalBaseStep + The base mesh step containing input files to this step + """ + + def __init__(self, test_case): + """ + Create a new step + + Parameters + ---------- + test_case : compass.ocean.tests.hurricane.files_for_e3sm.FilesForE3SM + The test case this step belongs to + """ + super().__init__(test_case, name='domain_files') + self.mesh = test_case.mesh + self.creation_date = test_case.creation_date + + def setup(self): + """ + Set up the step in the work directory, including downloading any + dependencies. + """ + super().setup() + + atm_grid = self.config.get('files_for_e3sm', 'atm_grid') + ocn_grid = self.mesh.mesh_name + creation_date = self.creation_date + + map_file_path = os.path.join( + self.test_case.steps['forcing_maps'].path, + f'map_{ocn_grid}_to_{atm_grid}_traave.{creation_date}.nc', + ) + self.add_input_file( + filename='map_ocn_to_atm_traave.nc', work_dir_target=map_file_path) + self.add_output_file( + filename=f'domain.lnd.{atm_grid}_{ocn_grid}.{creation_date}.nc') + self.add_output_file( + filename=f'domain.ocn.{atm_grid}_{ocn_grid}.{creation_date}.nc') + self.add_output_file( + filename=f'domain.ocn.{ocn_grid}.{creation_date}.nc') + + def run(self): + """ + Run this step of the test case + """ + super().run() + self._domain_files() + + def _domain_files(self): + """ + Create domain files + """ + + section = self.config['files_for_e3sm'] + domain_files_exe = section.get('domain_files_exe') + atm_grid = section.get('atm_grid') + ocn_grid = self.mesh.mesh_name + + logger = self.logger + logger.info( + f'Create domain files for {atm_grid}_{ocn_grid} grid pair.' + ) + + args = [ + 'python', domain_files_exe, + '--date-stamp', self.creation_date, + '-m', 'map_ocn_to_atm_traave.nc', + '-o', ocn_grid, + '-l', atm_grid, + ] + check_call(args, logger) + + logger.info(' Done.') diff --git a/compass/ocean/tests/hurricane/files_for_e3sm/forcing_maps.py b/compass/ocean/tests/hurricane/files_for_e3sm/forcing_maps.py new file mode 100644 index 0000000000..dbd0461561 --- /dev/null +++ b/compass/ocean/tests/hurricane/files_for_e3sm/forcing_maps.py @@ -0,0 +1,213 @@ +from mpas_tools.logging import check_call +from pyremap import MpasCellMeshDescriptor + +from compass.parallel import run_command +from compass.step import Step + + +class ForcingMaps(Step): + """ + A step for building mapping files for remapping forcing files + to a global MPAS-Ocean mesh + + Attributes + ---------- + mesh : compass.mesh.spherical.SphericalBaseStep + The base mesh step containing input files to this step + """ + + def __init__(self, test_case): + """ + Create a new step + + Parameters + ---------- + test_case : compass.ocean.tests.hurricane.files_for_e3sm.FilesForE3SM + The test case this step belongs to + """ + super().__init__(test_case, name='forcing_maps') + self.mesh = test_case.mesh + self.creation_date = test_case.creation_date + + def setup(self): + """ + Set up the step in the work directory, including downloading any + dependencies. + """ + super().setup() + + atm_grid = self.config.get('files_for_e3sm', 'atm_grid') + ocn_grid = self.mesh.mesh_name + mesh_path = self.mesh.steps['cull_mesh'].path + create_date = self.creation_date + + self.add_input_file( + filename='mesh.nc', work_dir_target=f'{mesh_path}/culled_mesh.nc') + self.add_output_file( + filename=f'map_{atm_grid}_to_{ocn_grid}_trbilin.{create_date}.nc') + self.add_output_file( + filename=f'map_{atm_grid}_to_{ocn_grid}_traave.{create_date}.nc') + self.add_output_file( + filename=f'map_{ocn_grid}_to_{atm_grid}_trbilin.{create_date}.nc') + self.add_output_file( + filename=f'map_{ocn_grid}_to_{atm_grid}_traave.{create_date}.nc') + + self._get_resources() + + def constrain_resources(self, available_resources): + """ + Constrain ``cpus_per_task`` and ``ntasks`` based on the number of + cores available to this step + + Parameters + ---------- + available_resources : dict + The total number of cores available to the step + """ + self._get_resources() + super().constrain_resources(available_resources) + + def run(self): + """ + Run this step of the test case + """ + super().run() + + atm_grid = self.config.get('files_for_e3sm', 'atm_grid') + ocn_grid = self.mesh.mesh_name + + self._scrip_file_gridded() + self._scrip_file_MPAS() + self._partition_scrip_file(atm_grid) + self._partition_scrip_file(ocn_grid) + self._create_weights(atm_grid, ocn_grid, 'trbilin') + self._create_weights(atm_grid, ocn_grid, 'traave') + self._create_weights(ocn_grid, atm_grid, 'trbilin') + self._create_weights(ocn_grid, atm_grid, 'traave') + + def _get_resources(self): + """ + Get resources + """ + section = self.config['hurricane'] + self.ntasks = section.getint('init_ntasks') + self.min_tasks = section.getint('init_min_tasks') + self.openmp_threads = section.getint('init_threads') + + def _scrip_file_gridded(self): + """ + Create gridded SCRIP file for atm + """ + logger = self.logger + grid_name = self.config.get('files_for_e3sm', 'atm_grid') + logger.info(f'Create gridded SCRIP file for {grid_name} grid') + + grids = { + 'T382': { + 'nlat': 576, + 'nlon': 1152, + 'lat_typ': 'gss', + 'lon_typ': 'grn_ctr', + }, + 'T574': { + 'nlat': 880, + 'nlon': 1760, + 'lat_typ': 'gss', + 'lon_typ': 'grn_ctr', + }, + } + + nlat, nlon, lat_typ, lon_typ = grids[grid_name].values() + args = [ + 'ncremap', + '-G', f'latlon={nlat},{nlon}#lat_typ={lat_typ}#lon_typ={lon_typ}', + '-g', f'{grid_name}.scrip.nc', + ] + check_call(args, logger) + + logger.info(' Done.') + + def _scrip_file_MPAS(self): + """ + Create SCRIP file from MPAS mesh file. + """ + grid_name = self.mesh.mesh_name + logger = self.logger + logger.info(f'Create MPAS SCRIP file for {grid_name} mesh') + + descriptor = MpasCellMeshDescriptor( + filename='mesh.nc', + mesh_name=grid_name, + ) + descriptor.to_scrip(f'{grid_name}.scrip.nc') + + logger.info(' Done.') + + def _partition_scrip_file(self, grid_name): + """ + Partition SCRIP file for parallel mbtempest use + """ + logger = self.logger + logger.info(f'Partition SCRIP file for {grid_name}') + + # Convert source SCRIP to mbtempest + args = [ + 'mbconvert', '-B', + f'{grid_name}.scrip.nc', + f'{grid_name}.scrip.h5m', + ] + # run in "parallel" with one task and one thread for Intel-MPI support + run_command(args, 1, 1, 1, self.config, logger) + + # Partition source SCRIP + args = [ + 'mbpart', f'{self.ntasks}', + '-z', 'RCB', + f'{grid_name}.scrip.h5m', + f'{grid_name}.scrip.p{self.ntasks}.h5m', + ] + # run in "parallel" with one task and one thread for Intel-MPI support + run_command(args, 1, 1, 1, self.config, logger) + + logger.info(' Done.') + + def _create_weights(self, src, tgt, method): + """ + Create mapping weights file using TempestRemap + """ + logger = self.logger + logger.info('Create weights file') + + if method not in ['traave', 'trbilin']: + raise ValueError(f'Unsupported regridding method {method}') + + src_file = f'{src}.scrip.p{self.ntasks}.h5m' + tgt_file = f'{tgt}.scrip.p{self.ntasks}.h5m' + map_file = f'map_{src}_to_{tgt}_{method}.{self.creation_date}.nc' + + # Build weights file + args = [ + 'mbtempest', '--type', '5', '--weights', + '--load', src_file, + '--load', tgt_file, + '--file', map_file, + '--method', 'fv', '--order', '1', + '--method', 'fv', '--order', '1', + ] + if method == 'trbilin': + args.extend(['--fvmethod', 'bilin']) + run_command( + args, self.cpus_per_task, self.ntasks, + self.openmp_threads, self.config, self.logger, + ) + + # Add attributes required by `generate_domain_files_E3SM.py` + args = [ + 'ncatted', '-O', + '-a', f'grid_file_src,global,o,c,{src_file}', + '-a', f'grid_file_dst,global,o,c,{tgt_file}', + map_file, + ] + check_call(args, logger) + + logger.info(' Done.') diff --git a/compass/ocean/tests/hurricane/hurricane.cfg b/compass/ocean/tests/hurricane/hurricane.cfg index 2f1678f65c..0089bd5111 100644 --- a/compass/ocean/tests/hurricane/hurricane.cfg +++ b/compass/ocean/tests/hurricane/hurricane.cfg @@ -8,3 +8,12 @@ convert_culled_mesh_to_cdf5 = False latitude_threshold = 43.0 # Maximum number of sweeps to search for land-locked cells sweep_count = 20 + +# config options related to regridding and domain files for E3SM configurations +[files_for_e3sm] + +# Grid for data atmosphere (currently T382 and T574 supported) +atm_grid = T574 + +# the relative or absolute path to the `generate_domain_files_E3SM.py` utility +domain_files_exe = ${paths:compass_branch}/E3SM-Project/tools/generate_domain_files/generate_domain_files_E3SM.py diff --git a/compass/version.py b/compass/version.py index fa98dd1c47..1452326dd3 100644 --- a/compass/version.py +++ b/compass/version.py @@ -1 +1 @@ -__version__ = '2.0.0-alpha.3' +__version__ = '2.0.0-alpha.4' diff --git a/deploy/pixi.toml.j2 b/deploy/pixi.toml.j2 index 88b9049075..66860b9653 100644 --- a/deploy/pixi.toml.j2 +++ b/deploy/pixi.toml.j2 @@ -40,6 +40,7 @@ mosaic = "*" mpas_tools = "{{ mpas_tools }}.*" nco = "*" netcdf4 = {version = "*", build = "nompi_*"} +numba = "*" numpy = ">=2.0,<3.0" {%- if platform == 'linux-64' %} otps = "{{ otps }}.*" diff --git a/docs/developers_guide/ocean/api.rst b/docs/developers_guide/ocean/api.rst index d91d1a76ac..7af90220f3 100644 --- a/docs/developers_guide/ocean/api.rst +++ b/docs/developers_guide/ocean/api.rst @@ -438,6 +438,16 @@ test cases and steps analysis.Analysis.plot_hwm analysis.Analysis.run + files_for_e3sm.FilesForE3SM + files_for_e3sm.FilesForE3SM.configure + files_for_e3sm.FilesForE3SM.run + files_for_e3sm.forcing_maps.ForcingMaps + files_for_e3sm.forcing_maps.ForcingMaps.setup + files_for_e3sm.forcing_maps.ForcingMaps.run + files_for_e3sm.domain_files.DomainFiles + files_for_e3sm.domain_files.DomainFiles.setup + files_for_e3sm.domain_files.DomainFiles.run + lts.mesh.lts_regions.LTSRegionsStep lts.mesh.lts_regions.LTSRegionsStep.setup lts.mesh.lts_regions.LTSRegionsStep.run diff --git a/docs/developers_guide/ocean/test_groups/hurricane.rst b/docs/developers_guide/ocean/test_groups/hurricane.rst index de81c346bf..60cbd813d0 100644 --- a/docs/developers_guide/ocean/test_groups/hurricane.rst +++ b/docs/developers_guide/ocean/test_groups/hurricane.rst @@ -276,3 +276,20 @@ The class :py:class:`compass.ocean.tests.hurricane.analysis.Analysis` defines a step to generate validation plots comparing sea surface height timeseries between modeled and observed data at several different stations. Both NOAA and USGS observations are plotted. + +.. _dev_ocean_hurricane_files_for_e3sm: + +files for e3sm test case +^^^^^^^^^^^^^^^^^^^^^^^^ +The ``files_for_e3sm`` test case builds the domain files and data atmosphere +regridding files needed to run MPAS-Ocean hurricane configurations in E3SM. + +forcing maps +"""""""""""" +The class :py:class:`compass.ocean.tests.hurricane.files_for_e3sm.forcing_maps.ForcingMaps` +builds the regridding files between the ocean mesh and the data atmosphere grid. + +domain files +"""""""""""" +The class :py:class:`compass.ocean.tests.hurricane.files_for_e3sm.domain_files.DomainFiles` +builds the domain files needed to run MPAS-Ocean hurricane configurations in E3SM. diff --git a/docs/users_guide/ocean/test_groups/hurricane.rst b/docs/users_guide/ocean/test_groups/hurricane.rst index a626d4b55f..d3dc5cd95a 100644 --- a/docs/users_guide/ocean/test_groups/hurricane.rst +++ b/docs/users_guide/ocean/test_groups/hurricane.rst @@ -3,8 +3,8 @@ hurricane ========= -The ``ocean/hurricane`` test group defines meshes, -initial conditions, forward simulations, and validation for global, +The ``ocean/hurricane`` test group defines meshes, initial conditions, forward +simulations, validation, and E3SM domain and mapping file creation for global, realistic ocean domains with regional refinement. These simulations are forced with time-varying atmospheric reanalysis data for tropical cyclone events and tides. The meshes contain refined regions in order @@ -110,6 +110,16 @@ Note that meshes and test cases may modify these options, as noted below. plot_station_dems = False + + # config options related to regridding and domain files for E3SM configurations + [files_for_e3sm] + + # Grid for data atmosphere (currently T382 and T574 supported) + atm_grid = T574 + + # the relative or absolute path to the `generate_domain_files_E3SM.py` utility + domain_files_exe = ${paths:compass_branch}/E3SM-Project/tools/generate_domain_files/generate_domain_files_E3SM.py + .. _ocean_hurricane_meshes: Meshes @@ -281,3 +291,26 @@ are used for the validation. .. image:: images/hurricane_subgrid_spatialerror.png :width: 800 px :align: center + +.. _ocean_hurricane_files_for_e3sm: + +files for e3sm test case +^^^^^^^^^^^^^^^^^^^^^^^^ +The ``files_for_e3sm`` test case builds the domain files and data atmosphere +regridding files needed to run MPAS-Ocean hurricane configurations in E3SM. +The LTS3 and FB-LTS meshes are not currently supported. + +forcing maps step +""""""""""""""""" +The ``forcing_maps`` step builds the regridding files between the ocean mesh +and the data atmosphere grid. Currently two data atmosphere grids are supported, +T382 and T574. These correspond to the NCEP Climate Forecast System Reanalysis +(CFSR) and Climate Forecast System Version 2 (CFSv2), respectively. The grid is +specified in the ``[files_for_e3sm]`` config options. + +domain files step +""""""""""""""""" +The ``domain_files`` step builds the ocean and atmosphere domain files needed +to define the ocean and data atmosphere grids in an E3SM configuration. The +step uses the domain file creator maintained in the tools directory of the +E3SM submodule.