Source code for decode.select

__all__ = ["by"]


# standard library
from collections.abc import Sequence
from datetime import datetime, timedelta
from logging import getLogger
from typing import Optional, TypeVar, Union


# dependencies
import xarray as xr


# constants
LOGGER = getLogger(__name__)


# type hints
T = TypeVar("T", bool, int, float, str, datetime, timedelta)
Multiple = Union[Sequence[T], T]




def by(
    dems: xr.DataArray,
    coord_name: str,
    /,
    include: Optional[Multiple[T]] = None,
    *,
    exclude: Optional[Multiple[T]] = None,
    min: Optional[T] = None,
    max: Optional[T] = None,
    sort: bool = False,
    as_dim: bool = False,
) -> xr.DataArray:
    """Select DEMS by values of a coordinate.

    Args:
        dems: DEMS DataArray to be selected.
        coord_name: Name of the coordinate for the selection.
        include: Coordinate values to be included.
            If not specified, all values are included.
        exclude: Coordinate values to be excluded.
            If not specified, any values are not excluded.
        min: Minimum selection bound (inclusive).
            If not specified, no bound is set.
        max: Maximum selection bound (exclusive).
            If not specified, no bound is set.
        sort: Whether to sort by the coordinate after selection.
        as_dim: Whether to use the coordinate as a dimension.

    Returns:
        Selected DEMS.

    Important:
        This function will be deprecated in a future release.
        For ``include`` and ``exclude`` options, use ``xarray.DataArray.isin``
        like ``dems.sel(time=dems.state.isin(['ON', 'OFF']))``.
        For ``min`` and ``max`` options, use ``ndtools.Range`` like
        ``dems.sel(chan=dems.frequency == Range(250e9, 300e9))``.

    """
    LOGGER.warning(
        "This function will be deprecated in a future release. "
        "For include and exclude options, use DataArray.isin like "
        "dems.sel(time=dems.state.isin(['ON', 'OFF'])). "
        "For min and max options, use ndtools.Range like "
        "dems.sel(chan=dems.frequency == Range(250e9, 300e9))."
    )
    coord = dems[coord_name]

    if not isinstance(coord, xr.DataArray):
        raise TypeError("Coordinate must be DataArray.")

    if not coord.ndim == 1:
        raise ValueError("Coordinate must be one-dimensional.")

    coord_dim = coord.dims[0]

    if include is not None:
        dems = dems.sel({coord_dim: dems[coord_name].isin(include)})

    if exclude is not None:
        dems = dems.sel({coord_dim: ~dems[coord_name].isin(exclude)})

    if min is not None:
        dems = dems.sel({coord_dim: dems[coord_name] >= min})  # type: ignore

    if max is not None:
        dems = dems.sel({coord_dim: dems[coord_name] < max})  # type: ignore

    if sort:
        dems = dems.sortby(coord_name)

    if as_dim:
        dems = dems.swap_dims({coord_dim: coord_name})

    return dems