No public description

PiperOrigin-RevId: 709298226

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)
 

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,
+    ),
 }