Prepare for hour_of_day and day_of_year

mllam_data_prep/derived_variables.py

@@ -63,10 +63,20 @@ def derive_variables(fp, derived_variables, chunking):
         kwargs = {v: ds_input[k] for k, v in required_variables.items()}
         func = get_derived_variable_function(function_name)
         derived_field = func(**kwargs)
-        derived_field = _return_dropped_coordinates(
-            derived_field, ds_input, required_coordinates, chunks
-        )
-        ds_subset[derived_field.name] = derived_field
+
+        # Some of the derived variables include two components, since
+        # they are cyclically encoded (cos and sin parts)
+        if isinstance(derived_field, xr.DataArray):
+            derived_field = _return_dropped_coordinates(
+                derived_field, ds_input, required_coordinates, chunks
+            )
+            ds_subset[derived_field.name] = derived_field
+        elif isinstance(derived_field, tuple):
+            for field in derived_field:
+                field = _return_dropped_coordinates(
+                    field, ds_input, required_coordinates, chunks
+                )
+                ds_subset[field.name] = field
 
     return ds_subset
 
@@ -173,62 +183,74 @@ def calculate_toa_radiation(lat, lon, time):
     return toa_radiation
 
 
-def derive_hour_of_day(ds):
+def calculate_hour_of_day(time):
     """
-    Derive hour of day features with a cyclic encoding
+    Function for calculating hour of day features with a cyclic encoding
 
     Parameters
     ----------
-    ds : xr.Dataset
-        The dataset with variables needed to derive hour of day
+    time : xr.DataArray or datetime object
+        Time
 
     Returns
     -------
-    ds: xr.Dataset
-        The dataset with hour of day added
+    hour_of_day_cos: xr.DataArray or float
+        cosine of the hour of day
+    hour_of_day_sin: xr.DataArray or float
+        sine of the hour of day
     """
     logger.info("Calculating hour of day")
 
     # Get the hour of the day
-    hour_of_day = ds.time.dt.hour
+    hour_of_day = time.dt.hour
 
     # Cyclic encoding of hour of day
     hour_of_day_cos, hour_of_day_sin = cyclic_encoding(hour_of_day, 24)
 
-    # Assign to the dataset
-    ds = ds.assign(hour_of_day_sin=hour_of_day_sin)
-    ds = ds.assign(hour_of_day_cos=hour_of_day_cos)
+    if isinstance(hour_of_day_cos, xr.DataArray):
+        # Add attributes
+        hour_of_day_cos.name = "hour_of_day_cos"
+
+    if isinstance(hour_of_day_sin, xr.DataArray):
+        # Add attributes
+        hour_of_day_sin.name = "hour_of_day_sin"
 
-    return ds
+    return hour_of_day_cos, hour_of_day_sin
 
 
-def derive_day_of_year(ds):
+def calculate_day_of_year(time):
     """
-    Derive day of year features with a cyclic encoding
+    Function for calculating day of year features with a cyclic encoding
 
     Parameters
     ----------
-    ds : xr.Dataset
-        The dataset with variables needed to derive day of year
+    time : xr.DataArray or datetime object
+        Time
 
     Returns
     -------
-    ds: xr.Dataset
-        The dataset with day of year added
+    day_of_year_cos: xr.DataArray or float
+        cosine of the day of year
+    day_of_year_sin: xr.DataArray or float
+        sine of the day of year
     """
     logger.info("Calculating day of year")
 
     # Get the day of year
-    day_of_year = ds.time.dt.dayofyear
+    day_of_year = time.dt.dayofyear
 
     # Cyclic encoding of day of year - use 366 to include leap years!
     day_of_year_cos, day_of_year_sin = cyclic_encoding(day_of_year, 366)
 
-    # Assign to the dataset
-    ds = ds.assign(day_of_year_sin=day_of_year_sin)
-    ds = ds.assign(day_of_year_cos=day_of_year_cos)
+    if isinstance(day_of_year_cos, xr.DataArray):
+        # Add attributes
+        day_of_year_cos.name = "day_of_year_cos"
+
+    if isinstance(day_of_year_sin, xr.DataArray):
+        # Add attributes
+        day_of_year_sin.name = "day_of_year_sin"
 
-    return ds
+    return day_of_year_cos, day_of_year_sin
 
 
 def cyclic_encoding(data_array, da_max):