From 0e95deca831deabd1caa879c817ff7922ba25057 Mon Sep 17 00:00:00 2001 From: bmooremaley Date: Tue, 3 Sep 2024 15:25:27 -0500 Subject: [PATCH 1/5] First draft adding cubed sphere interpolation to combine_topo --- .../tests/utility/combine_topo/__init__.py | 561 ++++++++++++------ 1 file changed, 379 insertions(+), 182 deletions(-) diff --git a/compass/ocean/tests/utility/combine_topo/__init__.py b/compass/ocean/tests/utility/combine_topo/__init__.py index 788d8da43..6fe641417 100644 --- a/compass/ocean/tests/utility/combine_topo/__init__.py +++ b/compass/ocean/tests/utility/combine_topo/__init__.py @@ -1,3 +1,7 @@ +import os +import subprocess +from datetime import datetime + import numpy as np import progressbar import pyproj @@ -14,7 +18,7 @@ class CombineTopo(TestCase): datasets """ - def __init__(self, test_group): + def __init__(self, test_group, target_grid): """ Create the test case @@ -23,28 +27,43 @@ def __init__(self, test_group): test_group : compass.ocean.tests.utility.Utility The test group that this test case belongs to """ - super().__init__(test_group=test_group, name='combine_topo') - self.add_step(Combine(test_case=self)) + subdir = os.path.join('combine_topo', target_grid) + + super().__init__( + test_group=test_group, name='combine_topo', subdir=subdir, + ) + + self.add_step(Combine(test_case=self, target_grid=target_grid)) class Combine(Step): """ A step for combining GEBCO 2023 with BedMachineAntarctica topography datasets + TODO: Document target grid throughout + + Attributes + ---------- + ... """ - def __init__(self, test_case): + def __init__(self, test_case, target_grid): """ Create a new step Parameters ---------- - test_case : compass.ocean.tests.utility.extrap_woa.ExtraWoa + test_case : compass.ocean.tests.utility.combine_topo.CombineTopo The test case this step belongs to + target_grid : `str`, either "lat_lon" or "cubed_sphere" + date_stamp : `datetime.datetime`, date when test was run """ - super().__init__(test_case, name='combine', ntasks=None, - min_tasks=None) + super().__init__( + test_case, name='combine', ntasks=None, min_tasks=None, + ) + self.target_grid = target_grid + self.date_stamp = None def setup(self): """ @@ -55,102 +74,71 @@ def setup(self): config = self.config section = config['combine_topo'] + + # Get input filenames and resolution antarctic_filename = section.get('antarctic_filename') - self.add_input_file(filename=antarctic_filename, - target=antarctic_filename, - database='bathymetry_database') global_filename = section.get('global_filename') - self.add_input_file(filename=global_filename, - target=global_filename, - database='bathymetry_database') - - cobined_filename = section.get('cobined_filename') - self.add_output_file(filename=cobined_filename) - + resolution = section.getint('resolution') + + # Datestamp + self.datestamp = datetime.now().strftime('%Y%m%d') + + # Build output filename + combined_filename = _get_combined_filename( + self.target_grid, self.datestamp, resolution, + antarctic_filename, global_filename, + ) + + # Add bathymetry data input files + self.add_input_file( + filename=antarctic_filename, + target=antarctic_filename, + database='bathymetry_database', + ) + self.add_input_file( + filename=global_filename, + target=global_filename, + database='bathymetry_database', + ) + self.add_output_file(filename=combined_filename) + + # Make tiles directory + os.makedirs('tiles', exist_ok=True) + + # Get ntasks and min_tasks self.ntasks = section.getint('ntasks') self.min_tasks = section.getint('min_tasks') + 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('combine_topo', 'ntasks') + self.min_tasks = config.getint('combine_topo', 'min_tasks') + super().constrain_resources(available_resources) + def run(self): """ Run this step of the test case """ - self._downsample_gebco() self._modify_bedmachine() - # we will work with bedmachine data at its original resolution - # self._downsample_bedmachine() + self._modify_gebco() + self._tile_gebco() + self._create_ne_scrip() + self._create_bedmachine_scrip() + self._get_map_gebco() + self._get_map_bedmachine() + self._remap_gebco() self._remap_bedmachine() self._combine() - def _downsample_gebco(self): - """ - Average GEBCO to 0.0125 degree grid. GEBCO is on 15" grid, so average - every 3x3 block of cells - """ - config = self.config - logger = self.logger - - section = config['combine_topo'] - in_filename = section.get('global_filename') - out_filename = 'GEBCO_2023_0.0125_degree.nc' - - gebco = xr.open_dataset(in_filename) - - nlon = gebco.sizes['lon'] - nlat = gebco.sizes['lat'] - - block = 3 - norm = 1.0 / block**2 - - nx = nlon // block - ny = nlat // block - - chunks = 2 - nxchunk = nlon // chunks - nychunk = nlat // chunks - - gebco = gebco.chunk({'lon': nxchunk, 'lat': nychunk}) - - bathymetry = np.zeros((ny, nx)) - - logger.info('Averaging GEBCO to 0.0125 degree grid') - widgets = [progressbar.Percentage(), ' ', - progressbar.Bar(), ' ', progressbar.ETA()] - bar = progressbar.ProgressBar(widgets=widgets, - maxval=chunks**2).start() - for ychunk in range(chunks): - for xchunk in range(chunks): - gebco_chunk = gebco.isel( - lon=slice(nxchunk * xchunk, nxchunk * (xchunk + 1)), - lat=slice(nychunk * ychunk, nychunk * (ychunk + 1))) - elevation = gebco_chunk.elevation.values - nxblock = nxchunk // block - nyblock = nychunk // block - bathy_block = np.zeros((nyblock, nxblock)) - for y in range(block): - for x in range(block): - bathy_block += elevation[y::block, x::block] - bathy_block *= norm - xmin = xchunk * nxblock - xmax = (xchunk + 1) * nxblock - ymin = ychunk * nyblock - ymax = (ychunk + 1) * nyblock - bathymetry[ymin:ymax, xmin:xmax] = bathy_block - bar.update(ychunk * chunks + xchunk + 1) - bar.finish() - - lon_corner = np.linspace(-180., 180., bathymetry.shape[1] + 1) - lat_corner = np.linspace(-90., 90., bathymetry.shape[0] + 1) - lon = 0.5 * (lon_corner[0:-1] + lon_corner[1:]) - lat = 0.5 * (lat_corner[0:-1] + lat_corner[1:]) - gebco_low = xr.Dataset({'bathymetry': (['lat', 'lon'], bathymetry)}, - coords={'lon': (['lon',], lon), - 'lat': (['lat',], lat)}) - gebco_low.attrs = gebco.attrs - gebco_low.lon.attrs = gebco.lon.attrs - gebco_low.lat.attrs = gebco.lat.attrs - gebco_low.bathymetry.attrs = gebco.elevation.attrs - gebco_low.to_netcdf(out_filename) - def _modify_bedmachine(self): """ Modify BedMachineAntarctica to compute the fields needed by MPAS-Ocean @@ -159,10 +147,9 @@ def _modify_bedmachine(self): logger.info('Modifying BedMachineAntarctica with MPAS-Ocean names') config = self.config + in_filename = config.get('combine_topo', 'antarctic_filename') + out_filename = in_filename.replace('.nc', '_mod.nc') - section = config['combine_topo'] - in_filename = section.get('antarctic_filename') - out_filename = 'BedMachineAntarctica-v3_mod.nc' bedmachine = xr.open_dataset(in_filename) mask = bedmachine.mask ice_mask = (mask != 0).astype(float) @@ -175,152 +162,362 @@ def _modify_bedmachine(self): bedmachine.ice_draft.attrs['units'] = 'meters' bedmachine['thickness'] = bedmachine.thickness + bedmachine['bathymetry'] = bedmachine['bathymetry'].where(ocean_mask, 0.) + bedmachine['ice_draft'] = bedmachine['ice_draft'].where(ocean_mask, 0.) + bedmachine['thickness'] = bedmachine['thickness'].where(ocean_mask, 0.) + bedmachine['ice_mask'] = ice_mask bedmachine['grounded_mask'] = grounded_mask bedmachine['ocean_mask'] = ocean_mask - varlist = ['bathymetry', 'ice_draft', 'thickness', 'ice_mask', - 'grounded_mask', 'ocean_mask'] + varlist = [ + 'bathymetry', 'ice_draft', 'thickness', + 'ice_mask', 'grounded_mask', 'ocean_mask', + ] bedmachine = bedmachine[varlist] bedmachine.to_netcdf(out_filename) logger.info(' Done.') - def _downsample_bedmachine(self): + def _modify_gebco(self): """ - Downsample bedmachine from 0.5 to 1 km grid + Modify GEBCO to include lon/lat bounds located at grid edges """ logger = self.logger - logger.info('Downsample BedMachineAntarctica from 500 m to 1 km') + logger.info('Adding bounds to GEBCO lat/lon') - in_filename = 'BedMachineAntarctica-v3_mod.nc' - out_filename = 'BedMachineAntarctica-v3_1k.nc' - bedmachine = xr.open_dataset(in_filename) - x = bedmachine.x.values - y = bedmachine.y.values - - nx = len(x) // 2 - ny = len(y) // 2 - x = 0.5 * (x[0:2 * nx:2] + x[1:2 * nx:2]) - y = 0.5 * (y[0:2 * ny:2] + y[1:2 * ny:2]) - bedmachine1k = xr.Dataset() - for field in bedmachine.data_vars: - in_array = bedmachine[field].values - out_array = np.zeros((ny, nx)) - for yoffset in range(2): - for xoffset in range(2): - out_array += 0.25 * \ - in_array[yoffset:2 * ny:2, xoffset:2 * nx:2] - da = xr.DataArray(out_array, dims=('y', 'x'), - coords={'x': (('x',), x), - 'y': (('y',), y)}) - bedmachine1k[field] = da - bedmachine1k[field].attrs = bedmachine[field].attrs - - bedmachine1k.to_netcdf(out_filename) + config = self.config + in_filename = config.get('combine_topo', 'global_filename') + out_filename = in_filename.replace('.nc', '_cf.nc') + + gebco = xr.open_dataset(in_filename) + lat = gebco.lat + lon = gebco.lon + dlat = gebco.lat.isel(lat=1) - gebco.lat.isel(lat=0) + dlon = gebco.lon.isel(lon=1) - gebco.lon.isel(lon=0) + lat_bnds = xr.concat([lat - 0.5 * dlat, lat + 0.5 * dlat], dim='bnds') + lon_bnds = xr.concat([lon - 0.5 * dlon, lon + 0.5 * dlon], dim='bnds') + gebco['lat_bnds'] = lat_bnds.transpose('lat', 'bnds') + gebco['lon_bnds'] = lon_bnds.transpose('lon', 'bnds') + gebco.lat.attrs['bounds'] = 'lat_bnds' + gebco.lon.attrs['bounds'] = 'lon_bnds' + + gebco.to_netcdf(out_filename) logger.info(' Done.') - def _remap_bedmachine(self): + def _tile_gebco(self): """ - Remap BedMachine Antarctica to GEBCO lat-lon grid + Tile GEBCO """ logger = self.logger - logger.info('Remap BedMachineAntarctica to GEBCO 1/80 deg grid') + logger.info('Tiling GEBCO data') - in_filename = 'BedMachineAntarctica-v3_mod.nc' - out_filename = 'BedMachineAntarctica_on_GEBCO_low.nc' - gebco_filename = 'GEBCO_2023_0.0125_degree.nc' + config = self.config + section = config['combine_topo'] + global_filename = section.get('global_filename') + lat_tiles = section.getint('lat_tiles') + lon_tiles = section.getint('lon_tiles') - projection = pyproj.Proj('+proj=stere +lat_ts=-71.0 +lat_0=-90 ' - '+lon_0=0.0 +k_0=1.0 +x_0=0.0 +y_0=0.0 ' - '+ellps=WGS84') + in_filename = global_filename.replace('.nc', '_cf.nc') - bedmachine = xr.open_dataset(in_filename) - x = bedmachine.x.values - y = bedmachine.y.values + gebco = xr.open_dataset(in_filename) + gebco = gebco.chunk({'lat': lat_tiles * lon_tiles}) + nlat = gebco.sizes['lat'] + nlon = gebco.sizes['lon'] + nlat_tile = nlat // lat_tiles + nlon_tile = nlon // lon_tiles + for lat_tile in range(lat_tiles): + for lon_tile in range(lon_tiles): + src_filename, _, _ = _get_tile_filenames(global_filename, resolution, method, lat_tile, lon_tile) + + lat_indices = np.arange( + lat_tile * nlat_tile - 1, (lat_tile + 1) * nlat_tile + 1 + ) + lon_indices = np.arange( + lon_tile * nlon_tile, (lon_tile + 1) * nlon_tile + 1 + ) + lat_indices = np.minimum(np.maximum(lat_indices, 0), nlat - 1) + lon_indices = np.mod(lon_indices, nlon) + ds_local = gebco.isel(lat=lat_indices, lon=lon_indices) + if lat_tile == 0: + # need to handle south pole + ds_local.lat.values[0] = -90. + if lat_tile == lat_tiles - 1: + # need to handle north pole + ds_local.lat.values[-1] = 90. + + write_job = ds_local.to_netcdf(out_filename, compute=False) + with ProgressBar(): + logger.info(f' writing {out_filename}') + write_job.compute() - in_descriptor = ProjectionGridDescriptor.create( - projection, x, y, 'BedMachineAntarctica_500m') + logger.info(' Done.') - out_descriptor = LatLonGridDescriptor.read(fileName=gebco_filename) + def _create_ne_scrip(self): + """ + Create SCRIP file for the NExxx (cubed-sphere) mesh + Reference: + https://acme-climate.atlassian.net/wiki/spaces/DOC/pages/872579110/ + Running+E3SM+on+New+Atmosphere+Grids + """ + logger = self.logger + logger.info('Creating NE scrip file') - mapping_filename = \ - 'map_BedMachineAntarctica_500m_to_GEBCO_0.0125deg_bilinear.nc' + resolution = self.resolution - remapper = Remapper(in_descriptor, out_descriptor, mapping_filename) - remapper.build_mapping_file(method='bilinear', mpiTasks=self.ntasks, - esmf_parallel_exec='srun', tempdir='.') - bedmachine = xr.open_dataset(in_filename) - bedmachine_on_gebco_low = remapper.remap(bedmachine) + args = [ + 'GenerateCSMesh', '--alt', '--res', f'{resolution}', + '--file', f'ne{resolution}.g', + ] + logger.info(f' Running: {" ".join(args)}') + subprocess.check_call(args) - for field in ['bathymetry', 'ice_draft', 'thickness']: - bedmachine_on_gebco_low[field].attrs['unit'] = 'meters' + args = [ + 'ConvertMeshToSCRIP', '--in', f'ne{resolution}.g', + '--out', f'ne{resolution}.scrip.nc', + ] + logger.info(f' Running: {" ".join(args)}') + subprocess.check_call(args) - bedmachine_on_gebco_low.to_netcdf(out_filename) logger.info(' Done.') - def _combine(self): + def _create_bedmachine_scrip(self): """ - Combine GEBCO with BedMachine Antarctica + Create SCRIP file for the BedMachineAntarctica-v3 bathymetry """ logger = self.logger - logger.info('Combine BedMachineAntarctica and GEBCO') + logger.info('Creating Bedmachine scrip file') + + in_filename = self.antarctic_filename + out_filename = in_filename.replace('.nc', '.scrip.nc') + + projection = pyproj.Proj( + '+proj=stere +lat_ts=-71.0 +lat_0=-90 +lon_0=0.0 ' + '+k_0=1.0 +x_0=0.0 +y_0=0.0 +ellps=WGS84' + ) + + bedmachine_descriptor = ProjectionGridDescriptor.read( + projection, in_filename, 'BedMachineAntarctica500m', + ) + bedmachine_descriptor.to_scrip(out_filename) + + logger.info(' Done.') + + def _get_map_gebco(self): + """ + Create mapping files for GEBCO tiles onto NE grid + """ + logger = self.logger + logger.info('Creating mapping files for GEBCO tiles') config = self.config section = config['combine_topo'] + method = section.get('method') + lat_tiles = section.getint('lat_tiles') + lon_tiles = section.getint('lon_tiles') + resolution = self.resolution + + global_name = self.global_filename.strip('.nc') + tile_prefix = f'tiles/{global_name}_tile' + map_prefix = f'tiles/map_{global_name}_to_ne{resolution}_{method}_tile' + ne_scrip_filename = f'ne{resolution}.scrip.nc' + + for lat_tile in range(self.lat_tiles): + for lon_tile in range(self.lon_tiles): + in_filename = f'{tile_prefix}_{lon_tile}_{lat_tile}.nc' + out_filename = f'{map_prefix}_{lon_tile}_{lat_tile}.nc' + + args = [ + 'srun', '-n', f'{self.ntasks}', + 'ESMF_RegridWeightGen', + '--source', in_filename, + '--destination', ne_scrip_filename, + '--weight', out_filename, + '--method', method, + '--netcdf4', + '--src_regional', + '--ignore_unmapped', + ] + logger.info(f' Running: {" ".join(args)}') + subprocess.check_call(args) - latmin = section.getfloat('latmin') - latmax = section.getfloat('latmax') - cobined_filename = section.get('cobined_filename') + logger.info(' Done.') - gebco_filename = 'GEBCO_2023_0.0125_degree.nc' - gebco = xr.open_dataset(gebco_filename) + def _get_map_bedmachine(self): + """ + Create BedMachine to NE mapping file + """ + logger = self.logger + logger.info('Creating BedMachine mapping file') - bedmachine_filename = 'BedMachineAntarctica_on_GEBCO_low.nc' - bedmachine = xr.open_dataset(bedmachine_filename) + config = self.config + method = config.get('combine_topo', 'method') + resolution = self.resolution + antarctic_name = self.antarctic_filename.strip('.nc') + in_filename = f'{antarctic_name}.scrip.nc' + out_filename = f'map_{antarctic_name}_to_ne{resolution}_{method}.nc' + ne_scrip_filename = f'ne{resolution}.scrip.nc' + + args = [ + 'srun', '-n', f'{self.ntasks}', + 'ESMF_RegridWeightGen', + '--source', in_filename, + '--destination', ne_scrip_filename, + '--weight', out_filename, + '--method', method, + '--netcdf4', + '--src_regional' + '--ignore_unmapped' + ] + logger.info(f' Running: {" ".join(args)}') + subprocess.check_call(args) + + logger.info(' Done.') + + def _remap_gebco(self): + """ + Remap GEBCO tiles to NE grid + """ + logger = self.logger + logger.info('Remapping GEBCO tiles to NE grid') + + global_name = self.global_filename.strip('.nc') + tile_prefix = f'tiles/{global_name}_tile' + ne_prefix = f'tiles/{global_name}_on_ne_tile' + map_prefix = f'tiles/map_{global_name}_to_ne{resolution}_{method}_tile' + + for lat_tile in range(self.lat_tiles): + for lon_tile in range(self.lon_tiles): + in_filename = f'{tile_prefix}_{lon_tile}_{lat_tile}.nc' + out_filename = f'{ne_prefix}_{lon_tile}_{lat_tile}.nc' + mapping_filename = f'{map_prefix}_{lon_tile}_{lat_tile}.nc' + + args = [ + 'ncremap', + '-m', mapping_filename, + '--vrb=1', + f'--renormalize={self.renorm_thresh}', + in_filename, + out_filename, + ] + logger.info(f' Running: {" ".join(args)}') + subprocess.check_call(args) + + logger.info(' Done.') + + def _remap_bedmachine(self): + """ + Remap BedMachine to NE grid + Still need: + bedmachine_filename, bedmachine_on_ne_filename, + mapping_filename, renorm_thresh + """ + logger = self.logger + logger.info('Remapping BedMachine to NE grid') + + args = [ + 'ncremap', + '-m', mapping_filename, + '--vrb=1', + f'--renormalize={renorm_thresh}', + '-R', '--rgr lat_nm=y --rgr lon_nm=x', + bedmachine_filename, + bedmachine_on_ne_filename, + ] + logger.info(f' Running: {" ".join(args)}') + subprocess.check_call(args) + + logger.info(' Done.') + + def _combine(self): + """ + """ + logger = self.logger + logger.info('Combine BedMachineAntarctica and GEBCO') combined = xr.Dataset() - alpha = (gebco.lat - latmin) / (latmax - latmin) + for lat_tile in range(lat_tiles): + for lon_tile in range(lon_tiles): + gebco_filename = f'{gebco_prefix}_{lon_tile}_{lat_tile}.nc' + gebco = xr.open_dataset(gebco_filename) + tile = gebco.elevation + tile = tile.where(tile.notnull(), 0.) + if 'bathymetry' in combined: + combined['bathymetry'] = combined.bathymetry + tile + else: + combined['bathymetry'] = tile + + bedmachine = xr.open_dataset(bedmachine_filename) + + alpha = (combined.lat - latmin) / (latmax - latmin) alpha = np.maximum(np.minimum(alpha, 1.0), 0.0) bedmachine_bathy = bedmachine.bathymetry valid = bedmachine_bathy.notnull() bedmachine_bathy = bedmachine_bathy.where(valid, 0.) - bedmachine_bathy = bedmachine_bathy.where(bedmachine_bathy < 0., 0.) - combined['bathymetry'] = \ - alpha * gebco.bathymetry.where(gebco.bathymetry < 0., 0.) + \ - (1.0 - alpha) * bedmachine_bathy - bathy_mask = xr.where(combined.bathymetry < 0., 1.0, 0.0) + combined['bathymetry'] = ( + alpha * combined.bathymetry + (1.0 - alpha) * bedmachine_bathy + ) + + mask = combined.bathymetry < 0. + combined['bathymetry'] = combined.bathymetry.where(mask, 0.) for field in ['ice_draft', 'thickness']: - combined[field] = bathy_mask * bedmachine[field] + combined[field] = bedmachine[field].astype(float) for field in ['bathymetry', 'ice_draft', 'thickness']: combined[field].attrs['unit'] = 'meters' - for field in ['ice_mask', 'grounded_mask', 'ocean_mask']: - combined[field] = bedmachine[field] - - combined['bathymetry_mask'] = bathy_mask - - fill = {'ice_draft': 0., 'thickness': 0., 'ice_mask': 0., - 'grounded_mask': 0., 'ocean_mask': bathy_mask} + fill = { + 'ice_mask': 0., + 'grounded_mask': 0., + 'ocean_mask': combined.bathymetry < 0. + } for field, fill_val in fill.items(): - valid = combined[field].notnull() - combined[field] = combined[field].where(valid, fill_val) + valid = bedmachine[field].notnull() + combined[field] = bedmachine[field].where(valid, fill_val) - combined['water_column'] = \ - combined['ice_draft'] - combined['bathymetry'] + combined['water_column'] = combined.ice_draft - combined.bathymetry combined.water_column.attrs['units'] = 'meters' + combined = combined.drop_vars(['x', 'y']) + + combined.to_netcdf(combined_filename) - combined.to_netcdf(cobined_filename) logger.info(' Done.') - diff = xr.Dataset() +def _get_combined_filename( + target_grid, datestamp, resolution, + antarctic_filename, global_filename, +): + """ + """ + + # Parse resolution + if target_grid == 'cubed_sphere': + resolution_str = f'ne{resolution}' + elif target_grid == 'lat_lon': + resolution_str = f'{resolution}_degree' + else: + raise ValueError('Unknown target grid ' + target_grid) + + # Build combined filename + combined_filename = '_'.join([ + antarctic_filename.strip('.nc'), global_filename.strip('.nc'), + resolution_str, datestamp.strftime('%Y%m%d.nc'), + ]) + + return combined_filename + +def _get_tile_filenames(global_filename, resolution, method, lat_tile, lon_tile, tiledir='tiles'): + """ + """ + + # Tiles + global_name = global_filename.strip('.nc') + src_filename = os.path.join(tiledir, f'{global_name}_tile_{lon_tile}_{lat_tile}.nc') + tgt_filename = os.path.join(tiledir, f'{global_name}_on_ne_tile_{lon_tile}_{lat_tile}.nc') + map_filename = os.path.join(tiledir, f'map_{global_name}_to_ne{resolution}_{method}_tile_{lon_tile}_{lat_tile}.nc') - diff['bathymetry'] = \ - gebco.bathymetry.where(gebco.bathymetry < 0, 0.) - \ - bedmachine.bathymetry - diff.to_netcdf('gebco_minus_bedmachine.nc') + return src_filename, tgt_filename, map_filename From 8c1a10325e95c4993942019d18f302b63d1a6039 Mon Sep 17 00:00:00 2001 From: bmooremaley Date: Mon, 23 Sep 2024 09:31:28 -0500 Subject: [PATCH 2/5] More development on implementing cubed sphere interpolation into combine_topo, nearing complete workflow for pull request --- compass/ocean/tests/utility/__init__.py | 2 +- .../tests/utility/combine_topo/__init__.py | 537 +++++++++--------- .../utility/combine_topo/combine_topo.cfg | 15 +- 3 files changed, 281 insertions(+), 273 deletions(-) diff --git a/compass/ocean/tests/utility/__init__.py b/compass/ocean/tests/utility/__init__.py index 483dcfada..e92218f20 100644 --- a/compass/ocean/tests/utility/__init__.py +++ b/compass/ocean/tests/utility/__init__.py @@ -16,6 +16,6 @@ def __init__(self, mpas_core): """ super().__init__(mpas_core=mpas_core, name='utility') - self.add_test_case(CombineTopo(test_group=self)) + self.add_test_case(CombineTopo(test_group=self, target_grid='lat_lon')) self.add_test_case(CullRestarts(test_group=self)) self.add_test_case(ExtrapWoa(test_group=self)) diff --git a/compass/ocean/tests/utility/combine_topo/__init__.py b/compass/ocean/tests/utility/combine_topo/__init__.py index 6fe641417..3c8c5a3be 100644 --- a/compass/ocean/tests/utility/combine_topo/__init__.py +++ b/compass/ocean/tests/utility/combine_topo/__init__.py @@ -3,10 +3,10 @@ from datetime import datetime import numpy as np -import progressbar import pyproj import xarray as xr -from pyremap import LatLonGridDescriptor, ProjectionGridDescriptor, Remapper +from dask.diagnostics import ProgressBar +from pyremap import ProjectionGridDescriptor from compass.step import Step from compass.testcase import TestCase @@ -57,13 +57,14 @@ def __init__(self, test_case, target_grid): test_case : compass.ocean.tests.utility.combine_topo.CombineTopo The test case this step belongs to target_grid : `str`, either "lat_lon" or "cubed_sphere" - date_stamp : `datetime.datetime`, date when test was run """ super().__init__( test_case, name='combine', ntasks=None, min_tasks=None, ) self.target_grid = target_grid - self.date_stamp = None + self.resolution = None + self.resolution_name = None + def setup(self): """ @@ -78,16 +79,20 @@ def setup(self): # Get input filenames and resolution antarctic_filename = section.get('antarctic_filename') global_filename = section.get('global_filename') - resolution = section.getint('resolution') - # Datestamp - self.datestamp = datetime.now().strftime('%Y%m%d') + # Parse resolution + if self.target_grid == 'cubed_sphere': + self.resolution = section.getint('resolution') + self.resolution_name = f'ne{self.resolution}' + elif self.target_grid == 'lat_lon': + self.resolution = section.getfloat('resolution') + self.resolution_name = f'{self.resolution:.4f}_degree' - # Build output filename - combined_filename = _get_combined_filename( - self.target_grid, self.datestamp, resolution, - antarctic_filename, global_filename, - ) + # Build combined filename + combined_filename = '_'.join([ + antarctic_filename.strip('.nc'), global_filename.strip('.nc'), + self.resolution_name, datetime.now().strftime('%Y%m%d.nc'), + ]) # Add bathymetry data input files self.add_input_file( @@ -102,13 +107,11 @@ def setup(self): ) self.add_output_file(filename=combined_filename) - # Make tiles directory - os.makedirs('tiles', exist_ok=True) - # Get ntasks and min_tasks self.ntasks = section.getint('ntasks') self.min_tasks = section.getint('min_tasks') + def constrain_resources(self, available_resources): """ Constrain ``cpus_per_task`` and ``ntasks`` based on the number of @@ -124,20 +127,49 @@ def constrain_resources(self, available_resources): self.min_tasks = config.getint('combine_topo', 'min_tasks') super().constrain_resources(available_resources) + def run(self): """ Run this step of the test case """ - self._modify_bedmachine() self._modify_gebco() - self._tile_gebco() - self._create_ne_scrip() - self._create_bedmachine_scrip() - self._get_map_gebco() - self._get_map_bedmachine() + self._modify_bedmachine() + self._create_gebco_tiles() + self._create_bedmachine_scrip_file() + self._create_target_scrip_file() self._remap_gebco() self._remap_bedmachine() - self._combine() + + + def _modify_gebco(self): + """ + Modify GEBCO to include lon/lat bounds located at grid edges + """ + logger = self.logger + logger.info('Adding bounds to GEBCO lat/lon') + + # Parse config + config = self.config + in_filename = config.get('combine_topo', 'global_filename') + out_filename = in_filename.replace('.nc', '_cf.nc') + + # Modify GEBCO + gebco = xr.open_dataset(in_filename) + lat = gebco.lat + lon = gebco.lon + dlat = lat.isel(lat=1) - lat.isel(lat=0) + dlon = lon.isel(lon=1) - lon.isel(lon=0) + lat_bnds = xr.concat([lat - 0.5 * dlat, lat + 0.5 * dlat], dim='bnds') + lon_bnds = xr.concat([lon - 0.5 * dlon, lon + 0.5 * dlon], dim='bnds') + gebco['lat_bnds'] = lat_bnds.transpose('lat', 'bnds') + gebco['lon_bnds'] = lon_bnds.transpose('lon', 'bnds') + gebco.lat.attrs['bounds'] = 'lat_bnds' + gebco.lon.attrs['bounds'] = 'lon_bnds' + + # Write modified GEBCO to netCDF + gebco.to_netcdf(out_filename) + logger.info(' Done.') + def _modify_bedmachine(self): """ @@ -146,157 +178,145 @@ def _modify_bedmachine(self): logger = self.logger logger.info('Modifying BedMachineAntarctica with MPAS-Ocean names') + # Parse config config = self.config in_filename = config.get('combine_topo', 'antarctic_filename') out_filename = in_filename.replace('.nc', '_mod.nc') + # Load BedMachine and get ice, ocean and grounded masks bedmachine = xr.open_dataset(in_filename) mask = bedmachine.mask ice_mask = (mask != 0).astype(float) - ocean_mask = (np.logical_or(mask == 0, mask == 3)).astype(float) + ocean_mask = np.logical_or(mask == 0, mask == 3).astype(float) grounded_mask = np.logical_or(np.logical_or(mask == 1, mask == 2), - mask == 4).astype(float) + mask == 4).astype(float) - bedmachine['bathymetry'] = bedmachine.bed - bedmachine['ice_draft'] = bedmachine.surface - bedmachine.thickness + # Add new variables and apply ocean mask + bedmachine['bathymetry'] = bedmachine.bed.where(ocean_mask, 0.) + bedmachine['thickness'] = bedmachine.thickness.where(ocean_mask, 0.) + bedmachine['ice_draft'] = \ + (bedmachine.surface - bedmachine.thickness).where(ocean_mask, 0.) bedmachine.ice_draft.attrs['units'] = 'meters' - bedmachine['thickness'] = bedmachine.thickness - - bedmachine['bathymetry'] = bedmachine['bathymetry'].where(ocean_mask, 0.) - bedmachine['ice_draft'] = bedmachine['ice_draft'].where(ocean_mask, 0.) - bedmachine['thickness'] = bedmachine['thickness'].where(ocean_mask, 0.) - bedmachine['ice_mask'] = ice_mask bedmachine['grounded_mask'] = grounded_mask bedmachine['ocean_mask'] = ocean_mask + # Remove all other variables varlist = [ 'bathymetry', 'ice_draft', 'thickness', 'ice_mask', 'grounded_mask', 'ocean_mask', ] - bedmachine = bedmachine[varlist] + # Write modified BedMachine to netCDF bedmachine.to_netcdf(out_filename) logger.info(' Done.') - def _modify_gebco(self): - """ - Modify GEBCO to include lon/lat bounds located at grid edges - """ - logger = self.logger - logger.info('Adding bounds to GEBCO lat/lon') - - config = self.config - in_filename = config.get('combine_topo', 'global_filename') - out_filename = in_filename.replace('.nc', '_cf.nc') - gebco = xr.open_dataset(in_filename) - lat = gebco.lat - lon = gebco.lon - dlat = gebco.lat.isel(lat=1) - gebco.lat.isel(lat=0) - dlon = gebco.lon.isel(lon=1) - gebco.lon.isel(lon=0) - lat_bnds = xr.concat([lat - 0.5 * dlat, lat + 0.5 * dlat], dim='bnds') - lon_bnds = xr.concat([lon - 0.5 * dlon, lon + 0.5 * dlon], dim='bnds') - gebco['lat_bnds'] = lat_bnds.transpose('lat', 'bnds') - gebco['lon_bnds'] = lon_bnds.transpose('lon', 'bnds') - gebco.lat.attrs['bounds'] = 'lat_bnds' - gebco.lon.attrs['bounds'] = 'lon_bnds' - - gebco.to_netcdf(out_filename) - logger.info(' Done.') - - def _tile_gebco(self): + def _create_gebco_tile(self, lon_tile, lat_tile): """ - Tile GEBCO + Create GEBCO tile """ logger = self.logger - logger.info('Tiling GEBCO data') + # Parse config config = self.config section = config['combine_topo'] - global_filename = section.get('global_filename') lat_tiles = section.getint('lat_tiles') lon_tiles = section.getint('lon_tiles') + global_name = section.get('global_filename').strip('.nc') - in_filename = global_filename.replace('.nc', '_cf.nc') + # Load GEBCO + gebco = xr.open_dataset(f'{global_name}_cf.nc') - gebco = xr.open_dataset(in_filename) - gebco = gebco.chunk({'lat': lat_tiles * lon_tiles}) + # Build lat and lon arrays for tile nlat = gebco.sizes['lat'] nlon = gebco.sizes['lon'] nlat_tile = nlat // lat_tiles nlon_tile = nlon // lon_tiles - for lat_tile in range(lat_tiles): - for lon_tile in range(lon_tiles): - src_filename, _, _ = _get_tile_filenames(global_filename, resolution, method, lat_tile, lon_tile) - - lat_indices = np.arange( - lat_tile * nlat_tile - 1, (lat_tile + 1) * nlat_tile + 1 - ) - lon_indices = np.arange( - lon_tile * nlon_tile, (lon_tile + 1) * nlon_tile + 1 - ) - lat_indices = np.minimum(np.maximum(lat_indices, 0), nlat - 1) - lon_indices = np.mod(lon_indices, nlon) - ds_local = gebco.isel(lat=lat_indices, lon=lon_indices) - if lat_tile == 0: - # need to handle south pole - ds_local.lat.values[0] = -90. - if lat_tile == lat_tiles - 1: - # need to handle north pole - ds_local.lat.values[-1] = 90. - - write_job = ds_local.to_netcdf(out_filename, compute=False) - with ProgressBar(): - logger.info(f' writing {out_filename}') - write_job.compute() - logger.info(' Done.') - - def _create_ne_scrip(self): + # Build tile latlon indices + lat_indices = [lat_tile * nlat_tile, (lat_tile + 1) * nlat_tile] + lon_indices = [lon_tile * nlon_tile, (lon_tile + 1) * nlon_tile] + if lat_tile == lat_tiles - 1: + lat_indices[1] = max([lat_indices[1], nlat]) + else: + lat_indices[1] += 1 + if lon_tile == lon_tiles - 1: + lon_indices[1] = max([lon_indices[1], nlon]) + else: + lon_indices[1] += 1 + lat_indices = np.arange(*lat_indices) + lon_indices = np.arange(*lon_indices) + + # Duplicate top and bottom rows to account for poles + if lat_tile == 0: + lat_indices = np.insert(lat_indices, 0, 0) + if lat_tile == lat_tiles - 1: + lat_indices = np.append(lat_indices, lat_indices[-1]) + + # Select tile from GEBCO + tile = gebco.isel(lat=lat_indices, lon=lon_indices) + if lat_tile == 0: + tile.lat.values[0] = -90. # Correct south pole + if lat_tile == lat_tiles - 1: + tile.lat.values[-1] = 90. # Correct north pole + + # Write tile to netCDF + out_filename = f'tiles/{global_name}_tile_{lon_tile}_{lat_tile}.nc' + write_job = tile.to_netcdf(out_filename, compute=False) + with ProgressBar(): + logger.info(f' writing {out_filename}') + write_job.compute() + + + def _create_gebco_tiles(self): """ - Create SCRIP file for the NExxx (cubed-sphere) mesh - Reference: - https://acme-climate.atlassian.net/wiki/spaces/DOC/pages/872579110/ - Running+E3SM+on+New+Atmosphere+Grids + Create GEBCO tiles. Wrapper around `_create_gebco_tile` """ logger = self.logger - logger.info('Creating NE scrip file') + logger.info('Creating GEBCO tiles') - resolution = self.resolution + # Make tiles directory + os.makedirs('tiles', exist_ok=True) - args = [ - 'GenerateCSMesh', '--alt', '--res', f'{resolution}', - '--file', f'ne{resolution}.g', - ] - logger.info(f' Running: {" ".join(args)}') - subprocess.check_call(args) + # Parse config + config = self.config + section = config['combine_topo'] + lat_tiles = section.getint('lat_tiles') + lon_tiles = section.getint('lon_tiles') - args = [ - 'ConvertMeshToSCRIP', '--in', f'ne{resolution}.g', - '--out', f'ne{resolution}.scrip.nc', - ] - logger.info(f' Running: {" ".join(args)}') - subprocess.check_call(args) + # Loop through tiles + for lat_tile in range(lat_tiles): + for lon_tile in range(lon_tiles): + + # Create GEBCO tile + self._create_gebco_tile(lon_tile, lat_tile) logger.info(' Done.') - def _create_bedmachine_scrip(self): + + def _create_bedmachine_scrip_file(self): """ - Create SCRIP file for the BedMachineAntarctica-v3 bathymetry + Create SCRIP file for BedMachineAntarctica data """ logger = self.logger - logger.info('Creating Bedmachine scrip file') + logger.info('Creating BedMachine SCRIP file') - in_filename = self.antarctic_filename - out_filename = in_filename.replace('.nc', '.scrip.nc') + # Parse config + config = self.config + section = config['combine_topo'] + antarctic_name = section.get('antarctic_filename').strip('.nc') + in_filename = f'{antarctic_name}_mod.nc' + out_filename = f'{antarctic_name}.scrip.nc' + # Define projection projection = pyproj.Proj( '+proj=stere +lat_ts=-71.0 +lat_0=-90 +lon_0=0.0 ' '+k_0=1.0 +x_0=0.0 +y_0=0.0 +ellps=WGS84' ) + # Create SCRIP file bedmachine_descriptor = ProjectionGridDescriptor.read( projection, in_filename, 'BedMachineAntarctica500m', ) @@ -304,220 +324,201 @@ def _create_bedmachine_scrip(self): logger.info(' Done.') - def _get_map_gebco(self): + + def _create_target_scrip_file(self): """ - Create mapping files for GEBCO tiles onto NE grid + Create SCRIP file for either the x.xxxx degree (lat-lon) mesh or + the NExxx (cubed-sphere) mesh, depending on the value of `self.target_grid` + References: + https://github.com/ClimateGlobalChange/tempestremap + https://acme-climate.atlassian.net/wiki/spaces/DOC/pages/872579110/ + Running+E3SM+on+New+Atmosphere+Grids """ logger = self.logger - logger.info('Creating mapping files for GEBCO tiles') - - config = self.config - section = config['combine_topo'] - method = section.get('method') - lat_tiles = section.getint('lat_tiles') - lon_tiles = section.getint('lon_tiles') - resolution = self.resolution - - global_name = self.global_filename.strip('.nc') - tile_prefix = f'tiles/{global_name}_tile' - map_prefix = f'tiles/map_{global_name}_to_ne{resolution}_{method}_tile' - ne_scrip_filename = f'ne{resolution}.scrip.nc' + logger.info('Creating EXODUS and SCRIP files') + + # Create EXODUS file + if self.target_grid == 'cubed_sphere': + args = [ + 'GenerateCSMesh', '--alt', '--res', f'{self.resolution}', + '--file', f'{self.resolution_name}.g', + ] + elif self.target_grid == 'lat_lon': + nlon = int(360 / self.resolution) + nlat = int(nlon / 2) + args = [ + 'GenerateRLLMesh', + '--lon', f'{nlon}', '--lat', f'{nlat}', + '--file', f'{self.resolution_name}.g', + ] + logger.info(f' Running: {" ".join(args)}') + subprocess.run(args, check=True) - for lat_tile in range(self.lat_tiles): - for lon_tile in range(self.lon_tiles): - in_filename = f'{tile_prefix}_{lon_tile}_{lat_tile}.nc' - out_filename = f'{map_prefix}_{lon_tile}_{lat_tile}.nc' - - args = [ - 'srun', '-n', f'{self.ntasks}', - 'ESMF_RegridWeightGen', - '--source', in_filename, - '--destination', ne_scrip_filename, - '--weight', out_filename, - '--method', method, - '--netcdf4', - '--src_regional', - '--ignore_unmapped', - ] - logger.info(f' Running: {" ".join(args)}') - subprocess.check_call(args) + # Create SCRIP file + args = [ + 'ConvertMeshToSCRIP', '--in', f'{self.resolution_name}.g', + '--out', f'{self.resolution_name}.scrip.nc', + ] + logger.info(f' Running: {" ".join(args)}') + subprocess.run(args, check=True) logger.info(' Done.') - def _get_map_bedmachine(self): + + def _create_weights(self, in_filename, out_filename): """ - Create BedMachine to NE mapping file + Create weights file for remapping to `self.target_grid` """ logger = self.logger - logger.info('Creating BedMachine mapping file') config = self.config method = config.get('combine_topo', 'method') - resolution = self.resolution - antarctic_name = self.antarctic_filename.strip('.nc') - in_filename = f'{antarctic_name}.scrip.nc' - out_filename = f'map_{antarctic_name}_to_ne{resolution}_{method}.nc' - ne_scrip_filename = f'ne{resolution}.scrip.nc' + # Generate weights file args = [ 'srun', '-n', f'{self.ntasks}', 'ESMF_RegridWeightGen', '--source', in_filename, - '--destination', ne_scrip_filename, + '--destination', f'{self.resolution_name}.scrip.nc', '--weight', out_filename, '--method', method, '--netcdf4', - '--src_regional' - '--ignore_unmapped' + '--src_regional', + '--ignore_unmapped', ] + logger.info(f' running: {" ".join(args)}') + subprocess.run(args, check=True) + + + def _remap_to_target(self, in_filename, mapping_filename, out_filename, default_dims=True): + """ + Remap to `self.target_grid` + """ + logger = self.logger + + config = self.config + section = config['combine_topo'] + renorm_thresh = section.getfloat('renorm_thresh') + + # Build command args + args = [ + 'ncremap', '-m', mapping_filename, '--vrb=1', + f'--renormalize={renorm_thresh}', + ] + + # Append non-default dimensions, if any + if not default_dims: + args = args + ['-R', '--rgr lat_nm=y --rgr lon_nm=x'] + + # Append input and output file names + args = args + [in_filename, out_filename] + + # Remap to target grid logger.info(f' Running: {" ".join(args)}') - subprocess.check_call(args) + subprocess.run(args, check=True) - logger.info(' Done.') def _remap_gebco(self): """ - Remap GEBCO tiles to NE grid + Remap GEBCO """ logger = self.logger - logger.info('Remapping GEBCO tiles to NE grid') + logger.info('Remapping GEBCO data') + + # Parse config + config = self.config + section = config['combine_topo'] + global_name = section.get('global_filename').strip('.nc') + method = section.get('method') + lat_tiles = section.getint('lat_tiles') + lon_tiles = section.getint('lon_tiles') + + # Create tile maps and remapped tiles + for lat_tile in range(lat_tiles): + for lon_tile in range(lon_tiles): - global_name = self.global_filename.strip('.nc') - tile_prefix = f'tiles/{global_name}_tile' - ne_prefix = f'tiles/{global_name}_on_ne_tile' - map_prefix = f'tiles/map_{global_name}_to_ne{resolution}_{method}_tile' + # File names + tile_suffix = f'tile_{lon_tile}_{lat_tile}.nc' + tile_filename = f'tiles/{global_name}_{tile_suffix}' + mapping_filename = f'tiles/map_{global_name}_to_{self.resolution_name}_{method}_{tile_suffix}' + remapped_filename = f'tiles/{global_name}_{self.resolution_name}_{tile_suffix}' - for lat_tile in range(self.lat_tiles): - for lon_tile in range(self.lon_tiles): - in_filename = f'{tile_prefix}_{lon_tile}_{lat_tile}.nc' - out_filename = f'{ne_prefix}_{lon_tile}_{lat_tile}.nc' - mapping_filename = f'{map_prefix}_{lon_tile}_{lat_tile}.nc' - - args = [ - 'ncremap', - '-m', mapping_filename, - '--vrb=1', - f'--renormalize={self.renorm_thresh}', - in_filename, - out_filename, - ] - logger.info(f' Running: {" ".join(args)}') - subprocess.check_call(args) + # Call remapping functions + self._create_weights(tile_filename, mapping_filename) + self._remap_to_target(tile_filename, mapping_filename, remapped_filename) logger.info(' Done.') + def _remap_bedmachine(self): """ - Remap BedMachine to NE grid - Still need: - bedmachine_filename, bedmachine_on_ne_filename, - mapping_filename, renorm_thresh + Remap BedMachineAntarctica """ logger = self.logger - logger.info('Remapping BedMachine to NE grid') + logger.info('Remapping BedMachine data') - args = [ - 'ncremap', - '-m', mapping_filename, - '--vrb=1', - f'--renormalize={renorm_thresh}', - '-R', '--rgr lat_nm=y --rgr lon_nm=x', - bedmachine_filename, - bedmachine_on_ne_filename, - ] - logger.info(f' Running: {" ".join(args)}') - subprocess.check_call(args) + # Parse config + config = self.config + section = config['combine_topo'] + antarctic_name = section.get('antarctic_filename').strip('.nc') + method = section.get('method') + + # File names + in_filename = f'{antarctic_name}_mod.nc' + scrip_filename = f'{antarctic_name}.scrip.nc' + mapping_filename = f'map_{antarctic_name}_to_{self.resolution_name}_{method}.nc' + remapped_filename = f'{antarctic_name}_{self.resolution_name}.nc' + + # Call remapping functions + self._create_weights(scrip_filename, mapping_filename) + self._remap_to_target( + in_filename, mapping_filename, remapped_filename, default_dims=False, + ) logger.info(' Done.') + def _combine(self): """ """ - logger = self.logger - logger.info('Combine BedMachineAntarctica and GEBCO') combined = xr.Dataset() + + # Combine remapped GEBCO tiles + tile_prefix = f'tiles/{global_name}_{self.resolution_name}_tile' for lat_tile in range(lat_tiles): for lon_tile in range(lon_tiles): - gebco_filename = f'{gebco_prefix}_{lon_tile}_{lat_tile}.nc' - gebco = xr.open_dataset(gebco_filename) - tile = gebco.elevation - tile = tile.where(tile.notnull(), 0.) + in_filename = f'{tile_prefix}_{lon_tile}_{lat_tile}.nc' + tile = xr.open_dataset(in_filename) + bathy = tile.elevation + mask = bathy.notnull() + bathy = bathy.where(mask, 0.) if 'bathymetry' in combined: - combined['bathymetry'] = combined.bathymetry + tile + combined['bathymetry'] = combined.bathymetry + bathy else: - combined['bathymetry'] = tile - - bedmachine = xr.open_dataset(bedmachine_filename) + combined['bathymetry'] = bathy + # Open BedMachine and blend into combined bathy with alpha factor + bedmachine = xr.open_dataset(antarctic_filename) alpha = (combined.lat - latmin) / (latmax - latmin) alpha = np.maximum(np.minimum(alpha, 1.0), 0.0) + bathy = bedmachine.bathymetry + valid = bathy.notnull() + bathy = bathy.where(valid, 0.) + combined['bathymetry'] = alpha * combined.bathymetry + (1.0 - alpha) * bathy + combined['bathymetry'] = combined.bathymetry.where(combined.bathymetry < 0, 0.) - bedmachine_bathy = bedmachine.bathymetry - valid = bedmachine_bathy.notnull() - bedmachine_bathy = bedmachine_bathy.where(valid, 0.) - - combined['bathymetry'] = ( - alpha * combined.bathymetry + (1.0 - alpha) * bedmachine_bathy - ) - - mask = combined.bathymetry < 0. - combined['bathymetry'] = combined.bathymetry.where(mask, 0.) - + # Handle remaining variables for field in ['ice_draft', 'thickness']: combined[field] = bedmachine[field].astype(float) for field in ['bathymetry', 'ice_draft', 'thickness']: combined[field].attrs['unit'] = 'meters' - - fill = { - 'ice_mask': 0., - 'grounded_mask': 0., - 'ocean_mask': combined.bathymetry < 0. - } - - for field, fill_val in fill.items(): + fill = {'ice_mask': 0., 'grounded_mask': 0., 'ocean_mask': combined['bathymetry'] < 0.} + for field in fill.keys(): valid = bedmachine[field].notnull() - combined[field] = bedmachine[field].where(valid, fill_val) - - combined['water_column'] = combined.ice_draft - combined.bathymetry + combined[field] = bedmachine[field].where(valid, fill[field]) + combined['water_column'] = combined['ice_draft'] - combined['bathymetry'] combined.water_column.attrs['units'] = 'meters' combined = combined.drop_vars(['x', 'y']) combined.to_netcdf(combined_filename) - - logger.info(' Done.') - -def _get_combined_filename( - target_grid, datestamp, resolution, - antarctic_filename, global_filename, -): - """ - """ - - # Parse resolution - if target_grid == 'cubed_sphere': - resolution_str = f'ne{resolution}' - elif target_grid == 'lat_lon': - resolution_str = f'{resolution}_degree' - else: - raise ValueError('Unknown target grid ' + target_grid) - - # Build combined filename - combined_filename = '_'.join([ - antarctic_filename.strip('.nc'), global_filename.strip('.nc'), - resolution_str, datestamp.strftime('%Y%m%d.nc'), - ]) - - return combined_filename - -def _get_tile_filenames(global_filename, resolution, method, lat_tile, lon_tile, tiledir='tiles'): - """ - """ - - # Tiles - global_name = global_filename.strip('.nc') - src_filename = os.path.join(tiledir, f'{global_name}_tile_{lon_tile}_{lat_tile}.nc') - tgt_filename = os.path.join(tiledir, f'{global_name}_on_ne_tile_{lon_tile}_{lat_tile}.nc') - map_filename = os.path.join(tiledir, f'map_{global_name}_to_ne{resolution}_{method}_tile_{lon_tile}_{lat_tile}.nc') - - return src_filename, tgt_filename, map_filename diff --git a/compass/ocean/tests/utility/combine_topo/combine_topo.cfg b/compass/ocean/tests/utility/combine_topo/combine_topo.cfg index 51bd2f2d1..8e0a7c03e 100644 --- a/compass/ocean/tests/utility/combine_topo/combine_topo.cfg +++ b/compass/ocean/tests/utility/combine_topo/combine_topo.cfg @@ -5,13 +5,20 @@ antarctic_filename = BedMachineAntarctica-v3.nc global_filename = GEBCO_2023.nc -# the name of the output topography file, to be copied to the bathymetry database -cobined_filename = BedMachineAntarctica_v3_and_GEBCO_2023_0.0125_degree_20240828.nc +resolution = 0.5 +method = bilinear + +# the threshold for masks below which interpolated variables are not renormalized +renorm_thresh = 1e-3 # the target and minimum number of MPI tasks to use in remapping -ntasks = 512 -min_tasks = 128 +ntasks = 1280 +min_tasks = 512 # latitudes between which the topography datasets get blended latmin = -62. latmax = -60. + +# the number of tiles in lat and lon for GEBCO remapping +lat_tiles = 3 +lon_tiles = 6 \ No newline at end of file From fcceb05750431b07a5d37a48d173949d7a13de9f Mon Sep 17 00:00:00 2001 From: bmooremaley Date: Wed, 25 Sep 2024 18:11:56 -0500 Subject: [PATCH 3/5] Finished working version of combine_topo test that includes lat_lon and cubed_sphere remapping cases, and successfully ran for 0.1 deg and ne30 target grids --- compass/ocean/tests/utility/__init__.py | 5 +- .../tests/utility/combine_topo/__init__.py | 308 +++++++++++------- .../utility/combine_topo/combine_topo.cfg | 6 +- 3 files changed, 193 insertions(+), 126 deletions(-) diff --git a/compass/ocean/tests/utility/__init__.py b/compass/ocean/tests/utility/__init__.py index e92218f20..78bf62196 100644 --- a/compass/ocean/tests/utility/__init__.py +++ b/compass/ocean/tests/utility/__init__.py @@ -16,6 +16,9 @@ def __init__(self, mpas_core): """ super().__init__(mpas_core=mpas_core, name='utility') - self.add_test_case(CombineTopo(test_group=self, target_grid='lat_lon')) + for target_grid in ['lat_lon', 'cubed_sphere']: + self.add_test_case( + CombineTopo(test_group=self, target_grid=target_grid), + ) self.add_test_case(CullRestarts(test_group=self)) self.add_test_case(ExtrapWoa(test_group=self)) diff --git a/compass/ocean/tests/utility/combine_topo/__init__.py b/compass/ocean/tests/utility/combine_topo/__init__.py index 3c8c5a3be..30a62bc5b 100644 --- a/compass/ocean/tests/utility/combine_topo/__init__.py +++ b/compass/ocean/tests/utility/combine_topo/__init__.py @@ -1,12 +1,13 @@ import os import subprocess +from glob import glob from datetime import datetime import numpy as np import pyproj import xarray as xr from dask.diagnostics import ProgressBar -from pyremap import ProjectionGridDescriptor +from pyremap import ProjectionGridDescriptor, get_lat_lon_descriptor from compass.step import Step from compass.testcase import TestCase @@ -26,6 +27,7 @@ def __init__(self, test_group, target_grid): ---------- test_group : compass.ocean.tests.utility.Utility The test group that this test case belongs to + target_grid : `str`, either "lat_lon" or "cubed_sphere" """ subdir = os.path.join('combine_topo', target_grid) @@ -41,11 +43,12 @@ class Combine(Step): """ A step for combining GEBCO 2023 with BedMachineAntarctica topography datasets - TODO: Document target grid throughout Attributes ---------- - ... + target_grid : `str`, either "lat_lon" or "cubed_sphere" + resolution : `float` degrees, or `int` NExxx + resolution_name: `str`, either x.xxxx_degrees or NExxx """ def __init__(self, test_case, target_grid): @@ -65,7 +68,6 @@ def __init__(self, test_case, target_grid): self.resolution = None self.resolution_name = None - def setup(self): """ Set up the step in the work directory, including downloading any @@ -80,12 +82,12 @@ def setup(self): antarctic_filename = section.get('antarctic_filename') global_filename = section.get('global_filename') - # Parse resolution + # Parse resolution and assign resolution attributes if self.target_grid == 'cubed_sphere': - self.resolution = section.getint('resolution') + self.resolution = section.getint('resolution_cubedsphere') self.resolution_name = f'ne{self.resolution}' elif self.target_grid == 'lat_lon': - self.resolution = section.getfloat('resolution') + self.resolution = section.getfloat('resolution_latlon') self.resolution_name = f'{self.resolution:.4f}_degree' # Build combined filename @@ -111,7 +113,6 @@ def setup(self): self.ntasks = section.getint('ntasks') self.min_tasks = section.getint('min_tasks') - def constrain_resources(self, available_resources): """ Constrain ``cpus_per_task`` and ``ntasks`` based on the number of @@ -127,19 +128,17 @@ def constrain_resources(self, available_resources): self.min_tasks = config.getint('combine_topo', 'min_tasks') super().constrain_resources(available_resources) - def run(self): """ Run this step of the test case """ self._modify_gebco() self._modify_bedmachine() - self._create_gebco_tiles() - self._create_bedmachine_scrip_file() self._create_target_scrip_file() self._remap_gebco() self._remap_bedmachine() - + self._combine() + self._cleanup() def _modify_gebco(self): """ @@ -170,7 +169,6 @@ def _modify_gebco(self): gebco.to_netcdf(out_filename) logger.info(' Done.') - def _modify_bedmachine(self): """ Modify BedMachineAntarctica to compute the fields needed by MPAS-Ocean @@ -212,10 +210,14 @@ def _modify_bedmachine(self): bedmachine.to_netcdf(out_filename) logger.info(' Done.') - def _create_gebco_tile(self, lon_tile, lat_tile): """ Create GEBCO tile + + Parameters + ---------- + lon_tile : `int`, tile number along lon dim + lat_tile : `int`, tile number along lat dim """ logger = self.logger @@ -225,6 +227,9 @@ def _create_gebco_tile(self, lon_tile, lat_tile): lat_tiles = section.getint('lat_tiles') lon_tiles = section.getint('lon_tiles') global_name = section.get('global_filename').strip('.nc') + out_filename = f'tiles/{global_name}_tile_{lon_tile}_{lat_tile}.nc' + + logger.info(f' creating {out_filename}') # Load GEBCO gebco = xr.open_dataset(f'{global_name}_cf.nc') @@ -263,45 +268,14 @@ def _create_gebco_tile(self, lon_tile, lat_tile): tile.lat.values[-1] = 90. # Correct north pole # Write tile to netCDF - out_filename = f'tiles/{global_name}_tile_{lon_tile}_{lat_tile}.nc' - write_job = tile.to_netcdf(out_filename, compute=False) - with ProgressBar(): - logger.info(f' writing {out_filename}') - write_job.compute() - - - def _create_gebco_tiles(self): - """ - Create GEBCO tiles. Wrapper around `_create_gebco_tile` - """ - logger = self.logger - logger.info('Creating GEBCO tiles') - - # Make tiles directory - os.makedirs('tiles', exist_ok=True) - - # Parse config - config = self.config - section = config['combine_topo'] - lat_tiles = section.getint('lat_tiles') - lon_tiles = section.getint('lon_tiles') - - # Loop through tiles - for lat_tile in range(lat_tiles): - for lon_tile in range(lon_tiles): - - # Create GEBCO tile - self._create_gebco_tile(lon_tile, lat_tile) - - logger.info(' Done.') - + tile.to_netcdf(out_filename) def _create_bedmachine_scrip_file(self): """ - Create SCRIP file for BedMachineAntarctica data + Create SCRIP file for BedMachineAntarctica data using pyremap """ logger = self.logger - logger.info('Creating BedMachine SCRIP file') + logger.info(' Creating BedMachine SCRIP file') # Parse config config = self.config @@ -316,58 +290,63 @@ def _create_bedmachine_scrip_file(self): '+k_0=1.0 +x_0=0.0 +y_0=0.0 +ellps=WGS84' ) - # Create SCRIP file + # Create SCRIP file using pyremap bedmachine_descriptor = ProjectionGridDescriptor.read( projection, in_filename, 'BedMachineAntarctica500m', ) bedmachine_descriptor.to_scrip(out_filename) - logger.info(' Done.') - - def _create_target_scrip_file(self): """ - Create SCRIP file for either the x.xxxx degree (lat-lon) mesh or - the NExxx (cubed-sphere) mesh, depending on the value of `self.target_grid` + Create SCRIP file for either the x.xxxx degree (lat-lon) mesh or the + NExxx (cubed-sphere) mesh, depending on the value of `self.target_grid` References: - https://github.com/ClimateGlobalChange/tempestremap https://acme-climate.atlassian.net/wiki/spaces/DOC/pages/872579110/ Running+E3SM+on+New+Atmosphere+Grids """ logger = self.logger - logger.info('Creating EXODUS and SCRIP files') + logger.info(f'Create SCRIP file for {self.resolution_name} mesh') + + out_filename = f'{self.resolution_name}.scrip.nc' - # Create EXODUS file + # Build cubed sphere SCRIP file using tempestremap if self.target_grid == 'cubed_sphere': + + # Create EXODUS file args = [ 'GenerateCSMesh', '--alt', '--res', f'{self.resolution}', '--file', f'{self.resolution_name}.g', ] - elif self.target_grid == 'lat_lon': - nlon = int(360 / self.resolution) - nlat = int(nlon / 2) + logger.info(f' Running: {" ".join(args)}') + subprocess.run(args, check=True) + + # Create SCRIP file args = [ - 'GenerateRLLMesh', - '--lon', f'{nlon}', '--lat', f'{nlat}', - '--file', f'{self.resolution_name}.g', + 'ConvertMeshToSCRIP', '--in', f'{self.resolution_name}.g', + '--out', out_filename, ] - logger.info(f' Running: {" ".join(args)}') - subprocess.run(args, check=True) + logger.info(f' Running: {" ".join(args)}') + subprocess.run(args, check=True) - # Create SCRIP file - args = [ - 'ConvertMeshToSCRIP', '--in', f'{self.resolution_name}.g', - '--out', f'{self.resolution_name}.scrip.nc', - ] - logger.info(f' Running: {" ".join(args)}') - subprocess.run(args, check=True) + # Build lat-lon SCRIP file using pyremap + elif self.target_grid == 'lat_lon': + descriptor = get_lat_lon_descriptor( + dLon=self.resolution, dLat=self.resolution, + ) + descriptor.to_scrip(out_filename) logger.info(' Done.') - def _create_weights(self, in_filename, out_filename): """ - Create weights file for remapping to `self.target_grid` + Create weights file for remapping to `self.target_grid`. Filenames + are passed as parameters so that the function can be applied to + GEBCO and BedMachine + + Parameters + ---------- + in_filename : `str`, source file name + out_filename : `str`, weights file name """ logger = self.logger @@ -389,10 +368,20 @@ def _create_weights(self, in_filename, out_filename): logger.info(f' running: {" ".join(args)}') subprocess.run(args, check=True) - - def _remap_to_target(self, in_filename, mapping_filename, out_filename, default_dims=True): + def _remap_to_target( + self, in_filename, mapping_filename, out_filename, default_dims=True, + ): """ - Remap to `self.target_grid` + Remap to `self.target_grid`. Filenames are passed as parameters so + that the function can be applied to GEBCO and BedMachine. + + Parameters + ---------- + in_filename : `str`, source file name + mapping_filename : `str`, weights file name + out_filename : `str`, remapped file name + default_dims : `bool`, default `True`, + if `False` specify non-default source dims y,x """ logger = self.logger @@ -402,25 +391,38 @@ def _remap_to_target(self, in_filename, mapping_filename, out_filename, default_ # Build command args args = [ - 'ncremap', '-m', mapping_filename, '--vrb=1', - f'--renormalize={renorm_thresh}', + 'ncremap', + '-m', mapping_filename, + '--vrb=1', f'--renormalize={renorm_thresh}', ] - # Append non-default dimensions, if any + # Add non-default gridding args + regridArgs = [] if not default_dims: - args = args + ['-R', '--rgr lat_nm=y --rgr lon_nm=x'] + regridArgs.extend([ + '--rgr lat_nm=y', + '--rgr lon_nm=x', + ]) + if self.target_grid == 'lat_lon': + regridArgs.extend([ + '--rgr lat_nm_out=lat', + '--rgr lon_nm_out=lon', + '--rgr lat_dmn_nm=lat', + '--rgr lon_dmn_nm=lon', + ]) + if len(regridArgs) > 0: + args.extend(['-R', ' '.join(regridArgs)]) # Append input and output file names - args = args + [in_filename, out_filename] + args.extend([in_filename, out_filename]) # Remap to target grid logger.info(f' Running: {" ".join(args)}') subprocess.run(args, check=True) - def _remap_gebco(self): """ - Remap GEBCO + Remap GEBCO to `self.target_grid` """ logger = self.logger logger.info('Remapping GEBCO data') @@ -433,6 +435,12 @@ def _remap_gebco(self): lat_tiles = section.getint('lat_tiles') lon_tiles = section.getint('lon_tiles') + # Make tiles directory + os.mkdir('tiles') + + # Initialize combined xarray.Dataset + gebco_remapped = xr.Dataset() + # Create tile maps and remapped tiles for lat_tile in range(lat_tiles): for lon_tile in range(lon_tiles): @@ -440,19 +448,42 @@ def _remap_gebco(self): # File names tile_suffix = f'tile_{lon_tile}_{lat_tile}.nc' tile_filename = f'tiles/{global_name}_{tile_suffix}' - mapping_filename = f'tiles/map_{global_name}_to_{self.resolution_name}_{method}_{tile_suffix}' - remapped_filename = f'tiles/{global_name}_{self.resolution_name}_{tile_suffix}' + mapping_filename = ( + f'tiles/map_{global_name}_to_{self.resolution_name}' + f'_{method}_{tile_suffix}' + ) + remapped_filename = ( + f'tiles/{global_name}_{self.resolution_name}_{tile_suffix}' + ) # Call remapping functions + self._create_gebco_tile(lon_tile, lat_tile) self._create_weights(tile_filename, mapping_filename) - self._remap_to_target(tile_filename, mapping_filename, remapped_filename) + self._remap_to_target( + tile_filename, mapping_filename, remapped_filename, + ) + + # Add tile to remapped global bathymetry + logger.info(f' adding {remapped_filename}') + bathy = xr.open_dataset(remapped_filename).elevation + bathy = bathy.where(bathy.notnull(), 0.) + if 'bathymetry' in gebco_remapped: + gebco_remapped['bathymetry'] = ( + gebco_remapped.bathymetry + bathy + ) + else: + gebco_remapped['bathymetry'] = bathy + + # Write tile to netCDF + out_filename = f'{global_name}_{self.resolution_name}.nc' + logger.info(f' writing {out_filename}') + gebco_remapped.to_netcdf(out_filename) logger.info(' Done.') - def _remap_bedmachine(self): """ - Remap BedMachineAntarctica + Remap BedMachineAntarctica to `self.target_grid` """ logger = self.logger logger.info('Remapping BedMachine data') @@ -466,59 +497,90 @@ def _remap_bedmachine(self): # File names in_filename = f'{antarctic_name}_mod.nc' scrip_filename = f'{antarctic_name}.scrip.nc' - mapping_filename = f'map_{antarctic_name}_to_{self.resolution_name}_{method}.nc' + mapping_filename = ( + f'map_{antarctic_name}_to_{self.resolution_name}_{method}.nc' + ) remapped_filename = f'{antarctic_name}_{self.resolution_name}.nc' # Call remapping functions + self._create_bedmachine_scrip_file() self._create_weights(scrip_filename, mapping_filename) self._remap_to_target( - in_filename, mapping_filename, remapped_filename, default_dims=False, + in_filename, mapping_filename, remapped_filename, + default_dims=False, ) logger.info(' Done.') - def _combine(self): """ + Combine remapped GEBCO and BedMachine data sets """ + logger = self.logger + logger.info('Combine BedMachineAntarctica and GEBCO') - combined = xr.Dataset() + # Parse config + config = self.config + section = config['combine_topo'] + global_name = section.get('global_filename').strip('.nc') + antarctic_name = section.get('antarctic_filename').strip('.nc') + latmin = section.getfloat('latmin') + latmax = section.getfloat('latmax') - # Combine remapped GEBCO tiles - tile_prefix = f'tiles/{global_name}_{self.resolution_name}_tile' - for lat_tile in range(lat_tiles): - for lon_tile in range(lon_tiles): - in_filename = f'{tile_prefix}_{lon_tile}_{lat_tile}.nc' - tile = xr.open_dataset(in_filename) - bathy = tile.elevation - mask = bathy.notnull() - bathy = bathy.where(mask, 0.) - if 'bathymetry' in combined: - combined['bathymetry'] = combined.bathymetry + bathy - else: - combined['bathymetry'] = bathy + # Initialize combined xarray.Dataset + in_filename = f'{global_name}_{self.resolution_name}.nc' + combined = xr.open_dataset(in_filename) - # Open BedMachine and blend into combined bathy with alpha factor - bedmachine = xr.open_dataset(antarctic_filename) + # Create blending factor alpha alpha = (combined.lat - latmin) / (latmax - latmin) alpha = np.maximum(np.minimum(alpha, 1.0), 0.0) + + # Open remapped BedMachine and blend into combined bathy + in_filename = f'{antarctic_name}_{self.resolution_name}.nc' + bedmachine = xr.open_dataset(in_filename) bathy = bedmachine.bathymetry - valid = bathy.notnull() - bathy = bathy.where(valid, 0.) - combined['bathymetry'] = alpha * combined.bathymetry + (1.0 - alpha) * bathy - combined['bathymetry'] = combined.bathymetry.where(combined.bathymetry < 0, 0.) + bathy = bathy.where(bathy.notnull(), 0.) + combined['bathymetry'] = ( + alpha * combined.bathymetry + (1.0 - alpha) * bathy + ) + ocean_mask = combined.bathymetry < 0. + combined['bathymetry'] = combined.bathymetry.where(ocean_mask, 0.) - # Handle remaining variables + # Add remaining Bedmachine variables to combined Dataset for field in ['ice_draft', 'thickness']: combined[field] = bedmachine[field].astype(float) - for field in ['bathymetry', 'ice_draft', 'thickness']: + combined['water_column'] = combined.ice_draft - combined.bathymetry + for field in ['bathymetry', 'ice_draft', 'thickness', 'water_column']: combined[field].attrs['unit'] = 'meters' - fill = {'ice_mask': 0., 'grounded_mask': 0., 'ocean_mask': combined['bathymetry'] < 0.} - for field in fill.keys(): - valid = bedmachine[field].notnull() - combined[field] = bedmachine[field].where(valid, fill[field]) - combined['water_column'] = combined['ice_draft'] - combined['bathymetry'] - combined.water_column.attrs['units'] = 'meters' - combined = combined.drop_vars(['x', 'y']) - - combined.to_netcdf(combined_filename) + + # Add Bedmachine masks to combined Dataset + masks = { + 'ice_mask': 0., + 'grounded_mask': 0., + 'ocean_mask': ocean_mask, + } + for field in masks: + combined[field] = bedmachine[field].where( + bedmachine[field].notnull(), masks[field], + ) + + # Drop x,y coordinates + #combined = combined.drop_vars(['x', 'y']) + + # Save combined bathy to NetCDF + combined.to_netcdf(self.outputs[0]) + + logger.info(' Done.') + + def _cleanup(self): + """ + Clean up work directory + """ + logger = self.logger + logger.info('Cleaning up work directory') + + # Remove PETxxx.RegridWeightGen.Log files + for f in glob('*.RegridWeightGen.Log'): + os.remove(f) + + logger.info(' Done.') \ No newline at end of file diff --git a/compass/ocean/tests/utility/combine_topo/combine_topo.cfg b/compass/ocean/tests/utility/combine_topo/combine_topo.cfg index 8e0a7c03e..8a2a34cc2 100644 --- a/compass/ocean/tests/utility/combine_topo/combine_topo.cfg +++ b/compass/ocean/tests/utility/combine_topo/combine_topo.cfg @@ -5,10 +5,12 @@ antarctic_filename = BedMachineAntarctica-v3.nc global_filename = GEBCO_2023.nc -resolution = 0.5 +# target resolution (degrees or NExxx) +resolution_latlon = 0.0125 +resolution_cubedsphere = 9000 method = bilinear -# the threshold for masks below which interpolated variables are not renormalized +# threshold for masks below which interpolated variables are not renormalized renorm_thresh = 1e-3 # the target and minimum number of MPI tasks to use in remapping From abd28866ce4f80e644c22beed9f84f5e40a22825 Mon Sep 17 00:00:00 2001 From: bmooremaley Date: Tue, 1 Oct 2024 15:28:19 -0500 Subject: [PATCH 4/5] Addressed pull request comments for combine-topo-to-cubedsphere --- .../tests/utility/combine_topo/__init__.py | 103 +++++++++--------- .../utility/combine_topo/combine_topo.cfg | 6 +- 2 files changed, 54 insertions(+), 55 deletions(-) diff --git a/compass/ocean/tests/utility/combine_topo/__init__.py b/compass/ocean/tests/utility/combine_topo/__init__.py index 30a62bc5b..903983c22 100644 --- a/compass/ocean/tests/utility/combine_topo/__init__.py +++ b/compass/ocean/tests/utility/combine_topo/__init__.py @@ -1,14 +1,14 @@ import os -import subprocess -from glob import glob from datetime import datetime +from glob import glob import numpy as np import pyproj import xarray as xr -from dask.diagnostics import ProgressBar +from mpas_tools.logging import check_call from pyremap import ProjectionGridDescriptor, get_lat_lon_descriptor +from compass.parallel import run_command from compass.step import Step from compass.testcase import TestCase @@ -187,13 +187,12 @@ def _modify_bedmachine(self): ice_mask = (mask != 0).astype(float) ocean_mask = np.logical_or(mask == 0, mask == 3).astype(float) grounded_mask = np.logical_or(np.logical_or(mask == 1, mask == 2), - mask == 4).astype(float) + mask == 4).astype(float) # Add new variables and apply ocean mask - bedmachine['bathymetry'] = bedmachine.bed.where(ocean_mask, 0.) - bedmachine['thickness'] = bedmachine.thickness.where(ocean_mask, 0.) - bedmachine['ice_draft'] = \ - (bedmachine.surface - bedmachine.thickness).where(ocean_mask, 0.) + bedmachine['bathymetry'] = bedmachine.bed + bedmachine['thickness'] = bedmachine.thickness + bedmachine['ice_draft'] = bedmachine.surface - bedmachine.thickness bedmachine.ice_draft.attrs['units'] = 'meters' bedmachine['ice_mask'] = ice_mask bedmachine['grounded_mask'] = grounded_mask @@ -317,16 +316,14 @@ def _create_target_scrip_file(self): 'GenerateCSMesh', '--alt', '--res', f'{self.resolution}', '--file', f'{self.resolution_name}.g', ] - logger.info(f' Running: {" ".join(args)}') - subprocess.run(args, check=True) + check_call(args, logger) # Create SCRIP file args = [ 'ConvertMeshToSCRIP', '--in', f'{self.resolution_name}.g', '--out', out_filename, ] - logger.info(f' Running: {" ".join(args)}') - subprocess.run(args, check=True) + check_call(args, logger) # Build lat-lon SCRIP file using pyremap elif self.target_grid == 'lat_lon': @@ -348,14 +345,11 @@ def _create_weights(self, in_filename, out_filename): in_filename : `str`, source file name out_filename : `str`, weights file name """ - logger = self.logger - config = self.config method = config.get('combine_topo', 'method') # Generate weights file args = [ - 'srun', '-n', f'{self.ntasks}', 'ESMF_RegridWeightGen', '--source', in_filename, '--destination', f'{self.resolution_name}.scrip.nc', @@ -365,8 +359,10 @@ def _create_weights(self, in_filename, out_filename): '--src_regional', '--ignore_unmapped', ] - logger.info(f' running: {" ".join(args)}') - subprocess.run(args, check=True) + run_command( + args, self.cpus_per_task, self.ntasks, + self.openmp_threads, config, self.logger, + ) def _remap_to_target( self, in_filename, mapping_filename, out_filename, default_dims=True, @@ -383,8 +379,6 @@ def _remap_to_target( default_dims : `bool`, default `True`, if `False` specify non-default source dims y,x """ - logger = self.logger - config = self.config section = config['combine_topo'] renorm_thresh = section.getfloat('renorm_thresh') @@ -417,8 +411,7 @@ def _remap_to_target( args.extend([in_filename, out_filename]) # Remap to target grid - logger.info(f' Running: {" ".join(args)}') - subprocess.run(args, check=True) + check_call(args, self.logger) def _remap_gebco(self): """ @@ -473,7 +466,7 @@ def _remap_gebco(self): ) else: gebco_remapped['bathymetry'] = bathy - + # Write tile to netCDF out_filename = f'{global_name}_{self.resolution_name}.nc' logger.info(f' writing {out_filename}') @@ -522,50 +515,56 @@ def _combine(self): # Parse config config = self.config section = config['combine_topo'] - global_name = section.get('global_filename').strip('.nc') - antarctic_name = section.get('antarctic_filename').strip('.nc') + sfx = f'_{self.resolution_name}.nc' + global_fname = section.get('global_filename').replace('.nc', sfx) + antarctic_fname = section.get('antarctic_filename').replace('.nc', sfx) latmin = section.getfloat('latmin') latmax = section.getfloat('latmax') - # Initialize combined xarray.Dataset - in_filename = f'{global_name}_{self.resolution_name}.nc' - combined = xr.open_dataset(in_filename) - - # Create blending factor alpha - alpha = (combined.lat - latmin) / (latmax - latmin) + # Load and mask GEBCO + gebco = xr.open_dataset(global_fname) + gebco_bathy = gebco.bathymetry + gebco_bathy = gebco_bathy.where(gebco_bathy.notnull(), 0.) + gebco_bathy = gebco_bathy.where(gebco_bathy < 0., 0.) + + # Load and mask BedMachine + bedmachine = xr.open_dataset(antarctic_fname) + bed_bathy = bedmachine.bathymetry + bed_bathy = bed_bathy.where(bed_bathy.notnull(), 0.) + bed_bathy = bed_bathy.where(bed_bathy < 0., 0.) + + # Blend data sets into combined data set + combined = xr.Dataset() + alpha = (gebco.lat - latmin) / (latmax - latmin) alpha = np.maximum(np.minimum(alpha, 1.0), 0.0) - - # Open remapped BedMachine and blend into combined bathy - in_filename = f'{antarctic_name}_{self.resolution_name}.nc' - bedmachine = xr.open_dataset(in_filename) - bathy = bedmachine.bathymetry - bathy = bathy.where(bathy.notnull(), 0.) combined['bathymetry'] = ( - alpha * combined.bathymetry + (1.0 - alpha) * bathy + alpha * gebco_bathy + (1.0 - alpha) * bed_bathy ) - ocean_mask = combined.bathymetry < 0. - combined['bathymetry'] = combined.bathymetry.where(ocean_mask, 0.) + bathy_mask = xr.where(combined.bathymetry < 0., 1.0, 0.0) # Add remaining Bedmachine variables to combined Dataset for field in ['ice_draft', 'thickness']: - combined[field] = bedmachine[field].astype(float) + combined[field] = bathy_mask * bedmachine[field] combined['water_column'] = combined.ice_draft - combined.bathymetry for field in ['bathymetry', 'ice_draft', 'thickness', 'water_column']: combined[field].attrs['unit'] = 'meters' - - # Add Bedmachine masks to combined Dataset - masks = { + + # Add masks + combined['bathymetry_mask'] = bathy_mask + for field in ['ice_mask', 'grounded_mask', 'ocean_mask']: + combined[field] = bedmachine[field] + + # Add fill values + fill_vals = { + 'ice_draft': 0., + 'thickness': 0., 'ice_mask': 0., 'grounded_mask': 0., - 'ocean_mask': ocean_mask, + 'ocean_mask': bathy_mask, } - for field in masks: - combined[field] = bedmachine[field].where( - bedmachine[field].notnull(), masks[field], - ) - - # Drop x,y coordinates - #combined = combined.drop_vars(['x', 'y']) + for field, fill_val in fill_vals.items(): + valid = combined[field].notnull() + combined[field] = combined[field].where(valid, fill_val) # Save combined bathy to NetCDF combined.to_netcdf(self.outputs[0]) @@ -583,4 +582,4 @@ def _cleanup(self): for f in glob('*.RegridWeightGen.Log'): os.remove(f) - logger.info(' Done.') \ No newline at end of file + logger.info(' Done.') diff --git a/compass/ocean/tests/utility/combine_topo/combine_topo.cfg b/compass/ocean/tests/utility/combine_topo/combine_topo.cfg index 8a2a34cc2..d9ec03d2d 100644 --- a/compass/ocean/tests/utility/combine_topo/combine_topo.cfg +++ b/compass/ocean/tests/utility/combine_topo/combine_topo.cfg @@ -7,7 +7,7 @@ global_filename = GEBCO_2023.nc # target resolution (degrees or NExxx) resolution_latlon = 0.0125 -resolution_cubedsphere = 9000 +resolution_cubedsphere = 3000 method = bilinear # threshold for masks below which interpolated variables are not renormalized @@ -15,7 +15,7 @@ renorm_thresh = 1e-3 # the target and minimum number of MPI tasks to use in remapping ntasks = 1280 -min_tasks = 512 +min_tasks = 1024 # latitudes between which the topography datasets get blended latmin = -62. @@ -23,4 +23,4 @@ latmax = -60. # the number of tiles in lat and lon for GEBCO remapping lat_tiles = 3 -lon_tiles = 6 \ No newline at end of file +lon_tiles = 6 From 1ed573a7d8c2881e86cb52251f4c093987814161 Mon Sep 17 00:00:00 2001 From: bmooremaley Date: Sun, 6 Oct 2024 14:36:28 -0500 Subject: [PATCH 5/5] Added scripfile output and updated docs for combine_topo/cubed_sphere --- .../tests/utility/combine_topo/__init__.py | 16 ++++++++++------ .../tests/utility/combine_topo/combine_topo.cfg | 2 +- .../ocean/test_groups/utility.rst | 13 +++++++++---- docs/users_guide/ocean/test_groups/utility.rst | 17 +++++++++++++---- 4 files changed, 33 insertions(+), 15 deletions(-) diff --git a/compass/ocean/tests/utility/combine_topo/__init__.py b/compass/ocean/tests/utility/combine_topo/__init__.py index 903983c22..227ac12ae 100644 --- a/compass/ocean/tests/utility/combine_topo/__init__.py +++ b/compass/ocean/tests/utility/combine_topo/__init__.py @@ -90,10 +90,13 @@ def setup(self): self.resolution = section.getfloat('resolution_latlon') self.resolution_name = f'{self.resolution:.4f}_degree' - # Build combined filename + # Build output filenames + datestamp = datetime.now().strftime('%Y%m%d') + scrip_filename = f'{self.resolution_name}_{datestamp}.scrip.nc' combined_filename = '_'.join([ - antarctic_filename.strip('.nc'), global_filename.strip('.nc'), - self.resolution_name, datetime.now().strftime('%Y%m%d.nc'), + antarctic_filename.strip('.nc'), + global_filename.strip('.nc'), + self.resolution_name, f'{datestamp}.nc', ]) # Add bathymetry data input files @@ -107,6 +110,7 @@ def setup(self): target=global_filename, database='bathymetry_database', ) + self.add_output_file(filename=scrip_filename) self.add_output_file(filename=combined_filename) # Get ntasks and min_tasks @@ -306,7 +310,7 @@ def _create_target_scrip_file(self): logger = self.logger logger.info(f'Create SCRIP file for {self.resolution_name} mesh') - out_filename = f'{self.resolution_name}.scrip.nc' + out_filename = self.outputs[0] # Build cubed sphere SCRIP file using tempestremap if self.target_grid == 'cubed_sphere': @@ -352,7 +356,7 @@ def _create_weights(self, in_filename, out_filename): args = [ 'ESMF_RegridWeightGen', '--source', in_filename, - '--destination', f'{self.resolution_name}.scrip.nc', + '--destination', self.outputs[0], '--weight', out_filename, '--method', method, '--netcdf4', @@ -567,7 +571,7 @@ def _combine(self): combined[field] = combined[field].where(valid, fill_val) # Save combined bathy to NetCDF - combined.to_netcdf(self.outputs[0]) + combined.to_netcdf(self.outputs[1]) logger.info(' Done.') diff --git a/compass/ocean/tests/utility/combine_topo/combine_topo.cfg b/compass/ocean/tests/utility/combine_topo/combine_topo.cfg index d9ec03d2d..6bcbad53c 100644 --- a/compass/ocean/tests/utility/combine_topo/combine_topo.cfg +++ b/compass/ocean/tests/utility/combine_topo/combine_topo.cfg @@ -15,7 +15,7 @@ renorm_thresh = 1e-3 # the target and minimum number of MPI tasks to use in remapping ntasks = 1280 -min_tasks = 1024 +min_tasks = 512 # latitudes between which the topography datasets get blended latmin = -62. diff --git a/docs/developers_guide/ocean/test_groups/utility.rst b/docs/developers_guide/ocean/test_groups/utility.rst index f1157ecbb..8c331f6b6 100644 --- a/docs/developers_guide/ocean/test_groups/utility.rst +++ b/docs/developers_guide/ocean/test_groups/utility.rst @@ -19,10 +19,15 @@ dataset. combine ~~~~~~~ The class :py:class:`compass.ocean.tests.utility.combine_topo.Combine` -defines a step for combining the datasets above. The GEBCO data is downsampled -to a 1/80 degree latitude-longitude grid to make later remapping to MPAS meshes -more manageable. The BedMachine data is remapped to this same mesh and the -two datasets are blended between 60 and 62 degrees south latitude. +defines a step for combining the datasets above. The GEBCO and BedMachine data +are remapped to a common global grid to make later remapping to MPAS meshes +more manageable, and the two datasets are blended between 60 and 62 degrees +south latitude. The GEBCO global dataset is divided into regional tiles prior +to remapping to improve performance. Two common global target grid options are +provided: a 1/80 degree latitude-longitude grid and an ne3000 cubed sphere +grid. These target grid options are selectable via the ``target_grid`` argument +in the :py:class:`compass.ocean.tests.utility.combine_topo.CombineTopo` test +case class. cull_restarts ------------- diff --git a/docs/users_guide/ocean/test_groups/utility.rst b/docs/users_guide/ocean/test_groups/utility.rst index a3cc2c027..9abfd128b 100644 --- a/docs/users_guide/ocean/test_groups/utility.rst +++ b/docs/users_guide/ocean/test_groups/utility.rst @@ -8,13 +8,22 @@ may create datasets for other test groups to use. It is partly designed to provide provenance for data processing that may be more complex than a single, short script. -combine_topo ------------- -The ``ocean/utility/combine_topo`` test case is used to combine the +combine_topo/lat_lon +-------------------- +The ``ocean/utility/combine_topo/lat_lon`` test case is used to combine the `BedMachine Antarctica v3 `_ dataset with the `GEBCO 2023 `_ dataset. The result is on a 1/80 degree latitude-longitude grid. This utility -is intended for provenance. It is intended to doucment the process for +is intended for provenance. It is intended to document the process for +producing the topography dataset used for E3SM ocean meshes. + +combine_topo/cubed_sphere +------------------------- +The ``ocean/utility/combine_topo/cubed_sphere`` test case is used to combine the +`BedMachine Antarctica v3 `_ +dataset with the `GEBCO 2023 `_ +dataset. The result is on an ne3000 cubed sphere grid. This utility +is intended for provenance. It is intended to document the process for producing the topography dataset used for E3SM ocean meshes. cull_restarts