如何在 numba 中创建 datetime64[D]

Question

我需要将日期传递给 numba 函数。

将它们作为 .astype('datetime64[D]') 传入效果很好。但我也需要在函数内部创建一个纪元日期。

import numpy as np
import pandas as pd
import numba
from numba import jit
from datetime import datetime, timedelta

def datetime_range(start, end, delta):
    current = start
    while current < end:
        yield current
        current += delta



@jit(nopython=True)
def myfunc(dts):
    epoch = np.datetime64('1970-01-01').astype('datetime64[D]')
    if epoch == dts[0]:
        n = 1
    return epoch


dts = [dt for dt in
       datetime_range(datetime(2016, 9, 1, 7), datetime(2016, 9, 2,7),
       timedelta(minutes=15))]

pandas_df = pd.DataFrame(index = dts)

res = myfunc(pandas_df.index.values.astype('datetime64[D]'))

print(res)

我得到错误：

numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Unknown attribute 'astype' of type datetime64[]

File "test5.py", line 17:
def myfunc(dts):
    epoch = np.datetime64('1970-01-01').astype('datetime64[D]')
    ^

During: typing of get attribute at C:/Users/PUser/PycharmProjects/pythonProjectTEST/test5.py (17)

File "test5.py", line 17:
def myfunc(dts):
    epoch = np.datetime64('1970-01-01').astype('datetime64[D]')
    ^

我怎样才能做到这一点

Answer 1

您的问题可能与this documented issue 和numba 有关。

---

第一个解决方法是在 jit 函数之外定义 epoch：

def myfunc(dts):
    @jit(nopython=True)
    def wrapper(dts, epoch):
        if epoch == dts[0]:
            n = 1
        return epoch
    epoch = np.datetime64('1970-01-01').astype('datetime64[D]')
    return wrapper(dts, epoch)

---

我想到的另一种 hacky 解决方案是将您的日期呈现为字符串，然后再将它们提供给 myfunc：

res = myfunc(np.datetime_as_string(pandas_df.index.values, unit='D'))

并在myfunc 内部定义epoch ='1970-01-01'。

您最终可以在此之后添加一个后处理步骤，将您的字符串转换回datetime64 或它们需要的任何内容。