rslearn.dataset.tile_utils

tile_utils

Low-level helpers for reading raster data from tile stores.

get_needed_band_sets_and_indexes

get_needed_band_sets_and_indexes(item: ItemType, bands: list[str], tile_store: TileStoreWithLayer) -> list[tuple[list[str], list[int], list[int]]]

Identify indexes of required bands in tile store.

Returns:

Type	Description
list[tuple[list[str], list[int], list[int]]]	A list for each tile-store layer that contains at least
list[tuple[list[str], list[int], list[int]]]	one requested band, a tuple: (src_bands, src_idx, dst_idx) where
list[tuple[list[str], list[int], list[int]]]	src_bands: the full band list for that layer,
list[tuple[list[str], list[int], list[int]]]	src_idx: indexes into src_bands of the bands that were requested,
list[tuple[list[str], list[int], list[int]]]	dst_idx: corresponding indexes in the requested bands list.

Source code in rslearn/dataset/tile_utils.py

def get_needed_band_sets_and_indexes(
    item: ItemType,
    bands: list[str],
    tile_store: TileStoreWithLayer,
) -> list[tuple[list[str], list[int], list[int]]]:
    """Identify indexes of required bands in tile store.

    Returns:
        A list for each tile-store layer that contains at least
        one requested band, a tuple: (src_bands, src_idx, dst_idx) where
        - src_bands: the full band list for that layer,
        - src_idx: indexes into src_bands of the bands that were requested,
        - dst_idx: corresponding indexes in the requested `bands` list.
    """
    wanted_band_indexes = {}
    for i, band in enumerate(bands):
        wanted_band_indexes[band] = i

    available_bands = tile_store.get_raster_bands(item)
    needed_band_sets_and_indexes = []

    for src_bands in available_bands:
        needed_src_indexes = []
        needed_dst_indexes = []
        for i, band in enumerate(src_bands):
            if band not in wanted_band_indexes:
                continue
            needed_src_indexes.append(i)
            needed_dst_indexes.append(wanted_band_indexes[band])
            del wanted_band_indexes[band]
        if len(needed_src_indexes) == 0:
            continue
        needed_band_sets_and_indexes.append(
            (src_bands, needed_src_indexes, needed_dst_indexes)
        )

    if len(wanted_band_indexes) > 0:
        return []

    return needed_band_sets_and_indexes

read_raster_window_from_tiles

read_raster_window_from_tiles(tile_store: TileStoreWithLayer, item: ItemType, bands: list[str], projection: Projection, bounds: PixelBounds, nodata_val: int | float | None, band_dtype: DTypeLike, remapper: Remapper | None = None, resampling: Resampling = bilinear, dst: RasterArray | None = None) -> RasterArray | None

Read an item's raster data from tiles into a window-aligned RasterArray.

Handles band mapping and spatial intersection internally. Uses first-valid nodata logic: only overwrites pixels where all bands equal the nodata value.

When nodata_val is None, every source pixel unconditionally overwrites the destination, and the destination is initialized as zeros.

Parameters:

Name	Type	Description	Default
tile_store	TileStoreWithLayer	the TileStore to read from.	required
item	ItemType	the item to read.	required
bands	list[str]	the requested band names (determines dst band order).	required
projection	Projection	the projection of the window.	required
bounds	PixelBounds	the pixel bounds of the window.	required
nodata_val	int | float | None	the scalar nodata value for the band set, or None.	required
band_dtype	DTypeLike	data type for the output array.	required
remapper	Remapper | None	optional remapper to apply on the source pixel values.	None
resampling	Resampling	how to resample pixels if re-projection is needed.	bilinear
dst	RasterArray | None	optional pre-allocated RasterArray to write into. If None, a new one is allocated with T from the item's tile store data and H/W from the given bounds.	None

Returns:

Type	Description
RasterArray | None	The dst RasterArray (allocated or provided), or None if the item has no
RasterArray | None	matching bands and dst was not provided.

Source code in rslearn/dataset/tile_utils.py

