From 8e5773f7870c0c8e9fd40c4ad42cfe3885741502 Mon Sep 17 00:00:00 2001 From: Stephan Rasp Date: Thu, 6 Feb 2025 00:17:42 -0800 Subject: [PATCH] No public description PiperOrigin-RevId: 723815173 --- scripts/compute_derived_variables.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/scripts/compute_derived_variables.py b/scripts/compute_derived_variables.py index 9341e1e..4931c9d 100644 --- a/scripts/compute_derived_variables.py +++ b/scripts/compute_derived_variables.py @@ -144,10 +144,21 @@ def _strip_offsets( def main(argv: list[str]) -> None: - derived_variables = { - variable_name: dvs.DERIVED_VARIABLE_DICT[variable_name] - for variable_name in DERIVED_VARIABLES.value - } + derived_variables = {} + for variable_name in DERIVED_VARIABLES.value: + # Remove suffix for precipitation accumulations + # E.g. total_precipitation_24hr_from_6hr should also be called + # total_precipitation_24hr + dv = dvs.DERIVED_VARIABLE_DICT[variable_name] + if ( + variable_name.startswith('total_precipitation_') + and '_from_' in variable_name + ): + variable_name = variable_name.split('_from_')[0] + assert ( + variable_name not in DERIVED_VARIABLES.value + ), 'Duplicate variable name after removing suffix.' + derived_variables[variable_name] = dv source_dataset, source_chunks = xbeam.open_zarr(INPUT_PATH.value)