tools

tools

Various utility functions for modifying xarray object

Functions

Name Description
align_lists Align items from several lists into a single list that respects the order in each sublist.
chunk Apply rechunking to a xr.Dataset ds along dimensions provided as kwargs
conform A method decorator which applies MapBlocksOutput.conform to the method output.
datetime Parse datetime (in isoformat) from ds attributes
drop_unused_dims Simple function to remove unused dimensions in a xarray.Dataset
getflag Return the binary flag with given name as a boolean array
getflags returns the flags in attributes of A as a dictionary {meaning: value}
haversine Calculate the great circle distance between two points (specified in
is_dask_backed Check if data variables in a Dataset are dask-backed.
is_numpy_backed Check if an xarray Dataset or DataArray is backed by numpy arrays
locate Locate lat0, lon0 within lat, lon (xr.DataArrays)
merge Merge DataArrays in ds along dimension dim.
only If iterable has only one item, return it.
raiseflag Raise a flag in DataArray A with name flag_name, value flag_value and condition
split Returns a Dataset where a given dimension is split into as many variables
str_to_bool Convert a string representation to a boolean value.
sub Creates a Dataset based on the conditions passed in parameters
sub_pt Creates a Dataset based on the circle specified in parameters
sub_rect Returns a Dataset based on the coordinates of the rectangle passed in parameters
trim_dims Trim the dimensions of Dataset A
wrap Wrap and reorder a cyclic dimension between vmin and vmax.
xr_filter Extracts a subset of the dataset where the condition is True, stacking the
xr_filter_decorator A decorator which applies the decorated function only where the condition is True.
xr_flat A method which flat a xarray.Dataset on a new dimension named ‘index’
xr_sample A method to extract a subset of sample from a flat xarray.Dataset
xr_unfilter Reconstructs the original dataset from a subset dataset where the condition is True,
xrcrop Crop a Dataset or DataArray along dimensions based on min/max values.

align_lists

tools.align_lists(lists)

Align items from several lists into a single list that respects the order in each sublist. Raises ValueError if the lists cannot be aligned (e.g., due to cycles).

Args: lists: List of lists, where each sublist defines a partial order.

Returns: A single list with items aligned according to the partial orders.

Example: align_lists([[‘x’], [‘x’, ‘y’], [‘y’, ‘z’]]) -> [‘x’, ‘y’, ‘z’] align_lists([[‘x’, ‘y’], [‘y’, ‘x’]]) -> ValueError

chunk

tools.chunk(ds, **kwargs)

Apply rechunking to a xr.Dataset ds along dimensions provided as kwargs

Works like ds.chunk but works also for Datasets with repeated dimensions.

conform

tools.conform(attrname, transpose=True)

A method decorator which applies MapBlocksOutput.conform to the method output.

The MapBlocksOutput should be an attribute attrname of the class.

datetime

tools.datetime(ds)

Parse datetime (in isoformat) from ds attributes

drop_unused_dims

tools.drop_unused_dims(ds)

Simple function to remove unused dimensions in a xarray.Dataset

getflag

tools.getflag(A, name)

Return the binary flag with given name as a boolean array

A: DataArray name: str

example: getflag(flags, ‘LAND’)

getflags

tools.getflags(A=None, meanings=None, masks=None, sep=None)

returns the flags in attributes of A as a dictionary {meaning: value}

Arguments:

provide either: A: Dataarray or: meanings: flag meanings ‘FLAG1 FLAG2’ masks: flag values [1, 2] sep: string separator

haversine

tools.haversine(lat1, lon1, lat2, lon2, radius=6371)

Calculate the great circle distance between two points (specified in decimal degrees) on a sphere of a given radius

Returns the distance in the same unit as radius (defaults to earth radius in km)

is_dask_backed

tools.is_dask_backed(data, strict=False)

Check if data variables in a Dataset are dask-backed.

Returns True if all variables are dask-backed, False if all are numpy-backed. In the mixed case (some dask, some numpy), returns True if strict=False (any variable is dask-backed), or raises ValueError if strict=True.

Args: data: The input Dataset to check. strict: If True, raise ValueError when the Dataset has a mix of dask-backed and numpy-backed variables. If False (default), return True if any variable is dask-backed.

Returns: True if any/all variables are dask-backed, False if all are numpy-backed.

is_numpy_backed

tools.is_numpy_backed(obj)

Check if an xarray Dataset or DataArray is backed by numpy arrays (as opposed to dask arrays or other backends).

For Datasets, recursively checks all data variables and coordinates.

Note: even when not dask-backed, DataArrays can be backed by ‘xarray.core.indexing.MemoryCachedArray’, which may be inefficient compared to np.ndarray.

Args: obj: An xarray Dataset or DataArray.

Returns: True if all data variables are numpy-backed, False otherwise.

locate

tools.locate(lat, lon, lat0, lon0, dist_min_km=None, verbose=False)

Locate lat0, lon0 within lat, lon (xr.DataArrays)

if dist_min_km is specified and if the minimal distance exceeds it, a ValueError is raised

returns a dictionary of the pixel coordinates

merge

tools.merge(ds, dim=None, varname=None, pattern='(.+)_(\\d+)', dtype=int)

Merge DataArrays in ds along dimension dim.

ds: xr.Dataset

dim: str or None name of the new or existing dimension if None, use the attribute split_dimension

varname: str or None name of the variable to create if None, detect variable name from regular expression

pattern: str Regular expression for matching variable names and coordinates if varname is None: First group represents the new variable name. Second group represents the coordinate value Ex: r’(.+)_()’ First group matches all characters. Second group matches digits. r’(+)()’ First group matches non-digit. Second group matches digits. if varname is not None: Match a single group representing the coordinate value

dtype: data type data type of the coordinate items

only

tools.only(iterable)

If iterable has only one item, return it. Otherwise raise a ValueError

raiseflag

tools.raiseflag(A, flag_name, flag_value, condition=None)

Raise a flag in DataArray A with name flag_name, value flag_value and condition The name and value of the flag is recorded in the attributes of A

Arguments:

A: DataArray of integers

flag_name: str Name of the flag flag_value: int Value of the flag condition: boolean array-like of same shape as A Condition to raise flag. If None, the flag values are unchanged ; the flag is simple registered in the attributes.

split

tools.split(d, dim, sep='_')

Returns a Dataset where a given dimension is split into as many variables

d: Dataset or DataArray

str_to_bool

tools.str_to_bool(value)

Convert a string representation to a boolean value.

Args: value: String value to convert. Case-insensitive comparison with ‘true’.

Returns: True if the lowercase value equals ‘true’, False otherwise.

Example: >>> str_to_bool(‘True’) True >>> str_to_bool(‘false’) False >>> str_to_bool(‘TRUE’) True

sub

tools.sub(ds, cond, drop_invalid=True, int_default_value=0)

Creates a Dataset based on the conditions passed in parameters

cond : a DataArray of booleans that defines which pixels are kept

drop_invalid, bool if True invalid pixels will be replace by nan for floats and int_default_value for other types

int_default_value, int for DataArrays of type int, this value is assigned on non-valid pixels

sub_pt

tools.sub_pt(ds, pt_lat, pt_lon, rad, drop_invalid=True, int_default_value=0)

Creates a Dataset based on the circle specified in parameters

pt_lat, pt_lon : Coordonates of the center of the point

rad : radius of the circle in km

drop_invalid, bool if True invalid pixels will be replace by nan for floats and int_default_value for other types

int_default_value, int for DataArrays of type int, this value is assigned on non-valid pixels

sub_rect

tools.sub_rect(
    ds,
    lat_min,
    lon_min,
    lat_max,
    lon_max,
    drop_invalid=True,
    int_default_value=0,
)

Returns a Dataset based on the coordinates of the rectangle passed in parameters

lat_min, lat_max, lon_min, lon_max : delimitations of the region of interest

drop_invalid, bool : if True, invalid pixels will be replace by nan for floats and int_default_value for other types

int_default_value, int : for DataArrays of type int, this value is assigned on non-valid pixels

trim_dims

tools.trim_dims(A)

Trim the dimensions of Dataset A

Rename all possible dimensions to avoid duplicate dimensions with same sizes Avoid any DataArray with duplicate dimensions

wrap

tools.wrap(ds, dim, vmin, vmax)

Wrap and reorder a cyclic dimension between vmin and vmax. The border value is duplicated at the edges. The period is (vmax-vmin)

Example: * Dimension [0, 359] -> [-180, 180] * Dimension [-180, 179] -> [-180, 180] * Dimension [0, 359] -> [0, 360]

Arguments:

ds: xarray.Dataset dim: str Name of the dimension to wrap vmin, vmax: float new values for the edges

xr_filter

tools.xr_filter(ds, condition, stackdim=None, transparent=False)

Extracts a subset of the dataset where the condition is True, stacking the condition dimensions. Equivalent to numpy’s boolean indexing, A[condition].

Parameters: ds (xr.Dataset): The input dataset. condition (xr.DataArray): A boolean DataArray indicating where the condition is True. stackdim (str, optional): The name of the new stacked dimension. If None, it will be determined automatically from the condition dimensions. transparent (bool, optional): whether to reassign the original dimension names to the Dataset (expanding with length-one dimensions).

Returns: xr.Dataset: A new dataset with the subset of data where the condition is True.

xr_filter_decorator

tools.xr_filter_decorator(
    argpos,
    condition,
    fill_value_float=np.nan,
    fill_value_int=0,
    transparent=False,
    stackdim=None,
)

A decorator which applies the decorated function only where the condition is True.

Args: argpos (int): Position index of the input dataset in the decorated function call. condition (Callable): A callable taking the Dataset as input and returning a boolean DataArray. fill_value_float (float, optional): Fill value for floating point data types. Default is np.nan. fill_value_int (int, optional): Fill value for integer data types. Default is 0 transparent (bool, optional): Whether to reassign the original dimension names to the Dataset (expanding with length-one dimensions). Default is False. stackdim (str | None, optional): The name of the new stacked dimension. If None, it will be determined automatically from the condition dimensions. Default is None.

Example: @xr_filter_decorator(0, lambda x: x.flags == 0) def my_func(ds: xr.Dataset) -> xr.Dataset: # my_func is applied only where ds.flags == 0 …

The decorator works by: 1. Extracting a subset of the dataset where the condition is True using xr_filter. 2. Applying the decorated function to the subset. 3. Reconstructing the original dataset from the subset using xr_unfilter.

Behavior with in-place vs. non-in-place modifications: - If the decorated function returns a Dataset (non-in-place), the decorator returns the unfiltered result. - If the decorated function returns None (in-place modification), the decorator updates the original input dataset in-place with the unfiltered modified subset.

NOTE: this decorator does not guarantee that the order of dimensions is maintained. When using this decorator with xr.apply_blocks, you may want to wrap your xr_filter_decorator decorated method with the conform decorator.

xr_flat

tools.xr_flat(ds)

A method which flat a xarray.Dataset on a new dimension named ‘index’

Args: ds (xr.Dataset): Dataset to flat

xr_sample

tools.xr_sample(ds, nb_sample, seed=None)

A method to extract a subset of sample from a flat xarray.Dataset

Args: ds (xr.Dataset): Input flat dataset nb_sample (int|float): Number or percentage of sample to extract seed (int, optional): Random seed to use. Defaults to None.

xr_unfilter

tools.xr_unfilter(
    sub,
    condition,
    stackdim=None,
    fill_value_float=np.nan,
    fill_value_int=0,
    transparent=False,
    ds_original=None,
)

Reconstructs the original dataset from a subset dataset where the condition is True, unstacking the condition dimensions.

Parameters: sub (xr.Dataset): The subset dataset where the condition is True. condition (xr.DataArray): A boolean DataArray indicating where the condition is True. stackdim (str, optional): The name of the stacked dimension. If None, it will be determined automatically from the condition dimensions. fill_value_float (float, optional): The fill value for floating point data types. Default is np.nan. fill_value_int (int, optional): The fill value for integer data types. Default is 0. transparent (bool, optional): whether to revert the transparent compatibility conversion applied in xrwhere. ds_original (xr.Dataset, optional): The original dataset before filtering. If provided, variables present in both sub and ds_original will preserve their original values where condition is False, instead of being filled with fill_value_float/fill_value_int. This is important for pass-through variables (e.g. flags) that should not be overwritten.

Returns: xr.DataArray: The reconstructed dataset with the specified dimensions unstacked.

xrcrop

tools.xrcrop(A, **kwargs)

Crop a Dataset or DataArray along dimensions based on min/max values.

For each dimension provided as kwarg, the min/max values along that dimension can be provided: - As a min/max tuple - As a DataArrat, for which the min/max are computed

Ex: crop dimensions latitude and longitude of gsw based on the min/max of ds.lat and ds.lon gsw = xrcrop( gsw, latitude=ds.lat, longitude=ds.lon, )

Note: the purpose of this function is to make it possible to .compute() the result of the cropped data, thus allowing to perform a sel over large arrays (otherwise extremely slow with dask based arrays).