def test_to_cftime_datetime(calendar, argument, expected_date_args):
date_type = get_date_type(calendar)
expected = date_type(*expected_date_args)
if isinstance(argument, tuple):
argument = date_type(*argument)
result = to_cftime_datetime(argument, calendar=calendar)
assert result == expected
def test_to_cftime_datetime_error_type_error():
with pytest.raises(TypeError):
to_cftime_datetime(1)
def test_to_cftime_datetime_error_no_calendar():
with pytest.raises(ValueError):
to_cftime_datetime('2000')
def test_to_cftime_datetime_error_type_error():
with pytest.raises(TypeError):
to_cftime_datetime(1)
def test_to_cftime_datetime_error_no_calendar():
with pytest.raises(ValueError):
to_cftime_datetime("2000")
def select_time(
da: Union[xr.DataArray, xr.Dataset],
drop: bool = False,
season: Union[str, Sequence[str]] = None,
month: Union[int, Sequence[int]] = None,
doy_bounds: Tuple[int, int] = None,
date_bounds: Tuple[str, str] = None,
) -> Union[xr.DataArray, xr.Dataset]:
"""Select entries according to a time period.
This conveniently improves xarray's :py:meth:`xarray.DataArray.where` and
:py:meth:`xarray.DataArray.sel` with fancier ways of indexing over time elements.
In addition to the data `da` and argument `drop`, only one of `season`, `month`,
`doy_bounds` or `date_bounds` may be passed.
Parameters
----------
da : xr.DataArray or xr.Dataset
Input data.
drop: boolean
Whether to drop elements outside the period of interest or
to simply mask them (default).
season: string or sequence of strings
One or more of 'DJF', 'MAM', 'JJA' and 'SON'.
month: integer or sequence of integers
Sequence of month numbers (January = 1 ... December = 12)
doy_bounds: 2-tuple of integers
The bounds as (start, end) of the period of interest expressed in day-of-year,
integers going from 1 (January 1st) to 365 or 366 (December 31st). If calendar
awareness is needed, consider using ``date_bounds`` instead.
Bounds are inclusive.
date_bounds: 2-tuple of strings
The bounds as (start, end) of the period of interest expressed as dates in the
month-day (%m-%d) format.
Bounds are inclusive.
Returns
-------
xr.DataArray or xr.Dataset
Selected input values. If ``drop=False``, this has the same length as ``da``
(along dimension 'time'), but with masked (NaN) values outside the period of
interest.
Examples
--------
Keep only the values of fall and spring.
>>> ds = open_dataset("ERA5/daily_surface_cancities_1990-1993.nc")
>>> ds.time.size
1461
>>> out = select_time(ds, drop=True, season=['MAM', 'SON'])
>>> out.time.size
732
Or all values between two dates (included).
>>> out = select_time(ds, drop=True, date_bounds=('02-29', '03-02'))
>>> out.time.values
array(['1990-03-01T00:00:00.000000000', '1990-03-02T00:00:00.000000000',
'1991-03-01T00:00:00.000000000', '1991-03-02T00:00:00.000000000',
'1992-02-29T00:00:00.000000000', '1992-03-01T00:00:00.000000000',
'1992-03-02T00:00:00.000000000', '1993-03-01T00:00:00.000000000',
'1993-03-02T00:00:00.000000000'], dtype='datetime64[ns]')
"""
N = sum(arg is not None for arg in [season, month, doy_bounds, date_bounds])
if N > 1:
raise ValueError(f"Only one method of indexing may be given, got {N}.")
if N == 0:
return da
def get_doys(start, end):
if start <= end:
return np.arange(start, end + 1)
return np.concatenate((np.arange(start, 367), np.arange(0, end + 1)))
if season is not None:
if isinstance(season, str):
season = [season]
mask = da.time.dt.season.isin(season)
elif month is not None:
if isinstance(month, int):
month = [month]
mask = da.time.dt.month.isin(month)
elif doy_bounds is not None:
mask = da.time.dt.dayofyear.isin(get_doys(*doy_bounds))
elif date_bounds is not None:
# This one is a bit trickier.
start, end = date_bounds
time = da.time
calendar = get_calendar(time)
if calendar not in uniform_calendars:
# For non-uniform calendars, we can't simply convert dates to doys
# conversion to all_leap is safe for all non-uniform calendar as it doesn't remove any date.
time = convert_calendar(time, "all_leap")
# values of time are the _old_ calendar
# and the new calendar is in the coordinate
calendar = "all_leap"
# Get doy of date, this is now safe because the calendar is uniform.
doys = get_doys(
to_cftime_datetime("2000-" + start, calendar).dayofyr,
to_cftime_datetime("2000-" + end, calendar).dayofyr,
)
mask = time.time.dt.dayofyear.isin(doys)
# Needed if we converted calendar, this puts back the correct coord
mask["time"] = da.time
return da.where(mask, drop=drop)
