diff --git a/cf_xarray/helpers.py b/cf_xarray/helpers.py index 94b024a7..436f5ebc 100644 --- a/cf_xarray/helpers.py +++ b/cf_xarray/helpers.py @@ -175,19 +175,79 @@ def bounds_to_vertices( f"Bounds format not understood. Got {bounds.dims} with shape {bounds.shape}." ) + core_dim_coords = { + dim: bounds.coords[dim].values for dim in core_dims if dim in bounds.coords + } + core_dim_orders = _get_core_dim_orders(core_dim_coords) + return xr.apply_ufunc( _bounds_helper, bounds, input_core_dims=[core_dims + [bounds_dim]], dask="parallelized", - kwargs={"n_core_dims": n_core_dims, "nbounds": nbounds, "order": order}, + kwargs={ + "n_core_dims": n_core_dims, + "nbounds": nbounds, + "order": order, + "core_dim_orders": core_dim_orders, + }, output_core_dims=[output_core_dims], dask_gufunc_kwargs=dict(output_sizes=output_sizes), output_dtypes=[bounds.dtype], ) -def _bounds_helper(values, n_core_dims, nbounds, order): +def _get_core_dim_orders(core_dim_coords: dict[str, np.ndarray]) -> dict[str, str]: + """ + Determine the order (ascending, descending, or mixed) of each core dimension + based on its coordinates. + + Repeated (equal) coordinates are ignored when determining the order. If all + coordinates are equal, the order is treated as "ascending". + + Parameters + ---------- + core_dim_coords : dict of str to np.ndarray + A dictionary mapping dimension names to their coordinate arrays. + + Returns + ------- + core_dim_orders : dict of str to str + A dictionary mapping each dimension name to a string indicating the order: + - "ascending": strictly increasing (ignoring repeated values) + - "descending": strictly decreasing (ignoring repeated values) + - "mixed": neither strictly increasing nor decreasing (ignoring repeated values) + """ + core_dim_orders = {} + + for dim, coords in core_dim_coords.items(): + diffs = np.diff(coords) + + # Handle datetime64 and timedelta64 safely for both numpy 1.26.4 and numpy 2 + if np.issubdtype(coords.dtype, np.datetime64) or np.issubdtype( + coords.dtype, np.timedelta64 + ): + # Cast to float64 for safe comparison + diffs_float = diffs.astype("float64") + nonzero_diffs = diffs_float[diffs_float != 0] + else: + zero = 0 + nonzero_diffs = diffs[diffs != zero] + + if nonzero_diffs.size == 0: + # All values are equal, treat as ascending + core_dim_orders[dim] = "ascending" + elif np.all(nonzero_diffs > 0): + core_dim_orders[dim] = "ascending" + elif np.all(nonzero_diffs < 0): + core_dim_orders[dim] = "descending" + else: + core_dim_orders[dim] = "mixed" + + return core_dim_orders + + +def _bounds_helper(values, n_core_dims, nbounds, order, core_dim_orders): if n_core_dims == 2 and nbounds == 4: # Vertices case (2D lat/lon) if order in ["counterclockwise", None]: @@ -211,11 +271,176 @@ def _bounds_helper(values, n_core_dims, nbounds, order): vertex_vals = np.block([[bot_left, bot_right], [top_left, top_right]]) elif n_core_dims == 1 and nbounds == 2: # Middle points case (1D lat/lon) - vertex_vals = np.concatenate((values[..., :, 0], values[..., -1:, 1]), axis=-1) + vertex_vals = _get_ordered_vertices(values, core_dim_orders) return vertex_vals +def _get_ordered_vertices( + bounds: np.ndarray, core_dim_orders: dict[str, str] +) -> np.ndarray: + """ + Convert a bounds array of shape (..., N, 2) or (N, 2) into a 1D array of vertices. + + This function reconstructs the vertices from a bounds array, handling both + monotonic and non-monotonic cases. + + Monotonic bounds (all values strictly increase or decrease when flattened): + - Concatenate the left endpoints (bounds[..., :, 0]) with the last right + endpoint (bounds[..., -1, 1]) to form the vertices. + + Non-monotonic bounds: + - Determine the order of the core dimension(s) ('ascending' or 'descending'). + - For ascending order: + - Use the minimum of each interval as the vertex. + - Use the maximum of the last interval as the final vertex. + - For descending order: + - Use the maximum of each interval as the vertex. + - Use the minimum of the last interval as the final vertex. + - Vertices are then sorted to match the coordinate direction. + + Features: + - Handles both ascending and descending bounds. + - Preserves repeated coordinates if present. + - Output shape is (..., N+1) or (N+1,). + + Parameters + ---------- + bounds : np.ndarray + Array of bounds, typically with shape (N, 2) or (..., N, 2). + core_dim_orders : dict[str, str] + Dictionary mapping core dimension names to their order ('ascending' or + 'descending'). Used for sorting the vertices. + + Returns + ------- + np.ndarray + Array of vertices with shape (..., N+1) or (N+1,). + """ + order = _get_order_of_core_dims(core_dim_orders) + + if _is_bounds_monotonic(bounds): + vertices = np.concatenate((bounds[..., :, 0], bounds[..., -1:, 1]), axis=-1) + else: + if order == "ascending": + endpoints = np.minimum(bounds[..., :, 0], bounds[..., :, 1]) + last_endpoint = np.maximum(bounds[..., -1, 0], bounds[..., -1, 1]) + elif order == "descending": + endpoints = np.maximum(bounds[..., :, 0], bounds[..., :, 1]) + last_endpoint = np.minimum(bounds[..., -1, 0], bounds[..., -1, 1]) + + vertices = np.concatenate( + [endpoints, np.expand_dims(last_endpoint, axis=-1)], axis=-1 + ) + + vertices = _sort_vertices(vertices, order) + + return vertices + + +def _is_bounds_monotonic(bounds: np.ndarray) -> bool: + """Check if the bounds are monotonic. + + Arrays are monotonic if all values are increasing or decreasing. This + functions ignores an intervals where consecutive values are equal, which + represent repeated coordinates. + + Parameters + ---------- + arr : np.ndarray + Numpy array to check, typically with shape (..., N, 2). + + Returns + ------- + bool + True if the flattened array is increasing or decreasing, False otherwise. + """ + # NOTE: Python 3.10 uses numpy 1.26.4. If the input is a datetime64 array, + # numpy 1.26.4 may raise: numpy.core._exceptions._UFuncInputCastingError: + # Cannot cast ufunc 'greater' input 0 from dtype(' 0) or np.all(nonzero_diffs < 0)) + + +def _get_order_of_core_dims(core_dim_orders: dict[str, str]) -> str: + """ + Determines the common order of core dimensions from a dictionary of + dimension orders. + + Parameters + ---------- + core_dim_orders : dict of str + A dictionary mapping dimension names to their respective order strings. + + Returns + ------- + order : str + The common order string shared by all core dimensions. + + Raises + ------ + ValueError + If the core dimension orders are not all aligned (i.e., not all values + are the same). + """ + orders = set(core_dim_orders.values()) + + if len(orders) != 1: + raise ValueError( + f"All core dimension orders must be aligned. Got orders: {core_dim_orders}" + ) + + order = next(iter(orders)) + + return order + + +def _sort_vertices(vertices: np.ndarray, order: str) -> np.ndarray: + """ + Sorts the vertices array along the last axis in ascending or descending order. + + Parameters + ---------- + vertices : np.ndarray + An array of vertices to be sorted. Sorting is performed along the last + axis. + order : str + The order in which to sort the vertices. Must be either "ascending" or + any other value for descending order. + + Returns + ------- + np.ndarray + The sorted array of vertices, with the same shape as the input. + + Examples + -------- + >>> import numpy as np + >>> vertices = np.array([[3, 1, 2], [6, 5, 4]]) + >>> _sort_vertices(vertices, "ascending") + array([[1, 2, 3], + [4, 5, 6]]) + >>> _sort_vertices(vertices, "descending") + array([[3, 2, 1], + [6, 5, 4]]) + """ + if order == "ascending": + new_vertices = np.sort(vertices, axis=-1) + else: + new_vertices = np.sort(vertices, axis=-1)[..., ::-1] + + return new_vertices + + def vertices_to_bounds( vertices: DataArray, out_dims: Sequence[str] = ("bounds", "x", "y") ) -> DataArray: diff --git a/cf_xarray/tests/test_helpers.py b/cf_xarray/tests/test_helpers.py index 1a26b1bc..c3c51d92 100644 --- a/cf_xarray/tests/test_helpers.py +++ b/cf_xarray/tests/test_helpers.py @@ -1,3 +1,4 @@ +import xarray as xr from numpy.testing import assert_array_equal from xarray.testing import assert_equal @@ -12,7 +13,7 @@ def test_bounds_to_vertices() -> None: - # 1D case + # 1D case (stricly monotonic, descending bounds) ds = airds.cf.add_bounds(["lon", "lat", "time"]) lat_c = cfxr.bounds_to_vertices(ds.lat_bounds, bounds_dim="bounds") assert_array_equal(ds.lat.values + 1.25, lat_c.values[:-1]) @@ -34,6 +35,59 @@ def test_bounds_to_vertices() -> None: lon_no = cfxr.bounds_to_vertices(rotds.lon_bounds, bounds_dim="bounds", order=None) assert_equal(lon_no, lon_ccw) + # 2D case (monotonicly increasing coords, non-monotonic bounds) + bounds_2d_desc = xr.DataArray( + [[50.5, 50.0], [51.0, 50.5], [51.0, 50.5], [52.0, 51.5], [52.5, 52.0]], + dims=("lat", "bounds"), + coords={"lat": [50.75, 50.75, 51.25, 51.75, 52.25]}, + ) + expected_vertices_2d_desc = xr.DataArray( + [50.0, 50.5, 50.5, 51.5, 52.0, 52.5], + dims=["lat_vertices"], + ) + vertices_2d_desc = cfxr.bounds_to_vertices(bounds_2d_desc, bounds_dim="bounds") + assert_equal(expected_vertices_2d_desc, vertices_2d_desc) + + # 3D case (non-monotonic bounds, monotonicly increasing coords) + bounds_3d = xr.DataArray( + [ + [ + [50.0, 50.5], + [50.5, 51.0], + [51.0, 51.5], + [51.5, 52.0], + [52.0, 52.5], + ], + [ + [60.0, 60.5], + [60.5, 61.0], + [61.0, 61.5], + [61.5, 62.0], + [62.0, 62.5], + ], + ], + dims=("extra", "lat", "bounds"), + coords={ + "extra": [0, 1], + "lat": [0, 1, 2, 3, 4], + "bounds": [0, 1], + }, + ) + expected_vertices_3d = xr.DataArray( + [ + [50.0, 50.5, 51.0, 51.5, 52.0, 52.5], + [60.0, 60.5, 61.0, 61.5, 62.0, 62.5], + ], + dims=("extra", "lat_vertices"), + coords={ + "extra": [0, 1], + }, + ) + vertices_3d = cfxr.bounds_to_vertices( + bounds_3d, bounds_dim="bounds", core_dims=["lat"] + ) + assert_equal(vertices_3d, expected_vertices_3d) + # Transposing the array changes the bounds direction ds = mollwds.transpose("x", "y", "x_vertices", "y_vertices", "bounds") lon_cw = cfxr.bounds_to_vertices(