rslearn 0.0.25__py3-none-any.whl → 0.0.27__py3-none-any.whl

rslearn/config/dataset.py

@@ -236,11 +236,9 @@ class BandSetConfig(BaseModel):
 
         warnings.warn(
             "`format = {'name': ...}` is deprecated; "
-            "use `{'class_path': '...', 'init_args': {...}}` instead.",
-            DeprecationWarning,
-        )
-        logger.warning(
-            "BandSet.format uses legacy format; support will be removed after 2026-03-01."
+            "use `{'class_path': '...', 'init_args': {...}}` instead. "
+            "Support will be removed after 2026-03-01.",
+            FutureWarning,
         )
 
         legacy_name_to_class_path = {
@@ -294,16 +292,6 @@ class SpaceMode(StrEnum):
     The duration of the sub-periods is controlled by another option in QueryConfig.
     """
 
-    COMPOSITE = "COMPOSITE"
-    """Creates one composite covering the entire window.
-
-    During querying all items intersecting the window are placed in one group.
-    The compositing_method in the rasterlayer config specifies how these items are reduced
-    to a single item (e.g MEAN/MEDIAN/FIRST_VALID) during materialization.
-    """
-
-    # TODO add PER_PERIOD_COMPOSITE
-
 
 class TimeMode(StrEnum):
     """Temporal  matching mode when looking up items corresponding to a window."""
@@ -353,6 +341,20 @@ class QueryConfig(BaseModel):
         default=timedelta(days=30),
         description="The duration of the periods, if the space mode is PER_PERIOD_MOSAIC.",
     )
+    mosaic_compositing_overlaps: int = Field(
+        default=1,
+        description="For MOSAIC and PER_PERIOD_MOSAIC modes, the number of overlapping items "
+        "wanted within each item group covering the window. Set to 1 for a single coverage "
+        "(default mosaic behavior), or higher for compositing multiple overlapping items."
+        "with mean or median compositing method.",
+    )
+    per_period_mosaic_reverse_time_order: bool = Field(
+        default=True,
+        description="For PER_PERIOD_MOSAIC mode, whether to return item groups in reverse "
+        "temporal order (most recent first). Set to False for chronological order (oldest first). "
+        "Default True is deprecated and will change to False with error if still unset or set True "
+        "after 2026-04-01.",
+    )
 
 
 class DataSourceConfig(BaseModel):
@@ -404,11 +406,9 @@ class DataSourceConfig(BaseModel):
 
         warnings.warn(
             "`Data source configuration {'name': ...}` is deprecated; "
-            "use `{'class_path': '...', 'init_args': {...}, ...}` instead.",
-            DeprecationWarning,
-        )
-        logger.warning(
-            "Data source configuration uses legacy format; support will be removed after 2026-03-01."
+            "use `{'class_path': '...', 'init_args': {...}, ...}` instead. "
+            "Support will be removed after 2026-03-01.",
+            FutureWarning,
         )
 
         # Split the dict into the base config that is in the pydantic model, and the
@@ -431,8 +431,9 @@ class DataSourceConfig(BaseModel):
             and "max_cloud_cover" in ds_init_args
         ):
             warnings.warn(
-                "Data source configuration specifies invalid 'max_cloud_cover' option.",
-                DeprecationWarning,
+                "Data source configuration specifies invalid 'max_cloud_cover' option."
+                "Support for ignoring this option will be removed after 2026-03-01.",
+                FutureWarning,
             )
             del ds_init_args["max_cloud_cover"]
 
@@ -449,7 +450,13 @@ class LayerType(StrEnum):
 
 
 class CompositingMethod(StrEnum):
-    """Method how to select pixels for the composite from corresponding items of a window."""
+    """Method how to select pixels for the composite from corresponding items of a window.
+
+    For MEAN and MEDIAN modes, mosaic_compositing_overlaps (in the QueryConfig) should
+    be set higher than 1 so that rslearn creates item groups during prepare that cover
+    the window with multiple overlaps. At each pixel/band, the mean and median can then
+    be computed across items in each group that cover that pixel.
+    """
 
     FIRST_VALID = "FIRST_VALID"
     """Select first valid pixel in order of corresponding items (might be sorted)"""

rslearn/data_sources/__init__.py

@@ -17,6 +17,7 @@ from .data_source import (
     ItemLookupDataSource,
     RetrieveItemDataSource,
 )
+from .direct_materialize_data_source import DirectMaterializeDataSource
 
 __all__ = (
     "DataSource",
@@ -24,5 +25,6 @@ __all__ = (
     "Item",
     "ItemLookupDataSource",
     "RetrieveItemDataSource",
+    "DirectMaterializeDataSource",
     "data_source_from_config",
 )

rslearn/data_sources/aws_landsat.py

@@ -9,32 +9,28 @@ import urllib.request
 import zipfile
 from collections.abc import Generator
 from datetime import datetime
-from typing import Any, BinaryIO
+from typing import BinaryIO
 
-import affine
 import boto3
 import dateutil.parser
 import fiona
 import fiona.transform
-import numpy.typing as npt
-import rasterio
 import shapely
 import shapely.geometry
 import tqdm
-from rasterio.enums import Resampling
 from upath import UPath
 
 import rslearn.data_sources.utils
-from rslearn.config import LayerConfig
 from rslearn.const import SHAPEFILE_AUX_EXTENSIONS, WGS84_PROJECTION
-from rslearn.dataset import Window
-from rslearn.dataset.materialize import RasterMaterializer
-from rslearn.tile_stores import TileStore, TileStoreWithLayer
+from rslearn.data_sources.direct_materialize_data_source import (
+    DirectMaterializeDataSource,
+)
+from rslearn.tile_stores import TileStoreWithLayer
 from rslearn.utils.fsspec import get_upath_local, join_upath, open_atomic
-from rslearn.utils.geometry import PixelBounds, Projection, STGeometry
+from rslearn.utils.geometry import STGeometry
 from rslearn.utils.grid_index import GridIndex
 
-from .data_source import DataSource, DataSourceContext, Item, QueryConfig
+from .data_source import DataSourceContext, Item, QueryConfig
 
 WRS2_GRID_SIZE = 1.0
 
@@ -79,7 +75,7 @@ class LandsatOliTirsItem(Item):
         )
 
 
-class LandsatOliTirs(DataSource, TileStore):
+class LandsatOliTirs(DirectMaterializeDataSource[LandsatOliTirsItem]):
     """A data source for Landsat 8/9 OLI-TIRS imagery on AWS.
 
     Specifically, uses the usgs-landsat S3 bucket maintained by USGS. The data includes
@@ -91,7 +87,7 @@ class LandsatOliTirs(DataSource, TileStore):
 
     bucket_name = "usgs-landsat"
     bucket_prefix = "collection02/level-1/standard/oli-tirs"
-    bands = ["B1", "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B9", "B10", "B11"]
+    BANDS = ["B1", "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B9", "B10", "B11"]
 
     wrs2_url = "https://d9-wret.s3.us-west-2.amazonaws.com/assets/palladium/production/s3fs-public/atoms/files/WRS2_descending_0.zip"  # noqa
     """URL to download shapefile specifying polygon of each (path, row)."""
@@ -110,6 +106,10 @@ class LandsatOliTirs(DataSource, TileStore):
                 SpaceMode.WITHIN.
             context: the data source context.
         """
+        # Each band is a separate single-band asset.
+        asset_bands = {band: [band] for band in self.BANDS}
+        super().__init__(asset_bands=asset_bands)
+
         # If context is provided, we join the directory with the dataset path,
         # otherwise we treat it directly as UPath.
         if context.ds_path is not None:
@@ -342,16 +342,43 @@ class LandsatOliTirs(DataSource, TileStore):
                 return item
         raise ValueError(f"item {name} not found")
 
-    def deserialize_item(self, serialized_item: Any) -> LandsatOliTirsItem:
+    # --- DirectMaterializeDataSource implementation ---
+
+    def get_asset_url(self, item_name: str, asset_key: str) -> str:
+        """Get the presigned URL to read the asset for the given item and asset key.
+
+        Args:
+            item_name: the name of the item.
+            asset_key: the key identifying which asset to get (the band name).
+
+        Returns:
+            the presigned URL to read the asset from.
+        """
+        # Get the item since it has the blob path.
+        item = self.get_item_by_name(item_name)
+
+        # For Landsat, the asset_key is the band name (e.g., "B1", "B2", etc.).
+        blob_key = item.blob_path + f"{asset_key}.TIF"
+        return self.client.generate_presigned_url(
+            "get_object",
+            Params={
+                "Bucket": self.bucket_name,
+                "Key": blob_key,
+                "RequestPayer": "requester",
+            },
+        )
+
+    # --- DataSource implementation ---
+
+    def deserialize_item(self, serialized_item: dict) -> LandsatOliTirsItem:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return LandsatOliTirsItem.deserialize(serialized_item)
 
     def retrieve_item(
         self, item: LandsatOliTirsItem
     ) -> Generator[tuple[str, BinaryIO], None, None]:
         """Retrieves the rasters corresponding to an item as file streams."""
-        for band in self.bands:
+        for band in self.BANDS:
             buf = io.BytesIO()
             self.bucket.download_fileobj(
                 item.blob_path + f"{band}.TIF",
@@ -376,7 +403,7 @@ class LandsatOliTirs(DataSource, TileStore):
             geometries: a list of geometries needed for each item
         """
         for item, cur_geometries in zip(items, geometries):
-            for band in self.bands:
+            for band in self.BANDS:
                 band_names = [band]
                 if tile_store.is_raster_ready(item.name, band_names):
                     continue
@@ -389,147 +416,3 @@ class LandsatOliTirs(DataSource, TileStore):
                         ExtraArgs={"RequestPayer": "requester"},
                     )
                     tile_store.write_raster_file(item.name, band_names, UPath(fname))
