Modified ocean/mesh/remap_topography to allow smoothing

compass/ocean/mesh/remap_topography.cfg

@@ -2,7 +2,14 @@
 [remap_topography]
 
 # the name of the topography file in the bathymetry database
-topo_filename = BedMachineAntarctica_v3_and_GEBCO_2023_0.0125_degree_20240828.nc
+#topo_filename = BedMachineAntarctica_v3_and_GEBCO_2023_ne3000_20241004.nc
+#src_scrip_filename = ne3000_20241004.scrip.nc
+topo_filename = BedMachineAntarctica_v3_and_GEBCO_2023_ne3000_20241004.nc
+src_scrip_filename = ne3000_20241004.scrip.nc
+
+# weight generator function:
+#    `tempest` for cubed-sphere bathy or `esmf` for latlon bathy
+weight_generator = tempest
 
 # variable names in topo_filename
 lon_var = lon
@@ -19,10 +26,11 @@ description = Bathymetry is from GEBCO 2023, combined with BedMachine
               Antarctica v3 around Antarctica.
 
 # the target and minimum number of MPI tasks to use in remapping
-ntasks = 4096
-min_tasks = 360
+ntasks = 1280
+min_tasks = 256
 
 # remapping method {'bilinear', 'neareststod', 'conserve'}
+# must use 'conserve' for tempestremap
 method = conserve
 
 # threshold of what fraction of an MPAS cell must contain ocean in order to
@@ -31,3 +39,10 @@ renorm_threshold = 0.01
 
 # the density of land ice from MALI (kg/m^3)
 ice_density = 910.0
+
+# smoothing parameters
+# no smoothing (required for esmf):
+#     expandDist = 0   [m]
+#     expandFactor = 1 [cell fraction]
+expandDist = 0
+expandFactor = 1

compass/ocean/mesh/remap_topography.py

@@ -4,8 +4,10 @@
 import xarray as xr
 from mpas_tools.cime.constants import constants
 from mpas_tools.io import write_netcdf
-from pyremap import LatLonGridDescriptor, MpasCellMeshDescriptor, Remapper
+from mpas_tools.logging import check_call
+from pyremap import MpasCellMeshDescriptor
 
+from compass.parallel import run_command
 from compass.step import Step
 
 
@@ -23,8 +25,10 @@ class RemapTopography(Step):
         The name of the MPAS mesh to include in the mapping file
     """
 
-    def __init__(self, test_case, base_mesh_step, name='remap_topography',
-                 subdir=None, mesh_name='MPAS_mesh'):
+    def __init__(
+        self, test_case, base_mesh_step, name='remap_topography', subdir=None,
+        mesh_name='MPAS_mesh',
+    ):
         """
         Create a new step
 
@@ -59,20 +63,30 @@ def setup(self):
         """
         super().setup()
         config = self.config
-        topo_filename = config.get('remap_topography', 'topo_filename')
+        section = config['remap_topography']
+        topo_filename = section.get('topo_filename')
+        src_scrip_filename = section.get('src_scrip_filename')
+
         self.add_input_file(
             filename='topography.nc',
             target=topo_filename,
-            database='bathymetry_database')
+            database='bathymetry_database',
+        )
+        self.add_input_file(
+            filename='source.scrip.nc',
+            target=src_scrip_filename,
+            database='bathymetry_database',
+        )
 
         base_path = self.base_mesh_step.path
         base_filename = self.base_mesh_step.config.get(
-            'spherical_mesh', 'mpas_mesh_filename')
+            'spherical_mesh', 'mpas_mesh_filename',
+        )
         target = os.path.join(base_path, base_filename)
         self.add_input_file(filename='base_mesh.nc', work_dir_target=target)
 
-        self.ntasks = config.getint('remap_topography', 'ntasks')
-        self.min_tasks = config.getint('remap_topography', 'min_tasks')
+        self.ntasks = section.getint('ntasks')
+        self.min_tasks = section.getint('min_tasks')
 
     def constrain_resources(self, available_resources):
         """
