fixed a bug in input argument type in timeaverage._get_time_method. U…

…pdated tests arguments for this function accordingly.
esm-tools · Nov 14, 2024 · b43b8bc · b43b8bc
1 parent fbeecbd
commit b43b8bc
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 68 deletions.
diff --git a/src/pymorize/timeaverage.py b/src/pymorize/timeaverage.py
@@ -13,8 +13,9 @@
 _split_by_chunks(dataset: xr.DataArray) -> Tuple[Dict, xr.DataArray]:
     Split a large dataset into sub-datasets for each chunk.
 
-_get_time_method(table_id: str) -> str:
-    Determine the time method based on the table_id string.
+_get_time_method(frequency: str) -> str:
+    Determine the time method based on the frequency string from
+    rule.data_request_variable.frequency.
 
 _frequency_from_approx_interval(interval: str) -> str:
     Convert an interval expressed in days to a frequency string.
@@ -83,28 +84,26 @@ def _split_by_chunks(dataset: xr.DataArray):
         yield (selection, dataset[selection])
 
 
-def _get_time_method(table_id: str) -> str:
+def _get_time_method(frequency: str) -> str:
     """
-    Determine the time method based on the table_id string.
+    Determine the time method based on the frequency string from CMIP6 table for
+    a specific variable (rule.data_request_variable.frequency).
 
-    This function checks the ending of the table_id string and returns a corresponding time method.
-    If the table_id ends with 'Pt', it returns 'INSTANTANEOUS'.
-    If the table_id ends with 'C' or 'CM', it returns 'CLIMATOLOGY'.
-    In all other cases, it returns 'MEAN'.
+    The type of time method influences how the data is processed for time averaging.
 
     Parameters
     ----------
-    table_id : str
-        The table_id string to check.
+    frequency : str
+        The frequency string from CMIP6 tables (example: "mon").
 
     Returns
     -------
     str
         The corresponding time method ('INSTANTANEOUS', 'CLIMATOLOGY', or 'MEAN').
     """
-    if table_id.endswith("Pt"):
+    if frequency.endswith("Pt"):
         return "INSTANTANEOUS"
-    if table_id.endswith("C") or table_id.endswith("CM"):
+    if frequency.endswith("C") or frequency.endswith("CM"):
         return "CLIMATOLOGY"
     return "MEAN"
 
@@ -225,7 +224,7 @@ def compute_average(da: xr.DataArray, rule):
     logger.debug(f"{approx_interval=} {frequency_str=}")
     # attach the frequency_str to rule, it is referenced when creating file name
     rule.frequency_str = frequency_str
-    time_method = _get_time_method(drv.table.table_id)
+    time_method = _get_time_method(drv.frequency)
     rule.time_method = time_method
     if time_method == "INSTANTANEOUS":
         ds = da.resample(time=frequency_str).first()
@@ -316,7 +315,5 @@ def compute_average(da: xr.DataArray, rule):
 time: mean
 time: mean grid_longitude: mean
 time: point
-""".strip().split(
-    "\n"
-)
+""".strip().split("\n")
 """list: cell_methods to ignore when calculating time averages"""
diff --git a/tests/unit/test_timeaverage.py b/tests/unit/test_timeaverage.py
@@ -5,61 +5,29 @@
 
 import pymorize.timeaverage
 
-TABLE_TIME_METHOD = {
-    "CMIP6_3hr": "MEAN",
-    "CMIP6_6hrLev": "MEAN",
-    "CMIP6_6hrPlev": "MEAN",
-    "CMIP6_6hrPlevPt": "INSTANTANEOUS",
-    "CMIP6_AERday": "MEAN",
-    "CMIP6_AERhr": "MEAN",
-    "CMIP6_AERmon": "MEAN",
-    "CMIP6_AERmonZ": "MEAN",
-    "CMIP6_Amon": "MEAN",
-    "CMIP6_CF3hr": "MEAN",
-    "CMIP6_CFday": "MEAN",
-    "CMIP6_CFmon": "MEAN",
-    "CMIP6_CFsubhr": "MEAN",
-    "CMIP6_CV": "MEAN",
-    "CMIP6_E1hr": "MEAN",
-    "CMIP6_E1hrClimMon": "CLIMATOLOGY",
-    "CMIP6_E3hr": "MEAN",
-    "CMIP6_E3hrPt": "INSTANTANEOUS",
-    "CMIP6_E6hrZ": "MEAN",
-    "CMIP6_Eday": "MEAN",
-    "CMIP6_EdayZ": "MEAN",
-    "CMIP6_Efx": "MEAN",
-    "CMIP6_Emon": "MEAN",
-    "CMIP6_EmonZ": "MEAN",
-    "CMIP6_Esubhr": "MEAN",
-    "CMIP6_Eyr": "MEAN",
-    "CMIP6_IfxAnt": "MEAN",
-    "CMIP6_IfxGre": "MEAN",
-    "CMIP6_ImonAnt": "MEAN",
-    "CMIP6_ImonGre": "MEAN",
-    "CMIP6_IyrAnt": "MEAN",
-    "CMIP6_IyrGre": "MEAN",
-    "CMIP6_LImon": "MEAN",
-    "CMIP6_Lmon": "MEAN",
-    "CMIP6_Oclim": "CLIMATOLOGY",
-    "CMIP6_Oday": "MEAN",
-    "CMIP6_Odec": "MEAN",
-    "CMIP6_Ofx": "MEAN",
-    "CMIP6_Omon": "MEAN",
-    "CMIP6_Oyr": "MEAN",
-    "CMIP6_SIday": "MEAN",
-    "CMIP6_SImon": "MEAN",
-    "CMIP6_coordinate": "MEAN",
-    "CMIP6_day": "MEAN",
-    "CMIP6_formula_terms": "MEAN",
-    "CMIP6_fx": "MEAN",
-    "CMIP6_grids": "MEAN",
-    "CMIP6_input_example": "MEAN",
+FREQUENCY_TIME_METHOD = {
+    "fx": "MEAN",
+    "1hr": "MEAN",
+    "3hr": "MEAN",
+    "6hr": "MEAN",
+    "day": "MEAN",
+    "mon": "MEAN",
+    "yr": "MEAN",
+    "dec": "MEAN",
+    "1hrPt": "INSTANTANEOUS",
+    "subhrPt": "INSTANTANEOUS",
+    "6hrPt": "INSTANTANEOUS",
+    "3hrPt": "INSTANTANEOUS",
+    "monPt": "INSTANTANEOUS",
+    "yrPt": "INSTANTANEOUS",
+    "1hrCM": "CLIMATOLOGY",
+    "monC": "CLIMATOLOGY",
 }
 
 
-@pytest.mark.parametrize("table_name, expected", TABLE_TIME_METHOD.items())
-def test__get_time_method(table_name, expected):
-    answer = pymorize.timeaverage._get_time_method(table_name)
+@pytest.mark.parametrize("frequency_name, expected", FREQUENCY_TIME_METHOD.items())
+def test__get_time_method(frequency_name, expected):
+    answer = pymorize.timeaverage._get_time_method(frequency_name)
     assert answer == expected