Support for global models

config/hourly-rollout-o96.yaml

@@ -0,0 +1,45 @@
+# yaml-language-server: $schema=../workflow/tools/config.schema.json
+description: |
+  Evaluate skill of COSMO-E emulator (M-1 forecaster).
+
+dates:
+  - 2022-02-03T00:00 # Some example
+
+runs:
+  - forecaster:
+      mlflow_id: 2bee8f6b7d0048089a5e25cb60def0e5
+      label: Hourly rollout O96 forecaster
+      steps: 0/120/6
+      config: resources/inference/configs/global-forecaster.yaml
+      disable_local_eccodes_definitions: true
+      extra_dependencies:
+        - git+https://github.com/ecmwf/anemoi-inference.git@main
+        - git+https://github.com/MeteoSwiss/anemoi-core.git@06e5533f3e8da37c44d887c42b67440b40286cb3#subdirectory=models
+
+baselines: []
+
+analysis:
+  label: ERA5
+  analysis_zarr: /scratch/mch/rradev/datasets/aifs-ea-an-oper-0001-mars-o96-1979-2022-1h-v4.zarr
+
+stratification:
+  regions: []
+  root: /scratch/mch/bhendj/regions/Prognoseregionen_LV95_20220517
+
+locations:
+  output_root: output/
+  mlflow_uri:
+    - https://servicedepl.meteoswiss.ch/mlstore
+    - https://mlflow.ecmwf.int
+
+profile:
+  executor: slurm
+  global_resources:
+    gpus: 16
+  default_resources:
+    slurm_partition: "postproc"
+    cpus_per_task: 1
+    mem_mb_per_cpu: 1800
+    runtime: "1h"
+    gpus: 0
+  jobs: 50

resources/inference/configs/global-forecaster.yaml

@@ -0,0 +1,12 @@
+input:
+  test:
+    use_original_paths: true
+
+allow_nans: true
+
+output:
+  grib:
+    path: grib/{date}{time}_{step:03}.grib
+    negative_step_mode: skip
+
+write_initial_state: true

src/data_input/__init__.py

@@ -1,21 +1,30 @@
+import yaml
 import logging
 import os
 import sys
 from datetime import datetime
 from pathlib import Path
 from typing import Iterable
+from functools import lru_cache
 
 eccodes_definition_path = Path(sys.prefix) / "share/eccodes-cosmo-resources/definitions"
 os.environ["ECCODES_DEFINITION_PATH"] = str(eccodes_definition_path)
 
-from meteodatalab import data_source, grib_decoder  # noqa: E402
-
 import numpy as np  # noqa: E402
 import xarray as xr  # noqa: E402
+import earthkit.data as ekd  # noqa: E402
 
 LOG = logging.getLogger(__name__)
 
 
+@lru_cache(maxsize=1)
+def earthkit_xarray_engine_profile() -> dict:
+    fn = Path(__file__).parent / "profile.yaml"
+    with open(fn) as f:
+        profile = yaml.safe_load(f)
+    return profile
+
+
 def load_analysis_data_from_zarr(
     analysis_zarr: Path, times: Iterable[datetime], params: list[str]
 ) -> xr.Dataset:
@@ -91,9 +100,16 @@ def load_fct_data_from_grib(
 ) -> xr.Dataset:
     """Load forecast data from GRIB files for a specific valid time."""
     files = sorted(grib_output_dir.glob("20*.grib"))
-    fds = data_source.FileDataSource(datafiles=files)
-    ds = grib_decoder.load(fds, {"param": params, "step": steps})
-    for var, da in ds.items():
+
+    profile = earthkit_xarray_engine_profile()
+    ds: xr.Dataset = (
+        ekd.from_source("file", files)
+        .sel(param=params, step=steps)
+        .to_xarray(profile=profile)
+    )
+    # fds = data_source.FileDataSource(datafiles=files)
+    # ds = grib_decoder.load(fds, {"param": params, "step": steps})
+    for var, da in ds.data_vars.items():
         if "z" in da.dims and da.sizes["z"] == 1:
             ds[var] = da.squeeze("z", drop=True)
         elif "z" in da.dims and da.sizes["z"] > 1:

src/data_input/profile.yaml

@@ -0,0 +1,20 @@
+ensure_dims: [z,number,step, forecast_reference_time]
+rename_dims: {
+   level: z,
+   number: eps,
+   step: lead_time,
+   forecast_reference_time: ref_time
+}
+variable_attrs:
+- cfName
+- name
+- units
+- typeOfLevel
+- levtype
+- paramId
+
+global_attrs:
+- Conventions: CF-1.8
+- institution: MeteoSwiss
+
+add_valid_time_coord: true

src/evalml/config.py

@@ -90,6 +90,11 @@ class RunConfig(BaseModel):
         description="Resource requirements for inference jobs (optional; defaults handled externally).",
     )
 
+    disable_local_eccodes_definitions: bool = Field(
+        False,
+        description="If true, the ECCODES_DEFINITION_PATH environment variable will not be set to the COSMO local definitions.",
+    )
+
     config: Dict[str, Any] | str
 
     @field_validator("steps")

workflow/rules/inference.smk

@@ -232,6 +232,9 @@ rule execute_inference:
         workdir=lambda wc: (
             OUT_ROOT / f"data/runs/{wc.run_id}/{wc.init_time}"
         ).resolve(),
+        disable_local_definitions=lambda wc: RUN_CONFIGS[wc.run_id].get(
+            "disable_local_eccodes_definitions", False
+        ),
     resources:
         slurm_partition=lambda wc: get_resource(wc, "slurm_partition", "short-shared"),
         cpus_per_task=lambda wc: get_resource(wc, "cpus_per_task", 24),
@@ -249,7 +252,10 @@ rule execute_inference:
 
         squashfs-mount {params.image_path}:/user-environment -- bash -c '
         source /user-environment/bin/activate
-        export ECCODES_DEFINITION_PATH=/user-environment/share/eccodes-cosmo-resources/definitions
+
+        if [ "{params.disable_local_definitions}" = "False" ]; then
+            export ECCODES_DEFINITION_PATH=/user-environment/share/eccodes-cosmo-resources/definitions
+        fi
 
         CMD_ARGS=()
 

workflow/tools/config.schema.json

@@ -191,6 +191,12 @@
           "default": null,
           "description": "Resource requirements for inference jobs (optional; defaults handled externally)."
         },
+        "disable_local_eccodes_definitions": {
+          "default": false,
+          "description": "If true, the ECCODES_DEFINITION_PATH environment variable will not be set to the COSMO local definitions.",
+          "title": "Disable Local Eccodes Definitions",
+          "type": "boolean"
+        },
         "config": {
           "anyOf": [
             {
@@ -372,6 +378,12 @@
           "default": null,
           "description": "Resource requirements for inference jobs (optional; defaults handled externally)."
         },
+        "disable_local_eccodes_definitions": {
+          "default": false,
+          "description": "If true, the ECCODES_DEFINITION_PATH environment variable will not be set to the COSMO local definitions.",
+          "title": "Disable Local Eccodes Definitions",
+          "type": "boolean"
+        },
         "config": {
           "anyOf": [
             {