Subsetting, Regions & Masks¶

Carve a dataset down to a region of interest — by bounding box, time window, named region, or boolean predicate — and add land / ocean / country masks. Named regions come from a small registry that also reads GeoJSON polygons.

Subsetting¶

SubsetBBox ¶

SubsetBBox(lon_bnds: tuple[float, float], lat_bnds: tuple[float, float], lon: str = 'lon', lat: str = 'lat')

Bases: Operator

Subset a dataset to a lon/lat bounding box.

Parameters:

Name	Type	Description	Default
`lon_bnds`	`tuple[float, float]`	`(min, max)` longitude bounds.	required
`lat_bnds`	`tuple[float, float]`	`(min, max)` latitude bounds.	required
`lon`	`str`	Name of the longitude coordinate.	`'lon'`
`lat`	`str`	Name of the latitude coordinate.	`'lat'`

Returns:

Type	Description
	The dataset restricted to the bounding box.

Source code in packages/xrtoolz-ops/src/xrops/geom/operators.py

def __init__(
    self,
    lon_bnds: tuple[float, float],
    lat_bnds: tuple[float, float],
    lon: str = "lon",
    lat: str = "lat",
):
    self.lon_bnds = tuple(lon_bnds)
    self.lat_bnds = tuple(lat_bnds)
    self.lon = lon
    self.lat = lat

SubsetTime ¶

SubsetTime(time_min: str, time_max: str, time: str = 'time')

Bases: Operator

Subset a dataset to an inclusive [time_min, time_max] window.

Parameters:

Name	Type	Description	Default
`time_min`	`str`	Start of the time window (ISO string).	required
`time_max`	`str`	End of the time window (ISO string).	required
`time`	`str`	Name of the time coordinate.	`'time'`

Returns:

Type	Description
	The dataset restricted to the time window.

Source code in packages/xrtoolz-ops/src/xrops/geom/operators.py

def __init__(self, time_min: str, time_max: str, time: str = "time"):
    self.time_min = time_min
    self.time_max = time_max
    self.time = time

SubsetToRegion ¶

SubsetToRegion(region: str | RegionSpec | Regions | dict[str, Any], *, lon: str = 'lon', lat: str = 'lat', validate: bool = True)

Bases: Operator

Subset a dataset to a named or custom geographic region.

Masks to the region and trims to its bounding box.

Parameters:

Name	Type	Description	Default
`region`	`str \| RegionSpec \| Regions \| dict[str, Any]`	Region specification — a registry name, a :class:`~xrops.geom._src.regions.RegionSpec`, a `regionmask.Regions`, or the dict form from `get_config()`.	required
`lon`	`str`	Name of the longitude coordinate.	`'lon'`
`lat`	`str`	Name of the latitude coordinate.	`'lat'`
`validate`	`bool`	Whether to validate coordinates before masking.	`True`

Returns:

Type	Description
	The dataset restricted to the region.

Source code in packages/xrtoolz-ops/src/xrops/geom/operators.py

def __init__(
    self,
    region: str | _regions.RegionSpec | regionmask.Regions | dict[str, Any],
    *,
    lon: str = "lon",
    lat: str = "lat",
    validate: bool = True,
):
    # Accept the dict form emitted by ``get_config`` so the standard
    # ``cls(**op.get_config())`` round-trip used by ApplyToEach works
    # for custom regions, not just registry strings.
    if isinstance(region, dict):
        region = _regions.region_from_dict(region)
    self.region = region
    self.lon = lon
    self.lat = lat
    self.validate = validate

SelectVariables ¶

SelectVariables(variables: str | Sequence[str])

Bases: Operator

Select a subset of data variables.

Parameters:

Name	Type	Description	Default
`variables`	`str \| Sequence[str]`	A variable name or sequence of names to keep.	required

Returns:

Type	Description
	The dataset containing only the selected variables.

Source code in packages/xrtoolz-ops/src/xrops/geom/operators.py

def __init__(self, variables: str | Sequence[str]):
    self.variables = [variables] if isinstance(variables, str) else list(variables)

Masks¶

AddLandMask ¶

AddLandMask(name: str = 'land_mask')

Bases: Operator

Add a boolean land mask derived from regionmask.

Parameters:

Name	Type	Description	Default
`name`	`str`	Name of the mask variable to add.	`'land_mask'`

Returns:

Type	Description
	The input dataset with a boolean land-mask variable.

Source code in packages/xrtoolz-ops/src/xrops/mask/operators.py

