Add natural levels for HRRR

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Set zarr chunks for lead time to size 1 in examples.
 - Updated HRRR tp to be hourly accumulated (Grib index 090)
+- Updated HRRR Lexicon to support natural levels
 
 ### Deprecated
 

earth2studio/data/hrrr.py

@@ -72,15 +72,21 @@ def fetch_dataarray(
         # Convert from Earth2Studio variable ID to HRRR id and modifier
         sfc_vars = {}
         prs_vars = {}
+        hybrid_vars = {}
         for var in variable:
             try:
                 hrrr_str, modifier = HRRRLexicon[var]
                 hrrr_name = hrrr_str.split("::")
 
                 if hrrr_name[0] == "sfc":
                     sfc_vars[var] = (f":{hrrr_name[1]}", modifier)
-                else:
+                elif hrrr_name[0] == "prs":
                     prs_vars[var] = (f":{hrrr_name[1]}", modifier)
+                elif hrrr_name[0] == "nat":
+                    hybrid_vars[var] = (f":{hrrr_name[1]}", modifier)
+                else:
+                    raise ValueError(f"HRRR name {hrrr_name[0]} not found.")
+
             except KeyError:  # noqa: PERF203
                 raise KeyError(f"variable id {var} not found in HRRR lexicon")
 
@@ -89,7 +95,9 @@ def fetch_dataarray(
 
         data_arrays = {}
         # Process surface and then pressure fields
-        for product, var_dict in zip(["sfc", "prs"], [sfc_vars, prs_vars]):
+        for product, var_dict in zip(
+            ["sfc", "prs", "nat"], [sfc_vars, prs_vars, hybrid_vars]
+        ):
             fh = FastHerbie(
                 time,
                 model="hrrr",

earth2studio/lexicon/hrrr.py

@@ -93,14 +93,25 @@ def build_vocab() -> dict[str, str]:
             975,
             1000,
         ]
+
         prs_names = ["UGRD", "VGRD", "HGT", "TMP", "RH", "SPFH"]
         e2s_id = ["u", "v", "z", "t", "r", "q"]
         prs_variables = {}
         for (id, variable) in zip(e2s_id, prs_names):
             for level in prs_levels:
                 prs_variables[f"{id}{level:d}"] = f"prs::{variable}:{level} mb"
 
-        return {**sfc_variables, **prs_variables}
+        hybrid_levels = list(range(51))
+        hybrid_names = ["UGRD", "VGRD", "HGT", "TMP", "SPFH", "PRES"]
+        e2s_hybrid_id = ["u_hl", "v_hl", "z_hl", "t_hl", "q_hl", "p_hl"]
+        hybrid_variables = {}
+        for (id, variable) in zip(e2s_hybrid_id, hybrid_names):
+            for level in hybrid_levels:
+                hybrid_variables[
+                    f"{id}{level:d}"
+                ] = f"nat::{variable}:{level} hybrid level"
+
+        return {**sfc_variables, **prs_variables, **hybrid_variables}
 
     VOCAB = build_vocab()
 

test/data/test_hrrr.py

@@ -37,7 +37,7 @@
         ],
     ],
 )
-@pytest.mark.parametrize("variable", ["t2m", ["u10m", "u100"]])
+@pytest.mark.parametrize("variable", ["t2m", ["u10m", "u100"], ["u_hl1"]])
 def test_hrrr_fetch(time, variable):
 
     ds = HRRR(cache=False)
@@ -115,7 +115,7 @@ def test_hrrr_fx_fetch(time, lead_time):
         np.array([np.datetime64("2024-01-01T00:00")]),
     ],
 )
-@pytest.mark.parametrize("variable", [["t2m", "sp"]])
+@pytest.mark.parametrize("variable", [["t2m", "sp", "t_hl10"]])
 @pytest.mark.parametrize("cache", [True, False])
 def test_hrrr_cache(time, variable, cache):
 

test/lexicon/test_hrrr_lexicon.py

@@ -21,7 +21,13 @@
 
 
 @pytest.mark.parametrize(
-    "variable", [["t2m"], ["u10m", "v200"], ["u80m", "z500", "q700"]]
+    "variable",
+    [
+        ["t2m"],
+        ["u10m", "v200"],
+        ["u80m", "z500", "q700"],
+        ["u_hl1", "v_hl4", "t_hl20", "p_hl30"],
+    ],
 )
 @pytest.mark.parametrize("device", ["cpu", "cuda:0"])
 def test_run_deterministic(variable, device):