rslearn 0.0.20__tar.gz → 0.0.21__tar.gz

{rslearn-0.0.20/rslearn.egg-info → rslearn-0.0.21}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.20
+Version: 0.0.21
 Summary: A library for developing remote sensing datasets and models
 Author: OlmoEarth Team
 License:                                  Apache License

{rslearn-0.0.20 → rslearn-0.0.21}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rslearn"
-version = "0.0.20"
+version = "0.0.21"
 description = "A library for developing remote sensing datasets and models"
 authors = [
     { name = "OlmoEarth Team" },

{rslearn-0.0.20 → rslearn-0.0.21}/rslearn/data_sources/climate_data_store.py

@@ -24,51 +24,55 @@ from rslearn.utils.geometry import STGeometry
 logger = get_logger(__name__)
 
 
-class ERA5LandMonthlyMeans(DataSource):
-    """A data source for ingesting ERA5 land monthly averaged data from the Copernicus Climate Data Store.
+class ERA5Land(DataSource):
+    """Base class for ingesting ERA5 land data from the Copernicus Climate Data Store.
 
     An API key must be passed either in the configuration or via the CDSAPI_KEY
     environment variable. You can acquire an API key by going to the Climate Data Store
     website (https://cds.climate.copernicus.eu/), registering an account and logging
-    in, and then
+    in, and then getting the API key from the user profile page.
 
     The band names should match CDS variable names (see the reference at
     https://confluence.ecmwf.int/display/CKB/ERA5-Land%3A+data+documentation). However,
     replace "_" with "-" in the variable names when specifying bands in the layer
     configuration.
 
-    This data source corresponds to the reanalysis-era5-land-monthly-means product.
-
-    All requests to the API will be for the whole globe. Although the API supports arbitrary
-    bounds in the requests, using the whole available area helps to reduce the total number of
-    requests.
+    By default, all requests to the API will be for the whole globe. To speed up ingestion,
+    we recommend specifying the bounds of the area of interest, in particular for hourly data.
     """
 
     api_url = "https://cds.climate.copernicus.eu/api"
-
-    # see: https://cds.climate.copernicus.eu/cdsapp#!/dataset/reanalysis-era5-land-monthly-means
-    DATASET = "reanalysis-era5-land-monthly-means"
-    PRODUCT_TYPE = "monthly_averaged_reanalysis"
     DATA_FORMAT = "netcdf"
     DOWNLOAD_FORMAT = "unarchived"
     PIXEL_SIZE = 0.1  # degrees, native resolution is 9km
 
     def __init__(
         self,
+        dataset: str,
+        product_type: str,
         band_names: list[str] | None = None,
         api_key: str | None = None,
+        bounds: list[float] | None = None,
         context: DataSourceContext = DataSourceContext(),
     ):
-        """Initialize a new ERA5LandMonthlyMeans instance.
+        """Initialize a new ERA5Land instance.
 
         Args:
+            dataset: the CDS dataset name (e.g., "reanalysis-era5-land-monthly-means").
+            product_type: the CDS product type (e.g., "monthly_averaged_reanalysis").
             band_names: list of band names to acquire. These should correspond to CDS
                 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.
+            bounds: optional bounding box as [min_lon, min_lat, max_lon, max_lat].
+                If not specified, the whole globe will be used.
             context: the data source context.
         """
+        self.dataset = dataset
+        self.product_type = product_type
+        self.bounds = bounds
+
         self.band_names: list[str]
         if context.layer_config is not None:
             self.band_names = []
@@ -134,8 +138,11 @@ class ERA5LandMonthlyMeans(DataSource):
                 # Collect Item list corresponding to the current month.
                 items = []
                 item_name = f"era5land_monthlyaveraged_{cur_date.year}_{cur_date.month}"
-                # Space is the whole globe.
-                bounds = (-180, -90, 180, 90)
+                # Use bounds if set, otherwise use whole globe
+                if self.bounds is not None:
+                    bounds = self.bounds
+                else:
+                    bounds = [-180, -90, 180, 90]
                 # Time is just the given month.
                 start_date = datetime(cur_date.year, cur_date.month, 1, tzinfo=UTC)
                 time_range = (
@@ -172,7 +179,9 @@ class ERA5LandMonthlyMeans(DataSource):
         # But the list of variables should include the bands we want in the correct
         # order. And we can distinguish those bands from other "variables" because they
         # will be 3D while the others will be scalars or 1D.
-        bands_data = []
+
+        band_arrays = []
+        num_time_steps = None
         for band_name in nc.variables:
             band_data = nc.variables[band_name]
             if len(band_data.shape) != 3:
@@ -182,18 +191,27 @@ class ERA5LandMonthlyMeans(DataSource):
             logger.debug(
                 f"NC file {nc_path} has variable {band_name} with shape {band_data.shape}"
             )
-            # Variable data is stored in a 3D array (1, height, width)
-            if band_data.shape[0] != 1:
+            # Variable data is stored in a 3D array (time, height, width)
+            # For hourly data, time is number of days in the month x 24 hours
+            if num_time_steps is None:
+                num_time_steps = band_data.shape[0]
+            elif band_data.shape[0] != num_time_steps:
                 raise ValueError(
-                    f"Bad shape for band {band_name}, expected 1 band but got {band_data.shape[0]}"
+                    f"Variable {band_name} has {band_data.shape[0]} time steps, "
+                    f"but expected {num_time_steps}"
                 )
-            bands_data.append(band_data[0, :, :])
+            # Original shape: (time, height, width)
+            band_array = np.array(band_data[:])
+            band_array = np.expand_dims(band_array, axis=1)
+            band_arrays.append(band_array)
 
-        array = np.array(bands_data)  # (num_bands, height, width)
-        if array.shape[0] != len(self.band_names):
-            raise ValueError(
-                f"Expected to get {len(self.band_names)} bands but got {array.shape[0]}"
-            )
+        # After concatenation: (time, num_variables, height, width)
+        stacked_array = np.concatenate(band_arrays, axis=1)
+
+        # After reshaping: (time x num_variables, height, width)
+        array = stacked_array.reshape(
+            -1, stacked_array.shape[2], stacked_array.shape[3]
+        )
 
         # Get metadata for the GeoTIFF
         lat = nc.variables["latitude"][:]
@@ -235,6 +253,58 @@ class ERA5LandMonthlyMeans(DataSource):
         ) as dst:
             dst.write(array)
 
+    def ingest(
+        self,
+        tile_store: TileStoreWithLayer,
+        items: list[Item],
+        geometries: list[list[STGeometry]],
+    ) -> None:
+        """Ingest items into the given tile store.
+
+        This method should be overridden by subclasses.
+
+        Args:
+            tile_store: the tile store to ingest into
+            items: the items to ingest
+            geometries: a list of geometries needed for each item
+        """
+        raise NotImplementedError("Subclasses must implement ingest method")
+
+
+class ERA5LandMonthlyMeans(ERA5Land):
+    """A data source for ingesting ERA5 land monthly averaged data from the Copernicus Climate Data Store.
+
+    This data source corresponds to the reanalysis-era5-land-monthly-means product.
+    """
+
+    def __init__(
+        self,
+        band_names: list[str] | None = None,
+        api_key: str | None = None,
+        bounds: list[float] | 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 "-". 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.
+            bounds: optional bounding box as [min_lon, min_lat, max_lon, max_lat].
+                If not specified, the whole globe will be used.
+            context: the data source context.
+        """
+        super().__init__(
+            dataset="reanalysis-era5-land-monthly-means",
+            product_type="monthly_averaged_reanalysis",
+            band_names=band_names,
+            api_key=api_key,
+            bounds=bounds,
+            context=context,
+        )
+
     def ingest(
         self,
         tile_store: TileStoreWithLayer,
@@ -256,25 +326,142 @@ class ERA5LandMonthlyMeans(DataSource):
                 continue
 
             # Send the request to the CDS API
-            # If area is not specified, the whole globe will be requested
+            if self.bounds is not None:
+                min_lon, min_lat, max_lon, max_lat = self.bounds
+                area = [max_lat, min_lon, min_lat, max_lon]
+            else:
+                area = [90, -180, -90, 180]  # Whole globe
+
             request = {
-                "product_type": [self.PRODUCT_TYPE],
+                "product_type": [self.product_type],
                 "variable": variable_names,
                 "year": [f"{item.geometry.time_range[0].year}"],  # type: ignore
                 "month": [
                     f"{item.geometry.time_range[0].month:02d}"  # type: ignore
                 ],
                 "time": ["00:00"],
+                "area": area,
+                "data_format": self.DATA_FORMAT,
+                "download_format": self.DOWNLOAD_FORMAT,
+            }
+            logger.debug(
+                f"CDS API request for year={request['year']} month={request['month']} area={area}"
+            )
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                local_nc_fname = os.path.join(tmp_dir, f"{item.name}.nc")
+                local_tif_fname = os.path.join(tmp_dir, f"{item.name}.tif")
+                self.client.retrieve(self.dataset, request, local_nc_fname)
+                self._convert_nc_to_tif(
+                    UPath(local_nc_fname),
+                    UPath(local_tif_fname),
+                )
+                tile_store.write_raster_file(
+                    item.name, self.band_names, UPath(local_tif_fname)
+                )
+
+
+class ERA5LandHourly(ERA5Land):
+    """A data source for ingesting ERA5 land hourly data from the Copernicus Climate Data Store.
+
+    This data source corresponds to the reanalysis-era5-land product.
+    """
+
+    def __init__(
+        self,
+        band_names: list[str] | None = None,
+        api_key: str | None = None,
+        bounds: list[float] | None = None,
+        context: DataSourceContext = DataSourceContext(),
+    ):
+        """Initialize a new ERA5LandHourly instance.
+
+        Args:
+            band_names: list of band names to acquire. These should correspond to CDS
+                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.
+            bounds: optional bounding box as [min_lon, min_lat, max_lon, max_lat].
+                If not specified, the whole globe will be used.
+            context: the data source context.
+        """
+        super().__init__(
+            dataset="reanalysis-era5-land",
+            product_type="reanalysis",
+            band_names=band_names,
+            api_key=api_key,
+            bounds=bounds,
+            context=context,
+        )
+
+    def ingest(
+        self,
+        tile_store: TileStoreWithLayer,
+        items: list[Item],
+        geometries: list[list[STGeometry]],
+    ) -> None:
+        """Ingest items into the given tile store.
+
+        Args:
+            tile_store: the tile store to ingest into
+            items: the items to ingest
+            geometries: a list of geometries needed for each item
+        """
+        # for CDS variable names, replace "-" with "_"
+        variable_names = [band.replace("-", "_") for band in self.band_names]
+
+        for item in items:
+            if tile_store.is_raster_ready(item.name, self.band_names):
+                continue
+
+            # Send the request to the CDS API
+            # If area is not specified, the whole globe will be requested
+            time_range = item.geometry.time_range
+            if time_range is None:
+                raise ValueError("Item must have a time range")
+
+            # For hourly data, request all days in the month and all 24 hours
+            start_time = time_range[0]
+
+            # Get all days in the month
+            year = start_time.year
+            month = start_time.month
+            # Get the last day of the month
+            if month == 12:
+                last_day = 31
+            else:
+                next_month = datetime(year, month + 1, 1, tzinfo=UTC)
+                last_day = (next_month - relativedelta(days=1)).day
+
+            days = [f"{day:02d}" for day in range(1, last_day + 1)]
+
+            # Get all 24 hours
+            hours = [f"{hour:02d}:00" for hour in range(24)]
+
+            if self.bounds is not None:
+                min_lon, min_lat, max_lon, max_lat = self.bounds
+                area = [max_lat, min_lon, min_lat, max_lon]
+            else:
+                area = [90, -180, -90, 180]  # Whole globe
+
+            request = {
+                "product_type": [self.product_type],
+                "variable": variable_names,
+                "year": [f"{year}"],
+                "month": [f"{month:02d}"],
+                "day": days,
+                "time": hours,
+                "area": area,
                 "data_format": self.DATA_FORMAT,
                 "download_format": self.DOWNLOAD_FORMAT,
             }
             logger.debug(
-                f"CDS API request for the whole globe for year={request['year']} month={request['month']}"
+                f"CDS API request for year={request['year']} month={request['month']} days={len(days)} hours={len(hours)} area={area}"
             )
             with tempfile.TemporaryDirectory() as tmp_dir:
                 local_nc_fname = os.path.join(tmp_dir, f"{item.name}.nc")
                 local_tif_fname = os.path.join(tmp_dir, f"{item.name}.tif")
-                self.client.retrieve(self.DATASET, request, local_nc_fname)
+                self.client.retrieve(self.dataset, request, local_nc_fname)
                 self._convert_nc_to_tif(
                     UPath(local_nc_fname),
                     UPath(local_tif_fname),

{rslearn-0.0.20 → rslearn-0.0.21}/rslearn/models/olmoearth_pretrain/model.py

@@ -159,20 +159,19 @@ class OlmoEarth(FeatureExtractor):
         that contains the distributed checkpoint. This is the format produced by
         pre-training runs in olmoearth_pretrain.
         """
-        # We avoid loading the train module here because it depends on running within
-        # olmo_core.
-        # Only pull in olmo_core when trying to load a distributed checkpoint to avoid dependency.
-        require_olmo_core("_load_model_from_checkpoint")
-        from olmo_core.distributed.checkpoint import load_model_and_optim_state
-
         with (checkpoint_upath / "config.json").open() as f:
             config_dict = json.load(f)
             model_config = Config.from_dict(config_dict["model"])
 
         model = model_config.build()
 
-        # Load the checkpoint.
+        # Load the checkpoint (requires olmo_core for distributed checkpoint loading).
         if not random_initialization:
+            require_olmo_core(
+                "_load_model_from_checkpoint with random_initialization=False"
+            )
+            from olmo_core.distributed.checkpoint import load_model_and_optim_state
+
             train_module_dir = checkpoint_upath / "model_and_optim"
             load_model_and_optim_state(str(train_module_dir), model)
             logger.info(f"loaded OlmoEarth encoder from {train_module_dir}")
@@ -242,7 +241,7 @@ class OlmoEarth(FeatureExtractor):
             present_modalities.append(modality)
             tensors = []
             for idx, inp in enumerate(context.inputs):
-                assert isinstance(inp, RasterImage)
+                assert isinstance(inp[modality], RasterImage)
                 tensors.append(inp[modality].image)
                 cur_timestamps = inp[modality].timestamps
                 if cur_timestamps is not None and len(cur_timestamps) > len(

{rslearn-0.0.20 → rslearn-0.0.21}/rslearn/train/dataset.py

@@ -205,8 +205,7 @@ def read_raster_layer_for_data_input(
     group_idx: int,
     layer_config: LayerConfig,
     data_input: DataInput,
-    layer_data: WindowLayerData | None,
-) -> tuple[torch.Tensor, tuple[datetime, datetime] | None]:
+) -> torch.Tensor:
     """Read a raster layer for a DataInput.
 
     This scans the available rasters for the layer at the window to determine which
@@ -219,11 +218,9 @@ def read_raster_layer_for_data_input(
         group_idx: the item group.
         layer_config: the layer configuration.
         data_input: the DataInput that specifies the bands and dtype.
-        layer_data: the WindowLayerData associated with this layer and window.
 
     Returns:
-        RasterImage containing raster data and the timestamp associated
-            with that data.
+        Raster data as a tensor.
     """
     # See what different sets of bands we need to read to get all the
     # configured bands.
@@ -294,34 +291,46 @@ def read_raster_layer_for_data_input(
             src[src_indexes, :, :].astype(data_input.dtype.get_numpy_dtype())
         )
 
-    # add the timestamp. this is a tuple defining the start and end of the time range.
-    time_range = None
-    if layer_data is not None:
-        item = Item.deserialize(layer_data.serialized_item_groups[group_idx][0])
-        if item.geometry.time_range is not None:
-            # we assume if one layer data has a geometry & time range, all of them do
-            time_ranges = [
-                (
-                    datetime.fromisoformat(
-                        Item.deserialize(
-                            layer_data.serialized_item_groups[group_idx][idx]
-                        ).geometry.time_range[0]  # type: ignore
-                    ),
-                    datetime.fromisoformat(
-                        Item.deserialize(
-                            layer_data.serialized_item_groups[group_idx][idx]
-                        ).geometry.time_range[1]  # type: ignore
-                    ),
-                )
-                for idx in range(len(layer_data.serialized_item_groups[group_idx]))
-            ]
-            # take the min and max
-            time_range = (
-                min([t[0] for t in time_ranges]),
-                max([t[1] for t in time_ranges]),
+    return image
+
+
+def read_layer_time_range(
+    layer_data: WindowLayerData | None, group_idx: int
+) -> tuple[datetime, datetime] | None:
+    """Extract the combined time range from all items in a layer data group.
+
+    Returns the min start time and max end time across all items, or None if
+    no items have time ranges.
+
+    Raises:
+        ValueError: If some items have time_range and others don't.
+    """
+    if layer_data is None:
+        return None
+
+    serialized_items = layer_data.serialized_item_groups[group_idx]
+    if not serialized_items:
+        return None
+
+    first_item = Item.deserialize(serialized_items[0])
+    if first_item.geometry.time_range is None:
+        return None
+
+    # If the first item has a time_range, all items must have one
+    time_ranges: list[tuple[datetime, datetime]] = []
+    for serialized_item in serialized_items:
+        item = Item.deserialize(serialized_item)
+        if item.geometry.time_range is None:
+            raise ValueError(
+                f"Item '{item.name}' has no time_range, but first item does. "
+                "All items in a group must consistently have or lack time_range."
             )
+        time_ranges.append(item.geometry.time_range)
 
-    return image, time_range
+    return (
+        min(tr[0] for tr in time_ranges),
+        max(tr[1] for tr in time_ranges),
+    )
 
 
 def read_data_input(
@@ -378,17 +387,17 @@ def read_data_input(
         time_ranges: list[tuple[datetime, datetime] | None] = []
         for layer_name, group_idx in layers_to_read:
             layer_config = dataset.layers[layer_name]
-            image, time_range = read_raster_layer_for_data_input(
+            image = read_raster_layer_for_data_input(
                 window,
                 bounds,
                 layer_name,
                 group_idx,
                 layer_config,
                 data_input,
-                # some layers (e.g. "label_raster") won't have associated
-                # layer datas
-                layer_datas[layer_name] if layer_name in layer_datas else None,
             )
+            # some layers (e.g. "label_raster") won't have associated layer datas
+            layer_data = layer_datas.get(layer_name)
+            time_range = read_layer_time_range(layer_data, group_idx)
             if len(time_ranges) > 0:
                 if type(time_ranges[-1]) is not type(time_range):
                     raise ValueError(

{rslearn-0.0.20 → rslearn-0.0.21/rslearn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.20
+Version: 0.0.21
 Summary: A library for developing remote sensing datasets and models
 Author: OlmoEarth Team
 License:                                  Apache License