带有日期时间的 Numba

Question

是否可以使用 Numba 加速以下功能？

@jit(nopython=True)
def loanDates(startDate, endDate, freq):
    startDate=pd.to_datetime(startDate)
    endDate=pd.to_datetime(endDate)
    if int(12/freq)==12/freq:

        FirstDate=pd.to_datetime(dt.datetime(startDate.year,startDate.month,1))

        ts=pd.date_range(start=FirstDate,end=endDate,freq=str(-int(12/freq))+"MS")
        ts=list(ts)
        for i in range(len(ts)):
            if startDate.day>ts[i].days_in_month:
                ts[i]=ts[i]+pd.DateOffset(days=ts[i].days_in_month-1)
            else:
                ts[i]=ts[i]+pd.DateOffset(days=startDate.day-1)
        ts=pd.DatetimeIndex(ts)

我收到这个错误：

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Untyped global name 'str': cannot determine Numba type of <class 'type'>

File "<ipython-input-58-3eac94417232>", line 16:
def loanDates(startDate, endDate, freq):
    <source elided>

        ts=pd.date_range(start=FirstDate,end=endDate,freq=str(-int(12/freq))+"MS")

据我了解，日期时间和 Numba 存在问题。我尝试使用而不是 pd.date_range 方法 numpy arange 但它没有满足我的需要：

x=np.arange(np.datetime64(sd), np.datetime64(ed),str(-int(12/fq))+"MS")

给出错误

ValueError: Could not convert object to NumPy timedelta

那么有没有可能使用 numba 来实现这种功能？

Answer 1

慢的部分是你的 for 循环，你不需要 Numba 来优化它。而不是这个：

    ts=list(ts)
    for i in range(len(ts)):
        if startDate.day>ts[i].days_in_month:
            ts[i]=ts[i]+pd.DateOffset(days=ts[i].days_in_month-1)
        else:
            ts[i]=ts[i]+pd.DateOffset(days=startDate.day-1)

做矢量化，像这样：

    overflow = startDate.day > ts.days_in_month # array of bools
    offsets = (ts.days_in_month-1).where(overflow, startDate.day-1)
    ts += pd.DateOffset(days=offsets)

或进一步简化：

    offsets = np.minimum(ts.days_in_month, startDate.day) - 1
    ts += pd.DateOffset(days=offsets)

这应该至少快 10 倍。

带有日期时间的 Numba

Numba with datetime

python

datetime

pandas

numba