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

rslearn/data_sources/__init__.py

@@ -10,40 +10,17 @@ Each source supports operations to lookup items that match with spatiotemporal
 geometries, and ingest those items.
 """
 
-import functools
-import importlib
-
-from upath import UPath
-
-from rslearn.config import LayerConfig
-from rslearn.log_utils import get_logger
-
-from .data_source import DataSource, Item, ItemLookupDataSource, RetrieveItemDataSource
-
-logger = get_logger(__name__)
-
-
-@functools.cache
-def data_source_from_config(config: LayerConfig, ds_path: UPath) -> DataSource:
-    """Loads a data source from config dict.
-
-    Args:
-        config: the LayerConfig containing this data source.
-        ds_path: the dataset root directory.
-    """
-    logger.debug("getting a data source for dataset at %s", ds_path)
-    if config.data_source is None:
-        raise ValueError("No data source specified")
-    name = config.data_source.name
-    module_name = ".".join(name.split(".")[:-1])
-    class_name = name.split(".")[-1]
-    module = importlib.import_module(module_name)
-    class_ = getattr(module, class_name)
-    return class_.from_config(config, ds_path)
-
+from .data_source import (
+    DataSource,
+    DataSourceContext,
+    Item,
+    ItemLookupDataSource,
+    RetrieveItemDataSource,
+)
 
 __all__ = (
     "DataSource",
+    "DataSourceContext",
     "Item",
     "ItemLookupDataSource",
     "RetrieveItemDataSource",

rslearn/data_sources/aws_landsat.py

@@ -25,7 +25,7 @@ from rasterio.enums import Resampling
 from upath import UPath
 
 import rslearn.data_sources.utils
-from rslearn.config import LayerConfig, RasterLayerConfig
+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
@@ -34,7 +34,7 @@ from rslearn.utils.fsspec import get_upath_local, join_upath, open_atomic
 from rslearn.utils.geometry import PixelBounds, Projection, STGeometry
 from rslearn.utils.grid_index import GridIndex
 
-from .data_source import DataSource, Item, QueryConfig
+from .data_source import DataSource, DataSourceContext, Item, QueryConfig
 
 WRS2_GRID_SIZE = 1.0
 
@@ -98,20 +98,25 @@ class LandsatOliTirs(DataSource, TileStore):
 
     def __init__(
         self,
-        config: RasterLayerConfig,
-        metadata_cache_dir: UPath,
+        metadata_cache_dir: str,
         sort_by: str | None = None,
+        context: DataSourceContext = DataSourceContext(),
     ) -> None:
         """Initialize a new LandsatOliTirs instance.
 
         Args:
-            config: configuration of this layer
-            metadata_cache_dir: directory to cache product metadata files.
+            metadata_cache_dir: directory to cache produtc metadata files.
             sort_by: can be "cloud_cover", default arbitrary order; only has effect for
                 SpaceMode.WITHIN.
+            context: the data source context.
         """
-        self.config = config
-        self.metadata_cache_dir = metadata_cache_dir
+        # 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:
+            self.metadata_cache_dir = join_upath(context.ds_path, metadata_cache_dir)
+        else:
+            self.metadata_cache_dir = UPath(metadata_cache_dir)
+
         self.sort_by = sort_by
 
         self.client = boto3.client("s3")
@@ -120,21 +125,6 @@ class LandsatOliTirs(DataSource, TileStore):
 
         self.wrs2_index: GridIndex | None = None
 
-    @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(f"data_source is required for config dict {config}")
-        d = config.data_source.config_dict
-        kwargs = dict(
-            config=config,
-            metadata_cache_dir=join_upath(ds_path, d["metadata_cache_dir"]),
-        )
-        if "sort_by" in d:
-            kwargs["sort_by"] = d["sort_by"]
-
-        return LandsatOliTirs(**kwargs)
-
     def _read_products(
         self, needed_year_pathrows: set[tuple[int, str, str]]
     ) -> Generator[LandsatOliTirsItem, None, None]:
@@ -536,7 +526,6 @@ class LandsatOliTirs(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,

rslearn/data_sources/aws_open_data.py

@@ -22,7 +22,6 @@ from rasterio.crs import CRS
 from upath import UPath
 
 import rslearn.data_sources.utils
-from rslearn.config import RasterLayerConfig
 from rslearn.const import SHAPEFILE_AUX_EXTENSIONS, WGS84_EPSG, WGS84_PROJECTION
 from rslearn.tile_stores import TileStoreWithLayer
 from rslearn.utils import GridIndex, Projection, STGeometry, daterange
@@ -32,6 +31,7 @@ from rslearn.utils.raster_format import get_raster_projection_and_bounds
 from .copernicus import get_harmonize_callback, get_sentinel2_tiles
 from .data_source import (
     DataSource,
+    DataSourceContext,
     Item,
     ItemLookupDataSource,
     QueryConfig,
@@ -83,16 +83,15 @@ class Naip(DataSource):
 
     def __init__(
         self,
-        config: RasterLayerConfig,
-        index_cache_dir: UPath,
+        index_cache_dir: str,
         use_rtree_index: bool = False,
         states: list[str] | None = None,
         years: list[int] | None = None,
+        context: DataSourceContext = DataSourceContext(),
     ) -> None:
         """Initialize a new Naip instance.
 
         Args:
