Source code for compass.ocean.mesh.remap_topography

import os
import pathlib

import numpy as np
import xarray as xr
from mpas_tools.cime.constants import constants
from mpas_tools.io import write_netcdf
from mpas_tools.logging import check_call
from pyremap import MpasCellMeshDescriptor

from compass.parallel import run_command
from compass.step import Step


[docs] class RemapTopography(Step): """ A step for remapping bathymetry and ice-shelf topography from a latitude-longitude grid to a global MPAS-Ocean mesh Attributes ---------- base_mesh_step : compass.mesh.spherical.SphericalBaseStep The base mesh step containing input files to this step mesh_name : str The name of the MPAS mesh to include in the mapping file """
[docs] def __init__( self, test_case, base_mesh_step, name='remap_topography', subdir=None, mesh_name='MPAS_mesh', ): """ Create a new step Parameters ---------- test_case : compass.ocean.tests.global_ocean.mesh.Mesh The test case this step belongs to base_mesh_step : compass.mesh.spherical.SphericalBaseStep The base mesh step containing input files to this step name : str, optional the name of the step subdir : str, optional the subdirectory for the step. The default is ``name`` mesh_name : str, optional The name of the MPAS mesh to include in the mapping file """ super().__init__(test_case, name=name, subdir=subdir, ntasks=None, min_tasks=None) self.base_mesh_step = base_mesh_step self.mesh_name = mesh_name self.add_output_file(filename='topography_remapped.nc')
[docs] def setup(self): """ Set up the step in the work directory, including downloading any dependencies. """ super().setup() config = self.config 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', ) 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', ) target = os.path.join(base_path, base_filename) self.add_input_file(filename='base_mesh.nc', work_dir_target=target) self.ntasks = section.getint('ntasks') self.min_tasks = section.getint('min_tasks')
[docs] 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 """ config = self.config self.ntasks = config.getint('remap_topography', 'ntasks') self.min_tasks = config.getint('remap_topography', 'min_tasks') super().constrain_resources(available_resources)
[docs] 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 logger.info('Partition SCRIP file') stem = pathlib.Path(in_filename).stem h5m_filename = f'{stem}.h5m' part_filename = f'{stem}.p{self.ntasks}.h5m' # Convert source SCRIP to mbtempest args = [ 'mbconvert', '-B', in_filename, h5m_filename, ] check_call(args, logger) # Partition source SCRIP args = [ 'mbpart', f'{self.ntasks}', '-z', 'RCB', h5m_filename, part_filename, ] 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') config = self.config method = config.get('remap_topography', 'method') if method != 'conserve': raise ValueError(f'Unsupported method {method} for TempestRemap') args = [ 'mbtempest', '--type', '5', '--load', f'source.scrip.p{self.ntasks}.h5m', '--load', f'target.scrip.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') 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', ] run_command( args, self.cpus_per_task, self.ntasks, self.openmp_threads, self.config, self.logger, ) logger.info(' Done.') def _remap_to_target(self): """ Remap combined bathymetry onto MPAS target mesh """ logger = self.logger logger.info('Remap to target') config = self.config method = config.get('remap_topography', 'method') # Build command args 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') 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': 'bed_elevation', 'thickness': 'landIceThkObserved', 'ice_mask': 'landIceFracObserved', 'grounded_mask': 'landIceGroundedFracObserved', 'ocean_mask': 'oceanFracObserved', 'bathymetry_mask': 'bathyFracObserved', } for in_var, out_var in rename.items(): ds_out[out_var] = ds_in[in_var] ds_out['landIceFloatingFracObserved'] = \ ds_out.landIceFracObserved - ds_out.landIceGroundedFracObserved # 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 norm = ds_out.bathyFracObserved valid = norm > renorm_threshold for var in ['bed_elevation', 'landIceThkObserved']: ds_out[var] = xr.where(valid, ds_out[var] / norm, 0.) thickness = ds_out.landIceThkObserved bed = ds_out.bed_elevation flotation_thickness = - (ocean_density / ice_density) * bed # 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.')