Merge pull request #154 from csiro-coasts/fix-rivers

Fix generated bounds for 'river' cells in CFGrid2D datasets

Author: mx-moth

docs/releases/development.rst

@@ -2,4 +2,5 @@
 Next release (in development)
 =============================
 
-* ...
+* Fix invalid geometry being generated for 'river' cells
+  in CFGrid2D datasets with no cell bounds (:pr:`154`).

src/emsarray/conventions/grid.py

@@ -497,6 +497,17 @@ def _get_or_make_bounds(self, coordinate: xarray.DataArray) -> xarray.DataArray:
         # On the edges where there are fewer 'surrounding' cells, the cell centres are used.
         # Edge and corner cells will be smaller than the surrounding cells because of this.
 
+        # Discard cell coordinates where the cell is bound by nan on both sides.
+        # This can occur when the grid tracks a river which might be only one cell wide.
+        coordinate_values = coordinate.values.copy()
+        nan_coordinates = numpy.isnan(coordinate_values)
+        j_pad = numpy.pad(nan_coordinates, ((1, 1), (0, 0)), constant_values=False)
+        j_bound_by_nan = j_pad[:-2, :] & j_pad[2:, :]
+        i_pad = numpy.pad(nan_coordinates, ((0, 0), (1, 1)), constant_values=False)
+        i_bound_by_nan = i_pad[:, :-2] & i_pad[:, 2:]
+        bound_by_nan = j_bound_by_nan | i_bound_by_nan
+        coordinate_values[bound_by_nan] = numpy.nan
+
         # numpy.nanmean will return nan for an all-nan column.
         # This is the exact behaviour that we want.
         # numpy emits a warning that can not be silenced when this happens,
@@ -505,7 +516,7 @@ def _get_or_make_bounds(self, coordinate: xarray.DataArray) -> xarray.DataArray:
             warnings.filterwarnings(
                 "ignore", "Mean of empty slice", category=RuntimeWarning)
             grid = numpy.nanmean([
-                numpy.pad(coordinate.values, pad, constant_values=numpy.nan)
+                numpy.pad(coordinate_values, pad, constant_values=numpy.nan)
                 for pad in itertools.product([(1, 0), (0, 1)], [(1, 0), (0, 1)])
             ], axis=0)
 
@@ -517,7 +528,8 @@ def _get_or_make_bounds(self, coordinate: xarray.DataArray) -> xarray.DataArray:
             ]
             for y in range(y_size)
         ])
-        # Any cell that has a `nan` in its bounds will be set to all nan
+
+        # Set nan bounds for all cells that have any `nan` in its bounds.
         cells_with_nans = numpy.isnan(bounds).any(axis=2)
         bounds[cells_with_nans] = numpy.nan
 

tests/conventions/test_cfgrid2d.py

@@ -24,8 +24,8 @@
 from emsarray.exceptions import NoSuchCoordinateError
 from emsarray.operations import geometry
 from tests.utils import (
-    DiagonalShocGrid, ShocGridGenerator, ShocLayerGenerator,
-    assert_property_not_cached
+    AxisAlignedShocGrid, DiagonalShocGrid, ShocGridGenerator,
+    ShocLayerGenerator, assert_property_not_cached, plot_geometry
 )
 
 
@@ -266,7 +266,7 @@ def test_polygons_no_bounds():
     assert polygons[20 * 9 + 3] is not None
 
 
-def test_polygons_with_bounds():
+def test_polygons_with_bounds() -> None:
     dataset = make_dataset(
         j_size=10, i_size=20,
         include_bounds=True, corner_size=3)
@@ -289,7 +289,7 @@ def test_polygons_with_bounds():
     assert polygons[20 * 9 + 3] is not None
 
 
-def test_holes():
+def test_holes() -> None:
     dataset = make_dataset(j_size=10, i_size=20, corner_size=5)
     only_polygons = dataset.ems.polygons[dataset.ems.mask]
 
@@ -313,15 +313,57 @@ def test_holes():
     assert poly.equals_exact(Polygon([(0.3, 1.9), (0.4, 1.8), (0.5, 1.9), (0.4, 2.0), (0.3, 1.9)]), 1e-6)
 
 
