diff --git a/lib/ants/analysis/__init__.py b/lib/ants/analysis/__init__.py index cddb441..4000a94 100644 --- a/lib/ants/analysis/__init__.py +++ b/lib/ants/analysis/__init__.py @@ -187,7 +187,7 @@ def standard_deviation(source, src_mean): return awm -def merge(primary_cube, alternate_cube, validity_polygon=None): +def merge(primary_cube, alternate_cube, validity_polygon=None, blending_distance=None): """ Merges data from the alternative cube into the primary cube. @@ -195,8 +195,15 @@ def merge(primary_cube, alternate_cube, validity_polygon=None): cube which lay outside the provided polygon, override the values of the primary at those locations. Containment is defined as any cell corner which lies within the polygon. "Within" explicitly does not include - those points which exactly lay on the polygon boundary. Where multiple - primary and alternate cubes are provided, then these are paired + those points which exactly lay on the polygon boundary. + + A blending between the sources can be applied by specifying the + ``blending_distance`` (for no blending, pass ``None``). A linear blending + between the primary and alternate sources will be applied in the region + immediately inside the polygon over the blending distance. + Beyond the blending distance, the alternate source is used. + + Where multiple primary and alternate cubes are provided, then these are paired appropriately where possible. Where these datasets are not defined on the same grid, the user should consider a regrid first to then utilise merge. @@ -220,6 +227,12 @@ def merge(primary_cube, alternate_cube, validity_polygon=None): stacked together with the primary_cube taking priority over alternate_cube in the case of an overlap. A runtime error will be raised if the primary_cube is wholly within the validity_polygon. + blending_distance : float + Distance over which blending between the primary and alternate sources + is applied. Note that this is in units of grid cells, not a physical distance. + If ``None``, no blending is applied, and there will be a hard edge between + the two sources. This option is only valid with a provided validity polygon, + and with a single level field. Returns ------- @@ -235,19 +248,51 @@ def merge(primary_cube, alternate_cube, validity_polygon=None): primary_cubes = ants.utils.cube.as_cubelist(primary_cube) alternate_cubes = ants.utils.cube.as_cubelist(alternate_cube) + if blending_distance: + _validate_args_with_blending( + primary_cubes, alternate_cubes, validity_polygon, blending_distance + ) + # Group (sort) cubes so they are ordered in a way suitable for merging. primary_cubes, alternate_cubes = ants.utils.cube.sort_cubes( primary_cubes, alternate_cubes ) result = iris.cube.CubeList([]) for src1, src2 in zip(primary_cubes, alternate_cubes): - nsource = _merge.merge(src1, src2, validity_polygon) + nsource = _merge.merge(src1, src2, validity_polygon, blending_distance) result.append(nsource) if isinstance(primary_cube, iris.cube.Cube): result = result[0] return result +def _validate_args_with_blending( + primary_cubes, alternate_cubes, validity_polygon, blending_distance +): + """Specific validation for merge arguments when blending is provided.""" + if validity_polygon is None: + raise ValueError( + "blending_distance can only be used with a validity_polygon. " + f"No polygon was provided, but got {blending_distance=}" + ) + + all_primary_single_level = all(map(ants.utils.cube.is_single_level, primary_cubes)) + if not all_primary_single_level: + raise ValueError( + "Blending is only supported for single level data sources. " + "The primary data source is not single level" + ) + + all_alternate_single_level = all( + map(ants.utils.cube.is_single_level, alternate_cubes) + ) + if not all_alternate_single_level: + raise ValueError( + "Blending is only supported for single level data sources. " + "The alternate data source is not single level" + ) + + def _flood_fill_neighbour_identify( shape, coords, seed_point, extended_neighbourhood, wraparound ): diff --git a/lib/ants/analysis/_merge.py b/lib/ants/analysis/_merge.py index 7a2d025..704f385 100644 --- a/lib/ants/analysis/_merge.py +++ b/lib/ants/analysis/_merge.py @@ -20,6 +20,7 @@ import numpy.lib.stride_tricks as stride import shapely from pykdtree.kdtree import KDTree +from scipy.ndimage import distance_transform_edt from shapely.geometry import Polygon from shapely.vectorized import contains @@ -485,7 +486,7 @@ def _unified_grid(cube, cube2): return merged_cube -def merge(primary_cube, alternate_cube, validity_polygon=None): +def merge(primary_cube, alternate_cube, validity_polygon=None, blending_distance=None): """ Merges data from the alternative cube into the primary cube. @@ -499,6 +500,12 @@ def merge(primary_cube, alternate_cube, validity_polygon=None): in the primary cube dataset, overrides that elements validity defined in the cases of a specified validity polygon. + A blending between the sources can be applied by specifying the + ``blending_distance`` (for no blending, pass ``None``). A linear blending + between the primary and alternate sources will be applied in the region + immediately inside the polygon over the blending distance. + Beyond the blending distance, the alternate source is used. + Parameters ---------- primary_cube : `~iris.cube.Cube` @@ -519,6 +526,11 @@ def merge(primary_cube, alternate_cube, validity_polygon=None): alternate_cube in the case of an overlap. If a validity polygon is provided and the entire primary_cube dataset is within the polygon then a Runtime error will be raised. + blending_distance : float + Distance over which blending between the primary and alternate sources + is applied. Note that this is in units of grid cells, not a physical distance. + If ``None``, no blending is applied, and there will be a hard edge between + the two sources. Raises ------ @@ -630,21 +642,31 @@ def _overwrite_data(cube, cube2): # Apply data which is outside of the polygon to the other. # Transpose the data (view) to allow broadcasting - pdata, pmask = horizontal_grid_reorder(merged_cube) - adata, amask = horizontal_grid_reorder(full_alternate_cube) - pdata[full_mask_outside] = adata[full_mask_outside] - pmask[full_mask_outside] = amask[full_mask_outside] + primary_data, primary_mask = horizontal_grid_reorder(merged_cube) + alternate_data, alternate_mask = horizontal_grid_reorder(full_alternate_cube) + if blending_distance: + is_circular = primary_cube.coord(axis="x").circular + primary_data[...] = blend_data( + from_array=alternate_data, + into_array=primary_data, + mask=~full_mask_outside, + blending_distance=blending_distance, + circular=is_circular, + ) + else: + primary_data[full_mask_outside] = alternate_data[full_mask_outside] + primary_mask[full_mask_outside] = alternate_mask[full_mask_outside] # Identify overlap priority using np.nan values (these are assigned # when source cells are beyond the extent of the target grid whilst # regridding). - pdata, pmask = horizontal_grid_reorder(merged_cube) - adata, amask = horizontal_grid_reorder(full_alternate_cube) - slices = [slice(None)] * pdata.ndim - slices[2:] = [0] * (pdata.ndim - 2) - nan_mask = np.isnan(pdata[tuple(slices)]) - pdata[nan_mask] = adata[nan_mask] - pmask[nan_mask] = amask[nan_mask] + primary_data, primary_mask = horizontal_grid_reorder(merged_cube) + alternate_data, alternate_mask = horizontal_grid_reorder(full_alternate_cube) + slices = [slice(None)] * primary_data.ndim + slices[2:] = [0] * (primary_data.ndim - 2) + nan_mask = np.isnan(primary_data[tuple(slices)]) + primary_data[nan_mask] = alternate_data[nan_mask] + primary_mask[nan_mask] = alternate_mask[nan_mask] # Clear some memory if we can np.ma.MaskedArray.shrink_mask(merged_cube.data) @@ -656,6 +678,127 @@ def _overwrite_data(cube, cube2): return merged_cube +def blend_data( + from_array: np.ndarray, + into_array: np.ndarray, + mask: np.ndarray, + blending_distance: float, + circular: bool = False, +): + """Blend two data sources across a specified blending distance. + + Returns an array with a weighted combination of data selected from the two + sources, as determined by the provided mask. + + This is calculated as follows: + + 1. For all points where ``mask == False``, use the "from" source + 2. For all points where ``mask == True``, determine the distance to the nearest + point in the "from" region. + 3. If this distance is greater than the blending distance, use the "into" source. + 4. If this distance is less than the blending distance, weight the two datasets + using a linear combination: blended = w * from_array + (1 - w) * into_array, + where w = distance / blending_distance. + + The following diagram illustrates the blending in one dimension, with a + blending_distance of 4. + + into ___________ + / + / + from ___________/ + + mask 0000000000011111111111111 + + Parameters + ---------- + from_array : np.ndarray + Source data to be blended from + into_array : np.ndarray + Source data to be blended into + mask : np.ndarray + A boolean mask identifying the two regions: False for the "from" source + region and True for the "into" source region. + blending_distance : float + Distance over which blending between the sources + is applied. Note that this is in units of grid cells, not a physical distance. + As such, this is resolution dependent. See notes for more detail. + + Returns + ------- + blended : nd.ndarray + The blended data + + Notes + ----- + The three arrays ``from_array``, ``into_array`` and ``mask`` + must have the same shape. + + This function uses :func:`scipy.ndimage.distance_transform_edt` to calculate + distances between points on the grid. As such, it has no knowledge of physical + distance or coordinate reference systems. + + Warning + ------- + This function does not support masked arrays. Passing a masked array may result + in unexpected behaviour. + """ + _validate_blend_args(from_array, into_array, mask, blending_distance) + + if circular: + # Pad either side of the domain to allow for wraparound in x + # Do not pad in y direction + pad_width_x = int(np.ceil(blending_distance)) + pad_width = [(0, 0), (pad_width_x, pad_width_x)] + mask = np.pad(mask, pad_width, mode="wrap") + + distance_into_region = distance_transform_edt(mask) + max_distance_into_region = distance_into_region.max() + if max_distance_into_region < blending_distance: + warnings.warn( + "All points within the region are within the blending distance. " + f"Specified {blending_distance=}, maximum distance into domain: " + f"{max_distance_into_region}" + ) + + if circular: + # Retrieve central slice of the padded distance array + npoints_x = from_array.shape[-1] + slice_x = slice(pad_width_x, pad_width_x + npoints_x) + distance_into_region = distance_into_region[:, slice_x] + + into_weight = np.clip(distance_into_region / blending_distance, 0.0, 1.0) + blended = (into_weight * into_array) + (1 - into_weight) * from_array + return blended + + +def _validate_blend_args(from_array, into_array, mask, blending_distance): + if blending_distance <= 0: + raise ValueError( + f"Invalid blending_distance: {blending_distance}. Must be greater than zero" + ) + if from_array.ndim != 2: + raise ValueError( + "Can only blend 2-dimensional data, got data with " + f"{from_array.ndim} dimensions" + ) + if from_array.shape != into_array.shape: + raise ValueError( + f"Cannot blend sources with different shapes: " + f"{from_array.shape} and {into_array.shape}" + ) + if from_array.shape != mask.shape: + raise ValueError( + "Cannot blend sources as mask shape is inconsistent with source shape. " + f"Source shape: {from_array.shape}, Mask shape: {mask.shape}" + ) + if blending_distance > min(from_array.shape) / 2: + raise ValueError( + f"Invalid {blending_distance=}: greater than half the domain size " + f"(shape={from_array.shape})" + ) + + def _spiral_wrapper( unresolved_mask, to_fill_mask, diff --git a/lib/ants/cli/ancil_fill_n_merge.py b/lib/ants/cli/ancil_fill_n_merge.py index 9bb7cea..cf83a4d 100755 --- a/lib/ants/cli/ancil_fill_n_merge.py +++ b/lib/ants/cli/ancil_fill_n_merge.py @@ -96,6 +96,7 @@ def main( end, netcdf_only, search_method, + blending_distance, ): """ Perform merge and fill operation on the provided sources. @@ -104,8 +105,13 @@ def main( to be provided, and may optionally have a ``polygon`` shapefile. The resulting data takes values from the ``primary_source`` within the ``polygon`` (or everywhere where valid data is present, if the ``polygon`` is not - provided), and values from the ``alternate_source`` everywhere else. See - :func:`ants.analysis.merge` for further details. + provided), and values from the ``alternate_source`` everywhere else. + A blending between the sources can be applied by specifying the + ``blending_distance`` (for no blending, pass ``None``). A linear blending + between the primary and alternate sources will be applied in the region + immediately outside the polygon over the blending distance. + Beyond the blending distance, the alternate source is used. + See :func:`ants.analysis.merge` for further details. The fill stage replaces missing data values with valid data, where missing is defined as data that is either masked or NaN. If a landseamask is @@ -141,6 +147,12 @@ def main( search_method : :obj:`str` Select the search method to be used when filling missing points. The methods currently supported are "spiral" and "kdtree". + blending_distance : float + Distance over which blending between the primary and alternate sources + is applied. Note that this is in units of grid cells, not a physical distance. + If ``None``, no blending is applied, and there will be a hard edge between + the two sources. + Returns ------- : :class:`~iris.cube.CubeList` @@ -163,7 +175,9 @@ def main( result = primary_cubes if alternate_cubes is not None: - result = ants.analysis.merge(primary_cubes, alternate_cubes, validity_polygon) + result = ants.analysis.merge( + primary_cubes, alternate_cubes, validity_polygon, blending_distance + ) if target_mask_filepath: ants.analysis.make_consistent_with_lsm(result, lbm, invert_mask, search_method) @@ -226,6 +240,11 @@ def _get_parser(): required=False, default="spiral", ) + blending_help = ( + "Distance over which blending between the primary and alternate sources " + "is applied. Note that this is in units of grid cells, not a physical distance." + ) + parser.add_argument("--blending-distance", type=float, help=blending_help) return parser @@ -251,6 +270,7 @@ def cli_interface(): end=args.end, netcdf_only=args.netcdf_only, search_method=args.search_method, + blending_distance=args.blending_distance, ) diff --git a/lib/ants/tests/analysis/merge/test_blend_data.py b/lib/ants/tests/analysis/merge/test_blend_data.py new file mode 100644 index 0000000..2606b1e --- /dev/null +++ b/lib/ants/tests/analysis/merge/test_blend_data.py @@ -0,0 +1,158 @@ +# (C) Crown Copyright, Met Office. All rights reserved. +# +# This file is part of ANTS and is released under the BSD 3-Clause license. +# See LICENSE.txt in the root of the repository for full licensing details. +import numpy as np +import pytest +from ants.analysis._merge import blend_data + + +class TestExceptions: + def test_invalid_blending_distance_0(self): + source1 = np.array([0]) + source2 = np.array([0]) + mask = np.array([0]) + blending_distance = 0 + expected_msg = "Invalid blending_distance: 0. Must be greater than zero" + with pytest.raises(ValueError, match=expected_msg): + blend_data(source1, source2, mask, blending_distance) + + def test_blending_distance_too_large(self): + source1 = np.zeros((10, 10)) + source2 = np.ones_like(source1) + mask = source1.copy() + blending_distance = 6 + expected_msg = ( + "Invalid blending_distance=6: greater than half the domain size " + r"\(shape=\(10, 10\)\)" + ) + with pytest.raises(ValueError, match=expected_msg): + blend_data(source1, source2, mask, blending_distance) + + def test_1D_fails(self): + source1 = np.zeros(3) + source2 = np.ones_like(source1) + mask = source1.copy() + blending_distance = 2 + expected_msg = "Can only blend 2-dimensional data, got data with 1 dimensions" + with pytest.raises(ValueError, match=expected_msg): + blend_data(source1, source2, mask, blending_distance) + + def test_3D_fails(self): + source1 = np.zeros((3, 3, 3)) + source2 = np.ones_like(source1) + mask = source1.copy() + blending_distance = 2 + expected_msg = "Can only blend 2-dimensional data, got data with 3 dimensions" + with pytest.raises(ValueError, match=expected_msg): + blend_data(source1, source2, mask, blending_distance) + + def test_different_source_shapes(self): + source1 = np.zeros((2, 3)) + source2 = np.ones((3, 2)) + mask = np.ones_like(source1, dtype=bool) + blending_distance = 1 + + expected_msg = ( + r"Cannot blend sources with different shapes: \(2, 3\) and \(3, 2\)" + ) + with pytest.raises(ValueError, match=expected_msg): + blend_data(source1, source2, mask, blending_distance) + + def test_different_source_and_mask_shapes(self): + source1 = np.zeros((2, 3)) + source2 = np.ones_like(source1) + mask = np.ones((3, 2), dtype=bool) + blending_distance = 1 + + expected_msg = ( + "Cannot blend sources as mask shape is inconsistent with source shape. " + r"Source shape: \(2, 3\), Mask shape: \(3, 2\)" + ) + with pytest.raises(ValueError, match=expected_msg): + blend_data(source1, source2, mask, blending_distance) + + def test_blending_covers_entire_region(self): + source1 = np.zeros((5, 5)) + source2 = np.ones_like(source1) + mask = np.array( + [ + [0, 0, 0, 0, 0], + [0, 1, 1, 1, 0], + [0, 1, 1, 1, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + ], + dtype=bool, + ) + blending_distance = 2.0 + expected_msg = ( + "All points within the region are within the blending distance. " + "Specified blending_distance=2.0, maximum distance into domain: 1.0" + ) + with pytest.warns(match=expected_msg): + blend_data(source1, source2, mask, blending_distance) + + +class TestFunctionality: + @pytest.fixture() + def source1(self): + return np.zeros((7, 7), dtype=np.float64) + + @pytest.fixture() + def source2(self): + return np.ones((7, 7), dtype=np.float64) + + @pytest.fixture() + def mask(self): + mask = np.array( + [ + [1, 1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1, 1], + [1, 1, 1, 0, 0, 1, 1], + [1, 1, 1, 0, 0, 0, 1], + [1, 1, 1, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0], + ], + dtype=bool, + ) + return mask + + def test_blending_2D(self, source1, source2, mask): + blending_distance = 2.5 + + expected = np.array( + [ + [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 0.89442719, 0.8, 0.8, 0.89442719, 1.0], + [1.0, 0.89442719, 0.56568542, 0.4, 0.4, 0.56568542, 0.89442719], + [1.0, 0.8, 0.4, 0.0, 0.0, 0.4, 0.56568542], + [1.0, 0.8, 0.4, 0.0, 0.0, 0.0, 0.4], + [1.0, 0.8, 0.4, 0.0, 0.0, 0.0, 0.0], + [1.0, 0.8, 0.4, 0.0, 0.0, 0.0, 0.0], + ] + ) + + blended = blend_data(source1, source2, mask, blending_distance) + + np.testing.assert_array_almost_equal(blended, expected) + + def test_blending_2D_circular(self, source1, source2, mask): + blending_distance = 2.5 + + expected = np.array( + [ + [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 0.89442719, 0.8, 0.8, 0.89442719, 1.0], + [1.0, 0.89442719, 0.56568542, 0.4, 0.4, 0.56568542, 0.89442719], + [0.89442719, 0.8, 0.4, 0.0, 0.0, 0.4, 0.56568542], + [0.56568542, 0.8, 0.4, 0.0, 0.0, 0.0, 0.4], + [0.4, 0.8, 0.4, 0.0, 0.0, 0.0, 0.0], + [0.4, 0.8, 0.4, 0.0, 0.0, 0.0, 0.0], + ] + ) + + blended = blend_data(source1, source2, mask, blending_distance, circular=True) + + np.testing.assert_array_almost_equal(blended, expected) diff --git a/lib/ants/tests/analysis/test_merge.py b/lib/ants/tests/analysis/test_merge.py index 1e2af77..1c6530e 100644 --- a/lib/ants/tests/analysis/test_merge.py +++ b/lib/ants/tests/analysis/test_merge.py @@ -7,6 +7,7 @@ import ants.tests import iris import numpy as np +import pytest from ants.analysis import merge @@ -40,7 +41,42 @@ def test_call_args(self): alternate_cube = self.generate_dummy_cube(shape=(4, 8)) with mock.patch("ants.analysis._merge.merge") as mock_method: - merge(primary_cube, alternate_cube, None) + merge(primary_cube, alternate_cube, None, None) - mock_method.assert_called_once_with(primary_cube, alternate_cube, None) + mock_method.assert_called_once_with(primary_cube, alternate_cube, None, None) self.assertFalse(self.mock_fill.called) + + def test_blending_distance_no_polygon(self): + primary_cube = self.generate_dummy_cube(shape=(4, 8)) + alternate_cube = self.generate_dummy_cube(shape=(4, 8)) + expected_msg = ( + "blending_distance can only be used with a validity_polygon. " + "No polygon was provided, but got blending_distance=1.0" + ) + with pytest.raises(ValueError, match=expected_msg): + merge( + primary_cube, + alternate_cube, + validity_polygon=None, + blending_distance=1.0, + ) + + def test_blending_distance_multi_level(self): + primary_cube = ants.tests.stock.simple_3d_time_varying() + alternate_cube = primary_cube.copy() + + # doesn't matter what the validity polygon is, as long as *something* is passed + validity_polygon = [[0, 0], [1, 1]] + + expected_msg = ( + "Blending is only supported for single level data sources. " + "The primary data source is not single level" + ) + + with pytest.raises(ValueError, match=expected_msg): + merge( + primary_cube, + alternate_cube, + validity_polygon=validity_polygon, + blending_distance=1.0, + ) diff --git a/lib/ants/tests/utils/cube/test_is_single_level.py b/lib/ants/tests/utils/cube/test_is_single_level.py new file mode 100644 index 0000000..87515dd --- /dev/null +++ b/lib/ants/tests/utils/cube/test_is_single_level.py @@ -0,0 +1,34 @@ +# (C) Crown Copyright, Met Office. All rights reserved. +# +# This file is part of ANTS and is released under the BSD 3-Clause license. +# See LICENSE.txt in the root of the repository for full licensing details. +import ants.tests.stock +from ants.utils.cube import is_single_level + + +def test_geodetic(): + cube = ants.tests.stock.geodetic((5, 5)) + assert is_single_level(cube) is True + + +def test_geodetic_transposed(): + cube = ants.tests.stock.geodetic((5, 5)) + cube.transpose() + assert is_single_level(cube) is True + + +def test_simple_4d_with_hybrid_height(): + cube = ants.tests.stock.simple_4d_with_hybrid_height() + assert is_single_level(cube) is False + + +def test_simple_3d_time_varying(): + cube = ants.tests.stock.simple_3d_time_varying() + assert is_single_level(cube) is False + + +def test_time_and_latitude(): + # construct cube to have time and latitude coordinates only + cube = ants.tests.stock.simple_3d_time_varying()[..., 0] + assert cube.ndim == 2 + assert is_single_level(cube) is False diff --git a/lib/ants/utils/cube.py b/lib/ants/utils/cube.py index b47a80e..076e919 100644 --- a/lib/ants/utils/cube.py +++ b/lib/ants/utils/cube.py @@ -1560,3 +1560,24 @@ def fetch_seed_index(cube, seed): xd = abs(x.points - seed[1]).argmin() yd = abs(y.points - seed[0]).argmin() return xd, yd + + +def is_single_level(cube: iris.cube.Cube) -> bool: + """Determine if a cube is defined on a single horizontal level. + + A cube is identified as single level if it is 2-dimensional, and those + dimensions correspond to the x and y axes (in any order). + + Parameters + ---------- + cube: iris.cube.Cube + The cube to check + + Returns + ------- + bool + Whether the cube is defined on a single horizontal level + """ + axes = {iris.util.guess_coord_axis(coord).lower() for coord in cube.dim_coords} + condition = cube.ndim == 2 and axes == {"x", "y"} + return condition