No public description

PiperOrigin-RevId: 698430292
google-research · Dec 2, 2024 · 3dc7e62 · 3dc7e62
1 parent 8308c9c
commit 3dc7e62
Showing 1 changed file with 11 additions and 0 deletions.
diff --git 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)