rslearn.dataset.window_data_storage.per_layer

per_layer

WindowDataStorage that packs all item groups into one raster file per layer.

Stores combined rasters at layers/{layer_name}/{bandset_dir}/... with a window_storage_meta.json sidecar recording each group's number of timesteps. Item groups are concatenated along the T axis, so reading back any individual group requires the sidecar to know which T-slice corresponds to which group.

Vector data is still stored per-item-group.

PerLayerStorage

Bases: WindowDataStorage

Storage that packs all item groups for a layer into one raster file.

Raster on-disk layout: layers/{layer_name}/{bandset_dir}/ with a window_storage_meta.json sidecar recording each group's number of timesteps. Item groups are concatenated along the T axis. This is more efficient when reading all groups at once but does not support concurrent or incremental writes.

Vector data is not packed by this storage; vector reads and writes fall through to :class:PerItemGroupStorage.

Source code in rslearn/dataset/window_data_storage/per_layer.py

class PerLayerStorage(WindowDataStorage):
    """Storage that packs all item groups for a layer into one raster file.

    Raster on-disk layout: ``layers/{layer_name}/{bandset_dir}/`` with a
    ``window_storage_meta.json`` sidecar recording each group's number of
    timesteps. Item groups are concatenated along the T axis. This is more
    efficient when reading all groups at once but does not support concurrent
    or incremental writes.

    Vector data is not packed by this storage; vector reads and writes fall
    through to :class:`PerItemGroupStorage`.
    """

    def __init__(self, window: Window) -> None:
        """Initialize the storage bound to a specific window."""
        super().__init__(window)
        # PerLayerStorage delegates vector ops to PerItemGroupStorage.
        self._per_item_group_storage = PerItemGroupStorage(window)

    @override
    def open_layer_writer(
        self,
        layer_name: str,
    ) -> LayerWriter:
        """Return a writer that buffers all groups before flushing."""
        return _PerLayerStorageLayerWriter(self.window, layer_name)

    @override
    def read_raster(
        self,
        layer_name: str,
        bands: list[str],
        raster_format: RasterFormat,
        projection: Projection | None = None,
        bounds: PixelBounds | None = None,
        group_idx: int = 0,
        resampling: Resampling = Resampling.bilinear,
    ) -> RasterArray:
        """Decode the combined raster and slice out the requested item group."""
        return self.read_rasters(
            layer_name,
            bands,
            [group_idx],
            raster_format,
            projection,
            bounds,
            resampling,
        )[0]

    @override
    def read_rasters(
        self,
        layer_name: str,
        bands: list[str],
        group_idxs: list[int],
        raster_format: RasterFormat,
        projection: Projection | None = None,
        bounds: PixelBounds | None = None,
        resampling: Resampling = Resampling.bilinear,
    ) -> list[RasterArray]:
        """Decode the combined raster once and slice out the requested groups."""
        proj = projection if projection is not None else self.window.projection
        bnds = bounds if bounds is not None else self.window.bounds
        raster_dir = _per_layer_raster_dir(self.window.window_root, layer_name, bands)
        meta = _PerLayerStorageMeta.read(raster_dir)
        for group_idx in group_idxs:
            if group_idx < 0 or group_idx >= meta.num_groups:
                raise ValueError(
                    f"PerLayerStorage: group_idx {group_idx} out of range "
                    f"[0, {meta.num_groups})"
                )
        combined = raster_format.decode_raster(raster_dir, proj, bnds, resampling)
        if combined.array.shape[1] != sum(meta.group_timestep_counts):
            raise ValueError(
                f"PerLayerStorage: combined raster T={combined.array.shape[1]} "
                f"does not match sidecar sum {sum(meta.group_timestep_counts)}"
            )
        out: list[RasterArray] = []
        for group_idx in group_idxs:
            t_slice = meta.t_slice_for_group(group_idx)
            sub_array = combined.array[:, t_slice, :, :].copy()
            sub_timestamps: list[tuple[datetime, datetime]] | None
            if combined.timestamps is None:
                sub_timestamps = None
            else:
                sub_timestamps = list(combined.timestamps[t_slice])
            out.append(
                RasterArray(
                    array=sub_array,
                    timestamps=sub_timestamps,
                    metadata=RasterMetadata(
                        nodata_value=combined.metadata.nodata_value
                    ),
                )
            )
        return out

    @override
    def read_vector(
        self,
        layer_name: str,
        vector_format: VectorFormat,
        projection: Projection | None = None,
        bounds: PixelBounds | None = None,
        group_idx: int = 0,
    ) -> list[Feature]:
        """Decode vector features per-item-group (we don't currently handle per-layer vector)."""
        return self._per_item_group_storage.read_vector(
            layer_name,
            vector_format,
            projection,
            bounds,
            group_idx=group_idx,
        )

open_layer_writer

open_layer_writer(layer_name: str) -> LayerWriter

Return a writer that buffers all groups before flushing.