-def test_bounds():
+def test_bounds() -> None:
     dataset = make_dataset(
         j_size=10, i_size=20, corner_size=5, include_bounds=True)
     # The corner cuts out some of the upper bounds
     assert_allclose(dataset.ems.bounds, (0, 0, 3, 2.5))
     assert_property_not_cached(dataset.ems, 'geometry')
 
 
-def test_face_centres():
+@pytest.mark.parametrize("include_bounds", [True, False], ids=["with-bounds", "without-bounds"])
+def test_bounds_river(
+    include_bounds: bool,
+    tmp_path: pathlib.Path,
+) -> None:
+    # Test calculated bounds for grids with one cell wide 'rivers'.
+    # The river cells have cell centres but insufficient information to
+    # construct any bounds for those cells.
+    dataset = make_dataset(
+        j_size=7, i_size=7, include_bounds=include_bounds, grid_type=AxisAlignedShocGrid)
+    # Clear out a corner on both sides, leaving a one cell wide river.
+    dataset['longitude'].values[:3, :3] = numpy.nan
+    dataset['latitude'].values[:3, :3] = numpy.nan
+    dataset['longitude'].values[-3:, :3] = numpy.nan
+    dataset['latitude'].values[-3:, :3] = numpy.nan
+    if include_bounds:
+        dataset['longitude_bounds'].values[:3, :3] = numpy.nan
+        dataset['latitude_bounds'].values[:3, :3] = numpy.nan
+        dataset['longitude_bounds'].values[-3:, :3] = numpy.nan
+        dataset['latitude_bounds'].values[-3:, :3] = numpy.nan
+
+    # Useful to see what the geometry looks like, not part of the test.
+    plot_geometry(dataset, tmp_path / 'river.png', extent=(-0.05, 0.75, -0.05, 0.75))
+
+    convention: ShocSimple = dataset.ems
+    polygons = convention.wind(xarray.DataArray(convention.polygons)).values
+    assert numpy.all(polygons[:, 3:] != None)  # noqa: E711
+    assert numpy.all(polygons[:3, :3] == None)  # noqa: E711
+    assert numpy.all(polygons[-3:, :3] == None)  # noqa: E711
+    river_polygons = polygons[3, :3]
+    if include_bounds:
+        assert numpy.all(river_polygons != None)  # noqa: E711
+    else:
+        assert numpy.all(river_polygons == None)  # noqa: E711
+
+    if not include_bounds:
+        # All the 'coast' cells should have the same minimum x coordinate,
+        # even the one next to the river.
+        min_x = [min(p.exterior.coords.xy[0]) for p in polygons[:, 3]]
+        numpy.testing.assert_allclose(min_x, 0.35)
+
+
+def test_face_centres() -> None:
     # SHOC simple face centres are taken directly from the coordinates,
     # not calculated from polygon centres.
     dataset = make_dataset(j_size=10, i_size=20, corner_size=3)
@@ -338,7 +380,7 @@ def test_face_centres():
             numpy.testing.assert_equal(face_centres[linear_index], [lon, lat])
 
 
-def test_selector_for_index():
+def test_selector_for_index() -> None:
     dataset = make_dataset(j_size=5, i_size=7)
     convention: ShocSimple = dataset.ems
 
@@ -348,13 +390,13 @@ def test_selector_for_index():
     assert selector == convention.selector_for_index(index)
 
 
-def test_make_geojson_geometry():
+def test_make_geojson_geometry() -> None:
     dataset = make_dataset(j_size=10, i_size=10, corner_size=3)
     out = json.dumps(geometry.to_geojson(dataset))
     assert isinstance(out, str)
 
 
-def test_ravel_index():
+def test_ravel_index() -> None:
     dataset = make_dataset(j_size=5, i_size=7)
     convention: ShocSimple = dataset.ems
 
@@ -365,15 +407,15 @@ def test_ravel_index():
         assert convention.wind_index(linear_index, grid_kind=CFGridKind.face) == index
 
 
-def test_grid_kinds():
+def test_grid_kinds() -> None:
     dataset = make_dataset(j_size=3, i_size=5)
     convention: ShocSimple = dataset.ems
 
     assert convention.grid_kinds == frozenset({CFGridKind.face})
     assert convention.default_grid_kind == CFGridKind.face
 
 
-def test_grid_kind_and_size():
+def test_grid_kind_and_size() -> None:
     dataset = make_dataset(j_size=5, i_size=7)
     convention: ShocSimple = dataset.ems
 