def __init__(self, name: str = "land_mask"):
    self.name = name

AddOceanMask ¶

AddOceanMask(ocean: str = 'global', name: str = 'ocean_mask')

Bases: Operator

Add a boolean ocean/basin mask derived from regionmask.

Parameters:

Name	Type	Description	Default
`ocean`	`str`	Ocean/basin selector (e.g. `"global"`).	`'global'`
`name`	`str`	Name of the mask variable to add.	`'ocean_mask'`

Returns:

Type	Description
	The input dataset with a boolean ocean-mask variable.

Source code in packages/xrtoolz-ops/src/xrops/mask/operators.py

def __init__(self, ocean: str = "global", name: str = "ocean_mask"):
    self.ocean = ocean
    self.name = name

AddCountryMask ¶

AddCountryMask(country: str, name: str = 'country_mask')

Bases: Operator

Add a boolean country mask derived from regionmask.

Parameters:

Name	Type	Description	Default
`country`	`str`	Country name/selector to mask.	required
`name`	`str`	Name of the mask variable to add.	`'country_mask'`

Returns:

Type	Description
	The input dataset with a boolean country-mask variable.

Source code in packages/xrtoolz-ops/src/xrops/mask/operators.py

def __init__(self, country: str, name: str = "country_mask"):
    self.country = country
    self.name = name

ApplyMask ¶

ApplyMask(mask, drop: bool = False)

Bases: Operator

Apply a boolean mask to a dataset.

Parameters:

Name	Type	Description	Default
`mask`		A boolean DataArray, or the name of a mask variable already in the dataset.	required
`drop`	`bool`	If `True`, drop fully-masked slices (`where(..., drop=True)`); otherwise keep them as NaN.	`False`

Returns:

Type	Description
	The masked dataset.

Source code in packages/xrtoolz-ops/src/xrops/mask/operators.py

def __init__(self, mask, drop: bool = False):
    self.mask = mask
    self.drop = drop

Functional primitives (Layer 0)¶

These pure functions back the operators above; each takes xr.DataArray/xr.Dataset and a dim: argument.

subset_bbox ¶

subset_bbox(ds: Dataset, lon_bnds: tuple[float, float], lat_bnds: tuple[float, float], lon: str = 'lon', lat: str = 'lat') -> Dataset

Keep points whose (lon, lat) fall inside a bounding box.

Works with both 1-D rectilinear grids (where sel with slices would also work) and 2-D lon/lat coordinate arrays, because it masks via where(..., drop=True).

Parameters:

Name	Type	Description	Default
`ds`	`Dataset`	Input dataset.	required
`lon_bnds`	`tuple[float, float]`	`(lon_min, lon_max)`.	required
`lat_bnds`	`tuple[float, float]`	`(lat_min, lat_max)`.	required
`lon`	`str`	Name of the longitude coordinate.	`'lon'`
`lat`	`str`	Name of the latitude coordinate.	`'lat'`

Returns:

Type	Description
`Dataset`	Dataset restricted to the bounding box.

Source code in packages/xrtoolz-products/src/xrproducts/types/_src/subset.py

def subset_bbox(
    ds: xr.Dataset,
    lon_bnds: tuple[float, float],
    lat_bnds: tuple[float, float],
    lon: str = "lon",
    lat: str = "lat",
) -> xr.Dataset:
    """Keep points whose ``(lon, lat)`` fall inside a bounding box.

    Works with both 1-D rectilinear grids (where ``sel`` with slices would
    also work) and 2-D lon/lat coordinate arrays, because it masks via
    ``where(..., drop=True)``.

    Args:
        ds: Input dataset.
        lon_bnds: ``(lon_min, lon_max)``.
        lat_bnds: ``(lat_min, lat_max)``.
        lon: Name of the longitude coordinate.
        lat: Name of the latitude coordinate.

    Returns:
        Dataset restricted to the bounding box.
    """
    lon_min, lon_max = lon_bnds
    lat_min, lat_max = lat_bnds
    mask = (
        (ds[lon] >= lon_min)
        & (ds[lon] <= lon_max)
        & (ds[lat] >= lat_min)
        & (ds[lat] <= lat_max)
    )
    return ds.where(mask, drop=True)

subset_time ¶

subset_time(ds: Dataset, time_min: str | datetime64 | Timestamp, time_max: str | datetime64 | Timestamp, time: str = 'time') -> Dataset

Restrict a dataset to the time range [time_min, time_max].

Parameters:

Name	Type	Description	Default
`ds`	`Dataset`	Input dataset.	required
`time_min`	`str \| datetime64 \| Timestamp`	Lower bound, any type xarray accepts in a `slice` (`str`, `numpy.datetime64`, `pandas.Timestamp`, ...).	required
`time_max`	`str \| datetime64 \| Timestamp`	Upper bound, same types as `time_min`.	required
`time`	`str`	Name of the time coordinate.	`'time'`

Returns:

Type	Description
`Dataset`	Dataset sliced along the time dimension.

Source code in packages/xrtoolz-products/src/xrproducts/types/_src/subset.py

def subset_time(
    ds: xr.Dataset,
    time_min: str | np.datetime64 | pd.Timestamp,
    time_max: str | np.datetime64 | pd.Timestamp,
    time: str = "time",
) -> xr.Dataset:
    """Restrict a dataset to the time range ``[time_min, time_max]``.

    Args:
        ds: Input dataset.
        time_min: Lower bound, any type xarray accepts in a ``slice``
            (``str``, ``numpy.datetime64``, ``pandas.Timestamp``, ...).
        time_max: Upper bound, same types as ``time_min``.
        time: Name of the time coordinate.

    Returns:
        Dataset sliced along the time dimension.
    """
    return ds.sel({time: slice(time_min, time_max)})

subset_to_region ¶

subset_to_region(ds: Dataset, region: str | RegionSpec | Regions, *, lon: str = 'lon', lat: str = 'lat', validate: bool = True) -> Dataset

Subset a Dataset to a named, custom, or polygon region.

Source code in packages/xrtoolz-ops/src/xrops/geom/_src/subset.py

def subset_to_region(
    ds: xr.Dataset,
    region: str | RegionSpec | regionmask.Regions,
    *,
    lon: str = "lon",
    lat: str = "lat",
    validate: bool = True,
) -> xr.Dataset:
    """Subset a Dataset to a named, custom, or polygon region."""
    import warnings

    if isinstance(region, str):
        region = resolve_region(region)
    if isinstance(region, RegionSpec):
        region = region.regions

    # regionmask emits a FutureWarning about the default mask method;
    # pin it explicitly so the choice is stable across regionmask
    # versions and the warning doesn't leak to callers.
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", FutureWarning)
        mask = region.mask(ds[lon], ds[lat], method="shapely").notnull()
    if validate and not bool(mask.any().item()):
        raise ValueError(
            "Region does not overlap dataset coordinates. "
            "Pass validate=False to allow empty results."
        )
    return ds.where(mask, drop=True)

subset_where ¶

subset_where(ds: Dataset, variable: str, min_val: float, max_val: float, drop: bool = True) -> Dataset

Keep points where variable is in [min_val, max_val].

Parameters:

Name	Type	Description	Default
`ds`	`Dataset`	Input dataset.	required
`variable`	`str`	Name of the variable used to build the mask.	required
`min_val`	`float`	Lower bound (inclusive).	required
`max_val`	`float`	Upper bound (inclusive).	required
`drop`	`bool`	Whether to drop cells outside the range (default `True`); if `False`, masked cells are replaced with NaN.	`True`

Returns:

Type	Description
`Dataset`	Dataset with cells outside the range masked or dropped.

Source code in packages/xrtoolz-products/src/xrproducts/types/_src/subset.py

def subset_where(
    ds: xr.Dataset,
    variable: str,
    min_val: float,
    max_val: float,
    drop: bool = True,
) -> xr.Dataset:
    """Keep points where ``variable`` is in ``[min_val, max_val]``.

    Args:
        ds: Input dataset.
        variable: Name of the variable used to build the mask.
        min_val: Lower bound (inclusive).
        max_val: Upper bound (inclusive).
        drop: Whether to drop cells outside the range (default ``True``);
            if ``False``, masked cells are replaced with NaN.

    Returns:
        Dataset with cells outside the range masked or dropped.
    """
    mask = (ds[variable] >= float(min_val)) & (ds[variable] <= float(max_val))
    return ds.where(mask, drop=drop)

select_variables ¶

select_variables(ds: Dataset, variables: str | Sequence[str]) -> Dataset

Return a dataset with only the requested variables.

Parameters:

Name	Type	Description	Default
`ds`	`Dataset`	Input dataset.	required
`variables`	`str \| Sequence[str]`	A single variable name or a sequence of names.	required

Returns:

Type	Description
`Dataset`	Dataset restricted to `variables`.

