time_shift added

earthkit/climate/aggregate/temporal.py

@@ -1,5 +1,6 @@
 import typing as T
 from copy import deepcopy
+from datetime import timedelta
 
 import numpy as np
 import xarray as xr
@@ -20,6 +21,8 @@ def daily_mean(
     ----------
     dataarray : xr.DataArray
         DataArray containing a `time` dimension.
+    time_dim : str
+        Name of the time dimension in the xarray object, default is `"time"`.
     **kwargs
         Keyword arguments to be passed to :func:`resample`.
 
@@ -43,6 +46,8 @@ def daily_max(
     ----------
     dataarray : xr.DataArray
         DataArray containing a `time` dimension.
+    time_dim : str
+        Name of the time dimension in the xarray object, default is `"time"`.
     **kwargs
         Keyword arguments to be passed to :func:`resample`.
 
@@ -57,6 +62,7 @@ def daily_max(
 def daily_min(
     dataarray: T.Union[xr.Dataset, xr.DataArray],
     time_dim: T.Union[str, None] = None,
+    time_shift: T.Union[None, timedelta, dict] = None,
     **kwargs,
 ):
     """
@@ -65,7 +71,11 @@ def daily_min(
     Parameters
     ----------
     dataarray : xr.DataArray
-        DataArray containing a `time` dimension.
+        DataArray containing a time dimension.
+    time_dim : (optional) str
+        Name of the time dimension in the xarray object, default is `"time"`.
+    time_shift : (optional) timedelta or dict
+        Any time shift to apply to the data prior to calculation, e.g. to change the local time zone.
     **kwargs
         Keyword arguments to be passed to :func:`resample`.
 
@@ -89,6 +99,8 @@ def monthly_mean(
     ----------
     dataarray : xr.DataArray
         DataArray containing a `time` dimension.
+    time_dim : str
+        Name of the time dimension in the xarray object, default is `"time"`.
     **kwargs
         Keyword arguments to be passed to :func:`resample`.
 
@@ -112,6 +124,8 @@ def monthly_max(
     ----------
     dataarray : xr.DataArray
         DataArray containing a `time` dimension.
+    time_dim : str
+        Name of the time dimension in the xarray object, default is `"time"`.
     **kwargs
         Keyword arguments to be passed to :func:`resample`.
 
@@ -135,6 +149,8 @@ def monthly_min(
     ----------
     dataarray : xr.DataArray
         DataArray containing a `time` dimension.
+    time_dim : str
+        Name of the time dimension in the xarray object, default is `"time"`.
     **kwargs
         Keyword arguments to be passed to :func:`resample`.
 

earthkit/climate/aggregate/tools.py

@@ -1,7 +1,10 @@
+from copy import deepcopy
 import functools
 import typing as T
+from datetime import timedelta
 
 import numpy as np
+import pandas as pd
 import xarray as xr
 
 #: Mapping from pandas frequency strings to xarray time groups
@@ -29,10 +32,25 @@ def wrapper(
         dataarray: T.Union[xr.Dataset, xr.DataArray],
         *args,
         time_dim: T.Union[str, None] = None,
+        time_shift: T.Union[None, timedelta, np.timedelta64, dict, str] = None,
         **kwargs,
     ):
         if time_dim is None:
             time_dim = get_dim_key(dataarray, "t")
+
+        if time_shift is not None:
+            # Create timedelta from dict
+            if isinstance(time_shift, str):
+                time_shift = pd.Timedelta(**time_shift)
+            elif isinstance(time_shift, dict):
+                time_shift = timedelta(**time_shift)
+
+            # Convert timedelta to timedelta64 (TODO: may need to be more robust here)
+            time_coord = dataarray.coords[time_dim] + pd.Timedelta(time_shift)
+            time_coord = time_coord.assign_attrs({"time_shift": f"{time_shift}"})
+
+            dataarray.assign_coords({time_dim: time_coord})
+
         return func(dataarray, *args, time_dim=time_dim, **kwargs)
 
     return wrapper