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_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)