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

rslearn/config/dataset.py

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

rslearn/data_sources/local_files.py

@@ -236,8 +236,8 @@ class RasterImporter(Importer):
                     "windows in the rslearn dataset. When using settings like "
                     "max_matches=1 and space_mode=MOSAIC, this may cause windows outside "
                     "the geometry’s valid bounds to be materialized from the global raster "
-                    "instead of a more appropriate source. Consider using COMPOSITE mode, "
-                    "or increasing max_matches if this behavior is unintended."
+                    "instead of a more appropriate source. Consider increasing max_matches"
+                    "if this behavior is unintended."
                 )
 
             if spec.name:

rslearn/data_sources/utils.py

@@ -1,7 +1,9 @@
 """Utilities shared by data sources."""
 
+import warnings
+from collections.abc import Callable
 from dataclasses import dataclass
-from datetime import UTC, datetime, timedelta
+from datetime import UTC, datetime
 from typing import TypeVar
 
 import shapely
@@ -40,13 +42,13 @@ class PendingMosaic:
     completed: bool = False
 
 
-def mosaic_matching(
+def _create_single_coverage_mosaics(
     window_geometry: STGeometry,
     items: list[ItemType],
     item_shps: list[shapely.Geometry],
-    max_matches: int,
+    max_mosaics: int,
 ) -> list[list[ItemType]]:
-    """Spatial item matching for mosaic space mode.
+    """Create mosaics where each mosaic covers the window geometry once.
 
     This attempts to piece together items into mosaics that fully cover the window
     geometry. If there are items leftover that only partially cover the window
@@ -56,15 +58,16 @@ def mosaic_matching(
         window_geometry: the geometry of the window.
         items: list of items.
         item_shps: the item shapes projected to the window's projection.
-        max_matches: the maximum number of matches (mosaics) to create.
+        max_mosaics: the maximum number of mosaics to create.
 
     Returns:
-        list of item groups, each one corresponding to a different mosaic.
+        list of item groups, each one corresponding to a different single-coverage
+        mosaic.
     """
     # To create mosaics, we iterate over the items in order, and add each item to
     # the first mosaic that the new item adds coverage to.
 
-    # max_matches could be very high if the user just wants us to create as many
+    # max_mosaics could be very high if the user just wants us to create as many
     # mosaics as possible, so we initialize the list here as empty and just add
     # more pending mosaics when it is necessary.
     pending_mosaics: list[PendingMosaic] = []
@@ -108,7 +111,7 @@ def mosaic_matching(
 
         # See if we can add a new mosaic based on this item. There must be room for
         # more mosaics, but the item must also intersect the requested geometry.
-        if len(pending_mosaics) >= max_matches:
+        if len(pending_mosaics) >= max_mosaics:
             continue
         intersect_area = item_shp.intersection(window_geometry.shp).area
         if (
@@ -127,18 +130,148 @@ def mosaic_matching(
     return [pending_mosaic.items for pending_mosaic in pending_mosaics]
 
 
-def per_period_mosaic_matching(
-    window_geometry: STGeometry,
-    item_list: list[ItemType],
-    period_duration: timedelta,
-    max_matches: int,
+def _consolidate_mosaics_by_overlaps(
+    mosaics: list[list[ItemType]],
+    overlaps: int,
+    max_groups: int,
+) -> list[list[ItemType]]:
+    """Consolidate single-coverage mosaics into groups based on desired overlaps.
+
+    Args:
+        mosaics: list of single-coverage mosaics (each mosaic is a list of items).
+        overlaps: the number of overlapping coverages wanted per group.
+        max_groups: the maximum number of groups to return.
+
+    Returns:
+        list of item groups, where each group contains items from multiple mosaics
+        to achieve the desired number of overlapping coverages.
+    """
+    if overlaps <= 0:
+        overlaps = 1
+
+    groups: list[list[ItemType]] = []
+    for i in range(0, len(mosaics), overlaps):
+        if len(groups) >= max_groups:
+            break
+        # Combine overlaps consecutive mosaics into one group
+        combined_items: list[ItemType] = []
+        for mosaic in mosaics[i : i + overlaps]:
+            combined_items.extend(mosaic)
+        if combined_items:
+            groups.append(combined_items)
+
+    return groups
+
+
+def match_with_space_mode_contains(
+    geometry: STGeometry,
+    items: list[ItemType],
+    item_shps: list[shapely.Geometry],
+    query_config: QueryConfig,
+) -> list[list[ItemType]]:
+    """Match items that fully contain the window geometry.
+
+    Args:
+        geometry: the window's geometry.
+        items: list of items.
+        item_shps: the item shapes projected to the window's projection.
+        query_config: the query configuration.
+
+    Returns:
+        list of matched item groups, where each group contains a single item.
+    """
+    groups: list[list[ItemType]] = []
+    for item, item_shp in zip(items, item_shps):
+        if not item_shp.contains(geometry.shp):
+            continue
+        groups.append([item])
+        if len(groups) >= query_config.max_matches:
+            break
+    return groups
+
+
+def match_with_space_mode_intersects(
+    geometry: STGeometry,
+    items: list[ItemType],
+    item_shps: list[shapely.Geometry],
+    query_config: QueryConfig,
+) -> list[list[ItemType]]:
+    """Match items that intersect any portion of the window geometry.
+
+    Args:
+        geometry: the window's geometry.
+        items: list of items.
+        item_shps: the item shapes projected to the window's projection.
+        query_config: the query configuration.
+
+    Returns:
+        list of matched item groups, where each group contains a single item.
+    """
+    groups: list[list[ItemType]] = []
+    for item, item_shp in zip(items, item_shps):
+        if not shp_intersects(item_shp, geometry.shp):
+            continue
+        groups.append([item])
+        if len(groups) >= query_config.max_matches:
+            break
+    return groups
+
+
+def match_with_space_mode_mosaic(
+    geometry: STGeometry,
+    items: list[ItemType],
+    item_shps: list[shapely.Geometry],
+    query_config: QueryConfig,
+) -> list[list[ItemType]]:
+    """Match items into mosaic groups that cover the window geometry.
+
+    Creates groups of items that together cover the window geometry. The number of
+    overlapping coverages in each group is controlled by mosaic_compositing_overlaps.
+
+    Args:
+        geometry: the window's geometry.
+        items: list of items.
+        item_shps: the item shapes projected to the window's projection.
+        query_config: the query configuration.
+
+    Returns:
+        list of matched item groups, where each group forms a mosaic covering the
+        window.
+    """
+    overlaps = query_config.mosaic_compositing_overlaps
+
+    # Calculate how many single-coverage mosaics we need to create.
+    # We need enough mosaics to consolidate into max_matches groups with the
+    # desired number of overlaps per group.
+    max_single_mosaics = query_config.max_matches * overlaps
+
+    # Create single-coverage mosaics
+    single_mosaics = _create_single_coverage_mosaics(
+        geometry, items, item_shps, max_single_mosaics
+    )
+
+    # Consolidate into groups based on overlaps
+    return _consolidate_mosaics_by_overlaps(
+        single_mosaics, overlaps, query_config.max_matches
+    )
+
+
+def match_with_space_mode_per_period_mosaic(
+    geometry: STGeometry,
+    items: list[ItemType],
+    item_shps: list[shapely.Geometry],
+    query_config: QueryConfig,
 ) -> list[list[ItemType]]:
     """Match items to the geometry with one mosaic per period.
 
     We divide the time range of the geometry into shorter periods. Within each period,
     we use the items corresponding to that period to create a mosaic. The returned item
-    groups include one group per period, starting from the most recent periods, up to
-    the provided max_matches.
+    groups include one group per period, up to the provided max_matches.
+
+    By default (reverse_time_order=True), groups are returned starting from the most
+    recent periods. When reverse_time_order=False, groups are returned in chronological
+    order (oldest first). reverse_time_order should always be set False, and
+    FutureWarning will be warned if it is not.
 
     The periods are also bounded to the window's time range, and aligned with the end
     of that time range, i.e. the most recent window is
@@ -159,42 +292,59 @@ def per_period_mosaic_matching(
     max_matches*period_duration is not equivalent to a longer window duration.
 
     Args:
-        window_geometry: the window geometry to match items to.
-        item_list: the list of items.
-        period_duration: the duration of one period.
-        max_matches: the number of per-period mosaics to create.
+        geometry: the window's geometry.
+        items: list of items.
+        item_shps: the item shapes projected to the window's projection (unused here)
+        query_config: the query configuration.
 
     Returns:
-        the matched item groups, where each group contains items that yield a
-            per-period mosaic.
+        list of matched item groups, where each group contains items that yield a
+        per-period mosaic.
     """
-    if window_geometry.time_range is None:
+    if geometry.time_range is None:
         raise ValueError(
             "all windows must have time range for per period mosaic matching"
         )
 
+    # Emit warning if per_period_mosaic_reverse_time_order is True (the default).
+    if query_config.per_period_mosaic_reverse_time_order:
+        warnings.warn(
+            "QueryConfig.per_period_mosaic_reverse_time_order defaults to True, which "
+            "returns item groups in reverse temporal order (most recent first) for "
+            "PER_PERIOD_MOSAIC mode. This default will change to False (chronological "
+            "order) after 2026-04-01. To silence this warning, explicitly set "
+            "per_period_mosaic_reverse_time_order=False.",
+            FutureWarning,
+            stacklevel=3,
+        )
+
+    period_duration = query_config.period_duration
+
     # For each period, we create an STGeometry with modified time range matching that
     # period, and use it with match_candidate_items_to_window to get a mosaic.
     cur_groups: list[list[ItemType]] = []
-    period_start = window_geometry.time_range[1] - period_duration
+    period_start = geometry.time_range[1] - period_duration
     while (
-        period_start >= window_geometry.time_range[0] and len(cur_groups) < max_matches
+        period_start >= geometry.time_range[0]
+        and len(cur_groups) < query_config.max_matches
     ):
         period_time_range = (
             period_start,
             period_start + period_duration,
         )
         period_start -= period_duration
-        period_geom = STGeometry(
-            window_geometry.projection, window_geometry.shp, period_time_range
-        )
+        period_geom = STGeometry(geometry.projection, geometry.shp, period_time_range)
 
         # We modify the QueryConfig here since caller should be asking for
         # multiple mosaics, but we just want one mosaic per period.
         period_groups = match_candidate_items_to_window(
             period_geom,
-            item_list,
-            QueryConfig(space_mode=SpaceMode.MOSAIC, max_matches=1),
+            items,
+            QueryConfig(
+                space_mode=SpaceMode.MOSAIC,
+                max_matches=1,
+                mosaic_compositing_overlaps=query_config.mosaic_compositing_overlaps,
+            ),
         )
 
         # There should be zero or one group depending on whether there were
@@ -204,9 +354,29 @@ def per_period_mosaic_matching(
             continue
         cur_groups.append(period_groups[0])
 
+    # Currently the item groups are in reverse chronologic order.
+    # Reverse it to correct chronological order if requested.
+    if not query_config.per_period_mosaic_reverse_time_order:
+        cur_groups.reverse()
+
     return cur_groups
 
 
+# Type alias for space mode handler functions
+SpaceModeHandler = Callable[
+    [STGeometry, list[ItemType], list[shapely.Geometry], QueryConfig],
+    list[list[ItemType]],
+]
+
+# Dict mapping SpaceMode values to their handler functions
+space_mode_handlers: dict[SpaceMode, SpaceModeHandler] = {
+    SpaceMode.CONTAINS: match_with_space_mode_contains,
+    SpaceMode.INTERSECTS: match_with_space_mode_intersects,
+    SpaceMode.MOSAIC: match_with_space_mode_mosaic,
+    SpaceMode.PER_PERIOD_MOSAIC: match_with_space_mode_per_period_mosaic,
+}
+
+
 def match_candidate_items_to_window(
     geometry: STGeometry, items: list[ItemType], query_config: QueryConfig
 ) -> list[list[ItemType]]:
@@ -262,43 +432,13 @@ def match_candidate_items_to_window(
                 item_geom = item_geom.to_projection(geometry.projection)
         item_shps.append(item_geom.shp)
 
-    if query_config.space_mode == SpaceMode.CONTAINS:
-        groups = []
-        for item, item_shp in zip(items, item_shps):
-            if not item_shp.contains(geometry.shp):
-                continue
-            groups.append([item])
-            if len(groups) >= query_config.max_matches:
-                break
-
-    elif query_config.space_mode == SpaceMode.INTERSECTS:
-        groups = []
-        for item, item_shp in zip(items, item_shps):
-            if not shp_intersects(item_shp, geometry.shp):
-                continue
-            groups.append([item])
-            if len(groups) >= query_config.max_matches:
-                break
-
-    elif query_config.space_mode == SpaceMode.MOSAIC:
-        groups = mosaic_matching(geometry, items, item_shps, query_config.max_matches)
-
-    elif query_config.space_mode == SpaceMode.PER_PERIOD_MOSAIC:
-        groups = per_period_mosaic_matching(
-            geometry, items, query_config.period_duration, query_config.max_matches
-        )
-
-    elif query_config.space_mode == SpaceMode.COMPOSITE:
-        group = []
-        for item, item_shp in zip(items, item_shps):
-            if not shp_intersects(item_shp, geometry.shp):
-                continue
-            group.append(item)
-        groups = [group]
-
-    else:
+    # Dispatch to the appropriate space mode handler
+    handler = space_mode_handlers.get(query_config.space_mode)
+    if handler is None:
         raise ValueError(f"invalid space mode {query_config.space_mode}")
 
+    groups = handler(geometry, items, item_shps, query_config)
+
     # Enforce minimum matches if set.
     if len(groups) < query_config.min_matches:
         logger.warning(

rslearn/dataset/materialize.py

@@ -236,7 +236,11 @@ def read_and_stack_raster_windows(
     band_dtype: npt.DTypeLike,
     resampling_method: Resampling = Resampling.bilinear,
 ) -> npt.NDArray[np.generic]:
-    """Create a stack of extent aligned raster windows.
+    """Create a stack of raster images, with one per item in the group.
+
+    We read the portion of each raster item corresponding to the window extent, and
+    stack the resulting images. This is used for the MEAN and MEDIAN compositing
+    methods to it can compute aggregate statistics across the stack.
 
     Args:
         group: Iterable of items (e.g., scene metadata objects) to read data from.

rslearn/models/clay/clay.py

@@ -105,7 +105,7 @@ class Clay(FeatureExtractor):
 
     def _resize_image(self, image: torch.Tensor, original_hw: int) -> torch.Tensor:
         """Resize the image to the input resolution."""
-        new_hw = self.patch_size if original_hw == 1 else DEFAULT_IMAGE_RESOLUTION
+        new_hw = PATCH_SIZE if original_hw == 1 else DEFAULT_IMAGE_RESOLUTION
         return F.interpolate(
             image, size=(new_hw, new_hw), mode="bilinear", align_corners=False
         )
@@ -123,7 +123,8 @@ class Clay(FeatureExtractor):
         device = param.device
 
         chips = torch.stack(
-            [inp[self.modality] for inp in context.inputs], dim=0
+            [inp[self.modality].single_ts_to_chw_tensor() for inp in context.inputs],
+            dim=0,
         )  # (B, C, H, W)
         if self.do_resizing:
             chips = self._resize_image(chips, chips.shape[2])
@@ -203,7 +204,6 @@ class ClayNormalize(Transform):
                 mean=means,
                 std=stds,
                 selectors=[modality],
-                num_bands=len(means),
             )
         self.normalizers = torch.nn.ModuleDict(normalizers)
 

rslearn/models/detr/detr.py

@@ -468,7 +468,10 @@ class Detr(Predictor):
 
         # Get image sizes.
         image_sizes = torch.tensor(
-            [[inp["image"].shape[2], inp["image"].shape[1]] for inp in context.inputs],
+            [
+                [inp["image"].image.shape[2], inp["image"].image.shape[1]]
+                for inp in context.inputs
+            ],
             dtype=torch.int32,
             device=features.device,
         )

rslearn/models/dinov3.py

@@ -159,7 +159,6 @@ class DinoV3Normalize(Transform):
         self.normalize = Normalize(
             [value * 255 for value in mean],
             [value * 255 for value in std],
-            num_bands=3,
         )
 
     def forward(

rslearn/models/olmoearth_pretrain/model.py

@@ -95,7 +95,9 @@ class OlmoEarth(FeatureExtractor):
         """
         if use_legacy_timestamps:
             warnings.warn(
-                "For new projects, don't use legacy timesteps.", DeprecationWarning
+                "For new projects, don't use legacy timesteps. "
+                "Support will be removed after 2026-04-01.",
+                FutureWarning,
             )
 
         if (

rslearn/models/pooling_decoder.py

@@ -124,6 +124,6 @@ class SegmentationPoolingDecoder(PoolingDecoder):
         """
         output_probs = super().forward(intermediates, context)
         # BC -> BCHW
-        h, w = context.inputs[0][self.image_key].shape[1:3]
+        h, w = context.inputs[0][self.image_key].image.shape[1:3]
         feat_map = output_probs.feature_vector[:, :, None, None].repeat([1, 1, h, w])
         return FeatureMaps([feat_map])

rslearn/models/prithvi.py

@@ -230,7 +230,6 @@ class PrithviNormalize(Transform):
         self.normalizer = Normalize(
             mean=config["mean"],
             std=config["std"],
-            num_bands=len(config["mean"]),
             selectors=[PrithviV2.INPUT_KEY],
         )