-
-    # The functions below are to emulate TileStore functionality so we can easily
-    # support materialization directly from the COGs.
-    def is_raster_ready(
-        self, layer_name: str, item_name: str, bands: list[str]
-    ) -> bool:
-        """Checks if this raster has been written to the store.
-
-        Args:
-            layer_name: the layer name or alias.
-            item_name: the item.
-            bands: the list of bands identifying which specific raster to read.
-
-        Returns:
-            whether there is a raster in the store matching the source, item, and
-                bands.
-        """
-        # Always ready since we access it on AWS bucket.
-        return True
-
-    def get_raster_bands(self, layer_name: str, item_name: str) -> list[list[str]]:
-        """Get the sets of bands that have been stored for the specified item.
-
-        Args:
-            layer_name: the layer name or alias.
-            item_name: the item.
-
-        Returns:
-            a list of lists of bands that are in the tile store (with one raster
-                stored corresponding to each inner list). If no rasters are ready for
-                this item, returns empty list.
-        """
-        return [[band] for band in self.bands]
-
-    def get_raster_bounds(
-        self, layer_name: str, item_name: str, bands: list[str], projection: Projection
-    ) -> PixelBounds:
-        """Get the bounds of the raster in the specified projection.
-
-        Args:
-            layer_name: the layer name or alias.
-            item_name: the item to check.
-            bands: the list of bands identifying which specific raster to read. These
-                bands must match the bands of a stored raster.
-            projection: the projection to get the raster's bounds in.
-
-        Returns:
-            the bounds of the raster in the projection.
-        """
-        item = self.get_item_by_name(item_name)
-        geom = item.geometry.to_projection(projection)
-        return (
-            int(geom.shp.bounds[0]),
-            int(geom.shp.bounds[1]),
-            int(geom.shp.bounds[2]),
-            int(geom.shp.bounds[3]),
-        )
-
-    def read_raster(
-        self,
-        layer_name: str,
-        item_name: str,
-        bands: list[str],
-        projection: Projection,
-        bounds: PixelBounds,
-        resampling: Resampling = Resampling.bilinear,
-    ) -> npt.NDArray[Any]:
-        """Read raster data from the store.
-
-        Args:
-            layer_name: the layer name or alias.
-            item_name: the item to read.
-            bands: the list of bands identifying which specific raster to read. These
-                bands must match the bands of a stored raster.
-            projection: the projection to read in.
-            bounds: the bounds to read.
-            resampling: the resampling method to use in case reprojection is needed.
-
-        Returns:
-            the raster data
-        """
-        # Landsat assets have single band per asset.
-        assert len(bands) == 1
-        band = bands[0]
-
-        # Get the item since it has the blob path.
-        item = self.get_item_by_name(item_name)
-
-        # Create pre-signed URL for rasterio access.
-        # We do this because accessing via URL is much faster since rasterio can use
-        # the URL directly.
-        blob_key = item.blob_path + f"{band}.TIF"
-        url = self.client.generate_presigned_url(
-            "get_object",
-            Params={
-                "Bucket": self.bucket_name,
-                "Key": blob_key,
-                "RequestPayer": "requester",
-            },
-        )
-
-        # Construct the transform to use for the warped dataset.
-        wanted_transform = affine.Affine(
-            projection.x_resolution,
-            0,
-            bounds[0] * projection.x_resolution,
-            0,
-            projection.y_resolution,
-            bounds[1] * projection.y_resolution,
-        )
-
-        with rasterio.open(url) as src:
-            with rasterio.vrt.WarpedVRT(
-                src,
-                crs=projection.crs,
-                transform=wanted_transform,
-                width=bounds[2] - bounds[0],
-                height=bounds[3] - bounds[1],
-                resampling=resampling,
-            ) as vrt:
-                return vrt.read()
-
-    def materialize(
-        self,
-        window: Window,
-        item_groups: list[list[LandsatOliTirsItem]],
-        layer_name: str,
-        layer_cfg: LayerConfig,
-    ) -> None:
-        """Materialize data for the window.
-
-        Args:
-            window: the window to materialize
-            item_groups: the items from get_items
-            layer_name: the name of this layer
-            layer_cfg: the config of this layer
-        """
-        RasterMaterializer().materialize(
-            TileStoreWithLayer(self, layer_name),
-            window,
-            layer_name,
-            layer_cfg,
-            item_groups,
-        )

