Search code examples
pythonnumpynumba

Unable to find out how to specify datetime dtype to Numba @guvectorize


I would like to handle array of datetime values in Numba. I tried so on a small example, but I am unable to have it working.

import numba as nb
from numba import guvectorize

@guvectorize('void(nb.types.NPDatetime("ns")[:],nb.types.NPDatetime("ns")[:])',
             '(m)->(m)')
def ts_copy(ts, ts_c):
    for idx, t in np.ndenumerate(ts):
        idx, = idx
        ts_c[idx] = t
        
ts = pd.date_range(start='2021/1/1 08:00', end='2021/1/1 18:00', freq='1H').to_numpy()
ts_c = np.empty(ts.size, dtype='datetime64[ns]')

res = ts_copy(ts, ts_c)

I get following error message:

Traceback (most recent call last):

  File "<ipython-input-39-1550d3aa0145>", line 3, in <module>
    def ts_copy(ts, ts_c):

  File "/home/pierre/anaconda3/lib/python3.8/site-packages/numba/np/ufunc/decorators.py", line 179, in wrap
    guvec.add(fty)

  File "/home/pierre/anaconda3/lib/python3.8/site-packages/numba/np/ufunc/ufuncbuilder.py", line 212, in add
    cres, args, return_type = _compile_element_wise_function(

  File "/home/pierre/anaconda3/lib/python3.8/site-packages/numba/np/ufunc/ufuncbuilder.py", line 144, in _compile_element_wise_function
    cres = nb_func.compile(sig, **targetoptions)

  File "/home/pierre/anaconda3/lib/python3.8/site-packages/numba/np/ufunc/ufuncbuilder.py", line 93, in compile
    return self._compile_core(sig, flags, locals)

  File "/home/pierre/anaconda3/lib/python3.8/site-packages/numba/np/ufunc/ufuncbuilder.py", line 124, in _compile_core
    args, return_type = sigutils.normalize_signature(sig)

  File "/home/pierre/anaconda3/lib/python3.8/site-packages/numba/core/sigutils.py", line 24, in normalize_signature
    parsed = _parse_signature_string(sig)

  File "/home/pierre/anaconda3/lib/python3.8/site-packages/numba/core/sigutils.py", line 14, in _parse_signature_string
    return eval(signature_str, {}, types.__dict__)

  File "<string>", line 1, in <module>

NameError: name 'nb' is not defined

Please, any idea? Thanks a lot for your help. Bests,


Solution

  • guvectorize does not work with datetime types. Convert to float timestamp.

    import numba as nb
    from numba import guvectorize
    import numpy as np
    import pandas as pd
    import arrow
    
    @guvectorize([(nb.types.float64[:],nb.types.float64[:])],'(m)->(m)')
    def ts_copy(ts, ts_c):
        for idx, t in np.ndenumerate(ts):
            idx, = idx
            ts_c[idx] = t
    
    
    ts = pd.date_range(start='2021/1/1 08:00', end='2021/1/1 18:00', freq='1H').to_numpy()
    ts_c = np.empty(ts.size, dtype='datetime64[ns]')
    
    def e(_list):
        return np.array([arrow.get(str(_)).float_timestamp for _ in _list])
    
    
    ts = e(ts)
    ts_c = e(ts_c)
    
    
    res = ts_copy(ts, ts_c)
    
    for each in res:
        print(arrow.get(each))