Source code in rslearn/dataset/window_data_storage/per_layer.py

@override
def open_layer_writer(
    self,
    layer_name: str,
) -> LayerWriter:
    """Return a writer that buffers all groups before flushing."""
    return _PerLayerStorageLayerWriter(self.window, layer_name)

read_raster

read_raster(layer_name: str, bands: list[str], raster_format: RasterFormat, projection: Projection | None = None, bounds: PixelBounds | None = None, group_idx: int = 0, resampling: Resampling = bilinear) -> RasterArray

Decode the combined raster and slice out the requested item group.

Source code in rslearn/dataset/window_data_storage/per_layer.py

@override
def read_raster(
    self,
    layer_name: str,
    bands: list[str],
    raster_format: RasterFormat,
    projection: Projection | None = None,
    bounds: PixelBounds | None = None,
    group_idx: int = 0,
    resampling: Resampling = Resampling.bilinear,
) -> RasterArray:
    """Decode the combined raster and slice out the requested item group."""
    return self.read_rasters(
        layer_name,
        bands,
        [group_idx],
        raster_format,
        projection,
        bounds,
        resampling,
    )[0]

read_rasters

read_rasters(layer_name: str, bands: list[str], group_idxs: list[int], raster_format: RasterFormat, projection: Projection | None = None, bounds: PixelBounds | None = None, resampling: Resampling = bilinear) -> list[RasterArray]

Decode the combined raster once and slice out the requested groups.

Source code in rslearn/dataset/window_data_storage/per_layer.py

@override
def read_rasters(
    self,
    layer_name: str,
    bands: list[str],
    group_idxs: list[int],
    raster_format: RasterFormat,
    projection: Projection | None = None,
    bounds: PixelBounds | None = None,
    resampling: Resampling = Resampling.bilinear,
) -> list[RasterArray]:
    """Decode the combined raster once and slice out the requested groups."""
    proj = projection if projection is not None else self.window.projection
    bnds = bounds if bounds is not None else self.window.bounds
    raster_dir = _per_layer_raster_dir(self.window.window_root, layer_name, bands)
    meta = _PerLayerStorageMeta.read(raster_dir)
    for group_idx in group_idxs:
        if group_idx < 0 or group_idx >= meta.num_groups:
            raise ValueError(
                f"PerLayerStorage: group_idx {group_idx} out of range "
                f"[0, {meta.num_groups})"
            )
    combined = raster_format.decode_raster(raster_dir, proj, bnds, resampling)
    if combined.array.shape[1] != sum(meta.group_timestep_counts):
        raise ValueError(
            f"PerLayerStorage: combined raster T={combined.array.shape[1]} "
            f"does not match sidecar sum {sum(meta.group_timestep_counts)}"
        )
    out: list[RasterArray] = []
    for group_idx in group_idxs:
        t_slice = meta.t_slice_for_group(group_idx)
        sub_array = combined.array[:, t_slice, :, :].copy()
        sub_timestamps: list[tuple[datetime, datetime]] | None
        if combined.timestamps is None:
            sub_timestamps = None
        else:
            sub_timestamps = list(combined.timestamps[t_slice])
        out.append(
            RasterArray(
                array=sub_array,
                timestamps=sub_timestamps,
                metadata=RasterMetadata(
                    nodata_value=combined.metadata.nodata_value
                ),
            )
        )
    return out

read_vector

read_vector(layer_name: str, vector_format: VectorFormat, projection: Projection | None = None, bounds: PixelBounds | None = None, group_idx: int = 0) -> list[Feature]

Decode vector features per-item-group (we don't currently handle per-layer vector).

Source code in rslearn/dataset/window_data_storage/per_layer.py

@override
def read_vector(
    self,
    layer_name: str,
    vector_format: VectorFormat,
    projection: Projection | None = None,
    bounds: PixelBounds | None = None,
    group_idx: int = 0,
) -> list[Feature]:
    """Decode vector features per-item-group (we don't currently handle per-layer vector)."""
    return self._per_item_group_storage.read_vector(
        layer_name,
        vector_format,
        projection,
        bounds,
        group_idx=group_idx,
    )

PerLayerStorageFactory

Bases: WindowDataStorageFactory

Factory that creates :class:PerLayerStorage instances.

Source code in rslearn/dataset/window_data_storage/per_layer.py

class PerLayerStorageFactory(WindowDataStorageFactory):
    """Factory that creates :class:`PerLayerStorage` instances."""

    @override
    def create(self, window: Window) -> PerLayerStorage:
        """Create a PerLayerStorage bound to the given window."""
        return PerLayerStorage(window)

create

create(window: Window) -> PerLayerStorage

Create a PerLayerStorage bound to the given window.

Source code in rslearn/dataset/window_data_storage/per_layer.py

@override
def create(self, window: Window) -> PerLayerStorage:
    """Create a PerLayerStorage bound to the given window."""
    return PerLayerStorage(window)