diff --git a/config/hourly-rollout-o96.yaml b/config/hourly-rollout-o96.yaml new file mode 100644 index 0000000..02f51f6 --- /dev/null +++ b/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 diff --git a/resources/inference/configs/global-forecaster.yaml b/resources/inference/configs/global-forecaster.yaml new file mode 100644 index 0000000..4c8d186 --- /dev/null +++ b/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 diff --git a/src/data_input/__init__.py b/src/data_input/__init__.py index 4ff4126..63c6459 100644 --- a/src/data_input/__init__.py +++ b/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: diff --git a/src/data_input/profile.yaml b/src/data_input/profile.yaml new file mode 100644 index 0000000..56d7517 --- /dev/null +++ b/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 diff --git a/src/evalml/config.py b/src/evalml/config.py index 8faa614..b347e8d 100644 --- a/src/evalml/config.py +++ b/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") diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index d0c66c7..31f59c9 100644 --- a/workflow/rules/inference.smk +++ b/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=() diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index 666e35b..a5dd7ba 100644 --- a/workflow/tools/config.schema.json +++ b/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": [ {