Source code in packages/xrtoolz-ops/src/xrops/geom/_src/subset.py

def select_variables(
    ds: xr.Dataset,
    variables: str | Sequence[str],
) -> xr.Dataset:
    """Return a dataset with only the requested variables.

    Args:
        ds: Input dataset.
        variables: A single variable name or a sequence of names.

    Returns:
        Dataset restricted to ``variables``.
    """
    names = [variables] if isinstance(variables, str) else list(variables)
    return ds.drop_vars([name for name in ds.data_vars if name not in names])

add_land_mask ¶

add_land_mask(ds: Dataset, name: str = 'land_mask') -> Dataset

Attach a 1/0 land mask (Natural Earth 110m) as a coordinate.

Parameters:

Name	Type	Description	Default
`ds`	`Dataset`	Input dataset with `lon` and `lat` coordinates.	required
`name`	`str`	Name of the mask coordinate to add.	`'land_mask'`

Returns:

Type	Description
`Dataset`	Dataset with a new `name` coordinate that is 1 over land and
`Dataset`	0 over ocean.

Source code in packages/xrtoolz-ops/src/xrops/mask/_src/masks.py

def add_land_mask(ds: xr.Dataset, name: str = "land_mask") -> xr.Dataset:
    """Attach a 1/0 land mask (Natural Earth 110m) as a coordinate.

    Args:
        ds: Input dataset with ``lon`` and ``lat`` coordinates.
        name: Name of the mask coordinate to add.

    Returns:
        Dataset with a new ``name`` coordinate that is 1 over land and
        0 over ocean.
    """
    regions = regionmask.defined_regions.natural_earth_v5_0_0.land_110
    mask = regions.mask_3D(ds).squeeze(drop=True).astype(np.int16)
    mask.attrs.update(
        standard_name="land_mask",
        long_name="Land Mask",
        source="Natural Earth 110m",
    )
    return _attach_mask(ds, mask, name)

add_ocean_mask ¶

add_ocean_mask(ds: Dataset, ocean: str = 'global', name: str = 'ocean_mask') -> Dataset

Attach a 1/0 ocean mask as a coordinate.

Parameters:

Name	Type	Description	Default
`ds`	`Dataset`	Input dataset with `lon` and `lat` coordinates.	required
`ocean`	`str`	Ocean basin name (e.g. `"indian"`, `"north atlantic"`) or `"global"` to union all Natural Earth ocean basins.	`'global'`
`name`	`str`	Name of the mask coordinate to add.	`'ocean_mask'`

Returns:

Type	Description
`Dataset`	Dataset with a new `name` coordinate that is 1 inside the
`Dataset`	requested ocean region(s) and 0 elsewhere.

Source code in packages/xrtoolz-ops/src/xrops/mask/_src/masks.py

def add_ocean_mask(
    ds: xr.Dataset,
    ocean: str = "global",
    name: str = "ocean_mask",
) -> xr.Dataset:
    """Attach a 1/0 ocean mask as a coordinate.

    Args:
        ds: Input dataset with ``lon`` and ``lat`` coordinates.
        ocean: Ocean basin name (e.g. ``"indian"``, ``"north atlantic"``)
            or ``"global"`` to union all Natural Earth ocean basins.
        name: Name of the mask coordinate to add.

    Returns:
        Dataset with a new ``name`` coordinate that is 1 inside the
        requested ocean region(s) and 0 elsewhere.
    """
    regions = regionmask.defined_regions.natural_earth_v5_0_0.ocean_basins_50
    mask_3d = regions.mask_3D(ds)

    if ocean == "global":
        mask = mask_3d.any(dim="region").astype(np.int16)
    else:
        selector = [i for i, n in enumerate(mask_3d.names.values) if n == ocean]
        if not selector:
            available = sorted(set(mask_3d.names.values))
            raise ValueError(
                f"Ocean basin {ocean!r} not found. Available: {available}."
            )
        mask = mask_3d.isel(region=selector[0]).astype(np.int16)

    mask.attrs.update(
        standard_name="ocean_mask",
        long_name="Ocean Mask",
        source="Natural Earth ocean_basins_50",
        basin=ocean,
    )
    return _attach_mask(ds, mask, name)

add_country_mask ¶

add_country_mask(ds: Dataset, country: str, name: str = 'country_mask') -> Dataset

Attach a 1/0 country mask as a coordinate.

Parameters:

