nested-eagle v1

src/Makefile

@@ -8,7 +8,12 @@ VERBOSE  = $(if $(filter undefined,$(origin DEBUG)),, --verbose)
 
 INSTALLDIR ?= conda
 SHELL      := $(shell which bash)
-NOW        := $(shell date -u +"%Y-%m-%dT%H:%M:%S")
+
+ifdef HINDCAST
+NOW := $(HINDCAST)
+else
+NOW := $(shell date -u +"%Y-%m-%dT%H:%M:%S")
+endif
 
 export NOW
 

src/config/base.yaml

@@ -62,7 +62,7 @@ training:
       export: false
       memory: 128G
       rundir: '{{ training.rundir }}'
-      walltime: '01:00:00'
+      walltime: '3:00:00'
     envcmds: !list '{{ val.training.envcmds }}'
     executable: '{{ app.launcher | default("") }} anemoi-training train --config-name=training'
   remove:
@@ -251,7 +251,7 @@ vx:
         cores: 1
         export: false
         partition: '{{ app.partitions.netaccess }}'
-        walltime: '00:30:00'
+        walltime: '3:00:00'
       envcmds:
         - source {{ val.conda }}/etc/profile.d/conda.sh
         - conda activate wxvx
@@ -280,7 +280,7 @@ vx:
             inittime: forecast_reference_time
             validtime: time
       leadtimes:
-        start: '{{ app.time.leadtime }}'
+        start: '6'
         step: '{{ app.time.step }}'
         stop: '{{ app.time.leadtime }}'
       meta:
@@ -294,26 +294,30 @@ vx:
           name: 2t
         gh:
           level_type: isobaricInhPa
-          levels: [500]
+          levels: [250, 500, 850]
           name: gh
         sp:
           level_type: surface
           name: sp
+        q:
+          level_type: isobaricInhPa
+          levels: [250, 500, 850]
+          name: q
         t:
           level_type: isobaricInhPa
-          levels: [850]
+          levels: [250, 500, 850]
           name: t
         u:
           level_type: isobaricInhPa
-          levels: [250]
+          levels: [250, 500, 850]
           name: u
         u10:
           level_type: heightAboveGround
           levels: [10]
           name: 10u
         v:
           level_type: isobaricInhPa
-          levels: [250]
+          levels: [250, 500, 850]
           name: v
         v10:
           level_type: heightAboveGround

src/config/nested.yaml

@@ -23,10 +23,10 @@ prewxvx:
       forecast_path: '{{ inference.anemoi.output_path }}'
       freq: '{{ app.time.freq }}'
       lead_time: !int '{{ app.time.leadtime }}'
-      lam_index: 64220
+      lam_index: 407040
       lcc_info:
-        n_x: 338
-        n_y: 190
+        n_x: 848
+        n_y: 480
       levels:
         - 250
         - 500
@@ -449,7 +449,7 @@ vx:
         wxvx:
           forecast:
             name: nested_global
-            path: '{{ meta.post_files }}/nested-global.{{ yyyymmdd[:4] }}-{{ yyyymmdd[4:6] }}-{{ yyyymmdd[6:8] }}T{{ hh }}.{{ "%02d" % fh }}h.nc'
+            path: '{{ meta.post_files }}/nested-global.{{ yyyymmdd[:4] }}-{{ yyyymmdd[4:6] }}-{{ yyyymmdd[6:8] }}T{{ hh }}.240h.nc'
           meta: !dict '{{ vx.common.wxvx.meta | replace("TRUTH", "grid2grid") | replace("EXTENT", "global") }}'
           paths: !dict '{{ vx.common.paths_grid | replace("TRUTH", "grid2grid") | replace("EXTENT", "global") }}'
 
@@ -461,7 +461,7 @@ vx:
             cores: 1
             export: false
             partition: '{{ app.partitions.netaccess }}'
-            walltime: '00:30:00'
+            walltime: '3:00:00'
             rundir: '{{ vx.grid2grid.lam.wxvx.rundir }}'
           envcmds: !list '{{ vx.common.execution.envcmds }}'
           executable: wxvx -c wxvx-grid2grid-lam.yaml -t plots -f
@@ -479,7 +479,7 @@ vx:
             coords: !dict '{{ vx.common.wxvx.forecast.coords }}'
             name: nested_lam
             # NB: Retain relative 'meta.post_files' reference in 'path' value!
-            path: '{{ meta.post_files }}/nested-lam.{{ yyyymmdd[:4] }}-{{ yyyymmdd[4:6] }}-{{ yyyymmdd[6:8] }}T{{ hh }}.{{ "%02d" % fh }}h.nc'
+            path: '{{ meta.post_files }}/nested-lam.{{ yyyymmdd[:4] }}-{{ yyyymmdd[4:6] }}-{{ yyyymmdd[6:8] }}T{{ hh }}.240h.nc'
             projection:
               a: 6371229
               b: 6371229
