diff --git a/ci/recipe/meta.yaml b/ci/recipe/meta.yaml index 0a2f267d..0ec0a624 100644 --- a/ci/recipe/meta.yaml +++ b/ci/recipe/meta.yaml @@ -1,5 +1,5 @@ {% set name = "pyremap" %} -{% set version = "2.3.0" %} +{% set version = "2.4.0" %} {% set python_min = "3.10" %} package: diff --git a/docs/developer_guide/api.md b/docs/developer_guide/api.md index 01aa15fb..26bf95ca 100644 --- a/docs/developer_guide/api.md +++ b/docs/developer_guide/api.md @@ -13,6 +13,8 @@ This page provides an auto-generated summary of the pyremap API. :toctree: generated/ + get_corners_1d + get_corners_2d interp_extrap_corner interp_extrap_corners_2d diff --git a/docs/mesh_descriptors/lat_lon_2d_grid_descriptor.md b/docs/mesh_descriptors/lat_lon_2d_grid_descriptor.md index 84a0c071..f84c43fc 100644 --- a/docs/mesh_descriptors/lat_lon_2d_grid_descriptor.md +++ b/docs/mesh_descriptors/lat_lon_2d_grid_descriptor.md @@ -15,6 +15,24 @@ The `LatLon2DGridDescriptor` class is used for grids where latitude and longitud - `read`: Reads a 2D latitude-longitude grid from a file. - `to_scrip`: Converts the grid to a SCRIP file. +## Grid-Cell Corners +As with {py:class}`LatLonGridDescriptor `, +`read()` uses the CF `bounds` of the latitude and longitude variables to find +grid-cell corners when they are available. For 2D coordinates, the bounds +give the 4 vertices of each cell: +``` +double lat(y, x) ; + lat:units = "degrees_north" ; + lat:bounds = "lat_bnds" ; +double lat_bnds(y, x, nv) ; +``` +CF does not say which vertex comes first or which direction the 4 vertices are +traversed in, so pyremap works this out from the bounds themselves. Both +latitude and longitude must have bounds, and neighboring cells must share +vertices, since the grid is described by 2D arrays of corners. Otherwise, +corners are interpolated and extrapolated from cell centers and a warning is +raised. + ## Example ```python from pyremap import LatLon2DGridDescriptor diff --git a/docs/mesh_descriptors/lat_lon_grid_descriptor.md b/docs/mesh_descriptors/lat_lon_grid_descriptor.md index 541b9fb2..55c8f8bc 100644 --- a/docs/mesh_descriptors/lat_lon_grid_descriptor.md +++ b/docs/mesh_descriptors/lat_lon_grid_descriptor.md @@ -16,6 +16,24 @@ The `LatLonGridDescriptor` class is used to describe a regular latitude-longitud - `create`: Creates a latitude-longitude grid programmatically. - `to_scrip`: Converts the grid to a SCRIP file. +## Grid-Cell Corners +Remapping needs the corners of each grid cell, not just the cell centers. +When `read()` is used, corners come from the CF `bounds` attribute of the +latitude and longitude variables if it is present: +``` +double lat(lat) ; + lat:units = "degrees_north" ; + lat:bounds = "lat_bnds" ; +double lat_bnds(lat, nbnd) ; +``` +The bounds must describe contiguous cells (the upper edge of each cell is the +lower edge of the next), since a 1D lat/lon grid is described by 1D arrays of +corners. If the `bounds` attribute is missing, points to a variable that is +not in the file, has the wrong shape, or describes cells with gaps or overlaps +between them, corners are instead interpolated between cell centers and +extrapolated at the ends of the grid, and a warning is raised in all but the +first of these cases. + ## Example ```python from pyremap import LatLonGridDescriptor diff --git a/docs/mesh_descriptors/projection_grid_descriptor.md b/docs/mesh_descriptors/projection_grid_descriptor.md index 0e2908e2..b7f6b59e 100644 --- a/docs/mesh_descriptors/projection_grid_descriptor.md +++ b/docs/mesh_descriptors/projection_grid_descriptor.md @@ -15,6 +15,13 @@ The `ProjectionGridDescriptor` class describes grids defined by map projections. - `create`: Creates a projection grid programmatically. - `to_scrip`: Converts the grid to a SCRIP file. +## Grid-Cell Corners +`read()` uses the CF `bounds` of the `x` and `y` variables to find the corners +of each grid cell in projection space when they are available and describe +contiguous cells. Otherwise, corners are interpolated between cell centers +and extrapolated at the ends of the grid. Corners are transformed from +projection space to latitude and longitude by `to_scrip()`. + ## Example ```python from pyremap import ProjectionGridDescriptor diff --git a/pyremap/descriptor/__init__.py b/pyremap/descriptor/__init__.py index 7cbd416a..6aeac94c 100644 --- a/pyremap/descriptor/__init__.py +++ b/pyremap/descriptor/__init__.py @@ -36,6 +36,12 @@ from pyremap.descriptor.projection_grid_descriptor import ( ProjectionGridDescriptor as ProjectionGridDescriptor, ) +from pyremap.descriptor.utility import ( + get_corners_1d as get_corners_1d, +) +from pyremap.descriptor.utility import ( + get_corners_2d as get_corners_2d, +) from pyremap.descriptor.utility import ( interp_extrap_corner as interp_extrap_corner, ) diff --git a/pyremap/descriptor/lat_lon_2d_grid_descriptor.py b/pyremap/descriptor/lat_lon_2d_grid_descriptor.py index c051b0b4..4b719b6c 100644 --- a/pyremap/descriptor/lat_lon_2d_grid_descriptor.py +++ b/pyremap/descriptor/lat_lon_2d_grid_descriptor.py @@ -18,7 +18,7 @@ from pyremap.descriptor.utility import ( add_history, expand_scrip, - interp_extrap_corners_2d, + get_corners_2d, round_res, unwrap_corners, ) @@ -87,6 +87,11 @@ def read( """ Read the lat-lon grid from a file with the given lat/lon var names. + Grid-cell corners come from the CF ``bounds`` of the latitude and + longitude variables if they are available and neighboring cells share + vertices. Otherwise, corners are interpolated and extrapolated from + the cell centers. + Parameters ---------- filename : str, optional @@ -127,9 +132,10 @@ def read( else: descriptor.units = 'radians' - # interp/extrap corners - descriptor.lon_corner = interp_extrap_corners_2d(descriptor.lon) - descriptor.lat_corner = interp_extrap_corners_2d(descriptor.lat) + # use CF bounds if available, otherwise interp/extrap corners + descriptor.lat_corner, descriptor.lon_corner = get_corners_2d( + ds, lat_var_name, lon_var_name + ) descriptor._set_coords( lat_var_name, diff --git a/pyremap/descriptor/lat_lon_grid_descriptor.py b/pyremap/descriptor/lat_lon_grid_descriptor.py index 4c0a8d19..48094fed 100644 --- a/pyremap/descriptor/lat_lon_grid_descriptor.py +++ b/pyremap/descriptor/lat_lon_grid_descriptor.py @@ -18,7 +18,7 @@ from pyremap.descriptor.utility import ( add_history, expand_scrip, - interp_extrap_corner, + get_corners_1d, round_res, unwrap_corners, ) @@ -121,6 +121,11 @@ def read( """ Read the lat-lon grid from a file with the given lat/lon var names. + Grid-cell corners come from the CF ``bounds`` of the latitude and + longitude variables if they are available and describe contiguous + cells. Otherwise, corners are interpolated and extrapolated from the + cell centers. + Parameters ---------- filename : str, optional @@ -157,9 +162,9 @@ def read( else: descriptor.units = 'radians' - # interp/extrap corners - descriptor.lon_corner = interp_extrap_corner(descriptor.lon) - descriptor.lat_corner = interp_extrap_corner(descriptor.lat) + # use CF bounds if available, otherwise interp/extrap corners + descriptor.lon_corner = get_corners_1d(ds, lon_var_name) + descriptor.lat_corner = get_corners_1d(ds, lat_var_name) descriptor._set_coords( lat_var_name, diff --git a/pyremap/descriptor/projection_grid_descriptor.py b/pyremap/descriptor/projection_grid_descriptor.py index 9a755326..308b1a9d 100644 --- a/pyremap/descriptor/projection_grid_descriptor.py +++ b/pyremap/descriptor/projection_grid_descriptor.py @@ -19,6 +19,7 @@ from pyremap.descriptor.utility import ( add_history, expand_scrip, + get_corners_1d, interp_extrap_corner, unwrap_corners, ) @@ -98,7 +99,11 @@ def read( """ Given a grid file with x and y coordinates defining the axes of the logically rectangular grid, read in the x and y coordinates and - interpolate/extrapolate to locate corners. + locate the corners. + + Corners come from the CF ``bounds`` of the x and y variables if they + are available and describe contiguous cells. Otherwise, corners are + interpolated and extrapolated from the cell centers. Parameters ---------- @@ -135,9 +140,9 @@ def read( ds[y_var_name].dims[0], ) - # interp/extrap corners - descriptor.x_corner = interp_extrap_corner(descriptor.x) - descriptor.y_corner = interp_extrap_corner(descriptor.y) + # use CF bounds if available, otherwise interp/extrap corners + descriptor.x_corner = get_corners_1d(ds, x_var_name) + descriptor.y_corner = get_corners_1d(ds, y_var_name) descriptor.history = add_history(ds=ds) return descriptor diff --git a/pyremap/descriptor/utility.py b/pyremap/descriptor/utility.py index 547f3c28..710e27e5 100644 --- a/pyremap/descriptor/utility.py +++ b/pyremap/descriptor/utility.py @@ -10,12 +10,213 @@ # https://raw.githubusercontent.com/MPAS-Dev/pyremap/main/LICENSE import sys +import warnings import numpy as np import pyproj.enums from pyproj import Transformer +def get_corners_1d(ds, var_name): + """ + Get the coordinates of grid-cell corners along a 1D coordinate, using the + CF ``bounds`` of the coordinate if they are available and usable, and + interpolating and extrapolating from cell centers if not. + + Parameters + ---------- + ds : xarray.Dataset + A dataset containing the coordinate variable ``var_name`` + + var_name : str + The name of the 1D coordinate variable + + Returns + ------- + corner : numpy.ndarray + A 1D array of corner coordinates with one more element than + ``ds[var_name]`` + """ + center = np.array(ds[var_name].values, float) + bounds = _get_cf_bounds(ds, var_name, shape=(len(center), 2)) + if bounds is not None: + corner = _corners_from_bounds_1d(bounds) + if corner is not None: + return corner + warnings.warn( + f'The CF bounds of {var_name} are not contiguous so corners ' + f'will be interpolated and extrapolated from cell centers ' + f'instead.', + stacklevel=2, + ) + + return interp_extrap_corner(center) + + +def get_corners_2d(ds, lat_var_name, lon_var_name): + """ + Get the coordinates of grid-cell corners for 2D latitude and longitude + coordinates, using the CF ``bounds`` of the coordinates if they are + available and usable, and interpolating and extrapolating from cell + centers if not. + + Parameters + ---------- + ds : xarray.Dataset + A dataset containing the coordinate variables + + lat_var_name, lon_var_name : str + The names of the 2D latitude and longitude variables + + Returns + ------- + lat_corner, lon_corner : numpy.ndarray + 2D arrays of corner coordinates with one more element than + ``ds[lat_var_name]`` and ``ds[lon_var_name]`` along each dimension + """ + lat = np.array(ds[lat_var_name].values, float) + lon = np.array(ds[lon_var_name].values, float) + shape = (lat.shape[0], lat.shape[1], 4) + lat_bounds = _get_cf_bounds(ds, lat_var_name, shape=shape) + lon_bounds = _get_cf_bounds(ds, lon_var_name, shape=shape) + if lat_bounds is not None and lon_bounds is not None: + corners = _corners_from_bounds_2d(lat_bounds, lon_bounds) + if corners is not None: + return corners + warnings.warn( + f'The CF bounds of {lat_var_name} and {lon_var_name} do not ' + f'share vertices between neighboring cells so corners will be ' + f'interpolated and extrapolated from cell centers instead.', + stacklevel=2, + ) + elif lat_bounds is not None or lon_bounds is not None: + warnings.warn( + f'Only one of {lat_var_name} and {lon_var_name} has usable CF ' + f'bounds so corners will be interpolated and extrapolated from ' + f'cell centers instead.', + stacklevel=2, + ) + + return interp_extrap_corners_2d(lat), interp_extrap_corners_2d(lon) + + +def _get_cf_bounds(ds, var_name, shape): + """ + Get the CF ``bounds`` of the given variable as a numpy array with the + expected shape, or ``None`` if they are not available + """ + bounds_name = ds[var_name].attrs.get('bounds') + if bounds_name is None: + return None + + if bounds_name not in ds: + warnings.warn( + f'{var_name} has a CF bounds attribute "{bounds_name}" but no ' + f'such variable is present in the dataset.', + stacklevel=3, + ) + return None + + bounds = np.array(ds[bounds_name].values, float) + if bounds.shape != shape: + warnings.warn( + f'The CF bounds variable {bounds_name} has shape ' + f'{bounds.shape}, not the expected {shape}.', + stacklevel=3, + ) + return None + + return bounds + + +def _bounds_tolerance(bounds): + """A tolerance for comparing bounds, based on the size of the cells""" + center = np.mean(bounds, axis=-1, keepdims=True) + scale = np.max(np.abs(bounds - center)) + return 1e-6 * scale + + +def _corners_from_bounds_1d(bounds): + """ + Convert CF bounds with shape ``(n, 2)`` to an array of ``n + 1`` corners, + or return ``None`` if the bounds are not contiguous (in which case they + cannot be described by a 1D array of corners) + """ + tol = _bounds_tolerance(bounds) + # bounds may be given in the direction of the coordinate or always from + # the lower to the upper edge, so try both orders + for flipped in [bounds, bounds[:, ::-1]]: + if np.all(np.abs(flipped[:-1, 1] - flipped[1:, 0]) <= tol): + return np.append(flipped[:, 0], flipped[-1, 1]) + + return None + + +def _corners_from_bounds_2d(lat_bounds, lon_bounds): + """ + Convert CF bounds with shape ``(ny, nx, 4)`` to arrays of corners with + shape ``(ny + 1, nx + 1)``, or return ``None`` if neighboring cells do not + share vertices (in which case they cannot be described by 2D arrays of + corners) + """ + # CF requires the vertices of each cell to be traversed in order around + # the cell but does not say which vertex comes first or which direction + # the traversal goes, so try each of the 8 possibilities. Each candidate + # gives the vertex indices of the lower-left, lower-right, upper-right and + # upper-left corners of the cell in index space. Requiring neighboring + # cells to share vertices picks out the right candidate except on grids + # too small to have neighbors in one or both directions, where the CF + # recommendation (anticlockwise starting from the lower left) is assumed. + candidates = [] + for base in ([0, 1, 2, 3], [0, 3, 2, 1]): + for shift in range(4): + candidates.append(base[shift:] + base[:shift]) + + tol = max(_bounds_tolerance(lat_bounds), _bounds_tolerance(lon_bounds)) + + for candidate in candidates: + if not all( + _vertices_are_shared(bounds, candidate, tol) + for bounds in [lat_bounds, lon_bounds] + ): + continue + + lower_left, lower_right, upper_right, upper_left = candidate + corners = [] + for bounds in [lat_bounds, lon_bounds]: + ny, nx = bounds.shape[0], bounds.shape[1] + corner = np.zeros((ny + 1, nx + 1)) + corner[:-1, :-1] = bounds[:, :, lower_left] + corner[:-1, -1] = bounds[:, -1, lower_right] + corner[-1, :-1] = bounds[-1, :, upper_left] + corner[-1, -1] = bounds[-1, -1, upper_right] + corners.append(corner) + + return corners[0], corners[1] + + return None + + +def _vertices_are_shared(bounds, candidate, tol): + """ + Whether neighboring cells share vertices if the 4 vertices of each cell + are in the order given by ``candidate`` (lower-left, lower-right, + upper-right and upper-left in index space) + """ + lower_left, lower_right, upper_right, upper_left = candidate + # neighbors in x share their left and right vertices while neighbors in y + # share their lower and upper vertices + shared = [ + (bounds[:, :-1, lower_right], bounds[:, 1:, lower_left]), + (bounds[:, :-1, upper_right], bounds[:, 1:, upper_left]), + (bounds[:-1, :, upper_left], bounds[1:, :, lower_left]), + (bounds[:-1, :, upper_right], bounds[1:, :, lower_right]), + ] + return all( + np.all(np.abs(first - second) <= tol) for first, second in shared + ) + + def interp_extrap_corner(in_field): """Interpolate/extrapolate a 1D field from grid centers to grid corners""" diff --git a/pyremap/version.py b/pyremap/version.py index 4b93182a..cae14223 100644 --- a/pyremap/version.py +++ b/pyremap/version.py @@ -1,2 +1,2 @@ -__version_info__ = (2, 3, 0) +__version_info__ = (2, 4, 0) __version__ = '.'.join(str(vi) for vi in __version_info__) diff --git a/tests/test_cf_bounds.py b/tests/test_cf_bounds.py new file mode 100644 index 00000000..1897eac1 --- /dev/null +++ b/tests/test_cf_bounds.py @@ -0,0 +1,356 @@ +# This software is open source software available under the BSD-3 license. +# +# Copyright (c) 2025 Triad National Security, LLC. All rights reserved. +# Copyright (c) 2025 Lawrence Livermore National Security, LLC. All rights +# reserved. +# Copyright (c) 2025 UT-Battelle, LLC. All rights reserved. +# +# Additional copyright and license information can be found in the LICENSE file +# distributed with this code, or at +# https://raw.githubusercontent.com/MPAS-Dev/pyremap/main/LICENSE +""" +Unit tests for using CF ``bounds`` to locate grid-cell corners. +""" + +import numpy as np +import pyproj +import pytest +import xarray as xr + +from pyremap import ( + LatLon2DGridDescriptor, + LatLonGridDescriptor, + ProjectionGridDescriptor, +) +from pyremap.descriptor import ( + get_corners_1d, + get_corners_2d, + interp_extrap_corner, + interp_extrap_corners_2d, +) + +# corners of a grid with cells that vary in size, so that the corners are not +# the same as those interpolated and extrapolated from the cell centers +LAT_CORNER = np.array([-90.0, -60.0, -10.0, 20.0, 30.0, 90.0]) +LON_CORNER = np.array([-180.0, -100.0, -30.0, 0.0, 45.0, 90.0, 180.0]) + + +def _centers(corner): + return 0.5 * (corner[:-1] + corner[1:]) + + +def _bounds_1d(corner): + """CF bounds of shape (n, 2) in the direction of the coordinate""" + return np.stack((corner[:-1], corner[1:]), axis=-1) + + +def _lat_lon_dataset(lat_bounds=None, lon_bounds=None): + """A 1D lat-lon dataset, optionally with CF bounds""" + lat = _centers(LAT_CORNER) + lon = _centers(LON_CORNER) + ds = xr.Dataset( + coords={ + 'lat': ('lat', lat, {'units': 'degrees_north'}), + 'lon': ('lon', lon, {'units': 'degrees_east'}), + } + ) + if lat_bounds is not None: + ds['lat_bnds'] = (('lat', 'nbnd'), lat_bounds) + ds.lat.attrs['bounds'] = 'lat_bnds' + if lon_bounds is not None: + ds['lon_bnds'] = (('lon', 'nbnd'), lon_bounds) + ds.lon.attrs['bounds'] = 'lon_bnds' + return ds + + +def _lat_lon_2d_dataset(order=(0, 1, 2, 3), lat_corner=None, lon_corner=None): + """ + A 2D lat-lon dataset with CF bounds, with the 4 vertices of each cell in + the order given by ``order`` (lower-left, lower-right, upper-right and + upper-left in index space) + """ + if lat_corner is None: + lat_corner = LAT_CORNER + if lon_corner is None: + lon_corner = LON_CORNER + lon_corner_2d, lat_corner_2d = np.meshgrid(lon_corner, lat_corner) + lat = 0.25 * ( + lat_corner_2d[:-1, :-1] + + lat_corner_2d[:-1, 1:] + + lat_corner_2d[1:, 1:] + + lat_corner_2d[1:, :-1] + ) + lon = 0.25 * ( + lon_corner_2d[:-1, :-1] + + lon_corner_2d[:-1, 1:] + + lon_corner_2d[1:, 1:] + + lon_corner_2d[1:, :-1] + ) + + ds = xr.Dataset() + for var_name, corner_2d, center in [ + ('lat2d', lat_corner_2d, lat), + ('lon2d', lon_corner_2d, lon), + ]: + # lower-left, lower-right, upper-right, upper-left + vertices = [ + corner_2d[:-1, :-1], + corner_2d[:-1, 1:], + corner_2d[1:, 1:], + corner_2d[1:, :-1], + ] + bounds = np.zeros((center.shape[0], center.shape[1], 4)) + for vertex_index, corner_index in enumerate(order): + bounds[:, :, vertex_index] = vertices[corner_index] + + units = 'degrees_north' if var_name == 'lat2d' else 'degrees_east' + ds[var_name] = (('y', 'x'), center, {'units': units}) + ds[f'{var_name}_bnds'] = (('y', 'x', 'nv'), bounds) + ds[var_name].attrs['bounds'] = f'{var_name}_bnds' + + return ds, lat_corner_2d, lon_corner_2d + + +def test_corners_1d_from_bounds(): + """CF bounds are used in place of interpolation and extrapolation""" + ds = _lat_lon_dataset( + lat_bounds=_bounds_1d(LAT_CORNER), lon_bounds=_bounds_1d(LON_CORNER) + ) + + np.testing.assert_allclose(get_corners_1d(ds, 'lat'), LAT_CORNER) + np.testing.assert_allclose(get_corners_1d(ds, 'lon'), LON_CORNER) + + # the point of the test data is that these are not the same as the + # interpolated and extrapolated corners + assert not np.allclose( + interp_extrap_corner(ds.lat.values), LAT_CORNER, atol=1e-10 + ) + + +def test_corners_1d_no_bounds(): + """Without CF bounds, corners are interpolated and extrapolated""" + ds = _lat_lon_dataset() + + np.testing.assert_allclose( + get_corners_1d(ds, 'lat'), interp_extrap_corner(ds.lat.values) + ) + + +def test_corners_1d_descending(): + """CF bounds are used for coordinates that decrease with index""" + lat_corner = LAT_CORNER[::-1] + lat = _centers(lat_corner) + ds = xr.Dataset(coords={'lat': ('lat', lat, {'units': 'degrees_north'})}) + ds['lat_bnds'] = (('lat', 'nbnd'), _bounds_1d(lat_corner)) + ds.lat.attrs['bounds'] = 'lat_bnds' + + np.testing.assert_allclose(get_corners_1d(ds, 'lat'), lat_corner) + + +def test_corners_1d_descending_min_max_bounds(): + """ + CF bounds are used for a descending coordinate whose bounds are always + given from the lower to the upper edge + """ + lat_corner = LAT_CORNER[::-1] + lat = _centers(lat_corner) + # each pair is [min, max] rather than in the direction of the coordinate + bounds = _bounds_1d(lat_corner)[:, ::-1] + ds = xr.Dataset(coords={'lat': ('lat', lat, {'units': 'degrees_north'})}) + ds['lat_bnds'] = (('lat', 'nbnd'), bounds) + ds.lat.attrs['bounds'] = 'lat_bnds' + + np.testing.assert_allclose(get_corners_1d(ds, 'lat'), lat_corner) + + +def test_corners_1d_noncontiguous_bounds(): + """Bounds with gaps between cells can't be used, so we fall back""" + bounds = _bounds_1d(LAT_CORNER) + # shrink each cell so neighbors no longer share an edge + center = np.mean(bounds, axis=-1, keepdims=True) + bounds = center + 0.9 * (bounds - center) + ds = _lat_lon_dataset(lat_bounds=bounds) + + with pytest.warns(UserWarning, match='not contiguous'): + corner = get_corners_1d(ds, 'lat') + + np.testing.assert_allclose(corner, interp_extrap_corner(ds.lat.values)) + + +def test_corners_1d_missing_bounds_variable(): + """A ``bounds`` attribute pointing at nothing is a warning, not an error""" + ds = _lat_lon_dataset() + ds.lat.attrs['bounds'] = 'lat_bnds' + + with pytest.warns(UserWarning, match='no such variable'): + corner = get_corners_1d(ds, 'lat') + + np.testing.assert_allclose(corner, interp_extrap_corner(ds.lat.values)) + + +def test_corners_1d_wrong_bounds_shape(): + """Bounds that aren't (n, 2) are a warning, not an error""" + ds = _lat_lon_dataset() + ds['lat_bnds'] = (('lat',), LAT_CORNER[:-1]) + ds.lat.attrs['bounds'] = 'lat_bnds' + + with pytest.warns(UserWarning, match='shape'): + corner = get_corners_1d(ds, 'lat') + + np.testing.assert_allclose(corner, interp_extrap_corner(ds.lat.values)) + + +@pytest.mark.parametrize( + 'order', + [ + (0, 1, 2, 3), # counterclockwise from the lower left + (1, 2, 3, 0), # counterclockwise from the lower right + (0, 3, 2, 1), # clockwise from the lower left + (2, 1, 0, 3), # clockwise from the upper right + ], +) +def test_corners_2d_from_bounds(order): + """ + 2D CF bounds are used whatever order the vertices of each cell are + traversed in + """ + ds, lat_corner_2d, lon_corner_2d = _lat_lon_2d_dataset(order=order) + + lat_corner, lon_corner = get_corners_2d(ds, 'lat2d', 'lon2d') + + np.testing.assert_allclose(lat_corner, lat_corner_2d) + np.testing.assert_allclose(lon_corner, lon_corner_2d) + + # again, these should not match what interpolation and extrapolation gives + assert not np.allclose( + interp_extrap_corners_2d(ds.lat2d.values), lat_corner_2d, atol=1e-10 + ) + + +def test_corners_2d_no_bounds(): + """Without CF bounds, 2D corners are interpolated and extrapolated""" + ds, _, _ = _lat_lon_2d_dataset() + ds = ds.drop_vars(['lat2d_bnds', 'lon2d_bnds']) + del ds.lat2d.attrs['bounds'] + del ds.lon2d.attrs['bounds'] + + lat_corner, lon_corner = get_corners_2d(ds, 'lat2d', 'lon2d') + + np.testing.assert_allclose( + lat_corner, interp_extrap_corners_2d(ds.lat2d.values) + ) + np.testing.assert_allclose( + lon_corner, interp_extrap_corners_2d(ds.lon2d.values) + ) + + +def test_corners_2d_unshared_vertices(): + """ + 2D bounds that neighboring cells don't share can't be described by corner + arrays, so we warn and fall back + """ + ds, _, _ = _lat_lon_2d_dataset() + bounds = ds.lat2d_bnds.values + center = np.mean(bounds, axis=-1, keepdims=True) + ds['lat2d_bnds'] = (ds.lat2d_bnds.dims, center + 0.9 * (bounds - center)) + + with pytest.warns(UserWarning, match='do not share vertices'): + lat_corner, lon_corner = get_corners_2d(ds, 'lat2d', 'lon2d') + + np.testing.assert_allclose( + lat_corner, interp_extrap_corners_2d(ds.lat2d.values) + ) + + +def test_corners_2d_bounds_on_one_coord_only(): + """Bounds on only one of the two coordinates are not enough""" + ds, _, _ = _lat_lon_2d_dataset() + ds = ds.drop_vars(['lon2d_bnds']) + del ds.lon2d.attrs['bounds'] + + with pytest.warns(UserWarning, match='Only one of'): + lat_corner, lon_corner = get_corners_2d(ds, 'lat2d', 'lon2d') + + np.testing.assert_allclose( + lat_corner, interp_extrap_corners_2d(ds.lat2d.values) + ) + + +def test_lat_lon_descriptor_honors_bounds(): + """``LatLonGridDescriptor.read()`` uses CF bounds""" + ds = _lat_lon_dataset( + lat_bounds=_bounds_1d(LAT_CORNER), lon_bounds=_bounds_1d(LON_CORNER) + ) + + descriptor = LatLonGridDescriptor.read(ds=ds, mesh_name='test') + + np.testing.assert_allclose(descriptor.lat_corner, LAT_CORNER) + np.testing.assert_allclose(descriptor.lon_corner, LON_CORNER) + + +def test_lat_lon_descriptor_scrip_from_bounds(tmp_path): + """The corners in the SCRIP file come from the CF bounds""" + ds = _lat_lon_dataset( + lat_bounds=_bounds_1d(LAT_CORNER), lon_bounds=_bounds_1d(LON_CORNER) + ) + + descriptor = LatLonGridDescriptor.read(ds=ds, mesh_name='test') + scrip_filename = str(tmp_path / 'scrip.nc') + descriptor.to_scrip(scrip_filename) + + with xr.open_dataset(scrip_filename) as ds_scrip: + corner_lat = ds_scrip.grid_corner_lat.values + corner_lon = ds_scrip.grid_corner_lon.values + + nlat = len(LAT_CORNER) - 1 + nlon = len(LON_CORNER) - 1 + assert corner_lat.shape == (nlat * nlon, 4) + # the first cell is bounded by the first two corners in each direction + np.testing.assert_allclose( + np.sort(np.unique(corner_lat[0, :])), LAT_CORNER[0:2] + ) + np.testing.assert_allclose( + np.sort(np.unique(corner_lon[0, :])), LON_CORNER[0:2] + ) + + +def test_lat_lon_2d_descriptor_honors_bounds(): + """``LatLon2DGridDescriptor.read()`` uses CF bounds""" + ds, lat_corner_2d, lon_corner_2d = _lat_lon_2d_dataset() + + descriptor = LatLon2DGridDescriptor.read( + ds=ds, lat_var_name='lat2d', lon_var_name='lon2d', mesh_name='test' + ) + + np.testing.assert_allclose(descriptor.lat_corner, lat_corner_2d) + np.testing.assert_allclose(descriptor.lon_corner, lon_corner_2d) + + +def test_projection_descriptor_honors_bounds(tmp_path): + """``ProjectionGridDescriptor.read()`` uses CF bounds on x and y""" + x_corner = 1e3 * np.array([-300.0, -100.0, 0.0, 50.0, 300.0]) + y_corner = 1e3 * np.array([-200.0, -50.0, 100.0, 400.0]) + + ds = xr.Dataset( + coords={ + 'x': ('x', _centers(x_corner), {'units': 'meters'}), + 'y': ('y', _centers(y_corner), {'units': 'meters'}), + } + ) + ds['x_bnds'] = (('x', 'nbnd'), _bounds_1d(x_corner)) + ds['y_bnds'] = (('y', 'nbnd'), _bounds_1d(y_corner)) + ds.x.attrs['bounds'] = 'x_bnds' + ds.y.attrs['bounds'] = 'y_bnds' + filename = str(tmp_path / 'proj_grid.nc') + ds.to_netcdf(filename) + + 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' + ) + descriptor = ProjectionGridDescriptor.read( + projection=projection, filename=filename, mesh_name='test' + ) + + np.testing.assert_allclose(descriptor.x_corner, x_corner) + np.testing.assert_allclose(descriptor.y_corner, y_corner)