rslearn 0.0.14__py3-none-any.whl → 0.0.16__py3-none-any.whl

rslearn/data_sources/planet_basemap.py

@@ -9,9 +9,9 @@ import requests
 import shapely
 from upath import UPath
 
-from rslearn.config import LayerConfig, QueryConfig, RasterLayerConfig
+from rslearn.config import QueryConfig
 from rslearn.const import WGS84_PROJECTION
-from rslearn.data_sources import DataSource, Item
+from rslearn.data_sources import DataSource, DataSourceContext, Item
 from rslearn.data_sources.utils import match_candidate_items_to_window
 from rslearn.log_utils import get_logger
 from rslearn.tile_stores import TileStoreWithLayer
@@ -68,21 +68,20 @@ class PlanetBasemap(DataSource):
 
     def __init__(
         self,
-        config: RasterLayerConfig,
         series_id: str,
         bands: list[str],
         api_key: str | None = None,
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Initialize a new Planet instance.
 
         Args:
-            config: the LayerConfig of the layer containing this data source
             series_id: the series of mosaics to use.
             bands: list of band names to use.
             api_key: optional Planet API key (it can also be provided via PL_API_KEY
                 environmnet variable).
+            context: the data source context
         """
-        self.config = config
         self.series_id = series_id
         self.bands = bands
 
@@ -94,26 +93,6 @@ class PlanetBasemap(DataSource):
         # Lazily load mosaics.
         self.mosaics: dict | None = None
 
-    @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "PlanetBasemap":
-        """Creates a new PlanetBasemap instance from a configuration dictionary."""
-        assert isinstance(config, RasterLayerConfig)
-        if config.data_source is None:
-            raise ValueError("data_source is required")
-        d = config.data_source.config_dict
-        kwargs = dict(
-            config=config,
-            series_id=d["series_id"],
-            bands=d["bands"],
-        )
-        optional_keys = [
-            "api_key",
-        ]
-        for optional_key in optional_keys:
-            if optional_key in d:
-                kwargs[optional_key] = d[optional_key]
-        return PlanetBasemap(**kwargs)
-
     def _load_mosaics(self) -> dict[str, STGeometry]:
         """Lazily load mosaics in the configured series_id from Planet API.
 

rslearn/data_sources/planetary_computer.py

@@ -18,9 +18,9 @@ import shapely
 from rasterio.enums import Resampling
 from upath import UPath
 
-from rslearn.config import LayerConfig, QueryConfig, RasterLayerConfig
+from rslearn.config import LayerConfig, QueryConfig
 from rslearn.const import WGS84_PROJECTION
-from rslearn.data_sources import DataSource, Item
+from rslearn.data_sources import DataSource, DataSourceContext, Item
 from rslearn.data_sources.utils import match_candidate_items_to_window
 from rslearn.dataset import Window
 from rslearn.dataset.materialize import RasterMaterializer
@@ -96,8 +96,9 @@ class PlanetaryComputer(DataSource, TileStore):
         sort_ascending: bool = True,
         timeout: timedelta = timedelta(seconds=10),
         skip_items_missing_assets: bool = False,
-        cache_dir: UPath | None = None,
+        cache_dir: str | None = None,
         max_items_per_client: int | None = None,
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Initialize a new PlanetaryComputer instance.
 
@@ -117,6 +118,7 @@ class PlanetaryComputer(DataSource, TileStore):
             max_items_per_client: number of STAC items to process before recreating
                 the client to prevent memory leaks from the resolved objects cache.
                 Defaults to DEFAULT_MAX_ITEMS_PER_CLIENT.
+            context: the data source context.
         """
         self.collection_name = collection_name
         self.asset_bands = asset_bands
@@ -125,46 +127,23 @@ class PlanetaryComputer(DataSource, TileStore):
         self.sort_ascending = sort_ascending
         self.timeout = timeout
         self.skip_items_missing_assets = skip_items_missing_assets
-        self.cache_dir = cache_dir
         self.max_items_per_client = (
             max_items_per_client or self.DEFAULT_MAX_ITEMS_PER_CLIENT
         )
 
-        if self.cache_dir is not None:
+        if cache_dir is not None:
+            if context.ds_path is not None:
+                self.cache_dir = join_upath(context.ds_path, cache_dir)
+            else:
+                self.cache_dir = UPath(cache_dir)
+
             self.cache_dir.mkdir(parents=True, exist_ok=True)
+        else:
+            self.cache_dir = None
 
         self.client: pystac_client.Client | None = None
         self._client_item_count = 0
 
-    @staticmethod
-    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "PlanetaryComputer":
-        """Creates a new PlanetaryComputer instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("config.data_source is required")
-        d = config.data_source.config_dict
-        kwargs: dict[str, Any] = dict(
-            collection_name=d["collection_name"],
-            asset_bands=d["asset_bands"],
-        )
-
-        if "timeout_seconds" in d:
-            kwargs["timeout"] = timedelta(seconds=d["timeout_seconds"])
-
-        if "cache_dir" in d:
-            kwargs["cache_dir"] = join_upath(ds_path, d["cache_dir"])
-
-        simple_optionals = [
-            "query",
-            "sort_by",
-            "sort_ascending",
-            "max_items_per_client",
-        ]
-        for k in simple_optionals:
-            if k in d:
-                kwargs[k] = d[k]
-
-        return PlanetaryComputer(**kwargs)
-
     def _load_client(
         self,
     ) -> pystac_client.Client:
@@ -545,7 +524,6 @@ class PlanetaryComputer(DataSource, TileStore):
             layer_name: the name of this layer
             layer_cfg: the config of this layer
         """
-        assert isinstance(layer_cfg, RasterLayerConfig)
         RasterMaterializer().materialize(
             TileStoreWithLayer(self, layer_name),
             window,
@@ -581,72 +559,48 @@ class Sentinel2(PlanetaryComputer):
 
     def __init__(
         self,
-        assets: list[str] | None = None,
         harmonize: bool = False,
+        assets: list[str] | None = None,
+        context: DataSourceContext = DataSourceContext(),
         **kwargs: Any,
     ):
         """Initialize a new Sentinel2 instance.
 
         Args:
-            assets: which assets in BANDS to ingest/materialize. None to ingest all
-                assets.
             harmonize: harmonize pixel values across different processing baselines,
                 see https://developers.google.com/earth-engine/datasets/catalog/COPERNICUS_S2_SR_HARMONIZED
+            assets: list of asset names to ingest, or None to ingest all assets. This
+                is only used if the layer config is missing from the context.
+            context: the data source context.
             kwargs: other arguments to pass to PlanetaryComputer.
         """
         self.harmonize = harmonize
 
-        if assets is None:
-            asset_bands = self.BANDS
-        else:
+        # Determine which assets we need based on the bands in the layer config.
+        if context.layer_config is not None:
+            asset_bands: dict[str, list[str]] = {}
+            for asset_key, band_names in self.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
+                    asset_bands[asset_key] = band_names
+                    break
+        elif assets is not None:
             asset_bands = {asset_key: self.BANDS[asset_key] for asset_key in assets}
+        else:
+            asset_bands = self.BANDS
 
         super().__init__(
             collection_name=self.COLLECTION_NAME,
             asset_bands=asset_bands,
             # Skip since all of the items should have the same assets.
             skip_items_missing_assets=True,
+            context=context,
             **kwargs,
         )
 
-    @staticmethod
-    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "Sentinel2":
-        """Creates a new Sentinel2 instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("config.data_source is required")
-        d = config.data_source.config_dict
-
-        # Determine the needed assets based on the band sets.
-        needed_assets: set[str] = set()
-        for asset_key, asset_bands in Sentinel2.BANDS.items():
-            for band_set in config.band_sets:
-                if not set(band_set.bands).intersection(set(asset_bands)):
-                    continue
-                needed_assets.add(asset_key)
-
-        kwargs: dict[str, Any] = dict(
-            assets=list(needed_assets),
-        )
-
-        if "timeout_seconds" in d:
-            kwargs["timeout"] = timedelta(seconds=d["timeout_seconds"])
-
-        if "cache_dir" in d:
-            kwargs["cache_dir"] = join_upath(ds_path, d["cache_dir"])
-
-        simple_optionals = [
-            "harmonize",
-            "query",
-            "sort_by",
-            "sort_ascending",
-            "max_items_per_client",
-        ]
-        for k in simple_optionals:
-            if k in d:
-                kwargs[k] = d[k]
-
-        return Sentinel2(**kwargs)
-
     def _get_product_xml(self, item: PlanetaryComputerItem) -> ET.Element:
         asset_url = planetary_computer.sign(item.asset_urls["product-metadata"])
         response = requests.get(asset_url, timeout=self.timeout.total_seconds())
@@ -779,55 +733,42 @@ class Sentinel1(PlanetaryComputer):
 
     def __init__(
         self,
-        band_names: list[str],
+        band_names: list[str] | None = None,
+        context: DataSourceContext = DataSourceContext(),
         **kwargs: Any,
     ):
         """Initialize a new Sentinel1 instance.
 
         Args:
-            band_names: list of bands to try to ingest.
+            band_names: list of bands to try to ingest, if the layer config is missing
+                from the context.
+            context: the data source context.
             kwargs: additional arguments to pass to PlanetaryComputer.
         """
+        # Get band names from the config if possible. If it isn't in the context, then
+        # we have to use the provided band names.
+        if context.layer_config is not None:
+            band_names = list(
+                {
+                    band
+                    for band_set in context.layer_config.band_sets
+                    for band in band_set.bands
+                }
+            )
+        if band_names is None:
+            raise ValueError(
+                "band_names must be set if layer config is not in the context"
+            )
+        # For Sentinel-1, the asset key should be the same as the band name (and all
+        # assets have one band).
         asset_bands = {band: [band] for band in band_names}
         super().__init__(
             collection_name=self.COLLECTION_NAME,
             asset_bands=asset_bands,
+            context=context,
             **kwargs,
         )
 
-    @staticmethod
-    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "Sentinel1":
-        """Creates a new  Sentinel1 instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("config.data_source is required")
-        d = config.data_source.config_dict
-        band_names: set[str] = set()
-        for band_set in config.band_sets:
-            for band in band_set.bands:
-                band_names.add(band)
-
-        kwargs: dict[str, Any] = dict(
-            band_names=list(band_names),
-        )
-
-        if "timeout_seconds" in d:
-            kwargs["timeout"] = timedelta(seconds=d["timeout_seconds"])
-
-        if "cache_dir" in d:
-            kwargs["cache_dir"] = join_upath(ds_path, d["cache_dir"])
-
-        simple_optionals = [
-            "query",
-            "sort_by",
-            "sort_ascending",
-            "max_items_per_client",
-        ]
-        for k in simple_optionals:
-            if k in d:
-                kwargs[k] = d[k]
-
-        return Sentinel1(**kwargs)
-
 
 class Naip(PlanetaryComputer):
     """A data source for NAIP data on Microsoft Planetary Computer.
@@ -840,42 +781,18 @@ class Naip(PlanetaryComputer):
 
     def __init__(
         self,
+        context: DataSourceContext = DataSourceContext(),
         **kwargs: Any,
     ):
         """Initialize a new Naip instance.
 
         Args:
-            band_names: list of bands to try to ingest.
+            context: the data source context.
             kwargs: additional arguments to pass to PlanetaryComputer.
         """
         super().__init__(
             collection_name=self.COLLECTION_NAME,
             asset_bands=self.ASSET_BANDS,
+            context=context,
             **kwargs,
         )
-
-    @staticmethod
-    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "Naip":
-        """Creates a new Naip instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("config.data_source is required")
-        d = config.data_source.config_dict
-        kwargs = {}
-
-        if "timeout_seconds" in d:
-            kwargs["timeout"] = timedelta(seconds=d["timeout_seconds"])
-
-        if "cache_dir" in d:
-            kwargs["cache_dir"] = join_upath(ds_path, d["cache_dir"])
-
-        simple_optionals = [
-            "query",
-            "sort_by",
-            "sort_ascending",
-            "max_items_per_client",
-        ]
-        for k in simple_optionals:
-            if k in d:
-                kwargs[k] = d[k]
-
-        return Naip(**kwargs)

rslearn/data_sources/usda_cdl.py

@@ -11,9 +11,9 @@ import requests.auth
 import shapely
 from upath import UPath
 
-from rslearn.config import QueryConfig, RasterLayerConfig
+from rslearn.config import QueryConfig
 from rslearn.const import WGS84_PROJECTION
-from rslearn.data_sources import DataSource, Item
+from rslearn.data_sources import DataSource, DataSourceContext, Item
 from rslearn.data_sources.utils import match_candidate_items_to_window
 from rslearn.log_utils import get_logger
 from rslearn.tile_stores import TileStoreWithLayer
@@ -60,39 +60,28 @@ class CDL(DataSource):
 
     def __init__(
         self,
-        band_name: str = "cdl",
         timeout: timedelta = timedelta(seconds=10),
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Initialize a new CDL instance.
 
         Args:
-            band_name: what to call the band.
             timeout: timeout for requests.
+            context: the data source context.
         """
-        self.band_name = band_name
         self.timeout = timeout
 
-    @staticmethod
-    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "CDL":
-        """Creates a new CDL instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("config.data_source is required")
-        d = config.data_source.config_dict
-
-        # Get the band name chosen by the user.
-        # There should be a single band set with a single band.
-        if len(config.band_sets) != 1:
-            raise ValueError("expected a single band set")
-        if len(config.band_sets[0].bands) != 1:
-            raise ValueError("expected band set to have a single band")
-        kwargs: dict[str, Any] = {
-            "band_name": config.band_sets[0].bands[0],
-        }
-
-        if "timeout_seconds" in d:
-            kwargs["timeout"] = timedelta(seconds=d["timeout_seconds"])
-
-        return CDL(**kwargs)
+        # Get the band name from the layer config, which should have a single band set
+        # with a single band. If the layer config is not available in the context, we
+        # default to "cdl".
+        if context.layer_config is not None:
+            if len(context.layer_config.band_sets) != 1:
+                raise ValueError("expected a single band set")
+            if len(context.layer_config.band_sets[0].bands) != 1:
+                raise ValueError("expected band set to have a single band")
+            self.band_name = context.layer_config.band_sets[0].bands[0]
+        else:
+            self.band_name = "cdl"
 
     def get_item_by_name(self, name: str) -> Item:
         """Gets an item by name.

rslearn/data_sources/usgs_landsat.py

@@ -18,9 +18,9 @@ import requests
 import shapely
 from upath import UPath
 
-from rslearn.config import QueryConfig, RasterLayerConfig
+from rslearn.config import QueryConfig
 from rslearn.const import WGS84_PROJECTION
-from rslearn.data_sources import DataSource, Item
+from rslearn.data_sources import DataSource, DataSourceContext, Item
 from rslearn.data_sources.utils import match_candidate_items_to_window
 from rslearn.tile_stores import TileStoreWithLayer
 from rslearn.utils import STGeometry
@@ -314,25 +314,24 @@ class LandsatOliTirs(DataSource):
 
     def __init__(
         self,
-        config: RasterLayerConfig,
         username: str,
         sort_by: str | None = None,
         password: str | None = None,
         token: str | None = None,
         timeout: timedelta = timedelta(seconds=10),
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Initialize a new LandsatOliTirs instance.
 
         Args:
-            config: the LayerConfig of the layer containing this data source
             username: EROS username
             sort_by: can be "cloud_cover", default arbitrary order; only has effect for
                 SpaceMode.WITHIN.
             password: EROS password (see M2MAPIClient).
             token: EROS application token (see M2MAPIClient).
             timeout: timeout for requests.
+            context: the data source context.
         """
-        self.config = config
         self.sort_by = sort_by
         self.timeout = timeout
 
@@ -340,30 +339,6 @@ class LandsatOliTirs(DataSource):
             username, password=password, token=token, timeout=timeout
         )
 
-    @staticmethod
-    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "LandsatOliTirs":
-        """Creates a new LandsatOliTirs instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("data_source is required")
-        d = config.data_source.config_dict
-
-        kwargs = dict(
-            config=config,
-            username=d["username"],
-            sort_by=d.get("sort_by"),
-        )
-
-        if "timeout_seconds" in d:
-            kwargs["timeout"] = timedelta(seconds=d["timeout_seconds"])
-
-        # Optional config.
-        for k in ["password", "token"]:
-            if k not in d:
-                continue
-            kwargs[k] = d[k]
-
-        return LandsatOliTirs(**kwargs)
-
     def _scene_metadata_to_item(self, result: dict[str, Any]) -> LandsatOliTirsItem:
         """Convert scene metadata from the API to a LandsatOliTirsItem."""
         metadata_dict = {}

rslearn/data_sources/utils.py

@@ -8,8 +8,11 @@ import shapely
 
 from rslearn.config import QueryConfig, SpaceMode, TimeMode
 from rslearn.data_sources import Item
+from rslearn.log_utils import get_logger
 from rslearn.utils import STGeometry, shp_intersects
 
+logger = get_logger(__name__)
+
 MOSAIC_MIN_ITEM_COVERAGE = 0.1
 """Minimum fraction of area that item should cover when adding it to a mosaic group."""
 
@@ -298,6 +301,12 @@ def match_candidate_items_to_window(
 
     # Enforce minimum matches if set.
     if len(groups) < query_config.min_matches:
+        logger.warning(
+            "Window rejected: found %d matches (required: %d) for time range %s",
+            len(groups),
+            query_config.min_matches,
+            geometry.time_range if geometry.time_range else "unlimited",
+        )
         return []
 
     return groups

rslearn/data_sources/worldcereal.py

@@ -6,18 +6,17 @@ import os
 import shutil
 import tempfile
 import zipfile
-from typing import Any
 
 import requests
 from fsspec.implementations.local import LocalFileSystem
 from upath import UPath
 
-from rslearn.config import DataSourceConfig, LayerConfig, QueryConfig, RasterLayerConfig
-from rslearn.data_sources.local_files import LocalFiles
+from rslearn.config import LayerType
+from rslearn.data_sources.local_files import LocalFiles, RasterItemSpec
 from rslearn.log_utils import get_logger
 from rslearn.utils.fsspec import get_upath_local, join_upath, open_atomic
 
-from .data_source import Item
+from .data_source import DataSourceContext, Item
 
 logger = get_logger(__name__)
 
@@ -237,71 +236,65 @@ class WorldCereal(LocalFiles):
 
     def __init__(
         self,
-        config: LayerConfig,
-        band: str,
-        worldcereal_dir: UPath,
+        worldcereal_dir: str,
+        band: str | None = None,
+        context: DataSourceContext = DataSourceContext(),
     ) -> None:
         """Create a new WorldCereal.
 
         Args:
-            config: configuration for this layer.
-            band: the worldcereal band being processed.
             worldcereal_dir: the directory to extract the WorldCereal GeoTIFF files. For
                 high performance, this should be a local directory; if the dataset is
                 remote, prefix with a protocol ("file://") to use a local directory
                 instead of a path relative to the dataset path.
+            band: the worldcereal band to process. This will only be used if the layer
+                config is missing from the context.
+            context: the data source context.
         """
-        self.band = band
-        tif_dir, tif_filepath = self.download_worldcereal_data(band, worldcereal_dir)
+        if context.ds_path is not None:
+            worldcereal_upath = join_upath(context.ds_path, worldcereal_dir)
+        else:
+            worldcereal_upath = UPath(worldcereal_dir)
+
+        if context.layer_config is not None:
+            if len(context.layer_config.band_sets) != 1:
+                raise ValueError("expected a single band set")
+            if len(context.layer_config.band_sets[0].bands) != 1:
+                raise ValueError("expected band set to have a single band")
+            self.band = context.layer_config.band_sets[0].bands[0]
+        elif band is not None:
+            self.band = band
+        else:
+            raise ValueError("band must be set if layer config is not in the context")
+
+        tif_dir, tif_filepath = self.download_worldcereal_data(
+            self.band, worldcereal_upath
+        )
         all_aezs: set[int] = self.all_aezs_from_tifs(tif_filepath)
 
         # now that we have all our aezs, lets match them to the bands
-        spec_dicts: list[dict] = []
+        item_specs: list[RasterItemSpec] = []
         for aez in all_aezs:
-            spec_dict: dict[str, Any] = {
+            item_spec = RasterItemSpec(
+                fnames=[],
+                bands=[],
                 # must be a str since we / with a posix path later
-                "name": str(aez),
-                "fnames": [],
-                "bands": [],
-            }
+                name=str(aez),
+            )
             aez_band_filepath = self.filepath_for_product_aez(tif_filepath, aez)
             if aez_band_filepath is not None:
-                spec_dict["fnames"].append(aez_band_filepath.absolute().as_uri())
-                spec_dict["bands"].append([band])
-            spec_dicts.append(spec_dict)
-        if len(spec_dicts) == 0:
-            raise ValueError(f"No AEZ files found for {band}")
-        # add this to the config
-        if config.data_source is not None:
-            if "item_specs" in config.data_source.config_dict:
-                logger.warning(
-                    "Overwriting item_specs in WorldCereal config.data_source"
-                )
-            config.data_source.config_dict["item_specs"] = spec_dicts
-        else:
-            config.data_source = DataSourceConfig(
-                name="rslearn.data_sources.WorldCereal",
-                query_config=QueryConfig.from_config({}),
-                config_dict={"item_specs": spec_dicts},
-            )
-
-        super().__init__(config, tif_dir)
-
-    @staticmethod
-    def from_config(config: LayerConfig, ds_path: UPath) -> "LocalFiles":
-        """Creates a new LocalFiles instance from a configuration dictionary."""
-        if config.data_source is None:
-            raise ValueError("LocalFiles data source requires a data source config")
-        d = config.data_source.config_dict
-        assert isinstance(config, RasterLayerConfig)
-        bandsets = config.band_sets
-        assert len(bandsets) == 1
-        assert len(bandsets[0].bands) == 1
-        band = bandsets[0].bands[0]
-        return WorldCereal(
-            config=config,
-            band=band,
-            worldcereal_dir=join_upath(ds_path, d["worldcereal_dir"]),
+                item_spec.fnames.append(aez_band_filepath.absolute().as_uri())
+                assert item_spec.bands is not None
+                item_spec.bands.append([self.band])
+            item_specs.append(item_spec)
+        if len(item_specs) == 0:
+            raise ValueError(f"No AEZ files found for {self.band}")
+
+        super().__init__(
+            src_dir=tif_dir,
+            raster_item_specs=item_specs,
+            layer_type=LayerType.RASTER,
+            context=context,
         )
 
     @staticmethod
@@ -441,7 +434,7 @@ class WorldCereal(LocalFiles):
         cache_fname = self.src_dir / f"{self.band}_summary.json"
         if not cache_fname.exists():
             logger.debug("cache at %s does not exist, listing items", cache_fname)
-            items = self.importer.list_items(self.config, self.src_dir)
+            items = self.importer.list_items(self.src_dir)
             serialized_items = [item.serialize() for item in items]
             with cache_fname.open("w") as f:
                 json.dump(serialized_items, f)