Merge pull request #943 from trhille/make_major_data_interpolation_optional

Add optional data interpolation for all meshes.

Add data interpolation option for all meshes. Interpolate data from high-resolution BedMachine and MEaSUREs data sets if interpolate_data = True and data paths are provided in .cfg file. Otherwise interpolation will be skipped. To preserve previous behavior, interpolation will be performed by default for antarctica and greenland, but not for regional meshes. This includes a significant refactor to reduce redundant code between all the different mesh_gen cases.

Author: matthewhoffman

compass/landice/mesh.py

@@ -2,6 +2,7 @@
 import re
 import sys
 import time
+import uuid
 from shutil import copyfile
 
 import jigsawpy
@@ -14,6 +15,7 @@
 from mpas_tools.mesh.conversion import convert, cull
 from mpas_tools.mesh.creation import build_planar_mesh
 from mpas_tools.mesh.creation.sort_mesh import sort_mesh
+from mpas_tools.scrip.from_mpas import scrip_from_mpas
 from netCDF4 import Dataset
 from scipy.interpolate import NearestNDInterpolator, interpn
 
@@ -636,13 +638,10 @@ def build_cell_width(self, section_name, gridded_dataset,
 
     f.close()
 
-    # Get bounds defined by user, or use bound of gridded dataset
-    bnds = [np.min(x1), np.max(x1), np.min(y1), np.max(y1)]
-    bnds_options = ['x_min', 'x_max', 'y_min', 'y_max']
-    for index, option in enumerate(bnds_options):
-        bnd = section.get(option)
-        if bnd != 'None':
-            bnds[index] = float(bnd)
+    # Get bounds defined by user, or use bounds from the gridded dataset.
+    bnds = get_mesh_config_bounding_box(
+        section,
+        default_bounds=[np.min(x1), np.max(x1), np.min(y1), np.max(y1)])
 
     geom_points, geom_edges = set_rectangular_geom_points_and_edges(*bnds)
 
@@ -1191,3 +1190,267 @@ def clean_up_after_interp(fname):
         data.variables['observedSurfaceVelocityUncertainty'][:] == 0.0)
     data.variables['observedSurfaceVelocityUncertainty'][0, mask[0, :]] = 1.0
     data.close()
