rslearn 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

rslearn/config/__init__.py

@@ -2,6 +2,7 @@
 
 from .dataset import (
     BandSetConfig,
+    CompositingMethod,
     DataSourceConfig,
     DType,
     LayerConfig,
@@ -10,7 +11,6 @@ from .dataset import (
     RasterFormatConfig,
     RasterLayerConfig,
     SpaceMode,
-    TileStoreConfig,
     TimeMode,
     VectorFormatConfig,
     VectorLayerConfig,
@@ -19,6 +19,7 @@ from .dataset import (
 
 __all__ = [
     "BandSetConfig",
+    "CompositingMethod",
     "DataSourceConfig",
     "DType",
     "LayerConfig",
@@ -27,7 +28,6 @@ __all__ = [
     "RasterFormatConfig",
     "RasterLayerConfig",
     "SpaceMode",
-    "TileStoreConfig",
     "TimeMode",
     "VectorFormatConfig",
     "VectorLayerConfig",

rslearn/config/dataset.py

@@ -1,5 +1,6 @@
 """Classes for storing configuration of a dataset."""
 
+import json
 from datetime import timedelta
 from enum import Enum
 from typing import Any
@@ -19,7 +20,11 @@ class DType(Enum):
     UINT8 = "uint8"
     UINT16 = "uint16"
     UINT32 = "uint32"
+    UINT64 = "uint64"
+    INT8 = "int8"
+    INT16 = "int16"
     INT32 = "int32"
+    INT64 = "int64"
     FLOAT32 = "float32"
 
     def get_numpy_dtype(self) -> npt.DTypeLike:
@@ -30,8 +35,16 @@ class DType(Enum):
             return np.uint16
         elif self == DType.UINT32:
             return np.uint32
+        elif self == DType.UINT64:
+            return np.uint64
+        elif self == DType.INT8:
+            return np.int8
+        elif self == DType.INT16:
+            return np.int16
         elif self == DType.INT32:
             return np.int32
+        elif self == DType.INT64:
+            return np.int64
         elif self == DType.FLOAT32:
             return np.float32
         raise ValueError(f"unable to handle numpy dtype {self}")
@@ -112,10 +125,12 @@ class BandSetConfig:
         self,
         config_dict: dict[str, Any],
         dtype: DType,
-        bands: list[str] | None = None,
+        bands: list[str],
         format: dict[str, Any] | None = None,
         zoom_offset: int = 0,
         remap: dict[str, Any] | None = None,
+        class_names: list[list[str]] | None = None,
+        nodata_vals: list[float] | None = None,
     ) -> None:
         """Creates a new BandSetConfig instance.
 
@@ -124,29 +139,43 @@ class BandSetConfig:
             dtype: the pixel value type to store tiles in
             bands: list of band names in this BandSetConfig
             format: the format to store tiles in, defaults to geotiff
-            zoom_offset: non-negative integer, store images at window resolution
-                divided by 2^(zoom_offset).
+            zoom_offset: store images at a resolution higher or lower than the window
+                resolution. This enables keeping source data at its native resolution,
+                either to save storage space (for lower resolution data) or to retain
+                details (for higher resolution data). If positive, store data at the
+                window resolution divided by 2^(zoom_offset) (higher resolution). If
+                negative, store data at the window resolution multiplied by
+                2^(-zoom_offset) (lower resolution).
             remap: config dict for Remapper to remap pixel values
+            class_names: optional list of names for the different possible values of
+                each band. The length of this list must equal the number of bands. For
+                example, [["forest", "desert"]] means that it is a single-band raster
+                where values can be 0 (forest) or 1 (desert).
+            nodata_vals: the nodata values for this band set. This is used during
+                materialization when creating mosaics, to determine which parts of the
+                source images should be copied.
         """
+        if class_names is not None and len(bands) != len(class_names):
+            raise ValueError(
+                f"the number of class lists ({len(class_names)}) does not match the number of bands ({len(bands)})"
+            )
+
         self.config_dict = config_dict
         self.bands = bands
-        self.format = format
         self.dtype = dtype
         self.zoom_offset = zoom_offset
         self.remap = remap
+        self.class_names = class_names
+        self.nodata_vals = nodata_vals
 
-        if not self.format:
+        if format is None:
             self.format = {"name": "geotiff"}
+        else:
+            self.format = format
 
     def serialize(self) -> dict[str, Any]:
-        """Serialize this BandSetConfig to a config dict, currently unused."""
-        return {
-            "bands": self.bands,
-            "format": self.format,
-            "dtype": self.dtype,
-            "zoom_offset": self.zoom_offset,
-            "remap": self.remap,
-        }
+        """Serialize this BandSetConfig to a config dict."""
+        return self.config_dict
 
     @staticmethod
     def from_config(config: dict[str, Any]) -> "BandSetConfig":
@@ -158,15 +187,16 @@ class BandSetConfig:
         kwargs = dict(
             config_dict=config,
             dtype=DType(config["dtype"]),
+            bands=config["bands"],
         )
-        for k in ["bands", "format", "zoom_offset", "remap"]:
+        for k in ["format", "zoom_offset", "remap", "class_names", "nodata_vals"]:
             if k in config:
                 kwargs[k] = config[k]
-        return BandSetConfig(**kwargs)
+        return BandSetConfig(**kwargs)  # type: ignore
 
     def get_final_projection_and_bounds(
-        self, projection: Projection, bounds: PixelBounds | None
-    ) -> tuple[Projection, PixelBounds | None]:
+        self, projection: Projection, bounds: PixelBounds
+    ) -> tuple[Projection, PixelBounds]:
         """Gets the final projection/bounds based on band set config.
 
         The band set config may apply a non-zero zoom offset that modifies the window's
@@ -187,11 +217,14 @@ class BandSetConfig:
             projection.x_resolution / (2**self.zoom_offset),
             projection.y_resolution / (2**self.zoom_offset),
         )
-        if bounds:
-            if self.zoom_offset > 0:
-                bounds = tuple(x * (2**self.zoom_offset) for x in bounds)
-            else:
-                bounds = tuple(x // (2 ** (-self.zoom_offset)) for x in bounds)
+        if self.zoom_offset > 0:
+            zoom_factor = 2**self.zoom_offset
+            bounds = tuple(x * zoom_factor for x in bounds)  # type: ignore
+        else:
+            bounds = tuple(
+                x // (2 ** (-self.zoom_offset))
+                for x in bounds  # type: ignore
+            )
         return projection, bounds
 
 
@@ -211,6 +244,22 @@ class SpaceMode(Enum):
     dataset.
     """
 
+    PER_PERIOD_MOSAIC = 4
+    """Create one mosaic per sub-period of the time range.
+
+    The duration of the sub-periods is controlled by another option in QueryConfig.
+    """
+
+    COMPOSITE = 5
+    """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(Enum):
     """Temporal  matching mode when looking up items corresponding to a window."""
@@ -222,10 +271,10 @@ class TimeMode(Enum):
     """Select items closest to the window time range, up to max_matches."""
 
     BEFORE = 3
-    """Select items before the start of the window time range, up to max_matches."""
+    """Select items before the end of the window time range, up to max_matches."""
 
     AFTER = 4
-    """Select items after the end of the window time range, up to max_matches."""
+    """Select items after the start of the window time range, up to max_matches."""
 
 
 class QueryConfig:
@@ -235,7 +284,9 @@ class QueryConfig:
         self,
         space_mode: SpaceMode = SpaceMode.MOSAIC,
         time_mode: TimeMode = TimeMode.WITHIN,
+        min_matches: int = 0,
         max_matches: int = 1,
+        period_duration: timedelta = timedelta(days=30),
     ):
         """Creates a new query configuration.
 
@@ -245,19 +296,29 @@ class QueryConfig:
         Args:
             space_mode: specifies how items should be matched with windows spatially
             time_mode: specifies how items should be matched with windows temporally
+            min_matches: the minimum number of item groups. If there are fewer than
+                this many matches, then no matches will be returned. This can be used
+                to prevent unnecessary data ingestion if the user plans to discard
+                windows that do not have a sufficient amount of data.
             max_matches: the maximum number of items (or groups of items, if space_mode
                 is MOSAIC) to match
+            period_duration: the duration of the periods, if the space mode is
+                PER_PERIOD_MOSAIC.
         """
         self.space_mode = space_mode
         self.time_mode = time_mode
+        self.min_matches = min_matches
         self.max_matches = max_matches
+        self.period_duration = period_duration
 
     def serialize(self) -> dict[str, Any]:
-        """Serialize this QueryConfig to a config dict, currently unused."""
+        """Serialize this QueryConfig to a config dict."""
         return {
             "space_mode": str(self.space_mode),
             "time_mode": str(self.time_mode),
+            "min_matches": self.min_matches,
             "max_matches": self.max_matches,
+            "period_duration": f"{self.period_duration.total_seconds()}s",
         }
 
     @staticmethod
@@ -267,11 +328,20 @@ class QueryConfig:
         Args:
             config: the config dict for this QueryConfig
         """
-        return QueryConfig(
-            space_mode=SpaceMode[config.get("space_mode", "MOSAIC")],
-            time_mode=TimeMode[config.get("time_mode", "WITHIN")],
-            max_matches=config.get("max_matches", 1),
-        )
+        kwargs: dict[str, Any] = dict()
+        if "space_mode" in config:
+            kwargs["space_mode"] = SpaceMode[config["space_mode"]]
+        if "time_mode" in config:
+            kwargs["time_mode"] = TimeMode[config["time_mode"]]
+        if "period_duration" in config:
+            kwargs["period_duration"] = timedelta(
+                seconds=pytimeparse.parse(config["period_duration"])
+            )
+        for k in ["min_matches", "max_matches"]:
+            if k not in config:
+                continue
+            kwargs[k] = config[k]
+        return QueryConfig(**kwargs)
 
 
 class DataSourceConfig:
@@ -307,16 +377,8 @@ class DataSourceConfig:
         self.ingest = ingest
 
     def serialize(self) -> dict[str, Any]:
-        """Serialize this DataSourceConfig to a config dict, currently unused."""
-        config_dict = self.config_dict.copy()
-        config_dict["name"] = self.name
-        config_dict["query_config"] = self.query_config.serialize()
-        config_dict["ingest"] = self.ingest
-        if self.time_offset:
-            config_dict["time_offset"] = str(self.time_offset)
-        if self.duration:
-            config_dict["duration"] = str(self.duration)
-        return config_dict
+        """Serialize this DataSourceConfig to a config dict."""
+        return self.config_dict
 
     @staticmethod
     def from_config(config: dict[str, Any]) -> "DataSourceConfig":
@@ -371,13 +433,40 @@ class LayerConfig:
         self.alias = alias
 
     def serialize(self) -> dict[str, Any]:
-        """Serialize this LayerConfig to a config dict, currently unused."""
+        """Serialize this LayerConfig to a config dict."""
         return {
             "layer_type": str(self.layer_type),
-            "data_source": self.data_source,
+            "data_source": self.data_source.serialize() if self.data_source else None,
             "alias": self.alias,
         }
 
+    def __hash__(self) -> int:
+        """Return a hash of this LayerConfig."""
+        return hash(json.dumps(self.serialize(), sort_keys=True))
+
+    def __eq__(self, other: Any) -> bool:
+        """Returns whether other is the same as this LayerConfig.
+
+        Args:
+            other: the other object to compare.
+        """
+        if not isinstance(other, LayerConfig):
+            return False
+        return self.serialize() == other.serialize()
+
+
+class CompositingMethod(Enum):
+    """Method how to select pixels for the composite from corresponding items of a window."""
+
+    FIRST_VALID = 1
+    """Select first valid pixel in order of corresponding items (might be sorted)"""
+
+    MEAN = 2
+    """Select per-pixel mean value of corresponding items of a window"""
+
+    MEDIAN = 3
+    """Select per-pixel median value of corresponding items of a window"""
+
 
 class RasterLayerConfig(LayerConfig):
     """Configuration of a raster layer."""
@@ -389,6 +478,7 @@ class RasterLayerConfig(LayerConfig):
         data_source: DataSourceConfig | None = None,
         resampling_method: Resampling = Resampling.bilinear,
         alias: str | None = None,
+        compositing_method: CompositingMethod = CompositingMethod.FIRST_VALID,
     ):
         """Initialize a new RasterLayerConfig.
 
@@ -398,10 +488,12 @@ class RasterLayerConfig(LayerConfig):
             data_source: optional DataSourceConfig if this layer is retrievable
             resampling_method: how to resample rasters (if needed), default bilinear resampling
             alias: alias for this layer to use in the tile store
+            compositing_method: how to compute pixel values in the composite of each windows items
         """
         super().__init__(layer_type, data_source, alias)
         self.band_sets = band_sets
         self.resampling_method = resampling_method
+        self.compositing_method = compositing_method
 
     @staticmethod
     def from_config(config: dict[str, Any]) -> "RasterLayerConfig":
@@ -422,7 +514,11 @@ class RasterLayerConfig(LayerConfig):
             ]
         if "alias" in config:
             kwargs["alias"] = config["alias"]
-        return RasterLayerConfig(**kwargs)
+        if "compositing_method" in config:
+            kwargs["compositing_method"] = CompositingMethod[
+                config["compositing_method"]
+            ]
+        return RasterLayerConfig(**kwargs)  # type: ignore
 
 
 class VectorLayerConfig(LayerConfig):
@@ -432,22 +528,28 @@ class VectorLayerConfig(LayerConfig):
         self,
         layer_type: LayerType,
         data_source: DataSourceConfig | None = None,
-        zoom_offset: int = 0,
         format: VectorFormatConfig = VectorFormatConfig("geojson"),
         alias: str | None = None,
+        class_property_name: str | None = None,
+        class_names: list[str] | None = None,
     ):
         """Initialize a new VectorLayerConfig.
 
         Args:
             layer_type: the LayerType (must be vector)
             data_source: optional DataSourceConfig if this layer is retrievable
-            zoom_offset: zoom offset at which to store the vector data
             format: the VectorFormatConfig, default storing as GeoJSON
             alias: alias for this layer to use in the tile store
+            class_property_name: optional metadata field indicating that the GeoJSON
+                features contain a property that corresponds to a class label, and this
+                is the name of that property.
+            class_names: the list of classes that the class_property_name property
+                could be set to.
         """
         super().__init__(layer_type, data_source, alias)
