From e2fcfb45006ab23d38850a26229c831adf5aaeeb Mon Sep 17 00:00:00 2001 From: Trevor Hillebrand Date: Tue, 18 Aug 2026 19:31:12 -0700 Subject: [PATCH 1/3] Mask raw MEaSUReS velocity to MALI hull instead of pre-extrapolating Add optional valid-data source masking to interp_gridded2mali via a new valid_mask_varnames argument. When provided, no-data source cells (where any listed variable is non-finite) are excluded from the ESMF remapping weights in addition to the destination-hull footprint mask, and the weight-gen call gains --norm_type fracarea and --extrap_method neareststod so partially covered destination cells are renormalized and any unmapped cells are filled from the nearest valid source. This lets the AIS mesh_gen case interpolate directly from the raw (unextrapolated) MEaSUReS velocity file, excluding no-data cells from the weights rather than relying on a slow offline extrapolation. BedMachine is unchanged (bed stays real, full-extent). Extrapolated-source cases are a no-op because their all-finite fields yield an all-ones mask that leaves the gated flags off. Update AIS user guide and framework developer docs accordingly. --- compass/landice/mesh.py | 71 ++++++++++++++++++- .../tests/antarctica/mesh_gen/mesh_gen.cfg | 2 +- docs/developers_guide/landice/framework.rst | 14 +++- .../landice/test_groups/antarctica.rst | 23 +++--- 4 files changed, 98 insertions(+), 12 deletions(-) diff --git a/compass/landice/mesh.py b/compass/landice/mesh.py index 2a2b86fd38..f8a9228ae3 100644 --- a/compass/landice/mesh.py +++ b/compass/landice/mesh.py @@ -1332,6 +1332,7 @@ def add_grid_imask_from_dst_scrip_hull(source_scrip, dest_scrip, source_crs='EPSG:4326', mesh_crs=None, hull_path=None, + valid_mask=None, logger=None): """ Create a new source SCRIP file with grid_imask set to 1 only for cells @@ -1373,6 +1374,11 @@ def add_grid_imask_from_dst_scrip_hull(source_scrip, dest_scrip, recomputed, which is useful when masking multiple source datasets against the same destination mesh. + valid_mask : numpy.ndarray or None, optional + 1D 0/1 (or boolean) mask over the source ``grid_size`` marking cells + with valid data. When provided, ``grid_imask`` is the intersection of + the hull footprint and this mask, so ESMF excludes no-data cells. + logger : logging.Logger, optional Logger for status messages; falls back to print if None. """ @@ -1417,6 +1423,13 @@ def _projected_centers(ds): src_pts = np.column_stack([xc, yc]) inside = hull_path.contains_points(src_pts).astype(np.int32) + if valid_mask is not None: + valid_mask = np.asarray(valid_mask).astype(np.int32).ravel() + if valid_mask.size != inside.size: + raise ValueError( + f'valid_mask size {valid_mask.size} does not match ' + f'source grid_size {inside.size}') + inside = inside * valid_mask _log(f'active source cells after masking: ' f'{inside.sum()} / {inside.size}') @@ -1564,9 +1577,43 @@ def _src_extent(filepath): _log(f'Warning: could not save mesh boundary plot: {exc}') +def _compute_source_valid_mask(source_file, varnames): + """ + Build a 1D valid-data mask (C-order, matching SCRIP ``grid_size``) that is + True only where every listed source variable has finite data. + + Parameters + ---------- + source_file : str + Path to the source gridded dataset. + + varnames : list of str + Source variable names whose combined finite footprint defines validity + (a cell is valid only where all listed variables are finite). + + Returns + ------- + numpy.ndarray + Boolean mask flattened in C-order over the source grid. + """ + # xarray decodes _FillValue to NaN, so isfinite drops no-data without + # discarding genuine zeros (e.g. stagnant ice). + with xarray.open_dataset(source_file) as ds: + valid = None + for name in varnames: + arr = ds[name].squeeze() + if arr.ndim != 2: + raise ValueError( + f"Expected a 2D field for '{name}' in {source_file}, " + f"got shape {arr.shape}") + finite = np.isfinite(arr.values) + valid = finite if valid is None else (valid & finite) + return valid.ravel() + + def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, nProcs, dest_file, proj, variables="all", - hull_path=None): + hull_path=None, valid_mask_varnames=None): """ Interpolate gridded dataset (e.g. MEASURES, BedMachine) onto a MALI mesh @@ -1600,6 +1647,12 @@ 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. + valid_mask_varnames : list of str or None, optional + Source variable names whose finite footprint defines valid data. When + provided, no-data source cells are excluded from the remapping weights + (and ESMF renormalizes/fills accordingly), removing the need for a + pre-extrapolated source raster. When None, all in-hull cells are used. + Returns ------- masked_source_scrip : str @@ -1637,6 +1690,12 @@ 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' + + valid_mask = None + if valid_mask_varnames is not None: + valid_mask = _compute_source_valid_mask(source_file, + valid_mask_varnames) + logger.info('masking source SCRIP to destination mesh footprint') add_grid_imask_from_dst_scrip_hull( source_scrip=source_scrip, @@ -1644,6 +1703,7 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, masked_source_scrip=masked_source_scrip, domain=proj, hull_path=hull_path, + valid_mask=valid_mask, logger=logger) # Generate remapping weights @@ -1659,6 +1719,11 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, "--dst_regional", "--src_regional", '--ignore_unmapped'] + # With no-data cells excluded, renormalize partially covered dst cells and + # fill any left unmapped from nearest source. + if valid_mask is not None and not valid_mask.all(): + args += ['--norm_type', 'fracarea', + '--extrap_method', 'neareststod'] check_call(args, logger=logger) # Perform actual interpolation using the weights @@ -1963,11 +2028,13 @@ def run_optional_interpolation( measures_vars = ['observedSurfaceVelocityX', 'observedSurfaceVelocityY', 'observedSurfaceVelocityUncertainty'] + # velocity vars are always named vx/vy 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, + valid_mask_varnames=['vx', 'vy']) # 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..29f7d50e2e 100644 --- a/docs/developers_guide/landice/framework.rst +++ b/docs/developers_guide/landice/framework.rst @@ -78,6 +78,16 @@ 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, ``valid_mask_varnames`` may be passed to also exclude no-data source +cells (those where any listed variable is not finite) from the remapping +weights. This lets a raw, unextrapolated source dataset be used directly instead +of a pre-extrapolated one: excluded cells cannot pollute ice-margin values, and +when any cells are excluded the ESMF call adds ``--norm_type fracarea`` (to +renormalize partially covered destination cells) and ``--extrap_method +neareststod`` (to fill any unmapped destination cells from the nearest valid +source). When ``valid_mask_varnames`` is ``None`` (the default), all +in-footprint cells are used, matching a pre-extrapolated source. + :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 +108,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..5eb02f16c6 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,15 @@ 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: no-data cells are excluded from the remapping weights via +source masking (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,7 +130,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. +boundary follows the actual domain shape rather than a simple convex hull. For +the MEaSUREs velocity dataset, no-data cells inside that boundary are also +excluded, and ``ESMF_RegridWeightGen`` is run with ``--norm_type fracarea`` and +``--extrap_method neareststod`` so partially covered destination cells are +renormalized and any left unmapped are filled from the nearest valid source. 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. From 11bc96e2566c5c93773661b8017e70027f94f5a4 Mon Sep 17 00:00:00 2001 From: Trevor Hillebrand Date: Thu, 20 Aug 2026 11:32:00 -0700 Subject: [PATCH 2/3] Fix ESMF conservative+extrapolation error in velocity interp ESMF forbids --extrap_method with conservative regridding, so the raw MEaSUReS velocity interpolation failed. Drop the invalid flag and instead fill destination cells left unmapped by the conservative pass with nearest active-source weights computed directly via a KD-tree and merged into the conservative weight file. ESMF_RegridWeightGen is still called only once, preserving the mesh_gen compute/memory profile. --- compass/landice/mesh.py | 122 ++++++++++++++++-- docs/developers_guide/landice/framework.rst | 15 ++- .../landice/test_groups/antarctica.rst | 11 +- 3 files changed, 130 insertions(+), 18 deletions(-) diff --git a/compass/landice/mesh.py b/compass/landice/mesh.py index f8a9228ae3..44ffca3f02 100644 --- a/compass/landice/mesh.py +++ b/compass/landice/mesh.py @@ -23,6 +23,7 @@ from scipy import ndimage from scipy.interpolate import interpn from scipy.ndimage import binary_dilation, distance_transform_edt +from scipy.spatial import cKDTree def mpas_flood_fill(seed_mask, grow_mask, cellsOnCell, nEdgesOnCell, @@ -1611,6 +1612,99 @@ def _compute_source_valid_mask(source_file, varnames): return valid.ravel() +def _scrip_centers_xyz(ds): + """ + Return SCRIP cell centers as unit-sphere Cartesian coordinates so nearest + neighbors are correct even over the poles (unlike raw lat/lon distance). + """ + lat = np.asarray(ds['grid_center_lat'].values, dtype=float).ravel() + lon = np.asarray(ds['grid_center_lon'].values, dtype=float).ravel() + units = ds['grid_center_lat'].attrs.get('units', 'degrees').lower() + if units.startswith('deg'): + lat = np.deg2rad(lat) + lon = np.deg2rad(lon) + cos_lat = np.cos(lat) + return np.column_stack([cos_lat * np.cos(lon), + cos_lat * np.sin(lon), + np.sin(lat)]) + + +def _fill_unmapped_weights(conserve_weights, masked_source_scrip, mali_scrip, + logger=None): + """ + Fill destination cells left unmapped by conservative regridding with + nearest active-source weights, merged into ``conserve_weights`` in place. + + ESMF forbids extrapolation with conservative methods, so any destination + cell with no valid source overlap is simply absent from the conservative + weight file. Rather than run a second (expensive) ``ESMF_RegridWeightGen`` + pass, the nearest active source cell for each unmapped destination cell is + found directly with a KD-tree and appended as a unit-weight row, producing + a single weight file usable by ``interpolate_to_mpasli_grid`` (which uses + ``S``/``col``/``row`` directly). + + Parameters + ---------- + conserve_weights : str + Path to the conservative ESMF weight file; overwritten with the merged + weights. + + masked_source_scrip : str + SCRIP file for the source grid whose ``grid_imask`` marks the active + (in-hull, valid-data) cells that may be used as nearest neighbors. + + mali_scrip : str + SCRIP file for the destination MALI mesh. + + logger : logging.Logger, optional + Logger for status messages; falls back to print if None. + """ + def _log(msg): + if logger is not None: + logger.info(msg) + else: + print(msg) + + with xarray.open_dataset(conserve_weights) as ds_w: + ds_w.load() + + n_b = int(ds_w.sizes['n_b']) + mapped = np.unique(ds_w['row'].values) # 1-based dst indices with weights + all_dst = np.arange(1, n_b + 1, dtype=ds_w['row'].dtype) + unmapped = all_dst[~np.isin(all_dst, mapped)] + + _log(f'filling {unmapped.size} unmapped destination cells from ' + f'nearest source') + if unmapped.size == 0: + return + + with xarray.open_dataset(masked_source_scrip) as ds_s: + src_xyz = _scrip_centers_xyz(ds_s) + active = np.flatnonzero( + np.asarray(ds_s['grid_imask'].values).ravel().astype(bool)) + with xarray.open_dataset(mali_scrip) as ds_d: + dst_xyz = _scrip_centers_xyz(ds_d) + + # Nearest active source cell for each unmapped destination cell. + tree = cKDTree(src_xyz[active]) + _, nn = tree.query(dst_xyz[unmapped - 1]) + nearest_src = active[nn] # 0-based source indices + + # ESMF stores col/row as 1-based; match that so the appended rows are + # applied identically to the conservative ones. + col = np.concatenate([ds_w['col'].values, + (nearest_src + 1).astype(ds_w['col'].dtype)]) + row = np.concatenate([ds_w['row'].values, unmapped]) + s = np.concatenate([ds_w['S'].values, + np.ones(unmapped.size, dtype=ds_w['S'].dtype)]) + + ds_out = ds_w.drop_vars(['S', 'col', 'row']) + ds_out['col'] = ('n_s', col) + ds_out['row'] = ('n_s', row) + ds_out['S'] = ('n_s', s) + ds_out.to_netcdf(conserve_weights) + + def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, nProcs, dest_file, proj, variables="all", hull_path=None, valid_mask_varnames=None): @@ -1649,9 +1743,11 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, valid_mask_varnames : list of str or None, optional Source variable names whose finite footprint defines valid data. When - provided, no-data source cells are excluded from the remapping weights - (and ESMF renormalizes/fills accordingly), removing the need for a - pre-extrapolated source raster. When None, all in-hull cells are used. + provided, no-data source cells are excluded from the remapping weights; + partially covered destination cells are renormalized (``fracarea``) and + any left unmapped are filled from a merged nearest-source weight file, + removing the need for a pre-extrapolated source raster. When None, all + in-hull cells are used. Returns ------- @@ -1708,6 +1804,9 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, # Generate remapping weights logger.info('generating gridded dataset -> MPAS weights') + # No-data source cells are excluded from the weights when a valid_mask is + # provided, so some destination cells may be left unmapped. + fill_unmapped = valid_mask is not None and not valid_mask.all() args = parallel_executable.split() + [ '-n', nProcs, 'ESMF_RegridWeightGen', @@ -1719,13 +1818,20 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, "--dst_regional", "--src_regional", '--ignore_unmapped'] - # With no-data cells excluded, renormalize partially covered dst cells and - # fill any left unmapped from nearest source. - if valid_mask is not None and not valid_mask.all(): - args += ['--norm_type', 'fracarea', - '--extrap_method', 'neareststod'] + if fill_unmapped: + # Renormalize partially covered dst cells by their covered fraction + # instead of diluting them with the excluded no-data area. + args += ['--norm_type', 'fracarea'] check_call(args, logger=logger) + if fill_unmapped: + # ESMF cannot extrapolate with conservative methods. Rather than a + # second (expensive) weight-gen pass, fill any dst cells left unmapped + # (no valid source overlap) with nearest active-source weights computed + # directly and merged into the conservative weights. + _fill_unmapped_weights(weights_filename, masked_source_scrip, + mali_scrip, logger=logger) + # Perform actual interpolation using the weights logger.info('calling interpolate_to_mpasli_grid') args = ['interpolate_to_mpasli_grid', diff --git a/docs/developers_guide/landice/framework.rst b/docs/developers_guide/landice/framework.rst index 29f7d50e2e..f3733e2e48 100644 --- a/docs/developers_guide/landice/framework.rst +++ b/docs/developers_guide/landice/framework.rst @@ -81,12 +81,15 @@ hull for each dataset. Optionally, ``valid_mask_varnames`` may be passed to also exclude no-data source cells (those where any listed variable is not finite) from the remapping weights. This lets a raw, unextrapolated source dataset be used directly instead -of a pre-extrapolated one: excluded cells cannot pollute ice-margin values, and -when any cells are excluded the ESMF call adds ``--norm_type fracarea`` (to -renormalize partially covered destination cells) and ``--extrap_method -neareststod`` (to fill any unmapped destination cells from the nearest valid -source). When ``valid_mask_varnames`` is ``None`` (the default), all -in-footprint cells are used, matching a pre-extrapolated source. +of a pre-extrapolated one: excluded cells cannot pollute ice-margin values. When +any cells are excluded the conservative weight-gen call adds ``--norm_type +fracarea`` (to renormalize partially covered destination cells). Because ESMF +cannot extrapolate with conservative methods, any destination cells left +unmapped (no valid source overlap) are then filled with nearest active-source +weights computed directly (via a KD-tree, so ``ESMF_RegridWeightGen`` is still +called only once) and merged into the conservative weights. When +``valid_mask_varnames`` is ``None`` (the default), all in-footprint cells are +used, matching a pre-extrapolated source. :py:func:`compass.landice.mesh.build_dst_scrip_hull()` builds a buffered concave boundary from a destination SCRIP file and returns it as a diff --git a/docs/users_guide/landice/test_groups/antarctica.rst b/docs/users_guide/landice/test_groups/antarctica.rst index 5eb02f16c6..25d3d728af 100644 --- a/docs/users_guide/landice/test_groups/antarctica.rst +++ b/docs/users_guide/landice/test_groups/antarctica.rst @@ -132,10 +132,13 @@ 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. For the MEaSUREs velocity dataset, no-data cells inside that boundary are also -excluded, and ``ESMF_RegridWeightGen`` is run with ``--norm_type fracarea`` and -``--extrap_method neareststod`` so partially covered destination cells are -renormalized and any left unmapped are filled from the nearest valid source. -This avoids unnecessary weight computation for the large portions of the +excluded, and the conservative ``ESMF_RegridWeightGen`` call uses ``--norm_type +fracarea`` so partially covered destination cells are renormalized. Since ESMF +cannot extrapolate with conservative methods, any destination cells left +unmapped are filled with nearest active-source weights computed directly and +merged into the conservative weights, so ``ESMF_RegridWeightGen`` is still run +only once. 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. From e61a2221d7827ececbc1b84ed664adfc606a231d Mon Sep 17 00:00:00 2001 From: Trevor Hillebrand Date: Mon, 24 Aug 2026 10:06:47 -0700 Subject: [PATCH 3/3] Extrapolate MEaSUReS velocity on source raster before remapping Replace the weight-level no-data exclusion + nearest fill with an on-the-fly nearest-value extrapolation of the raw velocity/error fields on the source raster (fast ndimage distance transform), then a plain conservative remap. This matches the smoothness of the old offline-extrapolated workflow and eliminates the single-pixel margin artifacts (spurious high speeds at the ice edge) produced by mesh-level nearest fill and fracarea normalization. Remove the now-unused valid_mask/fracarea/_fill_unmapped_weights machinery and promote _nearest_fill_from_valid to a shared module-level helper. --- compass/landice/mesh.py | 304 +++++++----------- docs/developers_guide/landice/framework.rst | 22 +- .../landice/test_groups/antarctica.rst | 19 +- 3 files changed, 131 insertions(+), 214 deletions(-) diff --git a/compass/landice/mesh.py b/compass/landice/mesh.py index 44ffca3f02..745c66953c 100644 --- a/compass/landice/mesh.py +++ b/compass/landice/mesh.py @@ -23,7 +23,6 @@ from scipy import ndimage from scipy.interpolate import interpn from scipy.ndimage import binary_dilation, distance_transform_edt -from scipy.spatial import cKDTree def mpas_flood_fill(seed_mask, grow_mask, cellsOnCell, nEdgesOnCell, @@ -977,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 @@ -997,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" @@ -1333,7 +1332,6 @@ def add_grid_imask_from_dst_scrip_hull(source_scrip, dest_scrip, source_crs='EPSG:4326', mesh_crs=None, hull_path=None, - valid_mask=None, logger=None): """ Create a new source SCRIP file with grid_imask set to 1 only for cells @@ -1375,11 +1373,6 @@ def add_grid_imask_from_dst_scrip_hull(source_scrip, dest_scrip, recomputed, which is useful when masking multiple source datasets against the same destination mesh. - valid_mask : numpy.ndarray or None, optional - 1D 0/1 (or boolean) mask over the source ``grid_size`` marking cells - with valid data. When provided, ``grid_imask`` is the intersection of - the hull footprint and this mask, so ESMF excludes no-data cells. - logger : logging.Logger, optional Logger for status messages; falls back to print if None. """ @@ -1424,13 +1417,6 @@ def _projected_centers(ds): src_pts = np.column_stack([xc, yc]) inside = hull_path.contains_points(src_pts).astype(np.int32) - if valid_mask is not None: - valid_mask = np.asarray(valid_mask).astype(np.int32).ravel() - if valid_mask.size != inside.size: - raise ValueError( - f'valid_mask size {valid_mask.size} does not match ' - f'source grid_size {inside.size}') - inside = inside * valid_mask _log(f'active source cells after masking: ' f'{inside.sum()} / {inside.size}') @@ -1578,10 +1564,17 @@ def _src_extent(filepath): _log(f'Warning: could not save mesh boundary plot: {exc}') -def _compute_source_valid_mask(source_file, varnames): +def _write_extrapolated_source(source_file, varnames, logger=None): """ - Build a 1D valid-data mask (C-order, matching SCRIP ``grid_size``) that is - True only where every listed source variable has finite data. + 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 ---------- @@ -1589,75 +1582,16 @@ def _compute_source_valid_mask(source_file, varnames): Path to the source gridded dataset. varnames : list of str - Source variable names whose combined finite footprint defines validity - (a cell is valid only where all listed variables are finite). - - Returns - ------- - numpy.ndarray - Boolean mask flattened in C-order over the source grid. - """ - # xarray decodes _FillValue to NaN, so isfinite drops no-data without - # discarding genuine zeros (e.g. stagnant ice). - with xarray.open_dataset(source_file) as ds: - valid = None - for name in varnames: - arr = ds[name].squeeze() - if arr.ndim != 2: - raise ValueError( - f"Expected a 2D field for '{name}' in {source_file}, " - f"got shape {arr.shape}") - finite = np.isfinite(arr.values) - valid = finite if valid is None else (valid & finite) - return valid.ravel() - - -def _scrip_centers_xyz(ds): - """ - Return SCRIP cell centers as unit-sphere Cartesian coordinates so nearest - neighbors are correct even over the poles (unlike raw lat/lon distance). - """ - lat = np.asarray(ds['grid_center_lat'].values, dtype=float).ravel() - lon = np.asarray(ds['grid_center_lon'].values, dtype=float).ravel() - units = ds['grid_center_lat'].attrs.get('units', 'degrees').lower() - if units.startswith('deg'): - lat = np.deg2rad(lat) - lon = np.deg2rad(lon) - cos_lat = np.cos(lat) - return np.column_stack([cos_lat * np.cos(lon), - cos_lat * np.sin(lon), - np.sin(lat)]) - - -def _fill_unmapped_weights(conserve_weights, masked_source_scrip, mali_scrip, - logger=None): - """ - Fill destination cells left unmapped by conservative regridding with - nearest active-source weights, merged into ``conserve_weights`` in place. - - ESMF forbids extrapolation with conservative methods, so any destination - cell with no valid source overlap is simply absent from the conservative - weight file. Rather than run a second (expensive) ``ESMF_RegridWeightGen`` - pass, the nearest active source cell for each unmapped destination cell is - found directly with a KD-tree and appended as a unit-weight row, producing - a single weight file usable by ``interpolate_to_mpasli_grid`` (which uses - ``S``/``col``/``row`` directly). - - Parameters - ---------- - conserve_weights : str - Path to the conservative ESMF weight file; overwritten with the merged - weights. - - masked_source_scrip : str - SCRIP file for the source grid whose ``grid_imask`` marks the active - (in-hull, valid-data) cells that may be used as nearest neighbors. - - mali_scrip : str - SCRIP file for the destination MALI mesh. + 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: @@ -1665,49 +1599,41 @@ def _log(msg): else: print(msg) - with xarray.open_dataset(conserve_weights) as ds_w: - ds_w.load() - - n_b = int(ds_w.sizes['n_b']) - mapped = np.unique(ds_w['row'].values) # 1-based dst indices with weights - all_dst = np.arange(1, n_b + 1, dtype=ds_w['row'].dtype) - unmapped = all_dst[~np.isin(all_dst, mapped)] + out_file = f'extrap_{os.path.basename(source_file)}' - _log(f'filling {unmapped.size} unmapped destination cells from ' - f'nearest source') - if unmapped.size == 0: - return + 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)) - with xarray.open_dataset(masked_source_scrip) as ds_s: - src_xyz = _scrip_centers_xyz(ds_s) - active = np.flatnonzero( - np.asarray(ds_s['grid_imask'].values).ravel().astype(bool)) - with xarray.open_dataset(mali_scrip) as ds_d: - dst_xyz = _scrip_centers_xyz(ds_d) - - # Nearest active source cell for each unmapped destination cell. - tree = cKDTree(src_xyz[active]) - _, nn = tree.query(dst_xyz[unmapped - 1]) - nearest_src = active[nn] # 0-based source indices - - # ESMF stores col/row as 1-based; match that so the appended rows are - # applied identically to the conservative ones. - col = np.concatenate([ds_w['col'].values, - (nearest_src + 1).astype(ds_w['col'].dtype)]) - row = np.concatenate([ds_w['row'].values, unmapped]) - s = np.concatenate([ds_w['S'].values, - np.ones(unmapped.size, dtype=ds_w['S'].dtype)]) - - ds_out = ds_w.drop_vars(['S', 'col', 'row']) - ds_out['col'] = ('n_s', col) - ds_out['row'] = ('n_s', row) - ds_out['S'] = ('n_s', s) - ds_out.to_netcdf(conserve_weights) + 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, valid_mask_varnames=None): + hull_path=None, extrap_varnames=None): """ Interpolate gridded dataset (e.g. MEASURES, BedMachine) onto a MALI mesh @@ -1741,13 +1667,13 @@ 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. - valid_mask_varnames : list of str or None, optional - Source variable names whose finite footprint defines valid data. When - provided, no-data source cells are excluded from the remapping weights; - partially covered destination cells are renormalized (``fracarea``) and - any left unmapped are filled from a merged nearest-source weight file, - removing the need for a pre-extrapolated source raster. When None, all - in-hull cells are used. + 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 ------- @@ -1757,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: @@ -1776,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'] @@ -1787,11 +1725,6 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, stem = os.path.splitext(source_scrip)[0] # strips .nc masked_source_scrip = f'{stem}_masked.nc' - valid_mask = None - if valid_mask_varnames is not None: - valid_mask = _compute_source_valid_mask(source_file, - valid_mask_varnames) - logger.info('masking source SCRIP to destination mesh footprint') add_grid_imask_from_dst_scrip_hull( source_scrip=source_scrip, @@ -1799,14 +1732,10 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, masked_source_scrip=masked_source_scrip, domain=proj, hull_path=hull_path, - valid_mask=valid_mask, logger=logger) # Generate remapping weights logger.info('generating gridded dataset -> MPAS weights') - # No-data source cells are excluded from the weights when a valid_mask is - # provided, so some destination cells may be left unmapped. - fill_unmapped = valid_mask is not None and not valid_mask.all() args = parallel_executable.split() + [ '-n', nProcs, 'ESMF_RegridWeightGen', @@ -1818,24 +1747,12 @@ def interp_gridded2mali(self, source_file, mali_scrip, parallel_executable, "--dst_regional", "--src_regional", '--ignore_unmapped'] - if fill_unmapped: - # Renormalize partially covered dst cells by their covered fraction - # instead of diluting them with the excluded no-data area. - args += ['--norm_type', 'fracarea'] check_call(args, logger=logger) - if fill_unmapped: - # ESMF cannot extrapolate with conservative methods. Rather than a - # second (expensive) weight-gen pass, fill any dst cells left unmapped - # (no valid source overlap) with nearest active-source weights computed - # directly and merged into the conservative weights. - _fill_unmapped_weights(weights_filename, masked_source_scrip, - mali_scrip, logger=logger) - # 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, @@ -1843,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 @@ -2134,13 +2055,14 @@ def run_optional_interpolation( measures_vars = ['observedSurfaceVelocityX', 'observedSurfaceVelocityY', 'observedSurfaceVelocityUncertainty'] - # velocity vars are always named vx/vy + # 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, - valid_mask_varnames=['vx', 'vy']) + extrap_varnames=['vx', 'vy', 'vErr']) # Diagnostic plot: show hull, MALI domain, and bounding boxes for # all source datasets that were interpolated. diff --git a/docs/developers_guide/landice/framework.rst b/docs/developers_guide/landice/framework.rst index f3733e2e48..4fec2e2477 100644 --- a/docs/developers_guide/landice/framework.rst +++ b/docs/developers_guide/landice/framework.rst @@ -78,18 +78,16 @@ 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, ``valid_mask_varnames`` may be passed to also exclude no-data source -cells (those where any listed variable is not finite) from the remapping -weights. This lets a raw, unextrapolated source dataset be used directly instead -of a pre-extrapolated one: excluded cells cannot pollute ice-margin values. When -any cells are excluded the conservative weight-gen call adds ``--norm_type -fracarea`` (to renormalize partially covered destination cells). Because ESMF -cannot extrapolate with conservative methods, any destination cells left -unmapped (no valid source overlap) are then filled with nearest active-source -weights computed directly (via a KD-tree, so ``ESMF_RegridWeightGen`` is still -called only once) and merged into the conservative weights. When -``valid_mask_varnames`` is ``None`` (the default), all in-footprint cells are -used, matching a pre-extrapolated source. +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 diff --git a/docs/users_guide/landice/test_groups/antarctica.rst b/docs/users_guide/landice/test_groups/antarctica.rst index 25d3d728af..75358747f0 100644 --- a/docs/users_guide/landice/test_groups/antarctica.rst +++ b/docs/users_guide/landice/test_groups/antarctica.rst @@ -104,9 +104,10 @@ to the Antarctic mesh. This takes care of the peculiarities of the current gridd compilation dataset (antarctica_8km_2020_10_20.nc), as well as using conservative remapping directly from the high-resolution BedMachineAntarctica and MeASUReS velocity datasets. The MEaSUREs velocity dataset is used directly in its raw -(lightly edited) form: no-data cells are excluded from the remapping weights via -source masking (see ``interp_gridded2mali()`` below), so no offline velocity -extrapolation is required. The BedMachine dataset still needs some pre-processing +(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 @@ -131,14 +132,10 @@ 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. For -the MEaSUREs velocity dataset, no-data cells inside that boundary are also -excluded, and the conservative ``ESMF_RegridWeightGen`` call uses ``--norm_type -fracarea`` so partially covered destination cells are renormalized. Since ESMF -cannot extrapolate with conservative methods, any destination cells left -unmapped are filled with nearest active-source weights computed directly and -merged into the conservative weights, so ``ESMF_RegridWeightGen`` is still run -only once. This avoids unnecessary weight computation for the large portions of -the +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.