+
+
+def get_optional_interp_datasets(section, logger):
+    """
+    Determine whether optional interpolation inputs are configured.
+
+    Parameters
+    ----------
+    section : configparser.SectionProxy
+        Config section containing optional interpolation options
+
+    logger : logging.Logger
+        Logger for status messages
+
+    Returns
+    -------
+    bedmachine_dataset : str or None
+        Path to BedMachine dataset if configured, otherwise ``None``
+
+    measures_dataset : str or None
+        Path to MEaSUREs dataset if configured, otherwise ``None``
+    """
+
+    def _specified(value):
+        return value is not None and str(value).strip().lower() not in [
+            '', 'none']
+
+    data_path = section.get('data_path', fallback=None)
+    bedmachine_filename = section.get('bedmachine_filename', fallback=None)
+    measures_filename = section.get('measures_filename', fallback=None)
+
+    use_bedmachine_interp = _specified(data_path) and \
+        _specified(bedmachine_filename)
+    use_measures_interp = _specified(data_path) and \
+        _specified(measures_filename)
+
+    if use_bedmachine_interp:
+        bedmachine_dataset = os.path.join(data_path, bedmachine_filename)
+    else:
+        bedmachine_dataset = None
+        logger.info('Skipping BedMachine interpolation because '
+                    '`data_path` and/or `bedmachine_filename` are '
+                    'not specified in config.')
+
+    if use_measures_interp:
+        measures_dataset = os.path.join(data_path, measures_filename)
+    else:
+        measures_dataset = None
+        logger.info('Skipping MEaSUREs interpolation because '
+                    '`data_path` and/or `measures_filename` are '
+                    'not specified in config.')
+
+    return bedmachine_dataset, measures_dataset
+
+
+def get_mesh_config_bounding_box(section, default_bounds=None):
+    """
+    Get bounding-box coordinates from a mesh config section.
+
+    Parameters
+    ----------
+    section : configparser.SectionProxy
+        Mesh config section containing ``x_min``, ``x_max``, ``y_min``,
+        and ``y_max``
+
+    default_bounds : list of float, optional
+        Default bounds in the form ``[x_min, x_max, y_min, y_max]`` to use
+        when config values are missing or set to ``None``
+
+    Returns
+    -------
+    bounding_box : list of float
+        Bounding box in the form ``[x_min, x_max, y_min, y_max]``
+    """
+
+    if default_bounds is None:
+        default_bounds = [None, None, None, None]
+
+    def _get_bound(option, default):
+        value = section.get(option, fallback=None)
+        if value is None or str(value).strip().lower() in ['', 'none']:
+            if default is None:
+                raise ValueError(
+                    f'Missing required config option `{option}` and no '
+                    'default was provided.')
+            return float(default)
+        return float(value)
+
+    return [
+        _get_bound('x_min', default_bounds[0]),
+        _get_bound('x_max', default_bounds[1]),
+        _get_bound('y_min', default_bounds[2]),
+        _get_bound('y_max', default_bounds[3])]
+
+
+def subset_gridded_dataset_to_bounds(
+        source_dataset, bounding_box, subset_tag, logger):
+    """
+    Subset a gridded source dataset to a bounding box.
+
+    Parameters
+    ----------
+    source_dataset : str
+        Path to source gridded dataset
+
+    bounding_box : list of float
+        Bounding box in the form ``[x_min, x_max, y_min, y_max]``
+
+    subset_tag : str
+        Tag to include in the subset filename
+
+    logger : logging.Logger
+        Logger for status messages
+
+    Returns
+    -------
+    subset_dataset : str
+        Path to subsetted gridded dataset written to the current directory
+    """
+
+    x_min, x_max, y_min, y_max = bounding_box
+    ds = xarray.open_dataset(source_dataset)
+
+    if 'x1' in ds and 'y1' in ds:
+        x_name = 'x1'
+        y_name = 'y1'
+    elif 'x' in ds and 'y' in ds:
+        x_name = 'x'
+        y_name = 'y'
+    else:
+        ds.close()
+        raise ValueError(
+            f'Could not find x/y coordinates in {source_dataset}. '
+            'Expected either x1/y1 or x/y.')
+
+    subset = ds.where(
+        (ds[x_name] >= x_min) & (ds[x_name] <= x_max) &
+        (ds[y_name] >= y_min) & (ds[y_name] <= y_max),
+        drop=True)
+
+    # Check for empty subset, handling possible mismatch
+    # between variable and dimension names
+    x_dim = x_name if x_name in subset.sizes else (
+        'x' if 'x' in subset.sizes else None)
+    y_dim = y_name if y_name in subset.sizes else (
+        'y' if 'y' in subset.sizes else None)
+    if x_dim is None or y_dim is None or subset.sizes[x_dim] == 0 or subset.sizes[y_dim] == 0:  # noqa
+        subset.close()
+        ds.close()
+        raise ValueError(
+            f'Bounding box {bounding_box} produced an empty subset for '
+            f'{source_dataset}. Dimension names in subset: '
+            f'{list(subset.sizes.keys())}')
+
+    base = os.path.splitext(os.path.basename(source_dataset))[0]
+    unique_id = uuid.uuid4().hex
+    subset_dataset = f'{base}_{subset_tag}_{unique_id}_subset.nc'
+    logger.info(f'Writing subset dataset: {subset_dataset}')
+    subset.to_netcdf(subset_dataset)
+
+    subset.close()
+    ds.close()
+    return subset_dataset
+
+
+def run_optional_interpolation(
+        self, mesh_filename, src_proj, parallel_executable, nProcs,
+        bedmachine_dataset=None, measures_dataset=None, subset_bounds=None):
+    """
+    Run optional interpolation from high-resolution BedMachine and MEaSUREs
+    datasets and perform some necessary cleanup. This can require many
+    more resources than the rest of the mesh generation process, so it is
+    usually desirable to skip this step when prototyping meshes. This step
+    is only run if `interpolate_data` is set to True in the config file
+    and the necessary dataset paths are provided.
+
+    Parameters
+    ----------
+    self : compass.step.Step
+        Step instance providing logger and context
+
+    mesh_filename : str
+        Destination MALI mesh file to interpolate to
+
+    src_proj : str
+        Source dataset projection for SCRIP generation
+
+    parallel_executable : str
+        Parallel launcher executable (e.g. ``srun``/``mpirun``)
+
+    nProcs : int or str
+        Number of processes for regridding weight generation
+
+    bedmachine_dataset : str, optional
+        BedMachine dataset path; if ``None`` this interpolation is skipped
+
+    measures_dataset : str, optional
+        MEaSUREs dataset path; if ``None`` this interpolation is skipped
+
+    subset_bounds : list of float, optional
+        Optional source-dataset subset bounds in the form
+        ``[x_min, x_max, y_min, y_max]``. If provided, BedMachine and
+        MEaSUREs datasets are subsetted before SCRIP generation and
+        interpolation.
+    """
+
+    logger = self.logger
+    do_optional_interp = bedmachine_dataset is not None or \
+        measures_dataset is not None
+    if not do_optional_interp:
+        return
+
+    if nProcs is None:
+        raise ValueError("nProcs must be provided as an int or str")
+    nProcs = str(nProcs)
+
+    subset_files = []
+
+    try:
+        if subset_bounds is not None:
+            if bedmachine_dataset is not None:
+                bedmachine_dataset = subset_gridded_dataset_to_bounds(
+                    bedmachine_dataset,
+                    subset_bounds,
+                    'bedmachine',
+                    logger)
+                subset_files.append(bedmachine_dataset)
+            if measures_dataset is not None:
+                measures_dataset = subset_gridded_dataset_to_bounds(
+                    measures_dataset,
+                    subset_bounds,
+                    'measures',
+                    logger)
+                subset_files.append(measures_dataset)
+
+        logger.info('creating scrip file for destination mesh')
+        mesh_base = os.path.splitext(mesh_filename)[0]
+        dst_scrip_file = f'{mesh_base}_scrip.nc'
+        scrip_from_mpas(mesh_filename, dst_scrip_file)
+
+        if bedmachine_dataset is not None:
+            interp_gridded2mali(self, bedmachine_dataset, dst_scrip_file,
+                                parallel_executable, nProcs,
+                                mesh_filename, src_proj, variables='all')
+
+        if measures_dataset is not None:
+            measures_vars = ['observedSurfaceVelocityX',
+                             'observedSurfaceVelocityY',
+                             'observedSurfaceVelocityUncertainty']
+            interp_gridded2mali(self, measures_dataset, dst_scrip_file,
+                                parallel_executable, nProcs,
+                                mesh_filename, src_proj,
+                                variables=measures_vars)
+
+        clean_up_after_interp(mesh_filename)
+    finally:
+        for subset_file in subset_files:
+            if os.path.exists(subset_file):
+                logger.info(f'Removing subset dataset: {subset_file}')
+                try:
+                    os.remove(subset_file)
+                except OSError as exc:
+                    logger.warning('Could not remove subset dataset '
+                                   f'{subset_file}: {exc}')