Python xarray.CFTimeIndex() Examples

def get_multiple_lead_cftime_shift_args(units, leads):
    """Returns ``CFTimeIndex.shift()`` offset increment for an arbitrary number of
    leads.

    Args:
        units (str): Units associated with the lead dimension. Must be one of
            years, seasons, months, weeks, pentads, days.
        leads (list, array, xr.DataArray of ints): Leads to return offset for.

    Returns:
        n (tuple of ints): Number of units to shift for ``leads``. ``value`` for
            ``CFTimeIndex.shift(value, str)``.
        freq (str): Pandas frequency alias. ``str`` for
            ``CFTimeIndex.shift(value, str)``.
    """
    n_freq_tuples = [get_lead_cftime_shift_args(units, lead) for lead in leads]
    n, freq = list(zip(*n_freq_tuples))
    return n, freq[0] 

def shift_cftime_singular(cftime, n, freq):
    """Shifts a singular ``cftime`` by the desired frequency.

    This directly pulls the ``shift`` method from ``CFTimeIndex`` in ``xarray``. This
    is useful if you need to shift a singular ``cftime`` by some offset, but are not
    working with a full ``CFTimeIndex``.

    Args:
        cftime (``cftime``): ``cftime`` object to shift.
        n (int): Number of steps to shift by.
        freq (str): Frequency string, per ``pandas`` convention.

    See:
    https://github.com/pydata/xarray/blob/master/xarray/coding/cftimeindex.py#L376.
    """
    if not isinstance(n, int):
        raise TypeError(f"'n' must be an int, got {n}.")
    if not isinstance(freq, str):
        raise TypeError(f"'freq' must be a str, got {freq}.")
    return cftime + n * to_offset(freq) 

def get_time_decoded(self, midpoint=True):
        """Return time decoded.
        """
        # to compute a time midpoint, we need a time_bound variable
        if midpoint and self.time_bound is None:
            raise ValueError('cannot compute time midpoint w/o time bounds')

        if midpoint:
            time_data = self.time_bound.mean(self.tb_dim)

        else:
            # if time has already been decoded and there's no year_offset,
            # just return the time as is
            if self.isdecoded(self.time):
                if self.year_offset is None:
                    return self.time

                # if we need to un-decode time to apply the year_offset,
                time_data = self.get_time_undecoded()

            # time has not been decoded
            else:
                time_data = self.time

        if self.year_offset is not None:
            time_data += cftime.date2num(
                datetime(int(self.year_offset), 1, 1),
                units=self.time_attrs['units'],
                calendar=self.time_attrs['calendar'],
            )
        time_out = self.time.copy()
        time_out.data = xr.CFTimeIndex(
            cftime.num2date(
                time_data,
                units=self.time_attrs['units'],
                calendar=self.time_attrs['calendar'],
                only_use_cftime_datetimes=True,
            )
        )
        return time_out 

def _same_verifs_alignment(init_lead_matrix, valid_inits, all_verifs, leads, n, freq):
    """Returns initializations and verification dates, maintaining a common verification
    window at all leads.

    See ``return_inits_and_verif_dates`` for descriptions of expected variables.
    """
    common_set_of_verifs = [
        i for i in all_verifs if (i == init_lead_matrix).any('time').all('lead')
    ]
    if not common_set_of_verifs:
        raise CoordinateError(
            'A common set of verification dates cannot be found for the '
            'initializations and verification data supplied. Change `alignment` to '
            "'same_inits' or 'maximize'."
        )
    # Force to CFTimeIndex for consistency with `same_inits`
    verif_dates = xr.concat(common_set_of_verifs, 'time').to_index()
    inits_that_verify_with_verif_dates = init_lead_matrix.isin(verif_dates)
    inits = {
        lead: valid_inits.where(
            inits_that_verify_with_verif_dates.sel(lead=lead), drop=True
        )
        for lead in leads
    }
    verif_dates = {lead: verif_dates for lead in leads}
    return inits, verif_dates 

