Merge pull request #82 from juaAI/feat/models-aifs-ens

feat(models): add hosted aifs_ens (additive)

Author: niasie

src/jua/weather/_model_meta.py

@@ -234,6 +234,14 @@ class ModelMetaInfo:
         special=((1, 0, 78),),
     ),
 )
+_MODEL_META_INFO[Models.AIFS_ENS] = ModelMetaInfo(
+    has_grid_access=True,
+    full_forecasted_hours=360,
+    has_statistics=True,
+    # 0.25 degree grid (1440x720 after south-pole drop) — matches the default
+    # num_lats=720 / num_lons=1440, so no override needed.
+    temporal_resolution=TemporalResolution(base=6),
+)
 _MODEL_META_INFO[Models.ECMWF_AIFS_ENSEMBLE] = ModelMetaInfo(
     forecast_name_mapping="ecmwf_aifs025_ensemble",
     full_forecasted_hours=360,

src/jua/weather/models.py

@@ -30,6 +30,7 @@ class Models(str, Enum):
     EPT2_HELIOS = "ept2_1_helios"
     EPT2_EUROPA = "ept2_1_europa"
     AIFS = "aifs"
+    AIFS_ENS = "aifs_ens"
     AURORA = "aurora"
     ECMWF_IFS_SINGLE = "ecmwf_ifs_single"
     ICON_EU = "icon_eu"

tests/functional/test_forecasts.py

@@ -36,6 +36,10 @@
 MODEL_SPECIFIC_FORECAST_DATES = {
     Models.EPT2_REASONING: datetime(2025, 11, 23, 0, 0, 0),
     Models.ICON_EU: datetime(2026, 2, 9, 0, 0, 0),
+    # AIFS ENS (ECMWF Open Data) backfill starts mid-2025 and is sparse at the
+    # edges; the default 2025-10-20 isn't present. Use a verified mid-history
+    # init time.
+    Models.AIFS_ENS: datetime(2026, 3, 3, 0, 0, 0),
 }
 
 SOLAR_ONLY_MODELS = {Models.EPT2_HELIOS}

tests/weather/test_aifs_ens_meta.py

@@ -0,0 +1,23 @@
+"""Regression tests for the hosted AIFS ENS model metadata.
+
+The hosted ClickHouse ``aifs_ens`` model is added alongside the existing
+open-meteo ``ecmwf_aifs025_ensemble`` (additive / expand phase). Unlike the
+point-only open-meteo model, the hosted model exposes full grid access plus
+ensemble statistics, so it must not regress into the "point queries only"
+behaviour gated by ``has_grid_access``.
+"""
+
+from jua.weather._model_meta import get_model_meta_info
+from jua.weather.models import Models
+
+
+def test_aifs_ens_is_fully_accessible() -> None:
+    meta = get_model_meta_info(Models.AIFS_ENS)
+    # Hosted ClickHouse grid model: grid/bbox slices must be allowed, not just points.
+    assert meta.has_grid_access is True
+    # 51-member ensemble: statistics available.
+    assert meta.has_statistics is True
+    # 0.25 degree global grid (1440x720 after south-pole drop).
+    assert (meta.num_lats, meta.num_lons) == (720, 1440)
+    # 0-360h forecast horizon.
+    assert meta.full_forecasted_hours == 360