Add monthly embedding experiment

one_off_projects/2026_02_23_monthly_embedding/README.md

@@ -0,0 +1,59 @@
+This was to verify that we can use rslearn to compute per-timestep embeddings with
+OlmoEarth. So we input a calendar year of data and get the embedding for each month
+(concatenated with the 768 embedding dimensions on the channel axis). The modalities
+get mixed in there too so we still need to do some pooling in post-processing.
+
+Create the dataset:
+
+```
+export DATASET_PATH=./dataset
+mkdir $DATASET_PATH
+cp one_off_projects/2026_02_23_monthly_embedding/config.json $DATASET_PATH/config.json
+rslearn dataset add_windows --root $DATASET_PATH --group default --utm --resolution 10 --window_size 1024 --src_crs EPSG:4326 --box=-122.255,47.589,-122.255,47.589 --start 2025-01-01T00:00:00+00:00 --end 2026-01-01T00:00:00+00:00 --name seattle
+```
+
+Materialize the data:
+
+```
+rslearn dataset prepare --root $DATASET_PATH --retry-max-attempts 5 --retry-backoff-seconds 5
+rslearn dataset materialize --root $DATASET_PATH --retry-max-attempts 5 --retry-backoff-seconds 5
+```
+
+Apply the model and get embeddings:
+
+```
+rslearn model predict --config one_off_projects/2026_02_23_monthly_embedding/model.yaml --data.init_args.path=$DATASET_PATH
+```
+
+Use this script to pool within each timestep and get one GeoTIFF per timestep:
+
+```python
+import os
+from pathlib import Path
+
+import rasterio
+from einops import rearrange
+
+num_bandsets = 3
+num_timesteps = 12
+embedding_dim = 768
+
+ds_path = Path(os.environ["DATASET_PATH"])
+print("Read embeddings")
+in_fname = next((ds_path / "windows" / "default" / "seattle" / "layers" / "embeddings").glob("*/geotiff.tif"))
+with rasterio.open(in_fname) as raster:
+    array = raster.read()
+    profile = raster.profile.copy()
+
+print("Get per-timestep embeddings")
+per_timestep = rearrange(
+    array, "(t s c) h w -> t c h w s", t=num_timesteps, s=num_bandsets, c=embedding_dim
+).mean(axis=4)
+
+for timestep_idx, timestep_embeddings in enumerate(per_timestep):
+    out_fname = in_fname.parent / f"geotiff_{timestep_idx}.tif"
+    print(f"Write to {out_fname}")
+    profile.update(count=embedding_dim, dtype=timestep_embeddings.dtype)
+    with rasterio.open(out_fname, "w", **profile) as dst:
+        dst.write(timestep_embeddings)
+```

one_off_projects/2026_02_23_monthly_embedding/config.json

@@ -0,0 +1,40 @@
+{
+    "layers": {
+      "sentinel2_l2a": {
+        "band_sets": [{
+            "bands": ["B01", "B02", "B03", "B04", "B05", "B06", "B07", "B08", "B8A", "B09", "B11", "B12"],
+            "dtype": "uint16"
+        }],
+        "data_source": {
+          "class_path": "rslearn.data_sources.planetary_computer.Sentinel2",
+          "init_args": {
+            "cache_dir": "cache/planetary_computer",
+            "harmonize": true,
+            "sort_by": "eo:cloud_cover"
+          },
+          "ingest": false,
+          "query_config": {
+            "max_matches": 12,
+            "period_duration": "30d",
+            "space_mode": "PER_PERIOD_MOSAIC"
+          }
+        },
+        "type": "raster"
+      },
+      "embeddings": {
+        "band_sets": [{
+            "dtype": "float32",
+            "num_bands": 27648,
+            "format": {
+                "class_path": "rslearn.utils.raster_format.GeotiffRasterFormat",
+                "init_args": {
+                    "geotiff_options": {
+                        "compress": null
+                    }
+                }
+            }
+        }],
+        "type": "raster"
+      }
+    }
+  }

one_off_projects/2026_02_23_monthly_embedding/model.yaml

@@ -0,0 +1,56 @@
+model:
+  class_path: rslearn.train.lightning_module.RslearnLightningModule
+  init_args:
+    model:
+      class_path: rslearn.models.singletask.SingleTaskModel
+      init_args:
+        encoder:
+          - class_path: rslearn.models.olmoearth_pretrain.model.OlmoEarth
+            init_args:
+              model_id: OLMOEARTH_V1_BASE
+              patch_size: 4
+              # Output BCHWN instead of pooling over timesteps/modalities.
+              token_pooling: false
+        decoder:
+          # Embedding will reshape BCHWN -> B(CN)HW
+          - class_path: rslearn.train.tasks.embedding.EmbeddingHead
+    optimizer:
+      class_path: rslearn.train.optimizer.AdamW
+data:
+  class_path: rslearn.train.data_module.RslearnDataModule
+  init_args:
+    path: ${DATASET_PATH}
+    inputs:
+      sentinel2_l2a:
+        data_type: "raster"
+        layers: ["sentinel2_l2a"]
+        # This is the band order expected by OlmoEarth.
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12", "B01", "B09"]
+        passthrough: true
+        dtype: FLOAT32
+        load_all_layers: true
+        load_all_item_groups: true
+    task:
+      class_path: rslearn.train.tasks.embedding.EmbeddingTask
+    batch_size: 8
+    num_workers: 32
+    predict_config:
+      transforms:
+        - class_path: rslearn.models.olmoearth_pretrain.norm.OlmoEarthNormalize
+          init_args:
+            band_names:
+              sentinel2_l2a: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12", "B01", "B09"]
+      load_all_crops: true
+      crop_size: 64
+      overlap_pixels: 8
+trainer:
+  callbacks:
+    - class_path: rslearn.train.prediction_writer.RslearnWriter
+      init_args:
+        path: placeholder
+        output_layer: embeddings
+        merger:
+          class_path: rslearn.train.prediction_writer.RasterMerger
+          init_args:
+            overlap_pixels: 2
+            downsample_factor: 4