@@ -382,7 +424,7 @@ def test_grid_kind_and_size():
     assert size == 5 * 7
 
 
-def test_ravel():
+def test_ravel() -> None:
     dataset = make_dataset(j_size=5, i_size=7)
     convention: ShocSimple = dataset.ems
 
@@ -394,7 +436,7 @@ def test_ravel():
         temp.values.reshape(temp.shape[:2] + (-1,)))
 
 
-def test_wind():
+def test_wind() -> None:
     dataset = make_dataset(j_size=5, i_size=7)
     convention: ShocSimple = dataset.ems
 
@@ -413,7 +455,7 @@ def test_wind():
         values.reshape(time_size, 5, 7))
 
 
-def test_drop_geometry(datasets: pathlib.Path):
+def test_drop_geometry(datasets: pathlib.Path) -> None:
     dataset = xarray.open_dataset(datasets / 'cfgrid2d.nc')
 
     dropped = dataset.ems.drop_geometry()
@@ -426,7 +468,7 @@ def test_drop_geometry(datasets: pathlib.Path):
     assert topology.longitude_name not in dropped.variables
 
 
-def test_values():
+def test_values() -> None:
     dataset = make_dataset(j_size=10, i_size=20, corner_size=5)
     eta = dataset.data_vars["eta"].isel(time=0)
     values = dataset.ems.ravel(eta)
@@ -439,7 +481,7 @@ def test_values():
 
 
 @pytest.mark.matplotlib
-def test_plot_on_figure():
+def test_plot_on_figure() -> None:
     # Not much to test here, mostly that it doesn't throw an error
     dataset = make_dataset(j_size=10, i_size=20)
     surface_temp = dataset.data_vars["temp"].isel(k=-1, time=0)

tests/utils.py

@@ -7,15 +7,18 @@
 from functools import cached_property
 from typing import Any
 
+import matplotlib.pyplot as plt
 import numpy
 import pytest
 import shapely
 import xarray
+from cartopy.mpl.geoaxes import GeoAxes
 from packaging.requirements import Requirement
 
 from emsarray.conventions.arakawa_c import (
     ArakawaCGridKind, c_mask_from_centres
 )
+from emsarray.types import Bounds, Pathish
 
 
 @contextlib.contextmanager
@@ -349,6 +352,14 @@ def y_back(self):
         return reduce_axes(self.y_grid, (False, True))
 
 
+class AxisAlignedShocGrid(ShocGridGenerator):
+    def make_x_grid(self, j: numpy.ndarray, i: numpy.ndarray) -> numpy.ndarray:
+        return 0.1 * i  # type: ignore
+
+    def make_y_grid(self, j: numpy.ndarray, i: numpy.ndarray) -> numpy.ndarray:
+        return 0.1 * j  # type: ignore
+
+
 class DiagonalShocGrid(ShocGridGenerator):
     def make_x_grid(self, j: numpy.ndarray, i: numpy.ndarray) -> numpy.ndarray:
         return 0.1 * (i + j)  # type: ignore
@@ -426,3 +437,28 @@ def only_versions(*requirements: str):
                 f'{requirement.name}=={version} does not satisfy {requirement}')
 
     return pytest.mark.skipif(len(invalid_versions) > 0, reason='\n'.join(invalid_versions))
+
+
+def plot_geometry(
+    dataset: xarray.Dataset,
+    out: Pathish,
+    *,
+    figsize: tuple[float, float] = (10, 10),
+    extent: Bounds | None = None,
+    title: str | None = None
+) -> None:
+    figure = plt.figure(layout='constrained', figsize=figsize)
+    axes: GeoAxes = figure.add_subplot(projection=dataset.ems.data_crs)
+    axes.set_aspect(aspect='equal', adjustable='datalim')
+    axes.gridlines(draw_labels=['left', 'bottom'], linestyle='dashed')
+
+    axes.add_collection(dataset.ems.make_poly_collection(
+        edgecolors='black', linewidth=0.5, facecolors='lightcyan'))
+    axes.scatter(dataset.ems.face_centres[:, 0], dataset.ems.face_centres[:, 1], c='red')
+
+    if title is not None:
+        axes.set_title(title)
+    if extent is not None:
+        axes.set_extent(extent)
+
+    figure.savefig(out)