diff --git a/lib/ants/decomposition.py b/lib/ants/decomposition.py index 1fa8246..2502e91 100644 --- a/lib/ants/decomposition.py +++ b/lib/ants/decomposition.py @@ -679,3 +679,493 @@ def _run(self, operation, args): bag = db.from_sequence(parameters) results = bag.starmap(operation).compute() return results + + +# ============================================================================= +# simple_split_and_process framework +# +# A simpler, more maintainable replacement for :func:`decompose`. All private +# helpers in this section are prefixed with ``_ssp_`` to make their ownership +# clear and to avoid confusion with the helpers used by the existing framework. +# ============================================================================= + + +def _ssp_compute_split_tuple_for_cube(cube, number_of_x_splits, number_of_y_splits): + """ + Build a split-count tuple aligned with the cube's dimension order. + + :func:`_mosaic_by_nsplits` requires one split count per cube dimension in + the same order as the cube's shape. This function maps the caller-supplied + *x* and *y* split counts onto the correct positions using the cube's + horizontal coordinate dimensions. + + Parameters + ---------- + cube : :class:`iris.cube.Cube` + The cube whose dimension order is used to place the split counts. + number_of_x_splits : int + How many pieces to split the x (longitude) dimension into. + number_of_y_splits : int + How many pieces to split the y (latitude) dimension into. + + Returns + ------- + tuple of int + A tuple with length equal to ``cube.ndim``. All non-horizontal + dimensions have a split count of 1 (no splitting). + + """ + # Start with 1 for every dimension (meaning "no split in this dimension"). + split_counts = np.ones(cube.ndim, dtype=int) + + x_coordinate, y_coordinate = ants.utils.cube.horizontal_grid(cube) + # coord_dims returns a tuple; for rectilinear grids each coordinate maps to + # exactly one dimension index. + x_dimension_index = cube.coord_dims(x_coordinate)[0] + y_dimension_index = cube.coord_dims(y_coordinate)[0] + + split_counts[x_dimension_index] = number_of_x_splits + split_counts[y_dimension_index] = number_of_y_splits + + return tuple(split_counts) + + +def _ssp_generate_cube_pieces(cube, number_of_x_splits, number_of_y_splits): + """ + Generator that yields sub-cubes produced by splitting the input cube. + + The cube is divided horizontally into ``number_of_x_splits * + number_of_y_splits`` pieces using :func:`_mosaic_by_nsplits`. Coordinate + bounds are guessed on the *full* cube before any splitting takes place; + guessing bounds on individual pieces can produce inconsistent results for + circular (global) coordinates due to floating-point representation of the + grid spacing. + + Parameters + ---------- + cube : :class:`iris.cube.Cube` + The cube to split. Modified in-place to add coordinate bounds if they + are not already present. + number_of_x_splits : int + Number of pieces along the x (longitude) dimension. + number_of_y_splits : int + Number of pieces along the y (latitude) dimension. + + Yields + ------ + :class:`iris.cube.Cube` + Sub-cube pieces in row-major order (y varies slowest). + + """ + # Guess bounds on the FULL cube before splitting. This is critical for + # circular (global) coordinates: guessing bounds on a sub-region can give + # a slightly different cell spacing than the full-resolution guess, which + # would make the pieces inconsistent with each other. + ants.utils.cube.guess_horizontal_bounds(cube) + + split_tuple = _ssp_compute_split_tuple_for_cube( + cube, number_of_x_splits, number_of_y_splits + ) + + for slice_tuple in _mosaic_by_nsplits(cube.shape, split_tuple): + yield cube[slice_tuple] + + +def _ssp_extract_source_region(source, target_piece, padding_cells): + """ + Extract the sub-region of *source* that spatially covers *target_piece*. + + This function is the owned extraction implementation for + :func:`simple_split_and_process`. It is deliberately kept separate from + the existing :func:`ants._constraints._extract_overlap` so that it can be + diagnosed and refined through the test battery without affecting other code. + + The function handles the following cases: + + * **Same coordinate convention** (both -180:180 or both 0:360): a single + contiguous slice of the source is returned. + * **Wrap-around extraction** (source is circular and the required region + straddles the source edge, e.g. UK region from a 0:360 source): two + slice groups are extracted, re-wrapped to be contiguous, and concatenated. + * **Coordinate convention mismatch** (source uses 0:360, target uses + -180:180 or vice versa): after extraction the source x-coordinates are + shifted by a whole number of modulus periods to align with the target's + x-coordinate range, so that iris can perform interpolation correctly. + * **Padding**: the extraction window is expanded by ``padding_cells`` + cells beyond the target bounds in both x and y, providing the stencil + context needed by the interpolation operation. + + Parameters + ---------- + source : :class:`iris.cube.Cube` + The full-resolution source cube. Bounds are added in-place if absent. + target_piece : :class:`iris.cube.Cube` + A sub-region of the target grid that defines the extraction window. + Bounds are added in-place if absent. + padding_cells : int + Number of source cells to include beyond the target bounds on each + side. A value of 1 is usually sufficient for linear interpolation. + Returns + ------- + :class:`iris.cube.Cube` + A sub-cube of *source* whose horizontal extent covers *target_piece* + (plus padding), with x-coordinates normalised to be compatible with + the target's coordinate range. + + Raises + ------ + RuntimeError + If :func:`ants.utils.cube.get_slices` returns more than two slice + groups (more than one wrap-around boundary is not currently supported). + + """ + # Ensure coordinate bounds are present on both cubes. We use BOUNDS (not + # coordinate points) to define the extraction window: this correctly + # captures source cells at the exact edges of target tile boundaries, + # where the cell centre may not lie within the target range but the cell + # edge does. + ants.utils.cube.guess_horizontal_bounds(source) + ants.utils.cube.guess_horizontal_bounds(target_piece) + + source_x_coordinate, source_y_coordinate = ants.utils.cube.horizontal_grid(source) + target_x_coordinate, target_y_coordinate = ants.utils.cube.horizontal_grid( + target_piece + ) + + # Determine the bounding box of the target piece using its coordinate + # BOUNDS, not its centre points. + target_x_minimum = float(np.min(target_x_coordinate.bounds)) + target_x_maximum = float(np.max(target_x_coordinate.bounds)) + target_y_minimum = float(np.min(target_y_coordinate.bounds)) + target_y_maximum = float(np.max(target_y_coordinate.bounds)) + + # Use ants.utils.cube.get_slices to identify which source cells overlap + # the target bounding box. get_slices is wrap-around aware: when the + # extraction window straddles the source edge it returns two slice groups + # (one for each side of the wrap-around boundary). + source_overlap_slices = ants.utils.cube.get_slices( + source, + ylim=[target_y_minimum, target_y_maximum], + xlim=[target_x_minimum, target_x_maximum], + pad_width=padding_cells, + ) + + if len(source_overlap_slices) > 2: + raise RuntimeError( + f"Source region extraction returned {len(source_overlap_slices)} " + "slice groups; at most 2 are supported (one wrap-around boundary)." + ) + + # Identify the dimension index for the x-coordinate so we can inspect + # individual slice objects from get_slices. + x_dimension_index = source.coord_dims(source_x_coordinate)[0] + + # Extract a sub-cube for each slice group. + extracted_cube_pieces = iris.cube.CubeList( + [source[slice_tuple] for slice_tuple in source_overlap_slices] + ) + + if len(extracted_cube_pieces) == 2: + # Two pieces arise when the extraction window wraps around the source + # edge (e.g. UK target from a 0:360 source crosses the 0/360 boundary). + # The two pieces have x-coordinates in disjoint ranges; we must + # re-wrap them to a common base so they are contiguous and can be + # concatenated by iris. + x_modulus = getattr(source_x_coordinate.units, "modulus", None) + if x_modulus is not None: + # Choose the contiguity base as the minimum x-bound of the + # highest-indexed slice group. The highest index corresponds to + # the "left" edge of the full extraction window in the periodic + # domain (e.g. the 350:360 cells when extracting the UK from a + # 0:360 source). + highest_start_index = max( + slice_tuple[x_dimension_index].start + for slice_tuple in source_overlap_slices + ) + x_contiguity_base = float( + source_x_coordinate.bounds[highest_start_index].min() + ) + + for piece in extracted_cube_pieces: + piece_x_coordinate = piece.coord(axis="x") + piece_x_coordinate.points = ants.utils.ndarray.wrap_lons( + piece_x_coordinate.points, x_contiguity_base, x_modulus + ) + piece_x_coordinate.bounds = ants.utils.ndarray.wrap_lons( + piece_x_coordinate.bounds, x_contiguity_base, x_modulus + ) + + # Sort pieces by their first x-coordinate value so that iris + # concatenate_cube sees them in ascending order. + extracted_cube_pieces = iris.cube.CubeList( + sorted( + extracted_cube_pieces, + key=lambda piece: piece.coord(axis="x").points[0], + ) + ) + + extracted_cube = extracted_cube_pieces.concatenate_cube() + + # Normalise the extracted source x-coordinates to be compatible with the + # target's x-coordinate range. + # + # When source and target use different longitude conventions (0:360 vs + # -180:180), the extracted source cells may have x-values like 347-367 + # while the target has x-values like -10 to 3. Without this step iris + # cannot find the source cells that bracket the target grid points. + # + # The fix: shift the extracted source x-coordinates by a whole number of + # modulus periods until the source centre-of-mass aligns with the target + # centre-of-mass. A whole-period shift is transparent to interpolation. + extracted_x_coordinate = extracted_cube.coord(axis="x") + x_modulus = getattr(extracted_x_coordinate.units, "modulus", None) + if x_modulus is not None: + extracted_x_centre = float(np.mean(extracted_x_coordinate.points)) + target_x_centre = (target_x_minimum + target_x_maximum) / 2.0 + number_of_periods_to_shift = round( + (target_x_centre - extracted_x_centre) / x_modulus + ) + if number_of_periods_to_shift != 0: + coordinate_shift = number_of_periods_to_shift * x_modulus + extracted_x_coordinate.points = ( + extracted_x_coordinate.points + coordinate_shift + ) + extracted_x_coordinate.bounds = ( + extracted_x_coordinate.bounds + coordinate_shift + ) + + ants.utils.cube.derive_circular_status(extracted_cube) + return extracted_cube + + +def _ssp_concatenate_result_pieces(result_pieces): + """ + Reassemble a list of per-piece operation results into a single cube. + + The pieces are expected to be spatially adjacent and non-overlapping, as + produced by applying an operation to the pieces generated by + :func:`_ssp_generate_cube_pieces`. + + After concatenation: + + * The circular attribute of the x-coordinate is re-derived, since it may + have been lost when the cube was split into regional pieces. + * The dtype of the result is checked against the first piece; if + concatenation has silently changed it (which should not happen but is + guarded against as a known past source of bugs) an explicit cast is + applied. + + Parameters + ---------- + result_pieces : list of :class:`iris.cube.Cube` + Ordered list of result cubes to assemble. Must be non-empty. + + Returns + ------- + :class:`iris.cube.Cube` + The fully reassembled result cube. + + """ + # Realise lazy data on every piece before inspecting dtype. Iris may + # return a cube whose .dtype property reflects the lazy (dask) graph's + # input dtype (e.g. int32 from the source array) rather than the dtype + # of the computed output. Accessing .data forces computation and ensures + # that .dtype subsequently returns the true output dtype. + for piece in result_pieces: + _ = piece.data # noqa: F841 — side-effect: realises lazy computation + + # Iris concatenate_cube requires all pieces to share the same dtype. + # Even after realisation, pieces may differ (e.g. some regrid pieces + # return float64, others int32 when source and target grids coincide + # exactly at some tiles). Promote all pieces to the common dtype before + # concatenating. + common_dtype = np.result_type(*[piece.dtype for piece in result_pieces]) + for piece in result_pieces: + if piece.dtype != common_dtype: + piece.data = piece.data.astype(common_dtype) + + assembled_cube = iris.cube.CubeList(result_pieces).concatenate_cube() + + # Restore the circular attribute on the x-coordinate where appropriate. + # The attribute is lost when a global cube is split into regional pieces + # (each piece is not global, so derive_circular_status removes it), but + # after reassembly the full extent is restored. + ants.utils.cube.derive_circular_status(assembled_cube) + + # Guard against silent dtype changes during concatenation. The iris + # concatenate operation should preserve dtype, but this explicit check + # ensures we catch any regression here rather than propagating an + # unexpected dtype change to the caller. + if assembled_cube.dtype != common_dtype: + assembled_cube.data = assembled_cube.data.astype(common_dtype) + + return assembled_cube + + +def _ssp_conform_result_piece_to_target_grid(result_piece, target_piece): + """ + Ensure a binary-operation result piece has the same grid shape as target. + + Some operations (notably area-weighted regridding) may collapse singleton + horizontal dimensions for small tiles (for example, returning shape + ``(nx,)`` instead of ``(1, nx)``). This prevents concatenation with + neighbouring tiles that retain both horizontal dimensions. + + This helper reshapes such results onto the target piece grid while + preserving the result metadata. + + Parameters + ---------- + result_piece : :class:`iris.cube.Cube` + Result cube returned by the binary operation. + target_piece : :class:`iris.cube.Cube` + Target tile used for the operation. + + Returns + ------- + :class:`iris.cube.Cube` + Result cube guaranteed to have ``target_piece.shape``. + + Raises + ------ + RuntimeError + If the result data cannot be reshaped onto ``target_piece.shape``. + + """ + if result_piece.shape == target_piece.shape: + return result_piece + + result_data = np.asarray(result_piece.data) + expected_size = int(np.prod(target_piece.shape)) + if result_data.size != expected_size: + raise RuntimeError( + "Binary operation result shape is incompatible with target piece: " + f"result shape {result_piece.shape}, target shape {target_piece.shape}." + ) + + reshaped_result_data = result_data.reshape(target_piece.shape) + conformed_result_piece = target_piece.copy(reshaped_result_data) + conformed_result_piece.metadata = result_piece.metadata + return conformed_result_piece + + +def simple_split_and_process( + operation, + source, + target=None, + number_of_x_splits=0, + number_of_y_splits=0, + padding_cells=1, +): + """ + Apply *operation* to *source* (and optionally *target*), optionally + splitting the data into smaller horizontal pieces first. + + This function is a simpler, more maintainable replacement for + :func:`decompose`. Key differences: + + * Split counts are passed directly as arguments rather than being read + from global configuration. + * No multiprocessing: all pieces are processed serially in a single + process. + * No temporary files or deferred data: results are kept in memory. + * Operates on single :class:`iris.cube.Cube` objects rather than + :class:`iris.cube.CubeList`. + + **Unary mode** (``target=None``): the *source* is split into pieces, the + *operation* is applied to each piece independently, and the pieces are + reassembled. + + **Binary mode** (``target`` provided): the *target* is split into pieces. + For each target piece the overlapping region of *source* is extracted + (plus ``padding_cells`` of context), the *operation* is applied to each + ``(source_piece, target_piece)`` pair, and the results are reassembled. + Splitting the *target* (rather than the source) ensures that the result + is identical to the no-split case: each target cell is interpolated from + exactly the same source cells regardless of how the decomposition is + configured. + + When ``number_of_x_splits == 0`` and ``number_of_y_splits == 0``, the + *operation* is applied to the full datasets without any splitting. + + Known correctness considerations (see private helpers for details): + + * Coordinate bounds are guessed on the full cubes *before* splitting to + avoid inconsistencies at circular/global boundaries. + * Source extraction uses the target piece *bounds* (not centre points) to + define the extraction window. + * Source x-coordinates are normalised after extraction to be compatible + with the target's x-coordinate range (handles 0:360 vs -180:180). + * Wrapping of circular (global) sources across the 0/360 or ±180 boundary + is handled by :func:`_ssp_extract_source_region`. + + Parameters + ---------- + operation : callable + The function to apply. Must have the signature + ``operation(source) -> result`` for unary mode or + ``operation(source, target) -> result`` for binary mode. + source : :class:`iris.cube.Cube` + The source dataset. + target : :class:`iris.cube.Cube`, optional + The target grid cube for binary operations. When provided, the + *target* is split into pieces and the *source* is subsetted to match + each piece. + number_of_x_splits : int, optional + Number of pieces to divide the x (longitude) dimension into. + Default is 0 (no splitting). + number_of_y_splits : int, optional + Number of pieces to divide the y (latitude) dimension into. + Default is 0 (no splitting). + padding_cells : int, optional + Number of source cells to include beyond the boundary of each target + piece during extraction. Provides the interpolation stencil context + needed at tile edges. Default is 1. + Returns + ------- + :class:`iris.cube.Cube` + The result of applying *operation* to the (optionally decomposed) + datasets. + + """ + # With zero splits in both dimensions, apply the operation directly to the + # full datasets without any decomposition overhead. + if number_of_x_splits == 0 and number_of_y_splits == 0: + if target is not None: + return operation(source, target) + else: + return operation(source) + + # Guess coordinate bounds on the full cubes before splitting. This must + # be done here (on the intact cubes) rather than inside the splitting + # loop: guessing bounds on a sub-region can produce slightly different + # spacing for circular grids, making the pieces inconsistent. + ants.utils.cube.guess_horizontal_bounds(source) + if target is not None: + ants.utils.cube.guess_horizontal_bounds(target) + + result_pieces = [] + + if target is not None: + # Binary path: iterate over target pieces and pair each with the + # overlapping region of the source. + for target_piece in _ssp_generate_cube_pieces( + target, number_of_x_splits, number_of_y_splits + ): + source_piece = _ssp_extract_source_region( + source, target_piece, padding_cells + ) + result_piece = operation(source_piece, target_piece) + result_piece = _ssp_conform_result_piece_to_target_grid( + result_piece, target_piece + ) + result_pieces.append(result_piece) + else: + # Unary path: iterate over source pieces and apply operation to each. + for source_piece in _ssp_generate_cube_pieces( + source, number_of_x_splits, number_of_y_splits + ): + result_piece = operation(source_piece) + result_pieces.append(result_piece) + + return _ssp_concatenate_result_pieces(result_pieces) diff --git a/lib/ants/tests/decomposition/test_simple_split_and_process.py b/lib/ants/tests/decomposition/test_simple_split_and_process.py new file mode 100644 index 0000000..ebfaa38 --- /dev/null +++ b/lib/ants/tests/decomposition/test_simple_split_and_process.py @@ -0,0 +1,577 @@ +# (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. +""" +Test battery for :func:`ants.decomposition.simple_split_and_process`. + +These tests verify the fundamental correctness invariant: decomposing data into +N pieces and applying an operation to those pieces yields the same result as +applying the operation to the full dataset in a single pass (0 splits). + +The test matrix covers: + + * Unary operations (single source, no target). + * Binary operations (source + target) across a range of geographic domains: + + - Global coarse source -> global finer target. + - Global coarse source -> UK regional target (straddles Greenwich Meridian). + - Global coarse source -> Australia regional target (southern hemisphere). + - Global coarse source -> New Zealand regional target (near International Dateline). + - Global coarse source -> Singapore regional target (small equatorial domain). + - Global coarse source -> Northern Greenland regional target (high latitude). + - Global source with 0:360 longitude -> UK target with -180:180 longitude + (tests coordinate convention mismatch handling). + +Each test is parametrized over four split configurations: + + * (2, 2): 2 x-splits and 2 y-splits. + * (3, 3): 3 x-splits and 3 y-splits. + * (1, 4): 4 y-splits only (no x-splits). + * (4, 1): 4 x-splits only (no y-splits). + +The result with 0 splits (no decomposition) is used as the ground truth. +""" + +import numpy +import numpy.testing +import pytest +import iris +import iris.analysis + +import ants.decomposition +import ants.tests.stock as stock + + +# --------------------------------------------------------------------------- +# Split configurations to parametrize over: (number_of_x_splits, number_of_y_splits) +# The id string is used in the pytest output to identify each configuration. +# --------------------------------------------------------------------------- +SPLIT_CONFIGURATIONS = [ + pytest.param(2, 2, id="2x2_splits"), + pytest.param(3, 3, id="3x3_splits"), + pytest.param(1, 4, id="1x4_splits"), + pytest.param(4, 1, id="4x1_splits"), +] + +# Mixed longitude convention decomposition (for example, 0:360 source and +# -180:180 target) is expected to agree with the non-decomposed reference up +# to machine precision. We enforce an explicit absolute tolerance here to +# avoid brittle exact-equality expectations in this numerically sensitive path. +MIXED_CONVENTION_ABSOLUTE_TOLERANCE = 1e-12 + + +# --------------------------------------------------------------------------- +# Stock cube factory functions +# +# Source cubes are deliberately coarser than target cubes so that the linear +# regrid operation is a genuine interpolation rather than a trivial identity. +# All shapes are kept small (single-digit or low tens of cells per dimension) +# to keep the test suite fast. +# --------------------------------------------------------------------------- + +def make_global_coarse_source(): + """ + Create a global coarse-resolution source cube using the -180:180 longitude + convention. + + The cube covers the full globe at approximately 10-degree resolution + (18 rows x 36 columns). Data values are sequential integers, suitable + for regridding tests. + """ + return stock.geodetic(shape=(18, 36)) + + +def make_global_coarse_source_0_to_360(): + """ + Create a global coarse-resolution source cube using the 0:360 longitude + convention. + + Identical to :func:`make_global_coarse_source` except the x-coordinate + spans 0 to 360 rather than -180 to 180. Used to exercise coordinate + convention mismatch between source and target during extraction. + """ + return stock.geodetic(shape=(18, 36), xlim=(0, 360)) + + +def make_global_fine_target(): + """ + Create a global fine-resolution target cube. + + The cube covers the full globe at approximately 5-degree resolution + (9 rows x 18 columns), finer than the 10-degree source but still small + enough for fast tests. + """ + return stock.geodetic(shape=(9, 18)) + + +def make_uk_target(): + """ + Create a target cube covering approximately the United Kingdom. + + The UK domain (49-61 N, -10 to 3 E) straddles the Greenwich Meridian. + A 0:360-convention source must wrap around the 0/360 boundary to cover + this region, making this an important test of extraction near zero longitude. + """ + return stock.geodetic(shape=(6, 7), ylim=(49, 61), xlim=(-10, 3)) + + +def make_australia_target(): + """ + Create a target cube covering approximately Australia. + + The domain (-45 to -10 N, 110 to 155 E) lies entirely in the southern + hemisphere, well away from the wrap-around boundaries. Both -180:180 and + 0:360 sources should yield this region without any wrapping. + """ + return stock.geodetic(shape=(7, 9), ylim=(-45, -10), xlim=(110, 155)) + + +def make_new_zealand_target(): + """ + Create a target cube covering approximately New Zealand. + + The domain (-48 to -33 N, 165 to 180 E) lies near the International + Dateline at 180 degrees east. This tests that extraction near the eastern + edge of the common -180:180 longitude range is handled correctly. + """ + return stock.geodetic(shape=(6, 6), ylim=(-48, -33), xlim=(165, 180)) + + +def make_singapore_target(): + """ + Create a target cube covering approximately the Singapore region. + + The domain (1 to 2 N, 103 to 105 E) is a small equatorial region. Used + to verify that decomposition works correctly for small regional domains + where individual split pieces may contain only 1 or 2 target grid cells. + """ + return stock.geodetic(shape=(4, 4), ylim=(1, 2), xlim=(103, 105)) + + +def make_northern_greenland_target(): + """ + Create a target cube covering approximately northern Greenland. + + The domain (75 to 85 N, -75 to -15 E) is at high latitude. Used to + verify that y-splitting behaves correctly in polar regions where grid + cells are narrow in the x-direction relative to their geographic extent. + """ + return stock.geodetic(shape=(4, 12), ylim=(75, 85), xlim=(-75, -15)) + + +# --------------------------------------------------------------------------- +# Operations used in tests +# --------------------------------------------------------------------------- + +def add_one_to_source(source): + """Unary operation: add 1.0 to every data value in the source cube.""" + return source + 1 + + +def regrid_source_to_target(source, target): + """ + Binary operation: regrid the source cube onto the target grid using + iris linear interpolation. + + :func:`iris.analysis.Linear` is used as a representative, well-understood + binary operation for validating the decomposition framework. + """ + return source.regrid(target, iris.analysis.Linear()) + + +def regrid_source_to_target_areaweighted(source, target): + """ + Binary operation: regrid the source cube onto the target grid using + iris area-weighted interpolation. + + This operation is sensitive to small target tiles that can collapse a + singleton horizontal dimension in intermediate pieces, so it is used to + validate decomposition piece-shape conformity before concatenation. + """ + return source.regrid(target, iris.analysis.AreaWeighted()) + + +# --------------------------------------------------------------------------- +# Assertion helpers +# --------------------------------------------------------------------------- + +def assert_decomposed_result_matches_reference(actual, expected): + """ + Assert that a decomposed result matches the reference (0-splits) result. + + Checks: + + * Data values are numerically equal within floating-point tolerance. + * Data dtype is identical. + * Cube metadata (name, units, coordinates) is identical. + + Parameters + ---------- + actual : :class:`iris.cube.Cube` + Result obtained by running with N splits. + expected : :class:`iris.cube.Cube` + Reference result obtained by running with 0 splits. + """ + numpy.testing.assert_array_almost_equal( + actual.data, + expected.data, + err_msg=( + "Data values differ between the decomposed result and the " + "reference (0-splits) result." + ), + ) + assert actual.data.dtype == expected.data.dtype, ( + f"dtype mismatch: decomposed result has dtype {actual.data.dtype!r}, " + f"but reference result has dtype {expected.data.dtype!r}." + ) + assert actual.metadata == expected.metadata, ( + f"Cube metadata mismatch:\n" + f" actual: {actual.metadata}\n" + f" expected: {expected.metadata}" + ) + + +# --------------------------------------------------------------------------- +# Test classes +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("number_of_x_splits,number_of_y_splits", SPLIT_CONFIGURATIONS) +class TestUnaryOperation: + """ + Tests that unary operations produce identical results with and without + splitting. + + A unary operation takes only a source cube and returns a result cube. + The source is split into pieces, the operation applied to each piece, and + the pieces reassembled. + """ + + def test_global_domain(self, number_of_x_splits, number_of_y_splits): + """ + Unary operation on a global coarse-resolution source should give the + same result whether run with 0 splits or N splits. + """ + source = make_global_coarse_source() + reference_result = ants.decomposition.simple_split_and_process( + add_one_to_source, + source, + ) + decomposed_result = ants.decomposition.simple_split_and_process( + add_one_to_source, + source, + number_of_x_splits=number_of_x_splits, + number_of_y_splits=number_of_y_splits, + ) + assert_decomposed_result_matches_reference(decomposed_result, reference_result) + + +@pytest.mark.parametrize("number_of_x_splits,number_of_y_splits", SPLIT_CONFIGURATIONS) +class TestBinaryOperation: + """ + Tests that binary operations (source regridded to target) produce identical + results with and without splitting, across a range of geographic domains. + + A binary operation takes a source cube and a target cube. In the + decomposed case the *target* is split into pieces, the source region + overlapping each target piece is extracted, the operation applied to each + (source piece, target piece) pair, and the results reassembled. + """ + + def _run_binary_test(self, source, target, number_of_x_splits, number_of_y_splits): + """ + Run the binary operation with 0 splits (reference) and with the given + split configuration, returning both results. + + Parameters + ---------- + source : :class:`iris.cube.Cube` + target : :class:`iris.cube.Cube` + number_of_x_splits : int + number_of_y_splits : int + + Returns + ------- + decomposed_result : :class:`iris.cube.Cube` + reference_result : :class:`iris.cube.Cube` + """ + reference_result = ants.decomposition.simple_split_and_process( + regrid_source_to_target, + source, + target=target, + ) + decomposed_result = ants.decomposition.simple_split_and_process( + regrid_source_to_target, + source, + target=target, + number_of_x_splits=number_of_x_splits, + number_of_y_splits=number_of_y_splits, + ) + return decomposed_result, reference_result + + def test_global_source_global_fine_target( + self, number_of_x_splits, number_of_y_splits + ): + """ + Binary operation: global coarse source -> global fine target. + + Tests that the decomposition preserves results for a straightforward + global-to-global regrid with no regional or wrap-around complications. + """ + source = make_global_coarse_source() + target = make_global_fine_target() + decomposed, reference = self._run_binary_test( + source, target, number_of_x_splits, number_of_y_splits + ) + assert_decomposed_result_matches_reference(decomposed, reference) + + + def test_global_source_uk_target(self, number_of_x_splits, number_of_y_splits): + """ + Binary operation: global coarse source -> UK regional target. + + The UK domain straddles the Greenwich Meridian (-10 to 3 E). + A -180:180 source must provide cells from both sides of 0 degrees. + """ + source = make_global_coarse_source() + target = make_uk_target() + decomposed, reference = self._run_binary_test( + source, target, number_of_x_splits, number_of_y_splits + ) + assert_decomposed_result_matches_reference(decomposed, reference) + def test_global_source_australia_target( + self, number_of_x_splits, number_of_y_splits + ): + """ + Binary operation: global coarse source -> Australia regional target. + + The domain is in the southern hemisphere (negative latitudes) and + entirely within a single contiguous longitude range with no wrapping + required. + """ + source = make_global_coarse_source() + target = make_australia_target() + decomposed, reference = self._run_binary_test( + source, target, number_of_x_splits, number_of_y_splits + ) + assert_decomposed_result_matches_reference(decomposed, reference) + def test_global_source_new_zealand_target( + self, number_of_x_splits, number_of_y_splits + ): + """ + Binary operation: global coarse source -> New Zealand target. + + The domain sits near the International Dateline (165 to 180 E). + This tests correct handling of the eastern edge of the standard + -180:180 longitude range. + """ + source = make_global_coarse_source() + target = make_new_zealand_target() + decomposed, reference = self._run_binary_test( + source, target, number_of_x_splits, number_of_y_splits + ) + assert_decomposed_result_matches_reference(decomposed, reference) + + def test_global_source_singapore_target( + self, number_of_x_splits, number_of_y_splits + ): + """ + Binary operation: global coarse source -> Singapore regional target. + + A small (4x4 cell) equatorial domain. With larger split counts some + individual target pieces will contain only 1 or 2 cells, exercising + the behaviour of iris linear regrid on very small sub-domains. + """ + source = make_global_coarse_source() + target = make_singapore_target() + decomposed, reference = self._run_binary_test( + source, target, number_of_x_splits, number_of_y_splits + ) + assert_decomposed_result_matches_reference(decomposed, reference) + + def test_global_source_northern_greenland_target( + self, number_of_x_splits, number_of_y_splits + ): + """ + Binary operation: global coarse source -> Northern Greenland target. + + A high-latitude domain (75 to 85 N). Tests that y-splitting near the + poles does not introduce numerical differences. + """ + source = make_global_coarse_source() + target = make_northern_greenland_target() + decomposed, reference = self._run_binary_test( + source, target, number_of_x_splits, number_of_y_splits + ) + assert_decomposed_result_matches_reference(decomposed, reference) + + def test_global_0_to_360_source_regional_target( + self, number_of_x_splits, number_of_y_splits + ): + """ + Binary operation: global 0:360 longitude source -> UK -180:180 target. + + Tests that the source region extraction correctly handles the case + where source and target use different x-coordinate conventions. The + UK target (−10 to 3 E) requires the 0:360 source to provide cells + from near 350-360 degrees, which must be re-mapped to the target's + negative-longitude range for iris to interpolate correctly. + """ + source = make_global_coarse_source_0_to_360() + target = make_uk_target() + decomposed, reference = self._run_binary_test( + source, target, number_of_x_splits, number_of_y_splits + ) + numpy.testing.assert_allclose( + decomposed.data, + reference.data, + rtol=0.0, + atol=MIXED_CONVENTION_ABSOLUTE_TOLERANCE, + err_msg=( + "Mixed longitude convention decomposition exceeded expected " + "machine-precision tolerance relative to the non-decomposed " + "reference." + ), + ) + assert decomposed.data.dtype == reference.data.dtype, ( + f"dtype mismatch: decomposed result has dtype {decomposed.data.dtype!r}, " + f"but reference result has dtype {reference.data.dtype!r}." + ) + assert decomposed.metadata == reference.metadata, ( + f"Cube metadata mismatch:\n" + f" decomposed: {decomposed.metadata}\n" + f" reference: {reference.metadata}" + ) + + +class TestConcatenateResultPiecesRegression: + """Regression tests for _ssp_concatenate_result_pieces dtype handling.""" + + def test_mixed_piece_dtypes_promoted_to_common_dtype(self): + """ + Mixed integer/float piece dtypes are promoted before concatenation. + + This protects against regressions in the final dtype guard path used + when assembling binary operation piece results. + """ + cube = stock.geodetic(shape=(6, 8)) + piece_left = cube[:, :4].copy() + piece_right = cube[:, 4:].copy() + + piece_left.data = piece_left.data.astype(numpy.int32) + piece_right.data = piece_right.data.astype(numpy.float64) + + expected_dtype = numpy.result_type(piece_left.dtype, piece_right.dtype) + assembled = ants.decomposition._ssp_concatenate_result_pieces( + [piece_left, piece_right] + ) + + assert assembled.dtype == expected_dtype + numpy.testing.assert_array_equal( + assembled.data, + cube.data.astype(expected_dtype), + ) + + +@pytest.mark.parametrize("number_of_x_splits,number_of_y_splits", SPLIT_CONFIGURATIONS) +class TestBinaryOperationAreaWeighted: + """ + Tests that area-weighted binary operations produce stable decomposed + results across supported domain/split combinations. + """ + + def _run_binary_test_areaweighted( + self, source, target, number_of_x_splits, number_of_y_splits + ): + reference_result = ants.decomposition.simple_split_and_process( + regrid_source_to_target_areaweighted, + source, + target=target, + ) + decomposed_result = ants.decomposition.simple_split_and_process( + regrid_source_to_target_areaweighted, + source, + target=target, + number_of_x_splits=number_of_x_splits, + number_of_y_splits=number_of_y_splits, + ) + return decomposed_result, reference_result + + def test_global_source_global_fine_target( + self, number_of_x_splits, number_of_y_splits + ): + source = make_global_coarse_source() + target = make_global_fine_target() + decomposed, reference = self._run_binary_test_areaweighted( + source, target, number_of_x_splits, number_of_y_splits + ) + assert_decomposed_result_matches_reference(decomposed, reference) + + def test_global_source_uk_target(self, number_of_x_splits, number_of_y_splits): + source = make_global_coarse_source() + target = make_uk_target() + decomposed, reference = self._run_binary_test_areaweighted( + source, target, number_of_x_splits, number_of_y_splits + ) + assert_decomposed_result_matches_reference(decomposed, reference) + + def test_global_source_australia_target( + self, number_of_x_splits, number_of_y_splits + ): + source = make_global_coarse_source() + target = make_australia_target() + decomposed, reference = self._run_binary_test_areaweighted( + source, target, number_of_x_splits, number_of_y_splits + ) + assert_decomposed_result_matches_reference(decomposed, reference) + + def test_global_source_new_zealand_target( + self, number_of_x_splits, number_of_y_splits + ): + source = make_global_coarse_source() + target = make_new_zealand_target() + decomposed, reference = self._run_binary_test_areaweighted( + source, target, number_of_x_splits, number_of_y_splits + ) + assert_decomposed_result_matches_reference(decomposed, reference) + + def test_global_source_singapore_target( + self, number_of_x_splits, number_of_y_splits + ): + source = make_global_coarse_source() + target = make_singapore_target() + decomposed, reference = self._run_binary_test_areaweighted( + source, target, number_of_x_splits, number_of_y_splits + ) + assert_decomposed_result_matches_reference(decomposed, reference) + + def test_global_source_northern_greenland_target( + self, number_of_x_splits, number_of_y_splits + ): + source = make_global_coarse_source() + target = make_northern_greenland_target() + decomposed, reference = self._run_binary_test_areaweighted( + source, target, number_of_x_splits, number_of_y_splits + ) + assert_decomposed_result_matches_reference(decomposed, reference) + + def test_global_0_to_360_source_regional_target( + self, number_of_x_splits, number_of_y_splits + ): + source = make_global_coarse_source_0_to_360() + target = make_uk_target() + decomposed, reference = self._run_binary_test_areaweighted( + source, target, number_of_x_splits, number_of_y_splits + ) + numpy.testing.assert_allclose( + decomposed.data, + reference.data, + rtol=0.0, + atol=MIXED_CONVENTION_ABSOLUTE_TOLERANCE, + err_msg=( + "Mixed longitude convention area-weighted decomposition " + "exceeded expected machine-precision tolerance." + ), + ) + assert decomposed.data.dtype == reference.data.dtype + assert decomposed.metadata == reference.metadata