@@ -94,64 +108,203 @@ def run(self):
         Run this step of the test case
         """
         config = self.config
+        weight_generator = config.get('remap_topography', 'weight_generator')
+
+        self._create_target_scrip_file()
+        if weight_generator == 'tempest':
+            self._partition_scrip_file('source.scrip.nc')
+            self._partition_scrip_file('target.scrip.nc')
+            self._create_weights_tempest()
+        elif weight_generator == 'esmf':
+            self._create_weights_esmf()
+        else:
+            msg = f'Unsupported weight generator function {weight_generator}'
+            raise ValueError(msg)
+        self._remap_to_target()
+        self._modify_remapped_bathymetry()
+
+    def _create_target_scrip_file(self):
+        """
+        Create target SCRIP file from MPAS mesh file.
+        """
+        logger = self.logger
+        logger.info('Create source SCRIP file')
+
+        config = self.config
+        section = config['remap_topography']
+        expandDist = section.getfloat('expandDist')
+        expandFactor = section.getfloat('expandFactor')
+
+        descriptor = MpasCellMeshDescriptor(
+            fileName='base_mesh.nc',
+            meshName=self.mesh_name,
+        )
+        descriptor.to_scrip(
+            'target.scrip.nc',
+            expandDist=expandDist,
+            expandFactor=expandFactor,
+        )
+
+        logger.info('  Done.')
+
+    def _partition_scrip_file(self, in_filename):
+        """
+        Partition SCRIP file for parallel mbtempest use
+        """
         logger = self.logger
-        parallel_executable = config.get('parallel', 'parallel_executable')
+        logger.info('Partition SCRIP file')
+
+        # Convert to NetCDF3 64-bit
+        args = [
+            'ncks', '-5',
+            in_filename,
+            in_filename.replace('.nc', '.64bit.nc'),
+        ]
+        check_call(args, logger)
+
+        # Convert source SCRIP to mbtempest
+        args = [
+            'mbconvert', '-B',
+            in_filename.replace('.nc', '.64bit.nc'),
+            in_filename.replace('.nc', '.64bit.h5m'),
+        ]
+        check_call(args, logger)
+
+        # Partition source SCRIP
+        args = [
+            'mbpart', f'{self.ntasks}',
+            '-z', 'RCB',
+            in_filename.replace('.nc', '.64bit.h5m'),
+            in_filename.replace('.nc', f'.64bit.p{self.ntasks}.h5m'),
+        ]
+        check_call(args, logger)
+
+        logger.info('  Done.')
+
+    def _create_weights_tempest(self):
+        """
+        Create mapping weights file using TempestRemap
+        """
+        logger = self.logger
+        logger.info('Create weights file')
 
-        lon_var = config.get('remap_topography', 'lon_var')
-        lat_var = config.get('remap_topography', 'lat_var')
+        config = self.config
         method = config.get('remap_topography', 'method')
-        renorm_threshold = config.getfloat('remap_topography',
-                                           'renorm_threshold')
-        ice_density = config.getfloat('remap_topography', 'ice_density')
-        ocean_density = constants['SHR_CONST_RHOSW']
-        g = constants['SHR_CONST_G']
+        if method != 'conserve':
+            raise ValueError(f'Unsupported method {method} for TempestRemap')
+
+        args = [
+            'mbtempest', '--type', '5',
+            '--load', f'source.scrip.64bit.p{self.ntasks}.h5m',
+            '--load', f'target.scrip.64bit.p{self.ntasks}.h5m',
+            '--file', f'map_source_to_target_{method}.nc',
+            '--weights', '--gnomonic',
+            '--boxeps', '1e-9',
+        ]
+
+        run_command(
+            args, self.cpus_per_task, self.ntasks,
+            self.openmp_threads, self.config, self.logger,
+        )
+
+        logger.info('  Done.')
+
+    def _create_weights_esmf(self):
+        """
+        Create mapping weights file using ESMF_RegridWeightGen
+        """
+        logger = self.logger
+        logger.info('Create weights file')
 
-        in_descriptor = LatLonGridDescriptor.read(fileName='topography.nc',
-                                                  lonVarName=lon_var,
-                                                  latVarName=lat_var)
+        config = self.config
+        method = config.get('remap_topography', 'method')
+
+        args = [
+            'ESMF_RegridWeightGen',
+            '--source', 'source.scrip.nc',
+            '--destination', 'target.scrip.nc',
+            '--weight', f'map_source_to_target_{method}.nc',
+            '--method', method,
+            '--netcdf4',
+            '--ignore_unmapped',
+        ]
 
-        in_mesh_name = in_descriptor.meshName
+        run_command(
+            args, self.cpus_per_task, self.ntasks,
+            self.openmp_threads, self.config, self.logger,
+        )
 
-        out_mesh_name = self.mesh_name
-        out_descriptor = MpasCellMeshDescriptor(fileName='base_mesh.nc',
-                                                meshName=self.mesh_name)
+        logger.info('  Done.')
+
+    def _remap_to_target(self):
+        """
+        Remap combined bathymetry onto MPAS target mesh
+        """
+        logger = self.logger
+        logger.info('Remap to target')
 
-        mapping_file_name = \
-            f'map_{in_mesh_name}_to_{out_mesh_name}_{method}.nc'
-        remapper = Remapper(in_descriptor, out_descriptor, mapping_file_name)
+        config = self.config
+        method = config.get('remap_topography', 'method')
 
-        remapper.build_mapping_file(method=method, mpiTasks=self.ntasks,
-                                    tempdir='.', logger=logger,
-                                    esmf_parallel_exec=parallel_executable)
+        # Build command args
+        # Unused options:
+        # -P mpas, handles some MPAS-specific index ordering, CF, etc...
+        # -C climatology, basically bypasses fill values
+        args = [
+            'ncremap',
+            '-m', f'map_source_to_target_{method}.nc',
+            '--vrb=1',
+            'topography.nc', 'topography_ncremap.nc',
+        ]
+        check_call(args, logger)
+
+        logger.info('  Done.')
+
+    def _modify_remapped_bathymetry(self):
+        """
+        Modify remapped bathymetry
+        """
+        logger = self.logger
+        logger.info('Modify remapped bathymetry')
 
-        remapper.remap_file(inFileName='topography.nc',
-                            outFileName='topography_ncremap.nc',
-                            logger=logger)
+        config = self.config
+        section = config['remap_topography']
+        renorm_threshold = section.getfloat('renorm_threshold')
+        ice_density = section.getfloat('ice_density')
+        ocean_density = constants['SHR_CONST_RHOSW']
+        g = constants['SHR_CONST_G']
 
         ds_in = xr.open_dataset('topography_ncremap.nc')
         ds_in = ds_in.rename({'ncol': 'nCells'})
-        ds_out = xr.Dataset()
-        rename = {'bathymetry_var': 'bed_elevation',
-                  'ice_thickness_var': 'landIceThkObserved',
-                  'ice_frac_var': 'landIceFracObserved',
-                  'grounded_ice_frac_var': 'landIceGroundedFracObserved',
-                  'ocean_frac_var': 'oceanFracObserved',
-                  'bathy_frac_var': 'bathyFracObserved'}
 
+        ds_out = xr.Dataset()
+        rename = {
+            'bathymetry_var': 'bed_elevation',
+            'ice_thickness_var': 'landIceThkObserved',
+            'ice_frac_var': 'landIceFracObserved',
+            'grounded_ice_frac_var': 'landIceGroundedFracObserved',
+            'ocean_frac_var': 'oceanFracObserved',
+            'bathy_frac_var': 'bathyFracObserved',
+        }
         for option, out_var in rename.items():
-            in_var = config.get('remap_topography', option)
+            in_var = section.get(option)
             ds_out[out_var] = ds_in[in_var]
 
         ds_out['landIceFloatingFracObserved'] = \
             ds_out.landIceFracObserved - ds_out.landIceGroundedFracObserved
 
-        # make sure fractions don't exceed 1
-        for var in ['landIceFracObserved', 'landIceGroundedFracObserved',
-                    'landIceFloatingFracObserved', 'oceanFracObserved',
-                    'bathyFracObserved']:
+        # Make sure fractions don't exceed 1
+        varNames = [
+            'landIceFracObserved',
+            'landIceGroundedFracObserved',
+            'landIceFloatingFracObserved',
+            'oceanFracObserved',
+            'bathyFracObserved',
+        ]
+        for var in varNames:
             ds_out[var] = np.minimum(ds_out[var], 1.)
 
-        # renormalize elevation variables
+        # Renormalize elevation variables
         norm = ds_out.bathyFracObserved
         valid = norm > renorm_threshold
         for var in ['bed_elevation', 'landIceThkObserved']:
@@ -163,13 +316,13 @@ def run(self):
         # not allowed to be thicker than the flotation thickness
         thickness = np.minimum(thickness, flotation_thickness)
         ds_out['landIceThkObserved'] = thickness
-
         ds_out['landIcePressureObserved'] = ice_density * g * thickness
 
         # compute the ice draft to be consistent with the land ice pressure
         # and using E3SM's density of seawater
-
         ds_out['landIceDraftObserved'] = \
             - (ice_density / ocean_density) * thickness
 
         write_netcdf(ds_out, 'topography_remapped.nc')
+
+        logger.info('  Done.')

conda/default.cfg

@@ -30,7 +30,7 @@ esmf = 8.6.1
 hdf5 = 1.14.3
 lapack = 3.9.1
 metis = 5.1.0
-moab = 5.5.1
+moab = master
 netcdf_c = 4.9.2
 netcdf_fortran = 4.6.1
 petsc = 3.19.1