Conditional indexing with object arrays

user_5717854 0
Viewed 186

I have a 1D numpy array of datetime objects, and I would like to grab all the objects which have a certain month. I can do this by iterating over the entire array:

    new_times = []
    for time in times:
        if time.month == 1:
            new_times.append(time)

but this isn't very efficient and doesn't seem very pythonic. If this was an array of integer months, I could do something like

    arr = np.array([1,1,2,3,4,5,6,1,7,4,8,1])
    new_arr = arr[np.where(arr == 1)]

but this method doesn't work when calling attributes of the array's objects. Is there a nice way of searching for all numpy array elements which have a certain object attribute?

numpy datetime arrays python
3 Answers

In general, I don't know you can do any better than using vectorize or frompyfunc (and then use standard numpy techniques for filtering etc.):

For example,

import datetime

M = np.array([datetime.datetime(1980,i,i) for i in range(1,4)])
M
# array([datetime.datetime(1980, 1, 1, 0, 0),
#        datetime.datetime(1980, 2, 2, 0, 0),
#        datetime.datetime(1980, 3, 3, 0, 0)], dtype=object)

import operator as op

np.frompyfunc(op.attrgetter("month"),1,1)(M)
# array([1, 2, 3], dtype=object)
np.vectorize(op.attrgetter("month"),otypes=(int,))(M)
# array([1, 2, 3])

In your specific case:

Mnp = M.astype("M8[M]")
Mnp - Mnp.astype("M8[Y]") + 1
# array([1, 2, 3], dtype='timedelta64[M]')
user_72073920

using a pandas.Series might be a good helper here,

from datetime import datetime
import numpy as np
import pandas as pd

a = np.array([datetime(2020, m, 1) for m in range(1, 5)])
s = pd.Series(a)

Now you can access the month as

s.dt.month
# 0    1
# 1    2
# 2    3
# 3    4
# dtype: int64

a[s.dt.month < 3]
# array([datetime.datetime(2020, 1, 1, 0, 0),
#        datetime.datetime(2020, 2, 1, 0, 0)], dtype=object)

This will work for an array of datetime.datetime objects as well as np.datetime64. But as always, convenience comes at a cost, so the list comprehension (in case of datetime.datetime), np.vectorize and np.frompyfunc will most likely run faster.

user_101974180

Using hpaulj's method here using a list comprehension is easy and fast:

>>> import numpy as np
>>> import datetime as dt
>>> a=np.array([dt.datetime(month=1,day=2,year=2020),dt.datetime(month=2,day=12,year=2019),dt.datetime(month=8,day=26,year=1952)])

>>> [e for e in a.tolist() if e.month<5]
[datetime.datetime(2020, 1, 2, 0, 0), datetime.datetime(2019, 2, 12, 0, 0)]

If you just want the indices:

>>> [i for i,e in enumerate(a.tolist()) if e.month<5]
[0, 1]
user_2986070
Related