def read_raster_window_from_tiles(
    tile_store: TileStoreWithLayer,
    item: ItemType,
    bands: list[str],
    projection: Projection,
    bounds: PixelBounds,
    nodata_val: int | float | None,
    band_dtype: npt.DTypeLike,
    remapper: Remapper | None = None,
    resampling: Resampling = Resampling.bilinear,
    dst: RasterArray | None = None,
) -> RasterArray | None:
    """Read an item's raster data from tiles into a window-aligned RasterArray.

    Handles band mapping and spatial intersection internally. Uses first-valid
    nodata logic: only overwrites pixels where all bands equal the nodata value.

    When nodata_val is None, every source pixel unconditionally overwrites the
    destination, and the destination is initialized as zeros.

    Args:
        tile_store: the TileStore to read from.
        item: the item to read.
        bands: the requested band names (determines dst band order).
        projection: the projection of the window.
        bounds: the pixel bounds of the window.
        nodata_val: the scalar nodata value for the band set, or ``None``.
        band_dtype: data type for the output array.
        remapper: optional remapper to apply on the source pixel values.
        resampling: how to resample pixels if re-projection is needed.
        dst: optional pre-allocated RasterArray to write into. If None, a new
            one is allocated with T from the item's tile store data and H/W
            from the given bounds.

    Returns:
        The dst RasterArray (allocated or provided), or None if the item has no
        matching bands and dst was not provided.
    """
    needed = get_needed_band_sets_and_indexes(item, bands, tile_store)
    if not needed:
        return dst

    for src_bands, src_indexes, dst_indexes in needed:
        src_bounds = tile_store.get_raster_bounds(item, src_bands, projection)
        intersection = (
            max(bounds[0], src_bounds[0]),
            max(bounds[1], src_bounds[1]),
            min(bounds[2], src_bounds[2]),
            min(bounds[3], src_bounds[3]),
        )
        if intersection[2] <= intersection[0] or intersection[3] <= intersection[1]:
            continue

        raster_array = tile_store.read_raster(
            item, src_bands, projection, intersection, resampling=resampling
        )
        src = raster_array.array  # (C_src, T, H_int, W_int)

        if dst is None:
            num_timesteps = src.shape[1]
            shape = (
                len(bands),
                num_timesteps,
                bounds[3] - bounds[1],
                bounds[2] - bounds[0],
            )
            if nodata_val is not None:
                dst_arr = np.full(shape, nodata_val, dtype=band_dtype)
            else:
                dst_arr = np.zeros(shape, dtype=band_dtype)
            dst = RasterArray(
                array=dst_arr,
                timestamps=raster_array.timestamps,
                metadata=RasterMetadata(nodata_value=nodata_val),
            )

        if src.shape[1] != dst.array.shape[1]:
            raise ValueError(
                f"Item {item.name!r} has T={src.shape[1]} in tile store but "
                f"dst has T={dst.array.shape[1]}"
            )

        dst_col = intersection[0] - bounds[0]
        dst_row = intersection[1] - bounds[1]

        out_crop = dst.array[
            :,
            :,
            dst_row : dst_row + src.shape[2],
            dst_col : dst_col + src.shape[3],
        ]  # (C, T, H_int, W_int) view

        # Build the first-valid mask: a pixel may be written only where it is
        # nodata in *every* destination band. Start from the first band's
        # nodata mask and logically-and (&=) each remaining band into it, so a
        # pixel that is already valid in any band stays False (untouched).
        if nodata_val is not None:
            mask = nodata_eq(out_crop[dst_indexes[0]], nodata_val)
            for dst_index in dst_indexes[1:]:
                mask &= nodata_eq(out_crop[dst_index], nodata_val)
        else:
            # nodata_val is None: every source pixel overwrites the destination.
            mask = None

        # Copy one band at a time. src[src_index] is a view (no copy), and
        # np.copyto writes straight into the destination window, so no
        # intersection-sized temporary array is allocated.
        for src_index, dst_index in zip(src_indexes, dst_indexes, strict=True):
            src_band = src[src_index]  # (T, H_int, W_int) view
            if remapper:
                src_band = remapper(src_band, band_dtype)
            if mask is not None:
                # where=mask restricts the write to first-valid pixels;
                # casting="unsafe" converts src_band's dtype to band_dtype
                # during the copy (e.g. uint16 -> float32).
                np.copyto(out_crop[dst_index], src_band, where=mask, casting="unsafe")
            else:
                np.copyto(out_crop[dst_index], src_band, casting="unsafe")

    return dst