Name	Type	Description	Default
`ds`	`Dataset`	Input dataset with `lon` and `lat` coordinates.	required
`country`	`str`	Country name or 3-letter abbreviation as recognised by Natural Earth 110m country regions.	required
`name`	`str`	Name of the mask coordinate to add.	`'country_mask'`

Returns:

Type	Description
`Dataset`	Dataset with a new `name` coordinate that is 1 inside
`Dataset`	`country` and 0 elsewhere.

Source code in packages/xrtoolz-ops/src/xrops/mask/_src/masks.py

def add_country_mask(
    ds: xr.Dataset,
    country: str,
    name: str = "country_mask",
) -> xr.Dataset:
    """Attach a 1/0 country mask as a coordinate.

    Args:
        ds: Input dataset with ``lon`` and ``lat`` coordinates.
        country: Country name or 3-letter abbreviation as recognised by
            Natural Earth 110m country regions.
        name: Name of the mask coordinate to add.

    Returns:
        Dataset with a new ``name`` coordinate that is 1 inside
        ``country`` and 0 elsewhere.
    """
    regions = regionmask.defined_regions.natural_earth_v5_0_0.countries_110
    mask_3d = regions.mask_3D(ds)
    names = list(mask_3d.names.values)
    abbrevs = list(mask_3d.abbrevs.values)

    if country in names:
        idx = names.index(country)
    elif country in abbrevs:
        idx = abbrevs.index(country)
    else:
        raise ValueError(
            f"Country {country!r} not found in Natural Earth 110m regions."
        )

    mask = mask_3d.isel(region=idx).astype(np.int16)
    mask.attrs.update(
        standard_name="country_mask",
        long_name="Country Mask",
        country=country,
        source="Natural Earth countries_110",
    )
    return _attach_mask(ds, mask, name)

apply_mask ¶

apply_mask(ds: Dataset, mask: DataArray | str, drop: bool = False) -> Dataset

Apply a boolean mask to every variable in the dataset.

Parameters:

Name	Type	Description	Default
`ds`	`Dataset`	Input dataset.	required
`mask`	`DataArray \| str`	Either a DataArray to mask by, or the name of a coordinate / variable in `ds` (e.g. `"ocean_mask"`).	required
`drop`	`bool`	If `True`, drop coordinates with no valid values.	`False`

Returns:

Type	Description
`Dataset`	Masked dataset.

Source code in packages/xrtoolz-ops/src/xrops/mask/_src/masks.py

def apply_mask(
    ds: xr.Dataset,
    mask: xr.DataArray | str,
    drop: bool = False,
) -> xr.Dataset:
    """Apply a boolean mask to every variable in the dataset.

    Args:
        ds: Input dataset.
        mask: Either a DataArray to mask by, or the name of a
            coordinate / variable in ``ds`` (e.g. ``"ocean_mask"``).
        drop: If ``True``, drop coordinates with no valid values.

    Returns:
        Masked dataset.
    """
    mask_da = ds[mask] if isinstance(mask, str) else mask
    return ds.where(mask_da.astype(bool), drop=drop)

Region registry¶

REGIONS is the built-in named-region table; RegionSpec is the frozen bounding-box descriptor. The helpers below resolve a region from a name, a dict, or a GeoJSON polygon.

REGIONS `module-attribute` ¶

REGIONS: dict[str, RegionSpec] = {'global': custom_region(id='global', display_name='Global', lat_min=-90.0, lat_max=90.0, lon_min=-180.0, lon_max=180.0, projection='Robinson'), 'north_atlantic': custom_region(id='north_atlantic', display_name='North Atlantic', lat_min=10.0, lat_max=65.0, lon_min=-80.0, lon_max=0.0), 'gulf_stream': custom_region(id='gulf_stream', display_name='Gulf Stream', lat_min=30.0, lat_max=45.0, lon_min=-80.0, lon_max=-50.0), 'kuroshio': custom_region(id='kuroshio', display_name='Kuroshio', lat_min=25.0, lat_max=45.0, lon_min=130.0, lon_max=180.0), 'mediterranean': custom_region(id='mediterranean', display_name='Mediterranean', lat_min=30.0, lat_max=46.0, lon_min=-6.0, lon_max=36.0), 'ibi': custom_region(id='ibi', display_name='Iberia-Biscay-Ireland (IBI)', lat_min=26.17, lat_max=56.08, lon_min=-19.08, lon_max=5.08)}

RegionSpec `dataclass` ¶

RegionSpec(id: str, display_name: str, regions: Regions, projection: str | None = None)

Named region pairing regionmask geometry with viz metadata.