python - 带有日期时间的 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 存在问题。我尝试使用 numpy arange 代替 pd.date_range 方法，但它没有满足我的需要:

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 倍。