-            config: the LayerConfig of the layer containing this data source.
             index_cache_dir: directory to cache index shapefiles.
             use_rtree_index: whether to create an rtree index to enable faster lookups
                 (default false)
@@ -100,9 +99,15 @@ class Naip(DataSource):
                 the search. If use_rtree_index is enabled, the rtree will only be
                 populated with data from these states.
             years: optional list of years to restrict the search
+            context: the data source context.
         """
-        self.config = config
-        self.index_cache_dir = index_cache_dir
+        # 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:
+            self.index_cache_dir = join_upath(context.ds_path, index_cache_dir)
+        else:
+            self.index_cache_dir = UPath(index_cache_dir)
+
         self.states = states
         self.years = years
 
@@ -119,22 +124,6 @@ class Naip(DataSource):
 
             self.rtree_index = get_cached_rtree(self.index_cache_dir, build_fn)
 
-    @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(f"data_source is required for config dict {config}")
-        d = config.data_source.config_dict
-        kwargs = dict(
-            config=config,
-            index_cache_dir=join_upath(ds_path, d["index_cache_dir"]),
-        )
-        simple_optionals = ["use_rtree_index", "states", "years"]
-        for k in simple_optionals:
-            if k in d:
-                kwargs[k] = d[k]
-        return Naip(**kwargs)
-
     def _download_manifest(self) -> UPath:
         """Download the manifest that enumerates files in the bucket.
 
