xarray slice data by multiple dimensions

Question

I am working with the forecast data from GFS. I have written the following function to extract a timeseries from an archive of forecasts:

def time_series_from_ensemble_archive(ensemble_archive, lead_time: int=0, ensemble_member: int=0):

    data = ensemble_archive

    data['age'] = (data.validityDateTime - data.as_of_datetime).astype(np.float)
    age_idx =  data['age']==lead_time

    return data[:, ensemble_member, :, :, :].values[age_idx.T, :, :]

It works as expected:

Here is the data:

Coordinates:
  * validityDateTime    (validityDateTime) datetime64[ns] 2017-10-01 ...
  * perturbationNumber  (perturbationNumber) int32 0 1 2 3 4 5 6 7 8 9 10 11 ...
  * lon                 (lon) float64 -119.0 -118.5 -118.0 -117.5 -117.0 ...
  * lat                 (lat) float64 45.5 45.0 44.5 44.0 43.5 43.0 42.5 ...
  * as_of_datetime      (as_of_datetime) datetime64[ns] 2017-10-01 ...
Attributes:
    name:                 2 metre temperature

And with my function:

temp_ts = time_series_from_ensemble_archive(data)
temp_ts.shape
(124, 10, 20)

type(temp_ts)
numpy.ndarray

However, I feel like it is not the most 'pythonic' or 'xarrayic' approach, and would be better to return another xarray object. Suggestions for improvement here? Could someone provide a solution using expand_dims or .sel methods?

user_3888719 · Answer

xarray provides a variety of ways to index and select data. You might try indexing with dimension names, e.g.:

# select using positional & boolean indices
return data[{
    'perturbationNumber': ensemble_member,
    'validityDateTime': (data['age'] == lead_time)}]

or if lead_time is actually a positional index, just

# select using positional indices
return data[{
    'perturbationNumber': ensemble_member,
    'validityDateTime': lead_time}]

If you'd like to provide the index labels rather than their positions, you can just use the .sel or .loc methods:

# select using labels
return data.sel(
    perturbationNumber=ensemble_member,
    validityDateTime=lead_time)

or

# select using labels and boolean indices
return data.loc[{
    'perturbationNumber': ensemble_member,
    'validityDateTime': (data['age'] == lead_time)}]

Calling da.values is the step which returns the numpy array backend of the xarray data. There's no reason why your code shouldn't work with the indices you provided indexing the actual xarray DataArray (without .values).

xarray slice data by multiple dimensions

1 Answers