Add files_for_e3sm test case to hurricane test group

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)

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)

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.')

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.')

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

compass/version.py

@@ -1 +1 @@
-__version__ = '2.0.0-alpha.3'
+__version__ = '2.0.0-alpha.4'

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 }}.*"