rslearn/data_sources/aws_open_data.py

@@ -318,9 +318,8 @@ class Naip(DataSource):
             groups.append(cur_groups)
         return groups
 
-    def deserialize_item(self, serialized_item: Any) -> NaipItem:
+    def deserialize_item(self, serialized_item: dict) -> NaipItem:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return NaipItem.deserialize(serialized_item)
 
     def ingest(
@@ -639,9 +638,8 @@ class Sentinel2(
                 return item
         raise ValueError(f"item {name} not found")
 
-    def deserialize_item(self, serialized_item: Any) -> Sentinel2Item:
+    def deserialize_item(self, serialized_item: dict) -> Sentinel2Item:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return Sentinel2Item.deserialize(serialized_item)
 
     def retrieve_item(

rslearn/data_sources/aws_sentinel1.py

@@ -2,7 +2,6 @@
 
 import os
 import tempfile
-from typing import Any
 
 import boto3
 from upath import UPath
@@ -78,9 +77,8 @@ class Sentinel1(DataSource, TileStore):
         """Gets an item by name."""
         return self.sentinel1.get_item_by_name(name)
 
-    def deserialize_item(self, serialized_item: Any) -> CopernicusItem:
+    def deserialize_item(self, serialized_item: dict) -> CopernicusItem:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return CopernicusItem.deserialize(serialized_item)
 
     def ingest(