[Feature]: Add BetaEarth to rs-embed

[Dinghye]

Problem

Good work from @mikonvergence !
https://github.com/asterisk-labs/beta-earth

Proposed solution

Example for betaearth

from __future__ import annotations

from typing import Any

import numpy as np

from ..core.embedding import Embedding
from ..core.errors import ModelError
from ..core.registry import register
from ..core.specs import OutputSpec, SensorSpec, SpatialSpec, TemporalSpec
from ..core.types import FetchResult
from ..providers import ProviderBase
from .base import EmbedderBase
from .meta_utils import build_meta


# ---- fixed model contract (do not expose via model_config) -----------------
_S2_COLLECTION = "COPERNICUS/S2_SR_HARMONIZED"
_S2_BANDS = ("B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B11", "B12")
_S1_COLLECTION = "COPERNICUS/S1_GRD_FLOAT"
_S1_BANDS = ("VV", "VH")
_DEM_COLLECTION = "COPERNICUS/DEM/GLO30"
_DEM_BAND = "DEM"

# ---- defaults (override via model_config) ----------------------------------
_S2_FRAMES_DEFAULT = 6
_S1_FRAMES_DEFAULT = 6


def _cfg(model_config: dict[str, Any] | None, key: str, default: Any) -> Any:
    return (model_config or {}).get(key, default)


def _fetch_dem_chw(provider: ProviderBase, *, spatial: SpatialSpec) -> np.ndarray:
    region = provider.get_region(spatial)
    sensor = SensorSpec(
        collection=_DEM_COLLECTION, bands=(_DEM_BAND,), scale_m=30, composite="mosaic"
    )
    img = provider.build_image(sensor=sensor, temporal=None, region=region)
    return provider.fetch_array_chw(
        image=img, bands=(_DEM_BAND,), region=region,
        scale_m=30, fill_value=0.0, collection=_DEM_COLLECTION,
    )


@register("betaearth")
class BetaEarthEmbedder(EmbedderBase):
    """ROI -> (S2 sequence + S1 sequence + COP-DEM GLO30) -> BetaEarth -> embedding."""

    def describe(self) -> dict[str, Any]:
        return {
            "type": "on_the_fly",
            "backend": ["provider"],
            "inputs": {
                "s2":  {"collection": _S2_COLLECTION,  "bands": list(_S2_BANDS)},
                "s1":  {"collection": _S1_COLLECTION,  "bands": list(_S1_BANDS)},
                "dem": {"collection": _DEM_COLLECTION, "bands": [_DEM_BAND]},
            },
            "temporal": {"mode": "range"},
            "output": ["pooled", "grid"],
            "defaults": {
                "s2_frames": _S2_FRAMES_DEFAULT,
                "s1_frames": _S1_FRAMES_DEFAULT,
                "scale_m": 10,
                "dem_scale_m": 30,
            },
            "model_config": {
                "s2_frames": {"type": "int", "default": _S2_FRAMES_DEFAULT},
                "s1_frames": {"type": "int", "default": _S1_FRAMES_DEFAULT},
            },
        }

    def fetch_input(
        self,
        provider: ProviderBase,
        *,
        spatial: SpatialSpec,
        temporal: TemporalSpec | None,
        sensor: SensorSpec | None,
        model_config: dict[str, Any] | None = None,
    ) -> FetchResult:
        if temporal is None:
            raise ModelError("betaearth requires TemporalSpec.range(start, end).")
        temporal.validate()

        s2_frames = int(_cfg(model_config, "s2_frames", _S2_FRAMES_DEFAULT))
        s1_frames = int(_cfg(model_config, "s1_frames", _S1_FRAMES_DEFAULT))

        s2_tchw = provider.fetch_multiframe_collection_raw_tchw(
            spatial=spatial, temporal=temporal,
            collection=_S2_COLLECTION, bands=_S2_BANDS, n_frames=s2_frames,
            scale_m=10, cloudy_pct=30, composite="median",
        )
        s1_tchw = provider.fetch_multiframe_collection_raw_tchw(
            spatial=spatial, temporal=temporal,
            collection=_S1_COLLECTION, bands=_S1_BANDS, n_frames=s1_frames,
            scale_m=10, cloudy_pct=None, composite="median",
        )
        dem_chw = _fetch_dem_chw(provider, spatial=spatial)

        bundle = {"s2": s2_tchw, "s1": s1_tchw, "dem": dem_chw}
        meta = {"s2": {"n_frames": s2_frames}, "s1": {"n_frames": s1_frames}}
        return FetchResult(data=bundle, meta=meta)

    def get_embedding(
        self,
        *,
        spatial: SpatialSpec,
        temporal: TemporalSpec | None,
        sensor: SensorSpec | None,
        output: OutputSpec,
        backend: str,
        device: str = "auto",
        input_chw: np.ndarray | dict[str, np.ndarray] | None = None,
        model_config: dict[str, Any] | None = None,
    ) -> Embedding:
        if input_chw is None:
            provider = self._get_provider(backend)
            result = self.fetch_input(
                provider, spatial=spatial, temporal=temporal,
                sensor=sensor, model_config=model_config,
            )
            bundle, fetch_meta = result.data, result.meta
        elif isinstance(input_chw, dict):
            bundle, fetch_meta = input_chw, {}
        else:
            raise ModelError("betaearth expects input_chw as dict[str, np.ndarray] (s2/s1/dem).")

        s2_tchw = bundle["s2"]   # [T, 10, H, W]
        s1_tchw = bundle["s1"]   # [T, 2,  H, W]
        dem_chw = bundle["dem"]  # [1, H, W]

        # Normalize: S2 /10000 clip; S1 log1p + 99-pct; DEM as BetaEarth expects.
        s2_tchw = np.clip(s2_tchw.astype(np.float32) / 10000.0, 0.0, 1.0)
        s1_log = np.log1p(np.maximum(s1_tchw.astype(np.float32), 0.0))
        denom = float(np.percentile(s1_log, 99)) or 1.0
        s1_tchw = np.clip(s1_log / denom, 0.0, 1.0).astype(np.float32)
        dem_chw = dem_chw.astype(np.float32)  # TODO: standardize as BetaEarth expects

        # TODO: replace with real BetaEarth forward pass
        vec = np.zeros(512, dtype=np.float32)

        meta = build_meta(
            model=self.model_name, kind="on_the_fly", backend=backend,
            source="multi[s2,s1,dem]",
            extra={"fetch": fetch_meta},
        )
        return Embedding(data=vec, meta=meta)

Example or context

emb = get_embedding(
    "betaearth",
    spatial=PointBuffer(121.5, 31.2, 2048),
    temporal=TemporalSpec.range("2022-06-01", "2022-09-01"),
    output=OutputSpec.pooled(),
    model_config={"s2_frames": 8, "s1_frames": 4},
)

Notes

• fetch_input returns FetchResult.data as a dict[str, np.ndarray] (s2, s1, dem) instead of a single CHW array — only get_embedding needs to know the shape, so this stays local to the model.
• model_config is plumbed end-to-end through get_embedding / get_embeddings_batch, so users can tune s2_frames / s1_frames without touching specs.
• export_batch works out of the box for the common case. If you need to persist the dict bundle to NPZ, override get_embeddings_batch_from_inputs and the export pipeline will call it.

⚠️ Check the GEE ↔ Planetary Computer correspondence before assuming the inputs match what BetaEarth was trained on.

Checks

• I searched existing issues and pull requests before opening this request.