From fba7ec3b79b2f626f9627aa61b5b8e3b88c66a56 Mon Sep 17 00:00:00 2001 From: Ian Langmore Date: Fri, 6 Dec 2024 16:07:45 -0800 Subject: [PATCH] --levels FLAG added to compute_probabilistic_climatological_forecasts. Allows selecting designated levels from input path. PiperOrigin-RevId: 703644807 --- .../compute_probabilistic_climatological_forecasts.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/scripts/compute_probabilistic_climatological_forecasts.py b/scripts/compute_probabilistic_climatological_forecasts.py index 4c9cf0b..aca59f1 100644 --- a/scripts/compute_probabilistic_climatological_forecasts.py +++ b/scripts/compute_probabilistic_climatological_forecasts.py @@ -100,6 +100,14 @@ 2020, help='Inclusive end year of --input_path to include.', ) +LEVELS = flags.DEFINE_list( + 'levels', + None, + help=( + 'Comma delimited list of pressure levels to compute spectra on. If' + ' empty, compute on all levels of --input_path' + ), +) VARIABLES = flags.DEFINE_list( 'variables', None, @@ -504,6 +512,9 @@ def main(argv: abc.Sequence[str]) -> None: if VARIABLES.value: input_ds = input_ds[VARIABLES.value] + if LEVELS.value: + input_ds = input_ds.sel(level=[int(l) for l in LEVELS.value]) + input_chunks = {k: v for k, v in input_chunks.items() if k in input_ds.dims} _check_input_spacing_and_time_flags(input_ds)