def _construct_init_lead_matrix(forecast, n, freq, leads):
    """Returns xr.DataArray of "real time" (init + lead) over all inits and leads.

    Arguments:
        forecast (``xarray object``): Prediction ensemble with ``init`` dim renamed to
            ``time`` and containing ``lead`` dim.
        n (tuple of ints): Number of units to shift for ``leads``. ``value`` for
            ``CFTimeIndex.shift(value, str)``.
        freq (str): Pandas frequency alias. ``str`` for
            ``CFTimeIndex.shift(value, str)``.
        leads (list, array, xr.DataArray of ints): Leads to return offset for.

    Returns:
        init_lead_matrix (``xr.DataArray``): DataArray with x=inits and y=lead with
            values corresponding to "real time", or ``init + lead`` over all inits and
            leads.
    """
    # Note that `init` is renamed to `time` in compute functions.
    init_lead_matrix = xr.concat(
        [
            xr.DataArray(
                shift_cftime_index(forecast, 'time', n, freq),
                dims=['time'],
                coords=[forecast['time']],
            )
            for n in n
        ],
        'lead',
    )
    init_lead_matrix['lead'] = leads
    return init_lead_matrix 

def get_lead_cftime_shift_args(units, lead):
    """Determines the date increment to use when adding the lead time to init time based
    on the units attribute.

    Args:
        units (str): Units associated with the lead dimension. Must be
            years, seasons, months, weeks, pentads, days.
        lead (int): Increment of lead being computed.

    Returns:
       n (int): Number of units to shift. ``value`` for
           ``CFTimeIndex.shift(value, str)``.
       freq (str): Pandas frequency alias. ``str`` for
           ``CFTimeIndex.shift(value, str)``.
    """
    lead = int(lead)

    d = {
        # Currently assumes yearly aligns with year start.
        'years': (lead, 'YS'),
        'seasons': (lead * 3, 'MS'),
        # Currently assumes monthly aligns with month start.
        'months': (lead, 'MS'),
        'weeks': (lead * 7, 'D'),
        'pentads': (lead * 5, 'D'),
        'days': (lead, 'D'),
    }

    try:
        n, freq = d[units]
    except KeyError:
        print(f'{units} is not a valid choice.')
        print(f'Accepted `units` values include: {d.keys()}')
    return n, freq 

def test_pandas_datetime_converted_to_cftime():
    """Tests that a pd.DatetimeIndex is converted to xr.CFTimeIndex."""
    inits = pd.date_range('1990', '2000', freq='YS')
    da = xr.DataArray(np.random.rand(len(inits)), dims='init', coords=[inits])
    new_inits = convert_time_index(da, 'init', '')
    assert isinstance(new_inits['init'].to_index(), xr.CFTimeIndex) 

def test_int64_converted_to_cftime():
    """Tests the xr.Int64Index is converted to xr.CFTimeIndex."""
    inits = np.arange(1990, 2000)
    da = xr.DataArray(np.random.rand(len(inits)), dims='init', coords=[inits])
    new_inits = convert_time_index(da, 'init', '')
    assert isinstance(new_inits['init'].to_index(), xr.CFTimeIndex) 

def test_float64_converted_to_cftime():
    """Tests the xr.Float64Index is converted to xr.CFTimeIndex."""
    inits = np.arange(1990, 2000) * 1.0
    da = xr.DataArray(np.random.rand(len(inits)), dims='init', coords=[inits])
    new_inits = convert_time_index(da, 'init', '')
    assert isinstance(new_inits['init'].to_index(), xr.CFTimeIndex) 

def test_convert_time_index_does_not_overwrite():
    """Tests that `convert_time_index` does not overwrite the original index."""
    inits = np.arange(1990, 2000)
    da = xr.DataArray(np.random.rand(len(inits)), dims='init', coords=[inits])
    new_inits = convert_time_index(da, 'init', '')
    assert isinstance(da.init.to_index(), pd.Int64Index)
    assert isinstance(new_inits.init.to_index(), xr.CFTimeIndex) 

