diff --git a/compass/landice/mesh.py b/compass/landice/mesh.py index 2a2b86fd38..745c66953c 100644 --- a/compass/landice/mesh.py +++ b/compass/landice/mesh.py @@ -976,6 +976,45 @@ def add_bedmachine_thk_to_ais_gridded_data(self, source_gridded_dataset, return gridded_dataset_with_bm_thk +def _nearest_fill_from_valid(field2d, valid_mask): + """ + Fill invalid cells in a 2D regular raster using the value from the nearest + valid cell on the same grid, via a fast Euclidean-distance transform. + + Parameters + ---------- + field2d : numpy.ndarray + 2D field to be filled. + + valid_mask : numpy.ndarray + Boolean mask where True marks valid cells. + + Returns + ------- + filled : numpy.ndarray + Copy of ``field2d`` with invalid cells filled from the nearest valid + cell. + """ + valid_mask = np.asarray(valid_mask, dtype=bool) + + if field2d.shape != valid_mask.shape: + raise ValueError('field2d and valid_mask must have the same shape') + + if not np.any(valid_mask): + raise ValueError('No valid cells available for nearest fill.') + + # distance_transform_edt maps each True (foreground) cell to the nearest + # False cell, so pass ~valid_mask to get nearest-valid indices. + nearest_inds = distance_transform_edt( + ~valid_mask, return_distances=False, return_indices=True) + + filled = np.array(field2d, copy=True) + invalid = ~valid_mask + filled[invalid] = field2d[nearest_inds[0, invalid], + nearest_inds[1, invalid]] + return filled + + def preprocess_ais_data(self, source_gridded_dataset, floodFillMask): """ Perform adjustments to gridded AIS datasets needed @@ -996,45 +1035,6 @@ def preprocess_ais_data(self, source_gridded_dataset, floodFillMask): """ logger = self.logger - def _nearest_fill_from_valid(field2d, valid_mask): - """ - Fill invalid cells in a 2D regular raster using the value from the - nearest valid cell on the same grid. - - Parameters - ---------- - field2d : numpy.ndarray - 2D field to be filled - valid_mask : numpy.ndarray - Boolean mask where True marks valid cells - - Returns - ------- - filled : numpy.ndarray - Copy of field2d with invalid cells filled - """ - valid_mask = np.asarray(valid_mask, dtype=bool) - - if field2d.shape != valid_mask.shape: - raise ValueError('field2d and valid_mask must have the same shape') - - if not np.any(valid_mask): - raise ValueError('No valid cells available for nearest fill.') - - # For EDT, foreground=True cells get mapped to nearest background=False - # cell when return_indices=True. So we pass ~valid_mask. - nearest_inds = distance_transform_edt( - ~valid_mask, return_distances=False, return_indices=True - ) - - filled = np.array(field2d, copy=True) - invalid = ~valid_mask - filled[invalid] = field2d[ - nearest_inds[0, invalid], - nearest_inds[1, invalid] - ] - return filled - # Apply floodFillMask to thickness field to help with culling file_with_flood_fill = \ f"{source_gridded_dataset.split('.')[:-1][0]}_floodFillMask.nc" @@ -1564,9 +1564,76 @@ def _src_extent(filepath): _log(f'Warning: could not save mesh boundary plot: {exc}') +def _write_extrapolated_source(source_file, varnames, logger=None): + """ + Write a copy of ``source_file`` with each listed variable nearest-filled + into its no-data (non-finite) cells using a fast Euclidean-distance + transform. + + Conservative remapping from this extrapolated copy has valid data across + the whole mesh footprint, avoiding the margin ramps and single-pixel fill + artifacts that arise when no-data cells are instead masked out of the + remapping weights. Only the listed variables and their coordinates are + written, keeping the temporary file small. + + Parameters + ---------- + source_file : str + Path to the source gridded dataset. + + varnames : list of str + Source variable names to extrapolate (e.g. ``['vx', 'vy', 'vErr']``). + Names not present in the source are skipped. + + logger : logging.Logger, optional + Logger for status messages; falls back to print if None. + + Returns + ------- + str + Path to the extrapolated source copy written to the current directory. + """ + def _log(msg): + if logger is not None: + logger.info(msg) + else: + print(msg) + + out_file = f'extrap_{os.path.basename(source_file)}' + + with xarray.open_dataset(source_file) as ds: + coord_vars = [c for c in ('x', 'y', 'x1', 'y1') if c in ds] + present = [v for v in varnames if v in ds] + for name in varnames: + if name not in ds: + _log(f" '{name}' not in source; skipping extrapolation") + ds_out = ds[present + coord_vars].load() + + for name in present: + da = ds_out[name] + values = np.asarray(da.values, dtype=float) + if values.ndim < 2: + raise ValueError( + f"Expected a >=2D field for '{name}', got shape " + f"{values.shape}") + ny, nx = values.shape[-2:] + planes = values.reshape(-1, ny, nx) + for k in range(planes.shape[0]): + valid = np.isfinite(planes[k]) + n_fill = int((~valid).sum()) + if n_fill == 0: + continue + _log(f" extrapolating '{name}': filling {n_fill} no-data cells") + planes[k] = _nearest_fill_from_valid(planes[k], valid) + ds_out[name] = (da.dims, planes.reshape(values.shape).astype(da.dtype)) + + ds_out.to_netcdf(out_file) + return out_file + + def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, nProcs, dest_file, proj, variables="all", - hull_path=None): + hull_path=None, extrap_varnames=None): """ Interpolate gridded dataset (e.g. MEASURES, BedMachine) onto a MALI mesh @@ -1600,6 +1667,14 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, avoids redundant I/O and computation when this function is called multiple times with the same destination mesh. + extrap_varnames : list of str or None, optional + Source variable names to extrapolate into their no-data regions (via a + fast nearest-value distance transform) before remapping. When provided, + a temporary extrapolated copy of the source is remapped instead of the + raw file, so conservative remapping has valid data across the whole + mesh footprint and margins are free of interpolation ramps and + single-pixel fill artifacts. When None, the source is remapped as-is. + Returns ------- masked_source_scrip : str @@ -1608,6 +1683,18 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, logger = self.logger + # Optionally extrapolate the source's no-data regions into a temporary + # copy so a plain conservative remap has valid data across the whole mesh + # footprint (avoids margin ramps and single-pixel fill artifacts). + extrap_file = None + if extrap_varnames is not None: + logger.info('extrapolating source no-data regions before remapping') + extrap_file = _write_extrapolated_source( + source_file, extrap_varnames, logger=logger) + interp_source = extrap_file + else: + interp_source = source_file + bare = os.path.splitext(os.path.basename(source_file))[0] match = re.search(r'(^.*[_-]v\d*[_-])+', bare) if match: @@ -1627,7 +1714,7 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, logger.info('creating scrip file for source dataset') # Note: writing scrip file to workdir args = ['create_scrip_file_from_planar_rectangular_grid', - '-i', source_file, + '-i', interp_source, '-s', source_scrip, '-p', proj, '-r', '2'] @@ -1637,6 +1724,7 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, # ESMF_RegridWeightGen only processes the cells that overlap the target. stem = os.path.splitext(source_scrip)[0] # strips .nc masked_source_scrip = f'{stem}_masked.nc' + logger.info('masking source SCRIP to destination mesh footprint') add_grid_imask_from_dst_scrip_hull( source_scrip=source_scrip, @@ -1664,7 +1752,7 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, # Perform actual interpolation using the weights logger.info('calling interpolate_to_mpasli_grid') args = ['interpolate_to_mpasli_grid', - '-s', source_file, + '-s', interp_source, '-d', dest_file, '-m', 'e', '-w', weights_filename, @@ -1672,6 +1760,10 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, check_call(args, logger=logger) + if extrap_file is not None and os.path.exists(extrap_file): + logger.info('removing temporary extrapolated source file') + os.remove(extrap_file) + return masked_source_scrip @@ -1963,11 +2055,14 @@ def run_optional_interpolation( measures_vars = ['observedSurfaceVelocityX', 'observedSurfaceVelocityY', 'observedSurfaceVelocityUncertainty'] + # extrapolate raw velocity/error into no-data gaps on the source + # raster so conservative remapping has clean margins interp_gridded2mali(self, measures_dataset, dst_scrip_file, parallel_executable, nProcs, mesh_filename, src_proj, variables=measures_vars, - hull_path=hull_path) + hull_path=hull_path, + extrap_varnames=['vx', 'vy', 'vErr']) # Diagnostic plot: show hull, MALI domain, and bounding boxes for # all source datasets that were interpolated. diff --git a/compass/landice/tests/antarctica/mesh_gen/mesh_gen.cfg b/compass/landice/tests/antarctica/mesh_gen/mesh_gen.cfg index 67847ecb87..a07ef1dfac 100644 --- a/compass/landice/tests/antarctica/mesh_gen/mesh_gen.cfg +++ b/compass/landice/tests/antarctica/mesh_gen/mesh_gen.cfg @@ -60,7 +60,7 @@ bedmachine_filename = NSIDC-0756_BedMachineAntarctica_19700101-20191001_V04.1_ed # filename of the MEASURES ice velocity dataset # (default value is for Perlmutter) -measures_filename = antarctica_ice_velocity_450m_v2_edits_extrap.nc +measures_filename = antarctica_ice_velocity_450m_v2_edits.nc # projection of the source datasets, according to the dictionary keys # create_scrip_file_from_planar_rectangular_grid from MPAS_Tools diff --git a/docs/developers_guide/landice/framework.rst b/docs/developers_guide/landice/framework.rst index 712eed33ba..4fec2e2477 100644 --- a/docs/developers_guide/landice/framework.rst +++ b/docs/developers_guide/landice/framework.rst @@ -78,6 +78,17 @@ pass a pre-built ``hull_path`` (from :py:func:`compass.landice.mesh.build_dst_scrip_hull()`) to avoid recomputing the hull for each dataset. +Optionally, ``extrap_varnames`` may be passed to extrapolate a raw, +unextrapolated source dataset on the fly. The listed variables are +nearest-filled into their no-data (non-finite) cells on the source raster using +a fast Euclidean-distance transform, written to a temporary copy that is +remapped in place of the original. This lets a plain conservative remap have +valid data across the whole mesh footprint, so margins are free of +interpolation ramps and single-pixel fill artifacts, without requiring an +offline pre-extrapolated source file. When ``extrap_varnames`` is ``None`` (the +default), the source is remapped as-is (appropriate for datasets that are +already gap-free or intentionally left unfilled, e.g. bed topography). + :py:func:`compass.landice.mesh.build_dst_scrip_hull()` builds a buffered concave boundary from a destination SCRIP file and returns it as a ``matplotlib.path.Path``. The boundary is constructed by rasterising @@ -98,7 +109,9 @@ fall within the concave boundary of the destination SCRIP footprint (plus a configurable buffer, default 50 km). The function projects both SCRIP files into a planar stereographic coordinate system (determined by the ``domain`` argument, e.g. ``'gis-gimp'`` or ``'ais-bedmap2'``) and uses a -``matplotlib.path.Path`` point-in-polygon test for efficiency. The resulting +``matplotlib.path.Path`` point-in-polygon test for efficiency. If an optional +``valid_mask`` array is supplied, ``grid_imask`` is further intersected with it +so no-data source cells are also excluded. The resulting masked SCRIP file is passed to ``ESMF_RegridWeightGen`` in place of the full-domain source SCRIP. diff --git a/docs/users_guide/landice/test_groups/antarctica.rst b/docs/users_guide/landice/test_groups/antarctica.rst index 941fb9d3ac..75358747f0 100644 --- a/docs/users_guide/landice/test_groups/antarctica.rst +++ b/docs/users_guide/landice/test_groups/antarctica.rst @@ -82,7 +82,7 @@ the mesh generation options are adjusted through the config file. # filename of the MEASURES ice velocity dataset # (default value is for Perlmutter) - measures_filename = antarctica_ice_velocity_450m_v2_edits_extrap.nc + measures_filename = antarctica_ice_velocity_450m_v2_edits.nc # projection of the source datasets, according to the dictionary keys # create_scrip_file_from_planar_rectangular_grid from MPAS_Tools @@ -103,12 +103,16 @@ The test case performs interpolation of observational data from gridded datasets to the Antarctic mesh. This takes care of the peculiarities of the current gridded compilation dataset (antarctica_8km_2020_10_20.nc), as well as using conservative remapping directly from the high-resolution BedMachineAntarctica and MeASUReS -velocity datasets. There is a fairly heavy degree of pre-processing done to get -the BedMachine and MeASUReS datasets ready to be used here. The pre-processing -includes renaming variables, setting reasonable _FillValue and missing_value -attributes, extrapolating fields to avoid interpolation ramps at ice margins, -updating mask values, and raising the bed topography at Lake Vostok to ensure -a flat ice surface there. +velocity datasets. The MEaSUREs velocity dataset is used directly in its raw +(lightly edited) form: its no-data gaps are extrapolated on the fly on the +source raster (a fast nearest-value distance transform, see +``interp_gridded2mali()`` below), so no offline velocity extrapolation is +required. The BedMachine dataset still needs some pre-processing +to be ready for use here, including renaming variables, setting reasonable +_FillValue and missing_value attributes, extrapolating thickness to avoid +interpolation ramps at ice margins, updating mask values, and raising the bed +topography at Lake Vostok to ensure a flat ice surface there. Bed topography is +never extrapolated, so it retains its true values everywhere. Those data files and processing scripts currently live here on Perlmutter: ``/global/cfs/cdirs/fanssie/standard_datasets/AIS_datasets``. @@ -127,8 +131,11 @@ case. Before running ``ESMF_RegridWeightGen``, the interpolation step automatically masks the source SCRIP file so that only cells overlapping a tight concave boundary around the destination mesh (plus a 50 km buffer) are active. The -boundary follows the actual domain shape rather than a simple convex hull. -This avoids unnecessary weight computation for the large portions of the +boundary follows the actual domain shape rather than a simple convex hull. For +the MEaSUREs velocity dataset, no-data gaps are additionally extrapolated on the +source raster (nearest-value distance transform) before a plain conservative +remap, so margins are free of interpolation ramps and fill artifacts. This +avoids unnecessary weight computation for the large portions of the BedMachine Antarctica domain that lie outside the target mesh and substantially reduces ESMF weight-generation time.