read_raster_windows

read_raster_windows(group: list[ItemType], bands: list[str], tile_store: TileStoreWithLayer, projection: Projection, bounds: PixelBounds, nodata_val: int | float | None, band_dtype: DTypeLike, remapper: Remapper | None = None, resampling_method: Resampling = bilinear) -> list[RasterArray]

Read each item in the group into a window-aligned RasterArray.

Each returned RasterArray has shape (C, T, H, W) where C = len(bands), T is determined by the item's data in the tile store, and H/W match the window bounds.

Parameters:

Name	Type	Description	Default
group	list[ItemType]	items to read. We create one RasterArray per item.	required
bands	list[str]	requested band names. The bands in the RasterArrays will appear in this order.	required
tile_store	TileStoreWithLayer	tile store containing raster data.	required
projection	Projection	target projection.	required
bounds	PixelBounds	target pixel bounds.	required
nodata_val	int | float | None	scalar nodata value for the band set, or None.	required
band_dtype	DTypeLike	output data type.	required
remapper	Remapper | None	optional remapper.	None
resampling_method	Resampling	resampling method.	bilinear

Returns:

Type	Description
list[RasterArray]	A list of RasterArrays, one per item that had matching bands.

Source code in rslearn/dataset/tile_utils.py

def read_raster_windows(
    group: list[ItemType],
    bands: list[str],
    tile_store: TileStoreWithLayer,
    projection: Projection,
    bounds: PixelBounds,
    nodata_val: int | float | None,
    band_dtype: npt.DTypeLike,
    remapper: Remapper | None = None,
    resampling_method: Resampling = Resampling.bilinear,
) -> list[RasterArray]:
    """Read each item in the group into a window-aligned RasterArray.

    Each returned RasterArray has shape (C, T, H, W) where C = len(bands),
    T is determined by the item's data in the tile store, and H/W match the
    window bounds.

    Args:
        group: items to read. We create one RasterArray per item.
        bands: requested band names. The bands in the RasterArrays will appear in
            this order.
        tile_store: tile store containing raster data.
        projection: target projection.
        bounds: target pixel bounds.
        nodata_val: scalar nodata value for the band set, or ``None``.
        band_dtype: output data type.
        remapper: optional remapper.
        resampling_method: resampling method.

    Returns:
        A list of RasterArrays, one per item that had matching bands.
    """
    results: list[RasterArray] = []
    for item in group:
        result = read_raster_window_from_tiles(
            tile_store=tile_store,
            item=item,
            bands=bands,
            projection=projection,
            bounds=bounds,
            nodata_val=nodata_val,
            band_dtype=band_dtype,
            remapper=remapper,
            resampling=resampling_method,
        )
        if result is not None:
            results.append(result)
    return results

mask_stacked_rasters

mask_stacked_rasters(stacked_rasters: NDArray[generic], nodata_val: int | float) -> MaskedArray

Mask stacked rasters where pixels equal the nodata value.

Parameters:

Name	Type	Description	Default
stacked_rasters	NDArray[generic]	numpy array of shape (num_items, num_bands, T, height, width) containing raster values for each item in the group.	required
nodata_val	int | float	scalar nodata value used to identify invalid pixels.	required

Returns:

Type	Description
MaskedArray	np.ma.MaskedArray with the same shape as stacked_rasters, where all
MaskedArray	pixels equal to the nodata value are masked.

Source code in rslearn/dataset/tile_utils.py

def mask_stacked_rasters(
    stacked_rasters: npt.NDArray[np.generic],
    nodata_val: int | float,
) -> np.ma.MaskedArray:
    """Mask stacked rasters where pixels equal the nodata value.

    Args:
        stacked_rasters: numpy array of shape (num_items, num_bands, T, height, width)
            containing raster values for each item in the group.
        nodata_val: scalar nodata value used to identify invalid pixels.

    Returns:
        np.ma.MaskedArray with the same shape as `stacked_rasters`, where all
        pixels equal to the nodata value are masked.
    """
    is_nodata = nodata_eq(stacked_rasters, nodata_val)
    return np.ma.masked_where(is_nodata, stacked_rasters)