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

rslearn/data_sources/earthdatahub.py

@@ -0,0 +1,425 @@
+"""Data sources backed by EarthDataHub-hosted datasets."""
+
+from __future__ import annotations
+
+import math
+import os
+import tempfile
+from datetime import UTC, datetime, timedelta
+from typing import Any, Literal
+
+import numpy as np
+import rasterio
+import shapely
+import xarray as xr
+from rasterio.transform import from_origin
+from upath import UPath
+
+from rslearn.config import QueryConfig, SpaceMode
+from rslearn.const import WGS84_EPSG, WGS84_PROJECTION
+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
+
+logger = get_logger(__name__)
+
+
+def _floor_to_utc_day(t: datetime) -> datetime:
+    """Return the UTC day boundary (00:00) for a datetime.
+
+    If `t` is naive, it is treated as UTC. If `t` is timezone-aware, it is converted
+    to UTC before flooring.
+    """
+    if t.tzinfo is None:
+        t = t.replace(tzinfo=UTC)
+    else:
+        t = t.astimezone(UTC)
+    return datetime(t.year, t.month, t.day, tzinfo=UTC)
+
+
+def _bounds_to_lon_ranges_0_360(
+    bounds: tuple[float, float, float, float],
+) -> list[tuple[float, float]]:
+    """Convert lon bounds to one or two non-wrapping ranges in [0, 360].
+
+    Expects non-wrapping bounds where min_lon <= max_lon.
+    """
+    min_lon, _, max_lon, _ = bounds
+
+    if min_lon >= 0 and max_lon >= 0:
+        return [(min_lon, max_lon)]
+    if min_lon < 0 and max_lon < 0:
+        return [(min_lon + 360.0, max_lon + 360.0)]
+
+    # Bounds cross 0 degrees (e.g. [-5, 5]) which wraps in the 0..360 convention.
+    return [(min_lon + 360.0, 360.0), (0.0, max_lon)]
+
+
+def _snap_bounds_outward(
+    bounds: tuple[float, float, float, float],
+    step_degrees: float,
+) -> tuple[float, float, float, float]:
+    """Snap lon/lat bounds outward to a fixed grid.
+
+    ERA5(-Land) datasets are provided on a regular 0.1° lat/lon grid. When a window is
+    much smaller than the grid spacing, its geographic bounds may not contain any grid
+    *centers*, which leads to empty `xarray.sel(..., slice(...))` selections.
+
+    Snapping outward ensures at least one grid cell is selected when there is any
+    overlap.
+    """
+    min_lon, min_lat, max_lon, max_lat = bounds
+    if min_lon > max_lon or min_lat > max_lat:
+        raise ValueError(f"invalid bounds: {bounds}")
+
+    min_lon = math.floor(min_lon / step_degrees) * step_degrees
+    min_lat = math.floor(min_lat / step_degrees) * step_degrees
+    max_lon = math.ceil(max_lon / step_degrees) * step_degrees
+    max_lat = math.ceil(max_lat / step_degrees) * step_degrees
+
+    # Ensure non-empty after snapping.
+    if max_lon == min_lon:
+        max_lon = min_lon + step_degrees
+    if max_lat == min_lat:
+        max_lat = min_lat + step_degrees
+
+    return (min_lon, min_lat, max_lon, max_lat)
+
+
+class ERA5LandDailyUTCv1(DataSource[Item]):
+    """ERA5-Land daily UTC (v1) hosted on EarthDataHub.
+
+    This data source reads from the EarthDataHub Zarr store and writes daily GeoTIFFs
+    into the dataset tile store.
+
+    Supported bands:
+    - d2m: 2m dewpoint temperature (units: K)
+    - e: evaporation (units: m of water equivalent)
+    - pev: potential evaporation (units: m)
+    - ro: runoff (units: m)
+    - sp: surface pressure (units: Pa)
+    - ssr: surface net short-wave (solar) radiation (units: J m-2)
+    - ssrd: surface short-wave (solar) radiation downwards (units: J m-2)
+    - str: surface net long-wave (thermal) radiation (units: J m-2)
+    - swvl1: volumetric soil water layer 1 (units: m3 m-3)
+    - swvl2: volumetric soil water layer 2 (units: m3 m-3)
+    - t2m: 2m temperature (units: K)
+    - tp: total precipitation (units: m)
+    - u10: 10m U wind component (units: m s-1)
+    - v10: 10m V wind component (units: m s-1)
+
+    Authentication:
+        EarthDataHub uses token-based auth. Configure your netrc file so HTTP clients
+        can attach the token automatically. On Linux and MacOS the netrc path is
+        `~/.netrc`.
+    """
+
+    DEFAULT_ZARR_URL = (
+        "https://data.earthdatahub.destine.eu/era5/era5-land-daily-utc-v1.zarr"
+    )
+    ALLOWED_BANDS = {
+        "d2m",
+        "e",
+        "pev",
+        "ro",
+        "sp",
+        "ssr",
+        "ssrd",
+        "str",
+        "swvl1",
+        "swvl2",
+        "t2m",
+        "tp",
+        "u10",
+        "v10",
+    }
+    PIXEL_SIZE_DEGREES = 0.1
+
+    def __init__(
+        self,
+        band_names: list[str] | None = None,
+        zarr_url: str = DEFAULT_ZARR_URL,
+        bounds: list[float] | None = None,
+        temperature_unit: Literal["celsius", "kelvin"] = "kelvin",
+        trust_env: bool = True,
+        context: DataSourceContext = DataSourceContext(),
+    ) -> None:
+        """Initialize a new ERA5LandDailyUTCv1 instance.
+
+        Args:
+            band_names: list of bands to ingest. If omitted and a LayerConfig is
+                provided via context, bands are inferred from that layer's band sets.
+            zarr_url: URL/path to the EarthDataHub Zarr store.
+            bounds: optional bounding box as [min_lon, min_lat, max_lon, max_lat]
+                in degrees (WGS84). For best performance, set bounds to your area of
+                interest.
+            temperature_unit: units to return for `t2m` ("celsius" or "kelvin").
+            trust_env: if True (default), allow the underlying HTTP client to read
+                environment configuration (including netrc) for auth/proxies.
+            context: rslearn data source context.
+        """
+        self.zarr_url = zarr_url
+        if bounds is not None:
+            if len(bounds) != 4:
+                raise ValueError(
+                    "ERA5LandDailyUTCv1 bounds must be [min_lon, min_lat, max_lon, max_lat] "
+                    f"(got {bounds!r})."
+                )
+            min_lon, min_lat, max_lon, max_lat = bounds
+            if min_lon > max_lon:
+                raise ValueError(
+                    "ERA5LandDailyUTCv1 does not yet support longitude ranges that cross the dateline "
+                    f"(got bounds min_lon={min_lon}, max_lon={max_lon})."
+                )
+            if min_lat > max_lat:
+                raise ValueError(
+                    "ERA5LandDailyUTCv1 bounds must have min_lat <= max_lat "
+                    f"(got bounds min_lat={min_lat}, max_lat={max_lat})."
+                )
+        self.bounds = bounds
+        self.temperature_unit = temperature_unit
+        self.trust_env = trust_env
+
+        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 not in self.band_names:
+                        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"
+            )
+
+        invalid_bands = [b for b in self.band_names if b not in self.ALLOWED_BANDS]
+        if invalid_bands:
+            raise ValueError(
+                f"unsupported ERA5LandDailyUTCv1 band(s): {invalid_bands}; "
+                f"supported: {sorted(self.ALLOWED_BANDS)}"
+            )
+
+        self._ds: Any | None = None
+
+    def _get_dataset(self) -> xr.Dataset:
+        """Open (and memoize) the backing ERA5-Land Zarr dataset."""
+        if self._ds is not None:
+            return self._ds
+
+        storage_options: dict[str, Any] | None = None
+        if self.zarr_url.startswith("http://") or self.zarr_url.startswith("https://"):
+            storage_options = {"client_kwargs": {"trust_env": self.trust_env}}
+
+        self._ds = xr.open_dataset(
+            self.zarr_url,
+            engine="zarr",
+            chunks=None,  # No dask
+            storage_options=storage_options,
+        )
+        return self._ds
+
+    def get_items(
+        self, geometries: list[STGeometry], query_config: QueryConfig
+    ) -> list[list[list[Item]]]:
+        """Get daily items intersecting the given geometries.
+
+        Returns one item per UTC day that intersects each requested geometry time
+        range.
+        """
+        if query_config.space_mode != SpaceMode.MOSAIC:
+            raise ValueError("expected mosaic space mode in the query configuration")
+
+        if self.bounds is not None:
+            min_lon, min_lat, max_lon, max_lat = self.bounds
+            item_shp = shapely.box(min_lon, min_lat, max_lon, max_lat)
+        else:
+            item_shp = shapely.box(-180, -90, 180, 90)
+
+        all_groups: list[list[list[Item]]] = []
+        for geometry in geometries:
+            if geometry.time_range is None:
+                raise ValueError("expected all geometries to have a time range")
+
+            start, end = geometry.time_range
+            cur_day = _floor_to_utc_day(start)
+            cur_groups: list[list[Item]] = []
+            while cur_day < end:
+                next_day = cur_day + timedelta(days=1)
+                item_name = f"era5land_dailyutc_v1_{cur_day.year:04d}{cur_day.month:02d}{cur_day.day:02d}"
+                item_geom = STGeometry(WGS84_PROJECTION, item_shp, (cur_day, next_day))
+                cur_groups.append([Item(item_name, item_geom)])
+                cur_day = next_day
+
+            all_groups.append(cur_groups)
+
+        return all_groups
+
+    def deserialize_item(self, serialized_item: Any) -> Item:
+        """Deserialize an `Item` previously produced by this data source."""
+        assert isinstance(serialized_item, dict)
+        return Item.deserialize(serialized_item)
+
+    def _get_effective_bounds(
+        self, geometries: list[STGeometry]
+    ) -> tuple[float, float, float, float]:
+        """Compute an effective WGS84 bounding box for ingestion.
+
+        If `self.bounds` is set, it is used as-is; otherwise, the bounds are derived
+        from the union of the provided geometries (after projecting each to WGS84).
+        """
+        if self.bounds is not None:
+            min_lon, min_lat, max_lon, max_lat = self.bounds
+            return (min_lon, min_lat, max_lon, max_lat)
+
+        min_lon = 180.0
+        min_lat = 90.0
+        max_lon = -180.0
+        max_lat = -90.0
+        for geom in geometries:
+            wgs84 = geom.to_projection(WGS84_PROJECTION)
+            b = wgs84.shp.bounds
+            min_lon = min(min_lon, b[0])
+            min_lat = min(min_lat, b[1])
+            max_lon = max(max_lon, b[2])
+            max_lat = max(max_lat, b[3])
+        return (min_lon, min_lat, max_lon, max_lat)
+
+    def _write_geotiff(
+        self,
+        tif_path: str,
+        lat: np.ndarray,
+        lon: np.ndarray,
+        band_arrays: list[np.ndarray],
+    ) -> None:
+        """Write a GeoTIFF with WGS84 georeferencing from ERA5(-Land) arrays.
+
+        Args:
+            tif_path: destination GeoTIFF path.
+            lat: 1D latitude coordinate (expected descending, north-to-south).
+            lon: 1D longitude coordinate (0..360 in the source dataset).
+            band_arrays: band arrays with shape (lat, lon), one per output band.
+        """
+        if lat.ndim != 1 or lon.ndim != 1:
+            raise ValueError("expected 1D latitude/longitude coordinates")
+        if len(band_arrays) == 0:
+            raise ValueError("expected at least one band array")
+
+        # Convert longitude to [-180, 180) and reorder so GeoTIFF coordinates match
+        # common WGS84 conventions and rslearn windows.
+        lon = ((lon + 180) % 360) - 180
+        lon_sort_idx = np.argsort(lon)
+        lon = lon[lon_sort_idx]
+        band_arrays = [a[:, lon_sort_idx] for a in band_arrays]
+
+        if len(lon) > 1:
+            dx = float(lon[1] - lon[0])
+        else:
+            dx = self.PIXEL_SIZE_DEGREES
+        if len(lat) > 1:
+            dy = float(abs(lat[1] - lat[0]))
+        else:
+            dy = self.PIXEL_SIZE_DEGREES
+
+        # ERA5-Land latitude is descending (north to south). This matches GeoTIFF row order.
+        if len(lat) > 1 and lat[1] > lat[0]:
+            raise ValueError("expected latitude coordinate to be descending")
+
+        west = float(lon.min() - dx / 2)
+        north = float(lat.max() + dy / 2)
+        transform = from_origin(west, north, dx, dy)
+        crs = f"EPSG:{WGS84_EPSG}"
+
+        array = np.stack(band_arrays, axis=0).astype(np.float32)
+        with rasterio.open(
+            tif_path,
+            "w",
+            driver="GTiff",
+            height=array.shape[1],
+            width=array.shape[2],
+            count=array.shape[0],
+            dtype=array.dtype,
+            crs=crs,
+            transform=transform,
+        ) as dst:
+            dst.write(array)
+
+    def ingest(
+        self,
+        tile_store: TileStoreWithLayer,
+        items: list[Item],
+        geometries: list[list[STGeometry]],
+    ) -> None:
+        """Ingest daily ERA5-Land rasters for the requested items/geometries.
+
+        For each item (one UTC day), this reads the corresponding slice from the
+        EarthDataHub Zarr store (subsetting by lat/lon for performance), then writes
+        a GeoTIFF into the dataset tile store.
+        """
+        ds = self._get_dataset()
+
+        for item, item_geoms in zip(items, geometries):
+            if tile_store.is_raster_ready(item.name, self.band_names):
+                continue
+
+            if item.geometry.time_range is None:
+                raise ValueError("expected item to have a time range")
+
+            day_start = _floor_to_utc_day(item.geometry.time_range[0])
+            day_str = f"{day_start.year:04d}-{day_start.month:02d}-{day_start.day:02d}"
+
+            bounds = _snap_bounds_outward(
+                self._get_effective_bounds(item_geoms),
+                step_degrees=self.PIXEL_SIZE_DEGREES,
+            )
+            lon_ranges_0_360 = _bounds_to_lon_ranges_0_360(bounds)
+            min_lat = bounds[1]
+            max_lat = bounds[3]
+
+            # Subset the dataset before computing, for performance.
+            # Latitude is descending in the dataset.
+            sel_kwargs_base: dict[str, Any] = dict(
+                valid_time=day_str,
+                latitude=slice(max_lat, min_lat),
+            )
+
+            band_arrays: list[np.ndarray] = []
+            lat: np.ndarray | None = None
+            lon: np.ndarray | None = None
+            for band in self.band_names:
+                if len(lon_ranges_0_360) == 1:
+                    da = ds[band].sel(
+                        **sel_kwargs_base,
+                        longitude=slice(lon_ranges_0_360[0][0], lon_ranges_0_360[0][1]),
+                    )
+                else:
+                    parts = [
+                        ds[band].sel(**sel_kwargs_base, longitude=slice(lo, hi))
+                        for (lo, hi) in lon_ranges_0_360
+                    ]
+                    da = xr.concat(parts, dim="longitude")
+
+                if band == "t2m" and self.temperature_unit == "celsius":
+                    da = da - 273.15
+
+                da = da.load()
+                if lat is None:
+                    lat = da["latitude"].to_numpy()
+                    lon = da["longitude"].to_numpy()
+                band_arrays.append(da.to_numpy())
+
+            assert lat is not None and lon is not None
+            if lat.size == 0 or lon.size == 0:
+                raise ValueError(
+                    f"ERA5LandDailyUTCv1 selection returned empty grid for item {item.name} "
+                    f"(bounds={bounds})"
+                )
+
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                local_tif_fname = os.path.join(tmp_dir, f"{item.name}.tif")
+                self._write_geotiff(local_tif_fname, lat, lon, band_arrays)
+                tile_store.write_raster_file(
+                    item.name, self.band_names, UPath(local_tif_fname)
+                )