@@ -460,51 +449,37 @@ class Sentinel2(
 
     def __init__(
         self,
-        config: RasterLayerConfig,
         modality: Sentinel2Modality,
-        metadata_cache_dir: UPath,
+        metadata_cache_dir: str,
         sort_by: str | None = None,
         harmonize: bool = False,
+        context: DataSourceContext = DataSourceContext(),
     ) -> None:
         """Initialize a new Sentinel2 instance.
 
         Args:
-            config: the LayerConfig of the layer containing this data source.
             modality: L1C or L2A.
             metadata_cache_dir: directory to cache product metadata files.
             sort_by: can be "cloud_cover", default arbitrary order; only has effect for
                 SpaceMode.WITHIN.
             harmonize: harmonize pixel values across different processing baselines,
                 see https://developers.google.com/earth-engine/datasets/catalog/COPERNICUS_S2_SR_HARMONIZED
+            context: the data source context.
         """  # noqa: E501
-        self.config = config
+        # 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:
+            self.metadata_cache_dir = join_upath(context.ds_path, metadata_cache_dir)
+        else:
+            self.metadata_cache_dir = UPath(metadata_cache_dir)
+
         self.modality = modality
-        self.metadata_cache_dir = metadata_cache_dir
         self.sort_by = sort_by
         self.harmonize = harmonize
 
         bucket_name = self.bucket_names[modality]
         self.bucket = boto3.resource("s3").Bucket(bucket_name)
 
-    @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("Sentinel2 data source requires a data source config")
-        d = config.data_source.config_dict
-        kwargs = dict(
-            config=config,
-            modality=Sentinel2Modality(d["modality"]),
-            metadata_cache_dir=join_upath(ds_path, d["metadata_cache_dir"]),
-        )
-
-        simple_optionals = ["sort_by", "harmonize"]
-        for k in simple_optionals:
-            if k in d:
-                kwargs[k] = d[k]
-
-        return Sentinel2(**kwargs)
-
     def _read_products(
         self, needed_cell_months: set[tuple[str, int, int]]
     ) -> Generator[Sentinel2Item, None, None]:

rslearn/data_sources/aws_sentinel1.py

@@ -7,7 +7,6 @@ from typing import Any
 import boto3
 from upath import UPath
 
-from rslearn.config import RasterLayerConfig
 from rslearn.data_sources.copernicus import (
     CopernicusItem,
     Sentinel1OrbitDirection,
@@ -19,7 +18,7 @@ from rslearn.log_utils import get_logger
 from rslearn.tile_stores import TileStore, TileStoreWithLayer
 from rslearn.utils.geometry import STGeometry
 
-from .data_source import DataSource, QueryConfig
+from .data_source import DataSource, DataSourceContext, QueryConfig
 
 WRS2_GRID_SIZE = 1.0
 
@@ -45,11 +44,13 @@ class Sentinel1(DataSource, TileStore):
     def __init__(
         self,
         orbit_direction: Sentinel1OrbitDirection | None = None,
+        context: DataSourceContext = DataSourceContext(),
     ) -> None:
         """Initialize a new Sentinel1 instance.
 
         Args:
             orbit_direction: optional orbit direction to filter by.
+            context: the data source context.
         """
         self.client = boto3.client("s3")
         self.bucket = boto3.resource("s3").Bucket(self.bucket_name)
@@ -59,18 +60,6 @@ class Sentinel1(DataSource, TileStore):
             orbit_direction=orbit_direction,
         )
 
-    @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(f"data_source is required for config dict {config}")
-        d = config.data_source.config_dict
-        kwargs: dict[str, Any] = {}
-        if "orbit_direction" in d:
-            d["orbit_direction"] = Sentinel1OrbitDirection[d["orbit_direction"]]
-
-        return Sentinel1(**kwargs)
-
     def get_items(
         self, geometries: list[STGeometry], query_config: QueryConfig
     ) -> list[list[list[CopernicusItem]]]:

rslearn/data_sources/climate_data_store.py

@@ -14,9 +14,9 @@ from dateutil.relativedelta import relativedelta
 from rasterio.transform import from_origin
 from upath import UPath
 
-from rslearn.config import QueryConfig, RasterLayerConfig, SpaceMode
+from rslearn.config import QueryConfig, SpaceMode
 from rslearn.const import WGS84_EPSG, WGS84_PROJECTION
-from rslearn.data_sources import DataSource, Item
+from rslearn.data_sources import DataSource, DataSourceContext, Item
 from rslearn.log_utils import get_logger
 from rslearn.tile_stores import TileStoreWithLayer
 from rslearn.utils.geometry import STGeometry
@@ -55,59 +55,40 @@ class ERA5LandMonthlyMeans(DataSource):
 
     def __init__(
         self,
-        band_names: list[str],
+        band_names: list[str] | None = None,
         api_key: str | None = None,
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Initialize a new ERA5LandMonthlyMeans instance.
 
         Args:
             band_names: list of band names to acquire. These should correspond to CDS
-                variable names but with "_" replaced with "-".
+                variable names but with "_" replaced with "-". This will only be used
+                if the layer config is missing from the context.
             api_key: the API key. If not set, it should be set via the CDSAPI_KEY
                 environment variable.
+            context: the data source context.
         """
-        self.band_names = band_names
+        self.band_names: list[str]
+        if context.layer_config is not None:
+            self.band_names = []
+            for band_set in context.layer_config.band_sets:
+                for band in band_set.bands:
+                    if band in self.band_names:
+                        continue
+                    self.band_names.append(band)
+        elif band_names is not None:
+            self.band_names = band_names
+        else:
+            raise ValueError(
+                "band_names must be set if layer_config is not in the context"
+            )
 
         self.client = cdsapi.Client(
             url=self.api_url,
             key=api_key,
         )
 
-    @staticmethod
-    def from_config(
-        config: RasterLayerConfig, ds_path: UPath
-    ) -> "ERA5LandMonthlyMeans":
-        """Creates a new ERA5LandMonthlyMeans instance from a configuration dictionary.
-
-        Args:
-            config: the LayerConfig of the layer containing this data source
-            ds_path: the path to the data source
-
-        Returns:
-            A new ERA5LandMonthlyMeans instance
-        """
-        if config.data_source is None:
-            raise ValueError("data_source is required")
-        d = config.data_source.config_dict
-
-        # Determine band names based on the configured band sets.
-        band_names = []
-        for band_set in config.band_sets:
-            for band in band_set.bands:
-                if band in band_names:
-                    continue
-                band_names.append(band)
-        kwargs: dict[str, Any] = dict(
-            band_names=band_names,
-        )
-
-        simple_optionals = ["api_key"]
-        for k in simple_optionals:
-            if k in d:
-                kwargs[k] = d[k]
-
-        return ERA5LandMonthlyMeans(**kwargs)
-
     def get_items(
         self, geometries: list[STGeometry], query_config: QueryConfig
     ) -> list[list[list[Item]]]:

rslearn/data_sources/copernicus.py

@@ -23,9 +23,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.data_source import DataSource, Item
+from rslearn.data_sources.data_source 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
@@ -306,6 +306,7 @@ class Copernicus(DataSource):
         sort_by: str | None = None,
         sort_desc: bool = False,
         timeout: float = 10,
+        context: DataSourceContext = DataSourceContext(),
     ):
         """Create a new Copernicus.
 
@@ -332,6 +333,7 @@ class Copernicus(DataSource):
             sort_desc: for sort_by, sort in descending order instead of ascending
                 order.
             timeout: timeout for requests.
+            context: the data source context.
         """
         self.glob_to_bands = glob_to_bands
         self.query_filter = query_filter
@@ -351,30 +353,6 @@ class Copernicus(DataSource):
                 self.username = os.environ["COPERNICUS_USERNAME"]
                 self.password = os.environ["COPERNICUS_PASSWORD"]
 
-    @staticmethod
-    def from_config(config: RasterLayerConfig, ds_path: UPath) -> "Copernicus":
-        """Creates a new Copernicus 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(
-            glob_to_bands=d["glob_to_bands"],
-        )
-
-        simple_optionals = [
-            "access_token",
-            "query_filter",
-            "order_by",
-            "sort_by",
-            "sort_desc",
-            "timeout",
-        ]
-        for k in simple_optionals:
-            if k in d:
-                kwargs[k] = d[k]
-
-        return Copernicus(**kwargs)
-
     def deserialize_item(self, serialized_item: Any) -> CopernicusItem:
         """Deserializes an item from JSON-decoded data."""
         assert isinstance(serialized_item, dict)
@@ -763,23 +741,43 @@ class Sentinel2(Copernicus):
 
     def __init__(
         self,
-        assets: list[str],
         product_type: Sentinel2ProductType,
         harmonize: bool = False,
+        assets: list[str] | None = None,
+        context: DataSourceContext = DataSourceContext(),
         **kwargs: Any,
     ):
         """Create a new Sentinel2.
 
         Args:
-            assets: list of assets corresponding to keys in BANDS, e.g. ["TCI", "B08"].
             product_type: desired product type, L1C or L2A.
             harmonize: harmonize pixel values across different processing baselines,
                 see https://developers.google.com/earth-engine/datasets/catalog/COPERNICUS_S2_SR_HARMONIZED
+            assets: the assets to download, or None to download all assets. This is
+                only used if the layer config is not in the context.
+            context: the data source context.
             kwargs: additional arguments to pass to Copernicus.
         """
         # Create glob to bands map.
+        # If the context is provided, we limit to needed assets based on the configured
+        # band sets.
+        if context.layer_config is not None:
+            needed_assets = []
+            for asset_key, asset_bands in Sentinel2.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(asset_bands)):
+                        continue
+                    needed_assets.append(asset_key)
+                    break
+        elif assets is not None:
+            needed_assets = assets
+        else:
+            needed_assets = list(Sentinel2.BANDS.keys())
+
         glob_to_bands = {}
-        for asset_key in assets:
+        for asset_key in needed_assets:
             band_names = self.BANDS[asset_key]
             glob_pattern = self.GLOB_PATTERNS[product_type][asset_key]
             glob_to_bands[glob_pattern] = band_names
@@ -788,46 +786,13 @@ class Sentinel2(Copernicus):
         query_filter = f"Attributes/OData.CSC.StringAttribute/any(att:att/Name eq 'productType' and att/OData.CSC.StringAttribute/Value eq '{quote(product_type.value)}')"
 
         super().__init__(
+            context=context,
             glob_to_bands=glob_to_bands,
             query_filter=query_filter,
             **kwargs,
         )
         self.harmonize = harmonize
 
-    @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 needed assets based on the configured 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),
-            product_type=Sentinel2ProductType[d["product_type"]],
-        )
-
-        simple_optionals = [
-            "harmonize",
-            "access_token",
-            "order_by",
-            "sort_by",
-            "sort_desc",
-            "timeout",
-        ]
-        for k in simple_optionals:
-            if k in d:
-                kwargs[k] = d[k]
-
-        return Sentinel2(**kwargs)
-
     # Override to support harmonization step.
     def _process_product_zip(
         self, tile_store: TileStoreWithLayer, item: CopernicusItem, local_zip_fname: str
@@ -922,6 +887,7 @@ class Sentinel1(Copernicus):
         product_type: Sentinel1ProductType,
         polarisation: Sentinel1Polarisation,
         orbit_direction: Sentinel1OrbitDirection | None = None,
+        context: DataSourceContext = DataSourceContext(),
         **kwargs: Any,
     ):
         """Create a new Sentinel1.
@@ -930,6 +896,7 @@ class Sentinel1(Copernicus):
             product_type: desired product type.
             polarisation: desired polarisation(s).
             orbit_direction: optional orbit direction to filter by.
+            context: the data source context.
             kwargs: additional arguments to pass to Copernicus.
         """
         # Create query filter based on the product type.
