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) diff --git a/weatherbench2/derived_variables.py b/weatherbench2/derived_variables.py index e7d8199..8d942c2 100644 --- a/weatherbench2/derived_variables.py +++ b/weatherbench2/derived_variables.py @@ -686,16 +686,16 @@ def interp_at_one_lat(da: xr.DataArray) -> xr.DataArray: class AggregatePrecipitationAccumulation(DerivedVariable): """Compute longer aggregation periods from existing shorter accumulations. - Note: This assumes a 6h raw time step and prediction_timedelta starting at 6h. - Attributes: accumulation_hours: Hours to accumulate precipitaiton over - raw_accumulation_name: Name of the 6hr accumulation + raw_accumulation_name: Name of the accumulation + raw_accumulation_hours: Hours of the raw accumulation lead_time_name: Name of lead_time dimension """ accumulation_hours: int raw_accumulation_name: str = 'total_precipitation_6hr' + raw_accumulation_hours: int = 6 lead_time_name: str = 'prediction_timedelta' @property @@ -711,7 +711,8 @@ def compute(self, dataset: xr.Dataset): # Compute aggregation steps steps = float( - np.timedelta64(self.accumulation_hours, 'h') / np.timedelta64(6, 'h') + np.timedelta64(self.accumulation_hours, 'h') + / np.timedelta64(self.raw_accumulation_hours, 'h') ) assert steps.is_integer(), 'Accumulation time must be multiple of timestep.' # Compute accumulation @@ -763,4 +764,10 @@ def compute(self, dataset: xr.Dataset): accumulation_hours=24, lead_time_name='prediction_timedelta', ), + 'total_precipitation_24hr_from_12hr': AggregatePrecipitationAccumulation( + accumulation_hours=24, + lead_time_name='prediction_timedelta', + raw_accumulation_name='total_precipitation_12hr', + raw_accumulation_hours=12, + ), }