-        self.zoom_offset = zoom_offset
         self.format = format
+        self.class_property_name = class_property_name
+        self.class_names = class_names
 
     @staticmethod
     def from_config(config: dict[str, Any]) -> "VectorLayerConfig":
@@ -456,42 +558,29 @@ class VectorLayerConfig(LayerConfig):
         Args:
             config: the config dict for this VectorLayerConfig
         """
-        kwargs = {"layer_type": LayerType(config["type"])}
+        kwargs: dict[str, Any] = {"layer_type": LayerType(config["type"])}
         if "data_source" in config:
             kwargs["data_source"] = DataSourceConfig.from_config(config["data_source"])
-        if "zoom_offset" in config:
-            kwargs["zoom_offset"] = config["zoom_offset"]
         if "format" in config:
             kwargs["format"] = VectorFormatConfig.from_config(config["format"])
-        if "alias" in config:
-            kwargs["alias"] = config["alias"]
-        return VectorLayerConfig(**kwargs)
 
-    def get_final_projection_and_bounds(
-        self, projection: Projection, bounds: PixelBounds | None
-    ) -> tuple[Projection, PixelBounds | None]:
-        """Gets the final projection/bounds based on zoom offset.
+        simple_optionals = [
+            "alias",
+            "class_property_name",
+            "class_names",
+        ]
+        for k in simple_optionals:
+            if k in config:
+                kwargs[k] = config[k]
 
