diff --git a/scripts/compute_derived_variables.py b/scripts/compute_derived_variables.py index bf20c8b..25b6aae 100644 --- a/scripts/compute_derived_variables.py +++ b/scripts/compute_derived_variables.py @@ -34,6 +34,7 @@ --job_name=compute-derived-variables-$USER ``` """ +import ast from absl import app from absl import flags import apache_beam as beam @@ -90,6 +91,14 @@ ' "total_precipitation_6hr" for backwards compatibility.' ), ) +RENAME_VARIABLES = flags.DEFINE_string( + 'rename_variables', + None, + help=( + 'Dictionary of variable to rename to standard names. E.g. {"2t":' + ' "2m_temperature"}' + ), +) WORKING_CHUNKS = flag_utils.DEFINE_chunks( 'working_chunks', '', @@ -143,6 +152,17 @@ def main(argv: list[str]) -> None: {RAW_TP_NAME.value: 'total_precipitation'} ) + rename_variables = ( + ast.literal_eval(RENAME_VARIABLES.value) + if RENAME_VARIABLES.value + else None + ) + if rename_variables: + source_dataset = source_dataset.rename(rename_variables) + source_chunks = { + rename_variables.get(k, k): v for k, v in source_chunks.items() + } + for var_name in PREEXISTING_VARIABLES_TO_REMOVE.value: if var_name in source_dataset: del source_dataset[var_name]