Skip to content

rslearn.utils.array

array

Array util functions.

nodata_eq

nodata_eq(array: NDArray[generic], nodata_value: int | float) -> NDArray[bool_]

NaN-aware element-wise equality between array and a nodata sentinel.

Equivalent to array == nodata_value but also matches NaN positions when the nodata sentinel is NaN.

Parameters:

Name Type Description Default
array NDArray[generic]

the data array (any shape).

required
nodata_value int | float

scalar nodata sentinel.

required

Returns:

Type Description
NDArray[bool_]

Boolean mask with the same shape as array; True where the value

NDArray[bool_]

equals (or is NaN-matching) the nodata sentinel.

Source code in rslearn/utils/array.py
def nodata_eq(
    array: npt.NDArray[np.generic],
    nodata_value: int | float,
) -> npt.NDArray[np.bool_]:
    """NaN-aware element-wise equality between *array* and a nodata sentinel.

    Equivalent to ``array == nodata_value`` but also matches NaN positions when
    the nodata sentinel is NaN.

    Args:
        array: the data array (any shape).
        nodata_value: scalar nodata sentinel.

    Returns:
        Boolean mask with the same shape as *array*; True where the value
        equals (or is NaN-matching) the nodata sentinel.
    """
    if isinstance(nodata_value, float) and math.isnan(nodata_value):
        return np.isnan(array)
    return array == nodata_value

unique_nodata_value

unique_nodata_value(values: Sequence[int | float]) -> int | float | None

Return the single unique value from values, or None if empty.

NaN-aware: all NaN entries are treated as equal.

Raises:

Type Description
ValueError

if values contains more than one distinct value.

Source code in rslearn/utils/array.py
def unique_nodata_value(values: Sequence[int | float]) -> int | float | None:
    """Return the single unique value from *values*, or None if empty.

    NaN-aware: all NaN entries are treated as equal.

    Raises:
        ValueError: if *values* contains more than one distinct value.
    """
    unique: list[int | float] = []
    for v in values:
        is_nan = isinstance(v, float) and math.isnan(v)
        if not any((math.isnan(u) if is_nan else u == v) for u in unique):
            unique.append(v)
    if not unique:
        return None
    if len(unique) > 1:
        raise ValueError(f"Expected exactly one unique nodata value, got {unique}")
    return unique[0]

copy_spatial_array

copy_spatial_array(src: Tensor | NDArray[Any], dst: Tensor | NDArray[Any], src_offset: tuple[int, int], dst_offset: tuple[int, int]) -> None

Copy image content from a source array onto a destination array.

The source and destination might be in the same coordinate system. Only the portion of the source array that overlaps in the coordinate system with the destination array will be copied, and other parts of the destination array will not be overwritten.

Supports arrays with any number of dimensions >= 2. The last two dimensions are treated as (H, W); any leading dimensions (e.g. C, or C and T) are preserved.

Parameters:

Name Type Description Default
src Tensor | NDArray[Any]

the source array (...HW).

required
dst Tensor | NDArray[Any]

the destination array (...HW).

required
src_offset tuple[int, int]

the (col, row) position of the top-left pixel of src in the coordinate system.

required
dst_offset tuple[int, int]

the (col, row) position of the top-left pixel of dst in the coordinate system.

required
Source code in rslearn/utils/array.py
def copy_spatial_array(
    src: "torch.Tensor | npt.NDArray[Any]",
    dst: "torch.Tensor | npt.NDArray[Any]",
    src_offset: tuple[int, int],
    dst_offset: tuple[int, int],
) -> None:
    """Copy image content from a source array onto a destination array.

    The source and destination might be in the same coordinate system. Only the portion
    of the source array that overlaps in the coordinate system with the destination
    array will be copied, and other parts of the destination array will not be
    overwritten.

    Supports arrays with any number of dimensions >= 2. The last two dimensions are
    treated as (H, W); any leading dimensions (e.g. C, or C and T) are preserved.

    Args:
        src: the source array (``...HW``).
        dst: the destination array (``...HW``).
        src_offset: the (col, row) position of the top-left pixel of src in the coordinate
            system.
        dst_offset: the (col, row) position of the top-left pixel of dst in the coordinate
            system.
    """
    if len(src.shape) < 2 or len(dst.shape) < 2:
        raise ValueError(f"src and dst must be at least 2-D, got shape {src.shape}")

    src_height, src_width = src.shape[-2:]
    dst_height, dst_width = dst.shape[-2:]
    # The top-left position within src that intersects with dst.
    src_col_offset = max(dst_offset[0] - src_offset[0], 0)
    src_row_offset = max(dst_offset[1] - src_offset[1], 0)
    # The top-left position within dst that intersects with src.
    # This is the position in dst of the same pixel as the one above in src.
    dst_col_offset = max(src_offset[0] - dst_offset[0], 0)
    dst_row_offset = max(src_offset[1] - dst_offset[1], 0)
    # Now compute how much of src we can copy.
    col_overlap = min(src_width - src_col_offset, dst_width - dst_col_offset)
    row_overlap = min(src_height - src_row_offset, dst_height - dst_row_offset)

    # Skip copy if there's no overlap.
    if col_overlap <= 0 or row_overlap <= 0:
        return

    dst[
        ...,
        dst_row_offset : dst_row_offset + row_overlap,
        dst_col_offset : dst_col_offset + col_overlap,
    ] = src[
        ...,
        src_row_offset : src_row_offset + row_overlap,
        src_col_offset : src_col_offset + col_overlap,
    ]