def read_netcdf(data_handle, domain=None, is_worker=False,
                start=None, stop=None, calendar='standard',
                var_dict=None) -> xr.Dataset:
    """Read in a NetCDF file"""
    if '*' in data_handle:
        ds = xr.open_mfdataset(data_handle)
    else:
        ds = xr.open_dataset(data_handle)

    if domain is not None:
        ds = ds.sel({k: domain[k]
                     for k in list(domain.dims.keys())
                     if k in list(ds.dims.keys())})
    else:
        dims_wo_coords = set(ds.dims) - set(ds.coords)
        for d in dims_wo_coords:
            if is_worker:
                logger = logging.getLogger('MetSim')
                logger.warning(
                    'Setting sequential coordinate on dimension {}'.format(d))
            ds[d] = np.arange(0, len(ds[d]))

    if 'time' in ds.coords:
        if isinstance(ds.indexes['time'], xr.CFTimeIndex):
            ds['time'] = ds.indexes['time'].to_datetimeindex()
        ds['time'] = (ds.indexes['time'] -
                      pd.Timedelta('11H59M59S')).round('D')

    if var_dict is not None:
        var_list = list(var_dict.keys())
        ds = ds[var_list]
        ds = ds.rename(var_dict)

    if start is not None or stop is not None:
        ds = ds.sel(time=slice(start, stop))
        dates = ds.indexes['time']
        ds['day_of_year'] = xr.Variable(('time', ), dates.dayofyear)

    return ds 

def convert_time_index(xobj, time_string, kind, calendar=HINDCAST_CALENDAR_STR):
    """Converts incoming time index to a standard xr.CFTimeIndex.

    Args:
        xobj (xarray object): Dataset or DataArray with a time dimension to convert.
        time_string (str): Name of time dimension.
        kind (str): Kind of object for error message.
        calendar (str): calendar to set time dimension to.

    Returns:
        Dataset or DataArray with converted time dimension. If incoming time index is
        ``xr.CFTimeIndex``, returns the same index. If ``pd.DatetimeIndex``, converts to
        ``cftime.ProlepticGregorian``. If ``pd.Int64Index`` or ``pd.Float64Index``,
        assumes annual resolution and returns year-start ``cftime.ProlepticGregorian``.

    Raises:
        ValueError: If ``time_index`` is not an ``xr.CFTimeIndex``, ``pd.Int64Index``,
            ``pd.Float64Index``, or ``pd.DatetimeIndex``.
    """
    xobj = xobj.copy()  # Ensures that main object index is not overwritten.
    time_index = xobj[time_string].to_index()

    if not isinstance(time_index, xr.CFTimeIndex):

        if isinstance(time_index, pd.DatetimeIndex):
            # Extract year, month, day strings from datetime.
            time_strings = [str(t) for t in time_index]
            split_dates = [d.split(' ')[0].split('-') for d in time_strings]

        # If Float64Index or Int64Index, assume annual and convert accordingly.
        elif isinstance(time_index, pd.Float64Index) | isinstance(
            time_index, pd.Int64Index
        ):
            warnings.warn(
                'Assuming annual resolution due to numeric inits. '
                'Change init to a datetime if it is another resolution.'
            )
            # TODO: What about decimal time? E.g. someone has 1955.5 or something?
            dates = [str(int(t)) + '-01-01' for t in time_index]
            split_dates = [d.split('-') for d in dates]
            if 'lead' in xobj.dims:
                # Probably the only case we can assume lead units, since `lead` does not
                # give us any information on this.
                xobj['lead'].attrs['units'] = 'years'

        else:
            raise ValueError(
                f'Your {kind} object must be pd.Float64Index, '
                'pd.Int64Index, xr.CFTimeIndex or '
                'pd.DatetimeIndex.'
            )
        # TODO: Account for differing calendars. Currently assuming `Gregorian`.
        cftime_dates = [
            getattr(cftime, calendar)(int(y), int(m), int(d))
            for (y, m, d) in split_dates
        ]
        time_index = xr.CFTimeIndex(cftime_dates)
        xobj[time_string] = time_index

    return xobj