Select xarray/pandas index based on specific months

Question

I have an xarray DataArray that I want to select the months April, May, June (similar to time.season=='JJA') for an entire time series.

Its structured like:

<xarray.DataArray 't2m' (time: 492, latitude: 81, longitude: 141)>

I have been previously selecting JJA by:

seasonal_data =temp_data.sel(time=temp_data['time.season']=='JJA')

I would like to do the same thing but with the months 'AMJ' instead. I can add any details that I might be missing.

Thanks

Answer 1

The easiest way to select custom months is to use boolean masks, e.g.,

def is_amj(month):
    return (month >= 4) & (month <= 6)

seasonal_data = temp_data.sel(time=is_amj(temp_data['time.month']))

Note that you need to use the bitwise operators like & or | because Python's built-ins and and or don't work on vectors. Also, you need the parentheses because bitwise operators have higher precedence than comparisons.