diff --git a/CHANGELOG.md b/CHANGELOG.md index 15266dbf..24e2a47b 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/earth2studio/data/hrrr.py b/earth2studio/data/hrrr.py index 01313ce1..56181992 100644 --- a/earth2studio/data/hrrr.py +++ b/earth2studio/data/hrrr.py @@ -72,6 +72,7 @@ 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] @@ -79,8 +80,13 @@ def fetch_dataarray( 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", diff --git a/earth2studio/lexicon/hrrr.py b/earth2studio/lexicon/hrrr.py index f9f0e9a7..ce3369a1 100644 --- a/earth2studio/lexicon/hrrr.py +++ b/earth2studio/lexicon/hrrr.py @@ -93,6 +93,7 @@ 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 = {} @@ -100,7 +101,17 @@ def build_vocab() -> dict[str, str]: 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() diff --git a/test/data/test_hrrr.py b/test/data/test_hrrr.py index e129ec9d..75709abd 100644 --- a/test/data/test_hrrr.py +++ b/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): diff --git a/test/lexicon/test_hrrr_lexicon.py b/test/lexicon/test_hrrr_lexicon.py index b03b6bee..e8e840e9 100644 --- a/test/lexicon/test_hrrr_lexicon.py +++ b/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):