@@ -945,31 +912,3 @@ class Sentinel1(Copernicus):
             query_filter=query_filter,
             **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
-
-        kwargs: dict[str, Any] = dict(
-            product_type=Sentinel1ProductType[d["product_type"]],
-            polarisation=Sentinel1Polarisation[d["polarisation"]],
-        )
-
-        if "orbit_direction" in d:
-            kwargs["orbit_direction"] = Sentinel1OrbitDirection[d["orbit_direction"]]
-
-        simple_optionals = [
-            "access_token",
-            "order_by",
-            "sort_by",
-            "sort_desc",
-            "timeout",
-        ]
-        for k in simple_optionals:
-            if k in d:
-                kwargs[k] = d[k]
-
-        return Sentinel1(**kwargs)

rslearn/data_sources/data_source.py

@@ -3,6 +3,8 @@
 from collections.abc import Generator
 from typing import Any, BinaryIO, Generic, TypeVar
 
+from upath import UPath
+
 from rslearn.config import LayerConfig, QueryConfig
 from rslearn.dataset import Window
 from rslearn.tile_stores import TileStoreWithLayer
@@ -127,3 +129,27 @@ class RetrieveItemDataSource(DataSource[ItemType]):
     ) -> Generator[tuple[str, BinaryIO], None, None]:
         """Retrieves the rasters corresponding to an item as file streams."""
         raise NotImplementedError
+
+
+class DataSourceContext:
+    """This context is passed to every data source.
+
+    When initializing data sources within rslearn, we always set the ds_path and
+    layer_config. However, for convenience (for users directly initializing the data
+    sources externally), each data source should allow for initialization when one or
+    both are missing.
+    """
+
+    def __init__(
+        self, ds_path: UPath | None = None, layer_config: LayerConfig | None = None
+    ):
+        """Create a new DataSourceContext.
+
+        Args:
+            ds_path: the path of the underlying dataset.
+            layer_config: the LayerConfig for the layer that the data source is for.
+        """
+        # We don't use dataclass here because otherwise jsonargparse will ignore our
+        # custom serializer/deserializer defined in rslearn.utils.jsonargparse.
+        self.ds_path = ds_path
+        self.layer_config = layer_config