@@ -499,7 +499,7 @@ vx:
       wxvx:
         wxvx:
           forecast:
-            path: '{{ meta.post_files }}/nested-global.{{ yyyymmdd[:4] }}-{{ yyyymmdd[4:6] }}-{{ yyyymmdd[6:8] }}T{{ hh }}.{{ "%02d" % fh }}h.nc'
+            path: '{{ meta.post_files }}/nested-global.{{ yyyymmdd[:4] }}-{{ yyyymmdd[4:6] }}-{{ yyyymmdd[6:8] }}T{{ hh }}.240h.nc'
           meta: !dict '{{ vx.common.wxvx.meta | replace("TRUTH", "grid2obs") | replace("EXTENT", "global") }}'
           paths: !dict '{{ vx.common.paths_obs | replace("TRUTH", "grid2obs") | replace("EXTENT", "global") }}'
     lam:
@@ -511,7 +511,7 @@ vx:
             export: false
             partition: '{{ app.partitions.netaccess }}'
             rundir: '{{ vx.grid2obs.lam.wxvx.rundir }}'
-            walltime: '00:30:00'
+            walltime: '3:00:00'
           envcmds: !list '{{ vx.common.execution.envcmds }}'
           executable: wxvx -c wxvx-grid2obs-lam.yaml -t plots -f
         name: grid2obs-lam
@@ -525,7 +525,7 @@ vx:
             coords: !dict '{{ vx.common.wxvx.forecast.coords }}'
             name: '{{ vx.grid2grid.lam.wxvx.wxvx.forecast.name }}'
             # NB: Retain relative 'meta.post_files' reference in 'path' value!
-            path: '{{ meta.post_files }}/nested-lam.{{ yyyymmdd[:4] }}-{{ yyyymmdd[4:6] }}-{{ yyyymmdd[6:8] }}T{{ hh }}.{{ "%02d" % fh }}h.nc'
+            path: '{{ meta.post_files }}/nested-lam.{{ yyyymmdd[:4] }}-{{ yyyymmdd[4:6] }}-{{ yyyymmdd[6:8] }}T{{ hh }}.240h.nc'
             projection: !dict '{{ vx.grid2grid.lam.wxvx.wxvx.forecast.projection }}'
           leadtimes: !dict '{{ vx.common.wxvx.leadtimes }}'
           meta: !dict '{{ vx.common.wxvx.meta | replace("TRUTH", "grid2obs") | replace("EXTENT", "lam") }}'

src/config/nrt.yaml

@@ -2,12 +2,12 @@
 
 app:
   nrt:
-    multistep_input_offset: !timedelta 0
+    multistep_input_offset: !timedelta 6
     now: !datetime '{{ "NOW" | env }}'
     ic_start: !datetime '{{ app.nrt.inference_start - app.nrt.multistep_input_offset }}'
     inference_start: !datetime '{{ app.nrt.now - app.nrt.offset - app.nrt.latency }}'
     latency: !timedelta 6
-    leadtime: !timedelta 48
+    leadtime: !timedelta 240
     offset: !timedelta '0:0:{{ (app.nrt.now.timestamp() % 21600) | int }}'
     stop: !datetime '{{ app.nrt.inference_start + app.nrt.leadtime }}'
   time:
@@ -16,17 +16,53 @@ app:
     start: !datetime '{{ app.nrt.ic_start }}'
     stop: !datetime '{{ app.nrt.stop }}'
     leadtime: !int '{{ (app.nrt.leadtime.total_seconds() / 3600) | int }}'
+    multistep_input: 2
   rundir: '{{ app.base }}/run/{{ app.experiment_name }}/nrt_inference/{{ app.time.inference_start.strftime("%Y/%m/%d/%H") }}'
 grids_and_meshes:
+  filenames:
+    gfs_target_grid: '{{ val.filenames.gfs_target_grid }}'
+    hrrr_target_grid: '{{ val.filenames.hrrr_target_grid }}'
+    latent_mesh: !remove
   rundir: '{{ app.base }}/run/{{ app.experiment_name }}/data'
+  conus_grid_resolution_km: 6
+  global_grid_resolution_deg: 0.25
+  latent_mesh_global_resolution_deg: !remove
+  latent_mesh_conus_coarsen_factor: !remove
 inference:
   anemoi:
-    input_dataset_kwargs:
+    checkpoint_path: '{{ app.base }}/inference-last.ckpt'
+    input_dataset_kwargs: &nrt-dataset-kwargs
       cutout:
         - dataset: '{{ val.data.rundir }}/hrrr.zarr'
-          trim_edge: [10, 11, 10, 11]
+          trim_edge: [25, 24, 25, 26]
         - dataset: '{{ val.data.rundir }}/gfs.zarr'
-  checkpoint_dir: '{{ app.base }}/run/{{ app.experiment_name }}/training/outputs/checkpoint'
+      adjust: all
+      min_distance_km: 6
+  checkpoint_dir: !remove
+  execution:
+    envcmds:
+      - source {{ val.conda }}/etc/profile.d/conda.sh
+      - conda activate anemoi
+      - set -ex
+      - export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH
+prewxvx:
+  common:
+    prewxvx:
+      lam_index: 407040
+      lcc_info:
+        n_x: 848
+        n_y: 480
+  global:
+    prewxvx:
+      eagle_tools:
+        anemoi_reference_dataset_kwargs:
+          <<: *nrt-dataset-kwargs
+        forecast_regrid_kwargs:
+          target_grid_path: '{{ grids_and_meshes.rundir }}/{{ val.filenames.gfs_target_grid }}'
+  lam:
+    prewxvx:
+      eagle_tools:
+        rename_curvilinear_coords_to_latlon: true
 training: !remove
 ufs2arco:
   mover:
@@ -50,6 +86,11 @@ ufs2arco:
     multistep_input: !int '{{ app.time.multistep_input }}'
     statistics_period: !remove
 val:
+  filenames:
+    gfs_target_grid: global_quarter_degree.nc
+    hrrr_target_grid: hrrr_6km.nc
+    latent_mesh: !remove
+    graph: !remove
   statsdelta: !remove
   variables:
     - gh
@@ -76,9 +117,7 @@ zarrs:
       ufs2arco:
         multisource: !remove
         source: !dict '{{ dict(ufs2arco.source, name="gfs_archive", slices=ufs2arco.slices) }}'
-        transforms:
-          horizontal_regrid:
-            target_grid_path: '{{ grids_and_meshes.rundir}}/{{ val.filenames.gfs_target_grid }}'
+        transforms: !remove
   hrrr:
     zarr:
       ufs2arco:

src/config/ursa.yaml

@@ -28,7 +28,7 @@ zarrs:
   common:
     execution:
       batchargs:
-        --mem-per-cpu: 3G
+        --mem-per-cpu: 20G
       envcmds: !list '{{ val.slurm_script + val.zarrs.envcmds }}'
       mpiargs:
         - "--export=ALL"

src/eagle/data/grids_and_meshes.py

@@ -54,20 +54,24 @@ def global_data_grid(self):
         The global grid, provisioned to the rundir.
         """
         res = self.config["global_grid_resolution_deg"]
-        if res == 0.25 or "gfs_target_grid" not in self.config["filenames"]:
-            yield self.taskname("global data grid (skipping)")
-            yield Asset(None, lambda: True)
-            yield None
-        else:
-            path = self.rundir / self.config["filenames"]["gfs_target_grid"]
-            yield self.taskname(f"global data grid {path}")
-            yield Asset(path, path.is_file)
-            yield None
-            path.parent.mkdir(parents=True, exist_ok=True)
-            ds = xesmf.util.grid_global(res, res, cf=True, lon1=360)
-            ds = ds.drop_vars("latitude_longitude")
-            ds = ds.sortby("lat", ascending=False)  # GFS goes north -> south
-            ds.to_netcdf(path)
+        if "gfs_target_grid" not in self.config["filenames"]:
+            return
+        path = self.rundir / self.config["filenames"]["gfs_target_grid"]
+        yield self.taskname(f"global data grid {path}")
+        yield Asset(path, path.is_file)
+        yield None
+        path.parent.mkdir(parents=True, exist_ok=True)
+        lon = np.arange(0, 360, res)
+        lat = np.arange(-90, 90, res)
+        ds = xr.Dataset(
+            coords={
+                "lon": ("lon", lon),
+                "lat": ("lat", lat),
+            }
+        )
+        ds = ds.sel(lat=slice(-89.9, 89.9))
+        ds = ds.sortby("lat", ascending=False)
+        ds.to_netcdf(path)
 
     @task
     def latent_mesh(self):

src/envs/anemoi.yaml

@@ -3,15 +3,21 @@ channels:
   - conda-forge
   - ufs-community
 dependencies:
-  - flash-attn 2.8.*
+  - flash-attn 2.7.*
   - mpi4py 4.1.*
-  - numpy 2.2.6.*
+  - numpy 1.26.4*
   - pip
   - python 3.12
+  - pytorch 2.6.*
+  - torch-geometric 2.6.*
+  - torchvision 0.21.*
   - setuptools 81.0.*
   - uwtools 2.16.*
   - pip:
-      - anemoi-inference==0.10.*
-      - anemoi-models==0.13.*
-      - anemoi-training==0.10.*
-      - eagle-tools==0.8.1
+    - anemoi-inference==0.9.1
+    - anemoi-models==0.12.0
+    - anemoi-datasets==0.5.30
+    - anemoi-graphs==0.8.2
+    - anemoi-transform==0.1.21
+    - anemoi-utils==0.4.42
+    - eagle-tools==0.8.1