rslearn/data_sources/eurocrops.py

@@ -5,7 +5,6 @@ import os
 import tempfile
 import zipfile
 from datetime import UTC, datetime, timedelta
-from typing import Any
 
 import fiona
 import requests
@@ -153,7 +152,7 @@ class EuroCrops(DataSource[EuroCropsItem]):
             groups.append(cur_groups)
         return groups
 
-    def deserialize_item(self, serialized_item: Any) -> EuroCropsItem:
+    def deserialize_item(self, serialized_item: dict) -> EuroCropsItem:
         """Deserializes an item from JSON-decoded data."""
         return EuroCropsItem.deserialize(serialized_item)
 

rslearn/data_sources/gcp_public_data.py

@@ -820,9 +820,8 @@ class Sentinel2(DataSource):
             groups.append(cur_groups)
         return groups
 
-    def deserialize_item(self, serialized_item: Any) -> Sentinel2Item:
+    def deserialize_item(self, serialized_item: dict) -> Sentinel2Item:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return Sentinel2Item.deserialize(serialized_item)
 
     def retrieve_item(

rslearn/data_sources/google_earth_engine.py

@@ -235,9 +235,8 @@ class GEE(DataSource, TileStore):
 
         return groups
 
-    def deserialize_item(self, serialized_item: Any) -> Item:
+    def deserialize_item(self, serialized_item: dict) -> Item:
         """Deserializes an item from JSON-decoded data."""
-        assert isinstance(serialized_item, dict)
         return Item.deserialize(serialized_item)
 
     def item_to_image(self, item: Item) -> ee.image.Image: