diff --git a/earthkit/climate/aggregate/temporal.py b/earthkit/climate/aggregate/temporal.py index 01d6ff3..07d64aa 100644 --- a/earthkit/climate/aggregate/temporal.py +++ b/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,7 +43,7 @@ def daily_max( **kwargs, ): """ - Calculate the daily max. + Calculate the daily maximum. Parameters ---------- @@ -48,6 +51,10 @@ def daily_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`. @@ -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`. diff --git a/earthkit/climate/aggregate/tools.py b/earthkit/climate/aggregate/tools.py index c7ab2ef..f26c83a 100644 --- a/earthkit/climate/aggregate/tools.py +++ b/earthkit/climate/aggregate/tools.py @@ -1,7 +1,5 @@ -from copy import deepcopy import functools import typing as T -from datetime import timedelta import numpy as np import pandas as pd @@ -32,7 +30,7 @@ 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: @@ -40,16 +38,16 @@ def wrapper( 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)