CODEWITHSUNDEEP
pythonpandasnumpypython-xarray
bchate
bchate

Reputation: 145

Select xarray/pandas index based on a list of months

I successfully implemented shoyer answer from Select xarray/pandas index based on specific months.

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

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

Unfortunately I need more flexibility in months selections (e.g. january and december, or february to november, or january,march,may,...). I imagine using a list of months.

I tried to modify the code as follows

sel_months = [1,3,5] #in the case of january,march and may

def to_keep(month):
    return (month in sel_months)

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

but I get the following message

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

and did not found a proper way to implement the a.any() recommendation.

Any help would be welcome.

Upvotes: 0

Views: 1352

Answers (3)

Gabe
Gabe

Reputation: 549

I modified shoyer's answer (from my question) and created this little function. t_months is a list like you have of months. t2m is likely the same structure as your temp_data. It works well for me, though there is likely a more elegant solution. 

def extract_months(month, start, end):
    return (month >= start) & (month <= end)


t2m_months = t2m_slice.sel(time=extract_months(t2m_slice['time.month'],t_months[0],t_months[-1]))

Upvotes: 0

bchate
bchate

Reputation: 145

Both COLDSPEED comments and whagi answer are working, but https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.in1d.html recommends to use np.isin, then I used np.isin(month, sel_months)

Upvotes: 0

whagi
whagi

Reputation: 21

I've been considering this problem lately, and I agree this method is kind of restrictive (sometimes you don't want exactly 3 months...). So adding more arguments to your function every time you need something different is not a good solution.

But, as you can see here there's also another way to do it, which is simply:

months = [1,2,3,4] # for example season = data.sel(time = np.in1d( data['time.month'], months))

And that will do.

Upvotes: 1

Related Questions