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

rslearn/data_sources/aws_sentinel2_element84.py

@@ -6,23 +6,20 @@ from collections.abc import Callable
 from datetime import timedelta
 from typing import Any
 
-import affine
 import numpy as np
 import numpy.typing as npt
 import rasterio
 import requests
-from rasterio.enums import Resampling
 from upath import UPath
 
-from rslearn.config import LayerConfig
+from rslearn.data_sources.direct_materialize_data_source import (
+    DirectMaterializeDataSource,
+)
 from rslearn.data_sources.stac import SourceItem, StacDataSource
-from rslearn.dataset import Window
-from rslearn.dataset.manage import RasterMaterializer
 from rslearn.log_utils import get_logger
-from rslearn.tile_stores import TileStore, TileStoreWithLayer
-from rslearn.utils import Projection, STGeometry
+from rslearn.tile_stores import TileStoreWithLayer
+from rslearn.utils import STGeometry
 from rslearn.utils.fsspec import join_upath
-from rslearn.utils.geometry import PixelBounds
 from rslearn.utils.raster_format import get_raster_projection_and_bounds
 
 from .data_source import (
@@ -32,7 +29,7 @@ from .data_source import (
 logger = get_logger(__name__)
 
 
-class Sentinel2(StacDataSource, TileStore):
+class Sentinel2(DirectMaterializeDataSource[SourceItem], StacDataSource):
     """A data source for Sentinel-2 L2A imagery on AWS from s3://sentinel-cogs.
 
     The S3 bucket has COGs so this data source supports direct materialization. It also
@@ -97,31 +94,36 @@ class Sentinel2(StacDataSource, TileStore):
             cache_upath.mkdir(parents=True, exist_ok=True)
 
         # Determine which assets we need based on the bands in the layer config.
-        self.asset_bands: dict[str, list[str]]
+        asset_bands: dict[str, list[str]]
         if context.layer_config is not None:
-            self.asset_bands = {}
+            asset_bands = {}
             for asset_key, band_names in self.ASSET_BANDS.items():
                 # See if the bands provided by this asset intersect with the bands in
                 # at least one configured band set.
                 for band_set in context.layer_config.band_sets:
                     if not set(band_set.bands).intersection(set(band_names)):
                         continue
-                    self.asset_bands[asset_key] = band_names
+                    asset_bands[asset_key] = band_names
                     break
         elif assets is not None:
-            self.asset_bands = {
+            asset_bands = {
                 asset_key: self.ASSET_BANDS[asset_key] for asset_key in assets
             }
         else:
-            self.asset_bands = self.ASSET_BANDS
+            asset_bands = dict(self.ASSET_BANDS)
+
+        # Initialize DirectMaterializeDataSource with asset_bands
+        DirectMaterializeDataSource.__init__(self, asset_bands=asset_bands)
 
-        super().__init__(
+        # Initialize StacDataSource
+        StacDataSource.__init__(
+            self,
             endpoint=self.STAC_ENDPOINT,
             collection_name=self.COLLECTION_NAME,
             query=query,
             sort_by=sort_by,
             sort_ascending=sort_ascending,
-            required_assets=list(self.asset_bands.keys()),
+            required_assets=list(asset_bands.keys()),
             cache_dir=cache_upath,
             properties_to_record=[self.HARMONIZE_PROPERTY_NAME],
         )
@@ -129,6 +131,42 @@ class Sentinel2(StacDataSource, TileStore):
         self.harmonize = harmonize
         self.timeout = timeout
 
+    # --- DirectMaterializeDataSource implementation ---
+
+    def get_asset_url(self, item_name: str, asset_key: str) -> str:
+        """Get the 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.
+
+        Returns:
+            the URL to read the asset from.
+        """
+        item = self.get_item_by_name(item_name)
+        return item.asset_urls[asset_key]
+
+    def get_read_callback(
+        self, item_name: str, asset_key: str
+    ) -> Callable[[npt.NDArray[Any]], npt.NDArray[Any]] | None:
+        """Return a callback to harmonize Sentinel-2 data if needed.
+
+        Args:
+            item_name: the name of the item being read.
+            asset_key: the key identifying which asset is being read.
+
+        Returns:
+            A callback function for harmonization, or None if not needed.
+        """
+        # Visual bands do not need harmonization.
+        if not self.harmonize or asset_key == "visual":
+            return None
+
+        item = self.get_item_by_name(item_name)
+        return self._get_harmonize_callback(item)
+
+    # --- Harmonization helpers ---
+
     def _get_harmonize_callback(
         self, item: SourceItem
     ) -> Callable[[npt.NDArray], npt.NDArray] | None:
@@ -223,152 +261,3 @@ class Sentinel2(StacDataSource, TileStore):
                     item.name,
                     asset_key,
                 )
-
-    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 wrap accesses to underlying API.
-        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 list(self.asset_bands.values())
-
-    def _get_asset_by_band(self, bands: list[str]) -> str:
-        """Get the name of the asset based on the band names."""
-        for asset_key, asset_bands in self.asset_bands.items():
-            if bands == asset_bands:
-                return asset_key
-
-        raise ValueError(f"no known asset with bands {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
-        """
-        asset_key = self._get_asset_by_band(bands)
-        item = self.get_item_by_name(item_name)
-        asset_url = item.asset_urls[asset_key]
-
-        # 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,
-        )
-
-        # Read from the raster under the specified projection/bounds.
-        with rasterio.open(asset_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:
-                raw_data = vrt.read()
-
-        # We can return the data now if harmonization is not needed.
-        if not self.harmonize or bands == self.ASSET_BANDS["visual"]:
-            return raw_data
-
-        # Otherwise we apply the harmonize_callback.
-        item = self.get_item_by_name(item_name)
-        harmonize_callback = self._get_harmonize_callback(item)
-
-        if harmonize_callback is None:
-            return raw_data
-
-        array = harmonize_callback(raw_data)
-        return array
-
-    def materialize(
-        self,
-        window: Window,
-        item_groups: list[list[SourceItem]],
-        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/climate_data_store.py

@@ -3,7 +3,6 @@
 import os
 import tempfile
 from datetime import UTC, datetime
-from typing import Any
 
 import cdsapi
 import netCDF4
@@ -160,9 +159,8 @@ class ERA5Land(DataSource):
 
         return all_groups
 
-    def deserialize_item(self, serialized_item: Any) -> Item:
+    def deserialize_item(self, serialized_item: dict) -> Item:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return Item.deserialize(serialized_item)
 
     def _convert_nc_to_tif(self, nc_path: UPath, tif_path: UPath) -> None:

rslearn/data_sources/copernicus.py

@@ -353,9 +353,8 @@ class Copernicus(DataSource):
                 self.username = os.environ["COPERNICUS_USERNAME"]
                 self.password = os.environ["COPERNICUS_PASSWORD"]
 
-    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 _get(self, path: str) -> dict[str, Any]:

rslearn/data_sources/data_source.py

@@ -76,7 +76,7 @@ class DataSource(Generic[ItemType]):
         """
         raise NotImplementedError
 
-    def deserialize_item(self, serialized_item: Any) -> ItemType:
+    def deserialize_item(self, serialized_item: dict) -> ItemType:
         """Deserializes an item from JSON-decoded data."""
         raise NotImplementedError
 

rslearn/data_sources/direct_materialize_data_source.py

@@ -0,0 +1,336 @@
+"""Base class for data sources that support direct materialization via TileStore."""
+
+from abc import abstractmethod
+from collections.abc import Callable
+from typing import Any, Generic
+
+import affine
+import numpy.typing as npt
+import rasterio
+import rasterio.vrt
+from rasterio.enums import Resampling
+
+from rslearn.config import LayerConfig
+from rslearn.data_sources.data_source import DataSource, ItemType
+from rslearn.dataset import Window
+from rslearn.dataset.materialize import RasterMaterializer
+from rslearn.tile_stores import TileStore, TileStoreWithLayer
+from rslearn.utils.geometry import PixelBounds, Projection
+
+
+class DirectMaterializeDataSource(DataSource[ItemType], TileStore, Generic[ItemType]):
+    """Base class for data sources that support direct materialization via TileStore.
+
+    This class provides common TileStore functionality for data sources that can read
+    raster data on-demand from remote sources (like cloud buckets or APIs) without
+    first ingesting into a local tile store.
+
+    Subclasses must implement:
+        - get_asset_url(): Get the URL for an asset given item name and bands
+        - get_item_by_name(): Get an item by its name
+
+    Subclasses may optionally override:
+        - get_raster_bands(): By default, we assume that items have all assets. If
+            items may have a subset of assets, override get_raster_bands to return
+            the sets of bands available for that item.
+        - get_read_callback(): Returns a callback to transform the raster array,
+            for post-processing like Sentinel-2 harmonization.
+    """
+
+    def __init__(self, asset_bands: dict[str, list[str]]):
+        """Initialize the DirectMaterializeDataSource.
+
+        Args:
+            asset_bands: mapping from asset key to the list of band names in that asset.
+        """
+        self.asset_bands = asset_bands
+
+    def _get_asset_key_by_bands(self, bands: list[str]) -> str:
+        """Get the asset key based on the band names.
+
+        Args:
+            bands: list of band names to look up.
+
+        Returns:
+            the asset key that provides those bands.
+
+        Raises:
+            ValueError: if no asset provides those bands.
+        """
+        for asset_key, asset_bands in self.asset_bands.items():
+            if bands == asset_bands:
+                return asset_key
+        raise ValueError(f"no known asset with bands {bands}")
+
+    # --- Methods that subclasses must implement ---
+
+    @abstractmethod
+    def get_asset_url(self, item_name: str, asset_key: str) -> str:
+        """Get the 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.
+
+        Returns:
+            the URL to read the asset from (must be readable by rasterio).
+        """
+        raise NotImplementedError
+
+    def get_item_by_name(self, name: str) -> ItemType:
+        """Get an item by its name.
+
+        Subclasses must implement this method, either directly or by inheriting from
+        a class that provides it (e.g., StacDataSource).
+
+        Args:
+            name: the name of the item to get.
+
+        Returns:
+            the item object.
+        """
+        raise NotImplementedError
+
+    # --- Optional hooks for subclasses ---
+
+    def get_read_callback(
+        self, item_name: str, asset_key: str
+    ) -> Callable[[npt.NDArray[Any]], npt.NDArray[Any]] | None:
+        """Return a callback to post-process raster data (e.g., harmonization).
+
+        Subclasses can override this to apply transformations to the raw raster data
+        after reading, such as harmonization for Sentinel-2 data.
+
+        Args:
+            item_name: the name of the item being read.
+            asset_key: the key identifying which asset is being read.
+
+        Returns:
+            A callback function that takes an array and returns a modified array,
+            or None if no post-processing is needed.
+        """
+        return None
+
+    # --- TileStore implementation ---
+
+    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.
+
+        For remote-backed tile stores, this always returns True since data is
+        read on-demand from the remote source.
+
+        Args:
+            layer_name: the layer name or alias.
+            item_name: the item.
+            bands: the list of bands identifying which specific raster to read.
+
+        Returns:
+            True, since data is always available from the remote source.
+        """
+        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.
+
+        By default, returns all band sets from the asset_bands configuration.
+        Subclasses can override this if not all items have all assets.
+
+        Args:
+            layer_name: the layer name or alias.
+            item_name: the item.
+
+        Returns:
+            a list of lists of bands available for this item.
+        """
+        return list(self.asset_bands.values())
+
+    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.
+            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_from_url(
+        self,
+        url: str,
+        projection: Projection,
+        bounds: PixelBounds,
+        resampling: Resampling,
+    ) -> npt.NDArray[Any]:
+        """Read raster data from a URL with reprojection.
+
+        This is the common logic for reading raster data from a URL and reprojecting
+        it to the target projection and bounds using rasterio's WarpedVRT.
+
+        Args:
+            url: the URL to read from (must be readable by rasterio).
+            projection: the projection to read in.
+            bounds: the bounds to read.
+            resampling: the resampling method to use.
+
+        Returns:
+            the raster data as a numpy array.
+        """
+        # 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 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.
+            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 as a numpy array.
+        """
+        # Get the asset key for the requested bands
+        asset_key = self._get_asset_key_by_bands(bands)
+
+        # Get the asset URL from the subclass
+        asset_url = self.get_asset_url(item_name, asset_key)
+
+        # Read the raster data
+        raw_data = self._read_raster_from_url(asset_url, projection, bounds, resampling)
+
+        # Apply any post-processing callback
+        callback = self.get_read_callback(item_name, asset_key)
+        if callback is not None:
+            raw_data = callback(raw_data)
+
+        return raw_data
+
+    def materialize(
+        self,
+        window: Window,
+        item_groups: list[list[ItemType]],
+        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,
+        )
+
+    # --- TileStore methods that are not supported ---
+
+    def write_raster(
+        self,
+        layer_name: str,
+        item_name: str,
+        bands: list[str],
+        projection: Projection,
+        bounds: PixelBounds,
+        array: npt.NDArray[Any],
+    ) -> None:
+        """Write raster data to the store.
+
+        This is not supported for remote-backed tile stores.
+        """
+        raise NotImplementedError(
+            "DirectMaterializeDataSource does not support writing raster data"
+        )
+
+    def write_raster_file(
+        self, layer_name: str, item_name: str, bands: list[str], fname: Any
+    ) -> None:
+        """Write raster data to the store.
+
+        This is not supported for remote-backed tile stores.
+        """
+        raise NotImplementedError(
+            "DirectMaterializeDataSource does not support writing raster files"
+        )
+
+    def is_vector_ready(self, layer_name: str, item_name: str) -> bool:
+        """Checks if this vector item has been written to the store.
+
+        This is not supported for remote-backed tile stores.
+        """
+        raise NotImplementedError(
+            "DirectMaterializeDataSource does not support vector operations"
+        )
+
+    def read_vector(
+        self,
+        layer_name: str,
+        item_name: str,
+        projection: Projection,
+        bounds: PixelBounds,
+    ) -> Any:
+        """Read vector data from the store.
+
+        This is not supported for remote-backed tile stores.
+        """
+        raise NotImplementedError(
+            "DirectMaterializeDataSource does not support vector operations"
+        )
+
+    def write_vector(self, layer_name: str, item_name: str, features: Any) -> None:
+        """Write vector data to the store.
+
+        This is not supported for remote-backed tile stores.
+        """
+        raise NotImplementedError(
+            "DirectMaterializeDataSource does not support vector operations"
+        )