Merge branch 'time_shift'

earthkit/climate/aggregate/temporal.py

@@ -1,6 +1,5 @@
 import typing as T
 from copy import deepcopy
-from datetime import timedelta
 
 import numpy as np
 import xarray as xr
@@ -23,6 +22,10 @@ def daily_mean(
         DataArray containing a `time` dimension.
     time_dim : str
         Name of the time dimension in the xarray object, default is `"time"`.
+    time_shift : (optional) timedelta or dict
+        A time shift to apply to the data prior to calculation, e.g. to change the local time zone.
+        It can be provided as any object that can be understood by `pandas.Timedelta`, a dictonary is passed
+        as kwargs to `pandas.Timedelta`
     **kwargs
         Keyword arguments to be passed to :func:`resample`.
 
@@ -40,14 +43,18 @@ def daily_max(
     **kwargs,
 ):
     """
-    Calculate the daily max.
+    Calculate the daily maximum.
 
     Parameters
     ----------
     dataarray : xr.DataArray
         DataArray containing a `time` dimension.
     time_dim : str
         Name of the time dimension in the xarray object, default is `"time"`.
+    time_shift : (optional) timedelta or dict
+        A time shift to apply to the data prior to calculation, e.g. to change the local time zone.
+        It can be provided as any object that can be understood by `pandas.Timedelta`, a dictonary is passed
+        as kwargs to `pandas.Timedelta`
     **kwargs
         Keyword arguments to be passed to :func:`resample`.
 
@@ -62,11 +69,10 @@ 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,
 ):
     """
-    Calculate the daily min.
+    Calculate the daily minimum.
 
     Parameters
     ----------
@@ -75,7 +81,9 @@ def daily_min(
     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.
+        A time shift to apply to the data prior to calculation, e.g. to change the local time zone.
+        It can be provided as any object that can be understood by `pandas.Timedelta`, a dictonary is passed
+        as kwargs to `pandas.Timedelta`
     **kwargs
         Keyword arguments to be passed to :func:`resample`.
 
@@ -86,6 +94,64 @@ def daily_min(
     return resample(dataarray, frequency="D", dim=time_dim, how="min", **kwargs)
 
 
+@tools.time_dim_decorator
+def daily_std(
+    dataarray: T.Union[xr.Dataset, xr.DataArray],
+    time_dim: T.Union[str, None] = None,
+    **kwargs,
+):
+    """
+    Calculate the daily standard deviation.
+
+    Parameters
+    ----------
+    dataarray : xr.DataArray
+        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
+        A time shift to apply to the data prior to calculation, e.g. to change the local time zone.
+        It can be provided as any object that can be understood by `pandas.Timedelta`, a dictonary is passed
+        as kwargs to `pandas.Timedelta`
+    **kwargs
+        Keyword arguments to be passed to :func:`resample`.
+
+    Returns
+    -------
+    xr.DataArray
+    """
+    return resample(dataarray, frequency="D", dim=time_dim, how="std", **kwargs)
+
+
+@tools.time_dim_decorator
+def daily_sum(
+    dataarray: T.Union[xr.Dataset, xr.DataArray],
+    time_dim: T.Union[str, None] = None,
+    **kwargs,
+):
+    """
+    Calculate the daily sum (accumulation).
+
+    Parameters
+    ----------
+    dataarray : xr.DataArray
+        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
+        A time shift to apply to the data prior to calculation, e.g. to change the local time zone.
+        It can be provided as any object that can be understood by `pandas.Timedelta`, a dictonary is passed
+        as kwargs to `pandas.Timedelta`
+    **kwargs
+        Keyword arguments to be passed to :func:`resample`.
+
+    Returns
+    -------
+    xr.DataArray
+    """
+    return resample(dataarray, frequency="D", dim=time_dim, how="sum", **kwargs)
+
+
 @tools.time_dim_decorator
 def monthly_mean(
     dataarray: T.Union[xr.Dataset, xr.DataArray],
@@ -101,6 +167,10 @@ def monthly_mean(
         DataArray containing a `time` dimension.
     time_dim : str
         Name of the time dimension in the xarray object, default is `"time"`.
+    time_shift : (optional) timedelta or dict
+        A time shift to apply to the data prior to calculation, e.g. to change the local time zone.
+        It can be provided as any object that can be understood by `pandas.Timedelta`, a dictonary is passed
+        as kwargs to `pandas.Timedelta`
     **kwargs
         Keyword arguments to be passed to :func:`resample`.
 
@@ -126,6 +196,10 @@ def monthly_max(
         DataArray containing a `time` dimension.
     time_dim : str
         Name of the time dimension in the xarray object, default is `"time"`.
+    time_shift : (optional) timedelta or dict
+        A time shift to apply to the data prior to calculation, e.g. to change the local time zone.
+        It can be provided as any object that can be understood by `pandas.Timedelta`, a dictonary is passed
+        as kwargs to `pandas.Timedelta`
     **kwargs
         Keyword arguments to be passed to :func:`resample`.
 
@@ -151,6 +225,10 @@ def monthly_min(
         DataArray containing a `time` dimension.
     time_dim : str
         Name of the time dimension in the xarray object, default is `"time"`.
+    time_shift : (optional) timedelta or dict
+        A time shift to apply to the data prior to calculation, e.g. to change the local time zone.
+        It can be provided as any object that can be understood by `pandas.Timedelta`, a dictonary is passed
+        as kwargs to `pandas.Timedelta`
     **kwargs
         Keyword arguments to be passed to :func:`resample`.
 

earthkit/climate/aggregate/tools.py

@@ -1,6 +1,5 @@
 import functools
 import typing as T
-from datetime import timedelta
 
 import numpy as np
 import pandas as pd
@@ -31,24 +30,24 @@ 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,
+        time_shift: T.Union[None, dict, str, pd.Timedelta] = 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):
+            if isinstance(time_shift, dict):
                 time_shift = pd.Timedelta(**time_shift)
-            elif isinstance(time_shift, dict):
-                time_shift = timedelta(**time_shift)
+            else:
+                time_shift = pd.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 = dataarray.coords[time_dim] + time_shift
             time_coord = time_coord.assign_attrs({"time_shift": f"{time_shift}"})
 
-            dataarray.assign_coords({time_dim: time_coord})
+            dataarray = dataarray.assign_coords({time_dim: time_coord})
 
         return func(dataarray, *args, time_dim=time_dim, **kwargs)