Source code for

# Copyright (c) 2018-2024 by xcube team and contributors
# Permissions are hereby granted under the terms of the MIT License:

import warnings
from typing import Collection, Optional, Tuple, Callable, Dict, Any, List, Mapping
from typing import Union

import cftime
import dask.array as da
import numpy as np
import pandas as pd
import xarray as xr

from xcube.core.gridmapping import GridMapping
from xcube.util.assertions import assert_given
from xcube.util.timeindex import ensure_time_index_compatible

Bbox = Tuple[float, float, float, float]
TimeRange = Union[
    Tuple[Optional[str], Optional[str]],
    Tuple[Optional[pd.Timestamp], Optional[pd.Timestamp]],

def select_subset(
    dataset: xr.Dataset,
    var_names: Optional[Collection[str]] = None,
    bbox: Optional[Bbox] = None,
    time_range: Optional[TimeRange] = None,
    grid_mapping: Optional[GridMapping] = None,
    """Create a subset from *dataset* given *var_names*,
    *bbox*, *time_range*.

    This is a high-level convenience function that may invoke

    * :func:`select_variables_subset`
    * :func:`select_spatial_subset`
    * :func:`select_temporal_subset`

        dataset: The dataset.
        var_names: Optional variable names.
        bbox: Optional bounding box in the dataset's CRS coordinate
        time_range: Optional time range
        grid_mapping: Optional dataset grid mapping.

        a subset of *dataset*, or unchanged *dataset* if no keyword-
        arguments are used.
    if var_names is not None:
        dataset = select_variables_subset(dataset, var_names=var_names)
    if bbox is not None:
        dataset = select_spatial_subset(
            dataset, xy_bbox=bbox, grid_mapping=grid_mapping
    if time_range is not None:
        dataset = select_temporal_subset(dataset, time_range=time_range)
    return dataset

[docs] def select_variables_subset( dataset: xr.Dataset, var_names: Optional[Collection[str]] = None ) -> xr.Dataset: """Select data variable from given *dataset* and create new dataset. Args: dataset: The dataset from which to select variables. var_names: The names of data variables to select. Returns: A new dataset. It is empty, if *var_names* is empty. It is *dataset*, if *var_names* is None. """ if var_names is None: return dataset dropped_variables = set(dataset.data_vars.keys()).difference(var_names) if not dropped_variables: return dataset return dataset.drop_vars(dropped_variables)
def select_spatial_subset( dataset: xr.Dataset, ij_bbox: Optional[Tuple[int, int, int, int]] = None, ij_border: int = 0, xy_bbox: Optional[Tuple[float, float, float, float]] = None, xy_border: float = 0.0, grid_mapping: Optional[GridMapping] = None, ) -> Optional[xr.Dataset]: """Select a spatial subset of *dataset* for the bounding box *ij_bbox* or *xy_bbox*. *ij_bbox* or *xy_bbox* must not be given both. Args: xy_bbox: The bounding box in x,y coordinates. xy_border: Border in units of the x,y coordinates. dataset: Source dataset. ij_bbox: Bounding box (i_min, i_min, j_max, j_max) in pixel coordinates. ij_border: Extra border added to *ij_bbox* in number of pixels grid_mapping: Optional dataset grid mapping. Returns: Spatial dataset subset """ if ij_bbox is None and xy_bbox is None: raise ValueError("One of ij_bbox and xy_bbox must be given") if ij_bbox and xy_bbox: raise ValueError("Only one of ij_bbox and xy_bbox can be given") if grid_mapping is None: grid_mapping = GridMapping.from_dataset(dataset) x_name, y_name = grid_mapping.xy_var_names x = dataset[x_name] y = dataset[y_name] if x.ndim == 1 and y.ndim == 1: # Hotfix für #981 and #985 if xy_bbox: if y.values[0] < y.values[-1]: ds = dataset.sel( **{ x_name: slice(xy_bbox[0] - xy_border, xy_bbox[2] + xy_border), y_name: slice(xy_bbox[1] - xy_border, xy_bbox[3] + xy_border), } ) else: ds = dataset.sel( **{ x_name: slice(xy_bbox[0] - xy_border, xy_bbox[2] + xy_border), y_name: slice(xy_bbox[3] + xy_border, xy_bbox[1] - xy_border), } ) return ds else: return dataset.isel( **{ x_name: slice(ij_bbox[0] - ij_border, ij_bbox[2] + ij_border), y_name: slice(ij_bbox[1] - ij_border, ij_bbox[3] + ij_border), } ) else: if xy_bbox: ij_bbox = grid_mapping.ij_bbox_from_xy_bbox( xy_bbox, ij_border=ij_border, xy_border=xy_border ) if ij_bbox[0] == -1: return None width, height = grid_mapping.size i_min, j_min, i_max, j_max = ij_bbox if i_min > 0 or j_min > 0 or i_max < width - 1 or j_max < height - 1: x_dim, y_dim = grid_mapping.xy_dim_names i_slice = slice(i_min, i_max + 1) j_slice = slice(j_min, j_max + 1) return dataset.isel({x_dim: i_slice, y_dim: j_slice}) return dataset def select_temporal_subset( dataset: xr.Dataset, time_range: TimeRange, time_name: str = "time" ) -> xr.Dataset: """Select a temporal subset from *dataset* given *time_range*. Args: dataset: The dataset. Must include time time_range: Time range given as two time stamps (start, end) that may be (ISO) strings or datetime objects. time_name: optional name of the time coordinate variable. Defaults to "time". Returns: """ assert_given(time_range, "time_range") time_name = time_name or "time" if time_name not in dataset: raise ValueError( f"cannot compute temporal subset: variable" f' "{time_name}" not found in dataset' ) time_1, time_2 = time_range time_1 = pd.to_datetime(time_1) if time_1 is not None else None time_2 = pd.to_datetime(time_2) if time_2 is not None else None if time_1 is None and time_2 is None: return dataset if time_2 is not None: delta = time_2 - time_2.floor("1D") if delta == pd.Timedelta("0 days 00:00:00"): time_2 += pd.Timedelta("1D") try: time_slice = ensure_time_index_compatible( dataset, slice(time_1, time_2), time_name ) return dataset.sel({time_name or "time": time_slice}) except TypeError: calendar = dataset.time.encoding.get("calendar") time_1 = cftime.datetime( time_1.year, time_1.month,, calendar=calendar ) time_2 = cftime.datetime( time_2.year, time_2.month,, calendar=calendar ) time_slice = ensure_time_index_compatible( dataset, slice(time_1, time_2), time_name ) return dataset.sel({time_name or "time": time_slice}) _PREDICATE_SIGNATURE = ( "predicate(" "slice_array: xr.DataArray, " "slice_info: Dict" ") -> bool" ) Predicate = Callable[[xr.DataArray, Dict[str, Any]], bool] def select_label_subset( dataset: xr.Dataset, dim: str, predicate: Union[Predicate, Mapping[str, Predicate]], use_dask: bool = False, ): """Select the labels in *dataset* along a given dimension *dim* using a predicate function *predicate* that is called for all variable slices for a current label. The *predicate* can also be provided as a mapping from variable names to dedicated predicate functions. The predicate function is called for all *dim* labels in *dataset* and for every variable that contains *dim*. If *predicate* returns False for any given label, that label will be dropped from dimension *dim*. Predicate functions are defined as follows::: def predicate(slice_array: xr.DataArray, slice_info: Dict) -> bool: ... Here, *slice_array* is a variable's array slice for the given label. The argument *slice_info* is a dictionary that contains the following keys: * var: str - name of the current variable. * dim: str - value of *dim*. * index: int - value for the current index within dimension *dim*. * label: Optional[xr.DataArray] - value for the current label within dimension *dim*. Note, the value of "label" will be None, if *dataset* does not contain a 1D-coordinate variable named *dim*. The following example selects only time labels from a 3-D (time, y, x) cube where the 2-D (y, x) images of variable "CHL" comprises more than 50% valid values::: >>> chl_data = np.random.random((5, 10, 20)) >>> chl_data = np.where(chl_data > 0.5, chl_data, np.nan) >>> ds = xr.Dataset({"CHL": (["time", "y", "x"], chl_data)}) >>> >>> def is_valid_slice(slice_array, slice_label): >>> return np.sum(np.isnan(slice_array)) / slice_array.size <= 0.5 >>> >>> ds_subset = select_label_subset(ds, "time", >>> predicate={"CHL": is_valid_slice}) Args: dataset: The dataset. dim: The name of the dimension from which to select the labels. predicate: The predicate function or a mapping from variable names to variable-specific predicate functions. use_dask: Whether to use a Dask graph that will compute the validity of labels in parallel. For a large number of labels, very complex Dask graphs will result (every label is a node) whose overhead may compensate the performance gain. Returns: A new dataset with labels along *dim* selected by the *predicate*. If all labels are selected, *dataset* is returned without change. """ if callable(predicate): predicate_lookup = { var_name: predicate for var_name, var in dataset.data_vars.items() if dim in var.dims } elif isinstance(predicate, Mapping): predicate_lookup = predicate for var_name, var_predicate in predicate_lookup.items(): if not callable(var_predicate): raise TypeError( f"predicate for variable {var_name!r}" f" must be callable with" f" signature {_PREDICATE_SIGNATURE}" ) else: raise TypeError( f"predicate" f" must be callable with" f" signature {_PREDICATE_SIGNATURE}" ) num_labels = dataset.dims[dim] valid_mask = [ _is_label_valid(dataset, predicate_lookup, dim, index) for index in range(num_labels) ] if use_dask: valid_mask = da.stack(valid_mask).compute() dropped_indexes = [i for i in range(num_labels) if not valid_mask[i]] if not dropped_indexes: return dataset return dataset.drop_isel({dim: dropped_indexes}) def _is_label_valid( dataset: xr.Dataset, predicate_lookup: Mapping[str, Predicate], dim: str, index: int ) -> da.Array: label = dataset[dim][index] if dim in dataset else None results: List[da.Array] = [] for var_name, var in dataset.data_vars.items(): if dim in var.dims: predicate = predicate_lookup.get(var_name) if predicate is not None: slice_array = var.isel({dim: index}) slice_info = dict(var=var_name, dim=dim, index=index, label=label) result = predicate(slice_array, slice_info) if isinstance(result, xr.DataArray): result = if isinstance(result, da.Array): results.append(result) else: results.append(da.from_array(result)) if len(results) == 0: return da.from_array(True) elif len(results) == 1: return results[0] else: return da.all(da.stack(results))