-        Args:
-            projection: the window's projection
-            bounds: the window's bounds (optional)
+        # The "zoom_offset" option was removed.
+        # We should change how we create configuration so we can error on all
+        # non-existing config options, but for now we make sure to raise error if
+        # zoom_offset is set since it is no longer supported.
+        if "zoom_offset" in config:
+            raise ValueError("unsupported zoom_offset option in vector layer config")
 
-        Returns:
-            tuple of updated projection and bounds with zoom offset applied
-        """
-        if self.zoom_offset == 0:
-            return projection, bounds
-        projection = Projection(
-            projection.crs,
-            projection.x_resolution / (2**self.zoom_offset),
-            projection.y_resolution / (2**self.zoom_offset),
-        )
-        if bounds:
-            if self.zoom_offset > 0:
-                bounds = tuple(x * (2**self.zoom_offset) for x in bounds)
-            else:
-                bounds = tuple(x // (2 ** (-self.zoom_offset)) for x in bounds)
-        return projection, bounds
+        return VectorLayerConfig(**kwargs)  # type: ignore
 
 
 def load_layer_config(config: dict[str, Any]) -> LayerConfig:
@@ -502,26 +591,3 @@ def load_layer_config(config: dict[str, Any]) -> LayerConfig:
     elif layer_type == LayerType.VECTOR:
         return VectorLayerConfig.from_config(config)
     raise ValueError(f"Unknown layer type {layer_type}")
-
-
-class TileStoreConfig:
-    """A configuration specifying a TileStore."""
-
-    def __init__(self, name: str, config_dict: dict[str, Any]) -> None:
-        """Create a new TileStoreConfig.
-
-        Args:
-            name: the tile store implementation name to use
-            config_dict: configuration options
-        """
-        self.name = name
-        self.config_dict = config_dict
-
-    @staticmethod
-    def from_config(config: dict[str, Any]) -> "TileStoreConfig":
-        """Create a TileStoreConfig from config dict.
-
-        Args:
-            config: the config dict for this TileStoreConfig
-        """
-        return TileStoreConfig(name=config["name"], config_dict=config)

rslearn/const.py

@@ -1,23 +1,17 @@
 """Constants."""
 
-from rasterio.crs import CRS
-
-from rslearn.utils import PixelBounds, Projection
-
-WGS84_EPSG = 4326
-"""The EPSG code for WGS-84."""
-
-WGS84_PROJECTION = Projection(CRS.from_epsg(WGS84_EPSG), 1, 1)
-"""The Projection for WGS-84 assuming 1 degree per pixel.
-
-This can be used to create STGeometry with shapes in longitude/latitude coordinates.
-"""
-
-WGS84_BOUNDS: PixelBounds = (-180, -90, 180, 90)
-"""The bounds of the WGS-84 projection."""
+from rslearn.utils.geometry import WGS84_BOUNDS, WGS84_EPSG, WGS84_PROJECTION
 
 TILE_SIZE = 512
 """Default tile size. TODO: remove this or move it elsewhere."""
 
 SHAPEFILE_AUX_EXTENSIONS = [".cpg", ".dbf", ".prj", ".sbn", ".sbx", ".shx", ".txt"]
 """Extensions of potential auxiliary files to .shp file."""
+
+__all__ = (
+    "WGS84_PROJECTION",
+    "WGS84_EPSG",
+    "WGS84_BOUNDS",
+    "TILE_SIZE",
+    "SHAPEFILE_AUX_EXTENSIONS",
+)

rslearn/data_sources/__init__.py

@@ -10,15 +10,20 @@ 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.
 
@@ -26,6 +31,9 @@ def data_source_from_config(config: LayerConfig, ds_path: UPath) -> DataSource:
         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]