Adopt forecast intervals including the end point (#71)

* Adopt forecast intervals including the end point

* Fix parsing

* Experiments work

* Update config/forecasters.yaml

* Align init times to availabiliy of COE

* run pre-commit

* Change README to COSMO-E availability

---------

Co-authored-by: Jonas Bhend <jonasbhend@users.noreply.github.com>
Co-authored-by: Jonas Bhend <jonas.bhend@meteoswiss.ch>

Author: dnerini

README.md

@@ -31,25 +31,25 @@ description: |
 dates:
   start: 2020-01-01T12:00
   end: 2020-01-10T00:00
-  frequency: 54h
+  frequency: 60h
 
 runs:
   - forecaster:
       mlflow_id: 2f962c89ff644ca7940072fa9cd088ec
       label: Stage D - N320 global grid with CERRA finetuning
-      steps: 0/126/6
+      steps: 0/120/6
   - forecaster:
       mlflow_id: d0846032fc7248a58b089cbe8fa4c511
       label: M-1 forecaster
-      steps: 0/126/6
+      steps: 0/120/6
 
 
 baselines:
   - baseline:
       baseline_id: COSMO-E
       label: COSMO-E
       root: /store_new/mch/msopr/ml/COSMO-E
-      steps: 0/126/6
+      steps: 0/120/6
 
 analysis:
   label: COSMO KENDA

config/forecasters-co1e.yaml

@@ -13,7 +13,7 @@ runs:
       mlflow_id: 2174c939c8844555a52843b71219d425
       label: Cosmo 1km + era5 N320, finetuned on cerra checkpoint, lam resolution 11
       config: resources/inference/configs/forecaster_no_trimedge_fromtraining.yaml
-      steps: 0/126/6
+      steps: 0/120/6
       inference_resources:
         gpu: 4
         tasks: 4
@@ -23,7 +23,7 @@ baselines:
       baseline_id: COSMO-1E
       label: COSMO-1E
       root: /scratch/mch/bhendj/COSMO-1E
-      steps: 0/126/6
+      steps: 0/33/6
 
 analysis:
   label: COSMO KENDA

config/forecasters.yaml

@@ -1,29 +1,24 @@
 # yaml-language-server: $schema=../workflow/tools/config.schema.json
 description: |
-  This is an experiment to do blabla.
+  Evaluate skill of COSMO-E emulator (M-1 forecaster).
 
 dates:
   start: 2020-01-01T12:00
   end: 2020-01-10T00:00
-  # end: 2020-03-30T00:00
-  frequency: 36h
+  frequency: 60h
 
 runs:
-  - forecaster:
-      mlflow_id: 2f962c89ff644ca7940072fa9cd088ec
-      label: Stage D - N320 global grid with CERRA finetuning
-      steps: 0/126/6
   - forecaster:
       mlflow_id: d0846032fc7248a58b089cbe8fa4c511
       label: M-1 forecaster
-      steps: 0/126/6
+      steps: 0/120/6
 
 baselines:
   - baseline:
       baseline_id: COSMO-E
       label: COSMO-E
       root: /store_new/mch/msopr/ml/COSMO-E
-      steps: 0/126/6
+      steps: 0/120/6
 
 analysis:
   label: COSMO KENDA
@@ -38,7 +33,7 @@ locations:
 profile:
   executor: slurm
   global_resources:
-    gpus: 15
+    gpus: 16
   default_resources:
     slurm_partition: "postproc"
     cpus_per_task: 1

config/interpolators.yaml

@@ -1,18 +1,18 @@
 # yaml-language-server: $schema=../workflow/tools/config.schema.json
 description: |
-  Stretched interpolator vs LAM interpolator.
+  Evaluate skill of SGM interpolator (M-2 interpolator).
 
 dates:
   start: 2020-01-01T12:00
   end: 2020-01-10T00:00
-  frequency: 54h
+  frequency: 60h
 
 runs:
   - interpolator:
-      mlflow_id: 9c18b90074214d769b8b383722fc5a06
+      mlflow_id: 8d1e0410ca7d4f74b368b3079878259a
       label: M-2 interpolator (KENDA)
-      steps: 0/121/1
-      config: resources/inference/configs/interpolator_from_test_data.yaml
+      steps: 0/120/1
+      config: resources/inference/configs/interpolator_from_test_data_stretched.yaml
       forecaster: null
       extra_dependencies:
         - git+https://github.com/ecmwf/anemoi-inference@14189907b4f4e3b204b7994f828831b8aa51e9b6
@@ -21,12 +21,12 @@ runs:
   - interpolator:
       mlflow_id: 8d1e0410ca7d4f74b368b3079878259a
       label: M-2 interpolator (M-1 forecaster)
-      steps: 0/121/1
+      steps: 0/120/1
       config: resources/inference/configs/interpolator_stretched.yaml
       forecaster:
         mlflow_id: d0846032fc7248a58b089cbe8fa4c511
         config: resources/inference/configs/forecaster_with_global.yaml
-        steps: 0/126/6
+        steps: 0/120/6
       extra_dependencies:
         - git+https://github.com/ecmwf/anemoi-inference@14189907b4f4e3b204b7994f828831b8aa51e9b6
         - torch-geometric==2.6.1
@@ -35,14 +35,14 @@ runs:
       mlflow_id: d0846032fc7248a58b089cbe8fa4c511
       label: M-1 forecaster
       config: resources/inference/configs/forecaster_with_global.yaml
-      steps: 0/126/6
+      steps: 0/120/6
 
 baselines:
   - baseline:
       baseline_id: COSMO-E-1h
       label: COSMO-E
       root: /scratch/mch/bhendj/COSMO-E
-      steps: 0/121/1
+      steps: 0/120/1
 
 analysis:
   label: COSMO KENDA
@@ -58,7 +58,7 @@ locations:
 profile:
   executor: slurm
   global_resources:
-    gpus: 15
+    gpus: 16
   default_resources:
     slurm_partition: "postproc"
     cpus_per_task: 1

pyproject.toml

@@ -50,4 +50,4 @@ markers = [
 packages = [
   "src/evalml",
   "src/verification"
-]
+]

src/evalml/config.py

@@ -74,8 +74,10 @@ class RunConfig(BaseModel):
         ...,
         description=(
             "Forecast lead times in hours, formatted as 'start/end/step'. "
-            "The range is half-open [start, end), meaning it includes the start  "
-            "but excludes the end. Example: '0/126/6' for lead times every 6 hours up to 120 hours."
+            "The range includes the start lead time and continues with the given step "
+            "until reaching or exceeding the end lead time. "
+            "Example: '0/120/6' for lead times every 6 hours up to 120 h, "
+            "or '0/33/6' up to 30 h."
         ),
     )
     extra_dependencies: List[str] = Field(
@@ -102,15 +104,13 @@ def validate_steps(cls, v: str) -> str:
         try:
             start, end, step = map(int, parts)
         except ValueError:
-            raise ValueError("Start, end, and interval must be integers.")
-        if start >= end:
-            raise ValueError(f"Start ({start}) must be less than end ({end}).")
-        if step <= 0:
-            raise ValueError(f"Interval ({step}) must be a positive integer.")
-        if (end - start) % step != 0:
+            raise ValueError("Start, end, and step must be integers.")
+        if start > end:
             raise ValueError(
-                f"The step ({step}) must evenly divide the range ({end - start})."
+                f"Start ({start}) must be less than or equal to end ({end})."
             )
+        if step <= 0:
+            raise ValueError(f"Step ({step}) must be a positive integer.")
         return v
 
 

workflow/rules/data.smk

@@ -18,15 +18,15 @@ if "extract_cosmoe" in config.get("include-optional-rules", []):
             runtime="24h",
         params:
             year_postfix=lambda wc: f"FCST{wc.year}",
-            lead_time="0/126/6",
+            steps="0/120/6",
         log:
             OUT_ROOT / "logs/extract-cosmoe-fcts-{year}.log",
         shell:
             """
             python workflow/scripts/extract_baseline_fct.py \
                 --archive_dir {input.archive}/{params.year_postfix} \
                 --output_store {output.fcts} \
-                --lead_time {params.lead_time} \
+                --steps {params.steps} \
                     > {log} 2>&1
             """
 
@@ -45,14 +45,14 @@ if "extract_cosmo1e" in config.get("include-optional-rules", []):
             runtime="24h",
         params:
             year_postfix=lambda wc: f"FCST{wc.year}",
-            lead_time="0/34/1",
+            steps="0/33/1",
         log:
             OUT_ROOT / "logs/extract-cosmo1e-fcts-{year}.log",
         shell:
             """
             python workflow/scripts/extract_baseline_fct.py \
                 --archive_dir {input.archive}/{params.year_postfix} \
                 --output_store {output.fcts} \
-                --lead_time {params.lead_time} \
+                --steps {params.steps} \
                     > {log} 2>&1
             """

workflow/rules/inference.smk

@@ -121,7 +121,6 @@ def get_resource(wc, field: str, default):
 def get_leadtime(wc):
     """Get the lead time from the run config."""
     start, end, step = RUN_CONFIGS[wc.run_id]["steps"].split("/")
-    end = int(end) - int(step)  # make inclusive
     return f"{end}h"
 
 

workflow/rules/verif.smk

@@ -22,7 +22,7 @@ rule verif_metrics_baseline:
         analysis_zarr=config["analysis"].get("analysis_zarr"),
     params:
         baseline_label=lambda wc: BASELINE_CONFIGS[wc.baseline_id].get("label"),
-        baseline_steps=lambda wc: BASELINE_CONFIGS[wc.baseline_id].get("steps"),
+        baseline_steps=lambda wc: BASELINE_CONFIGS[wc.baseline_id]["steps"],
         analysis_label=config["analysis"].get("label"),
     output:
         OUT_ROOT / "data/baselines/{baseline_id}/{init_time}/verif.nc",
@@ -38,7 +38,7 @@ rule verif_metrics_baseline:
             --analysis_zarr {input.analysis_zarr} \
             --baseline_zarr {input.baseline_zarr} \
             --reftime {wildcards.init_time} \
-            --lead_time "{params.baseline_steps}" \
+            --steps "{params.baseline_steps}" \
             --baseline_label "{params.baseline_label}" \
             --analysis_label "{params.analysis_label}" \
             --output {output} > {log} 2>&1
@@ -79,7 +79,7 @@ rule verif_metrics:
         uv run {input.script} \
             --grib_output_dir {input.grib_output} \
             --analysis_zarr {input.analysis_zarr} \
-            --lead_time "{params.fcst_steps}" \
+            --steps "{params.fcst_steps}" \
             --fcst_label "{params.fcst_label}" \
             --analysis_label "{params.analysis_label}" \
             --output {output} > {log} 2>&1

workflow/scripts/extract_baseline.py

@@ -42,7 +42,7 @@ def check_reftime_consistency(tarfiles: list[Path], delta_h: int = 12):
 
 
 def extract(
-    tar: Path, lead_time: list[int], run_id: str, params: list[str]
+    tar: Path, lead_times: list[int], run_id: str, params: list[str]
 ) -> xr.Dataset:
     LOG.info(f"Extracting fields from {tar}.")
     reftime = reftime_from_tarfile(tar)
@@ -54,7 +54,7 @@ def extract(
         raise ValueError("Currently only COSMO-E and COSMO-1E are supported.")
     tar_archive = tarfile.open(tar, mode="r:*")
     out = ekd.SimpleFieldList()
-    for lt in lead_time:
+    for lt in lead_times:
         filename = f"{tar.stem}/grib/{gribname}{lt:03}_{run_id}"
         LOG.info(f"Extracting {filename}.")
         stream = tar_archive.extractfile(filename)
@@ -79,23 +79,19 @@ def extract(
 class ScriptConfig(Namespace):
     archive_dir: Path
     output_store: Path
-    lead_time: int
+    steps: list[int]
     run_id: str
     params: list[str]
 
 
-def _parse_lead_time(lead_time: str) -> int:
-    # check that lead_time is in the format "start/stop/step"
-    if "/" not in lead_time:
-        raise ValueError(
-            f"Expected lead_time in format 'start/stop/step', got '{lead_time}'"
-        )
-    if len(lead_time.split("/")) != 3:
-        raise ValueError(
-            f"Expected lead_time in format 'start/stop/step', got '{lead_time}'"
-        )
-
-    return list(range(*map(int, lead_time.split("/"))))
+def _parse_steps(steps: str) -> int:
+    # check that steps is in the format "start/stop/step"
+    if "/" not in steps:
+        raise ValueError(f"Expected steps in format 'start/stop/step', got '{steps}'")
+    if len(steps.split("/")) != 3:
+        raise ValueError(f"Expected steps in format 'start/stop/step', got '{steps}'")
+    start, end, step = map(int, steps.split("/"))
+    return list(range(start, end + 1, step))
 
 
 def main(cfg: ScriptConfig):
@@ -135,7 +131,7 @@ def main(cfg: ScriptConfig):
 
     for i in indices:
         file = tarfiles[i]
-        ds = extract(file, cfg.lead_time, cfg.run_id, cfg.params)
+        ds = extract(file, cfg.steps, cfg.run_id, cfg.params)
 
         LOG.info(f"Extracted: {ds}")
 
@@ -167,7 +163,7 @@ def main(cfg: ScriptConfig):
         help="Path to the output zarr store.",
     )
 
-    parser.add_argument("--lead_time", type=_parse_lead_time, default="0/126/6")
+    parser.add_argument("--steps", type=_parse_steps, default="0/120/6")
 
     parser.add_argument("--run_id", type=str, default="000")
 
@@ -193,10 +189,10 @@ def main(cfg: ScriptConfig):
 python workflow/scripts/extract_baseline_fct.py \
     --archive_dir /archive/mch/msopr/osm/COSMO-E/FCST20 \
     --output_store /store_new/mch/msopr/ml/COSMO-E/FCST20.zarr \
-    --lead_time 0/126/6
+    --steps 0/120/6
 
 python workflow/scripts/extract_baseline_fct.py \
     --archive_dir /archive/mch/s83/osm/from_GPFS/COSMO-1E/FCST20 \
     --output_store /store_new/mch/msopr/ml/COSMO-1E/FCST20.zarr \
-    --lead_time 0/34/1
+    --steps 0/33/1
 """