how to vectorize this function to pandas or polars

Hello everyone I need to replicate the below behavior using vectorized operations. (currently written in cython and it is still too slow as I call it over many groups). However, I am not sure if this behavior is possible. Any help would be appreciated.

python version / pseudocode:

def calc_pulses(np.ndarray time, double pgt, double dt):
    curr_pulse = 1
    res = np.empty(len(time))
    begin = time[0]
    for k in range(len(time)):
        if time[k] - begin <= pgt:
            res[k] = curr_pulse
        elif time[k] - begin > pgt+dt:
            begin = time[k]
            curr_pulse += 1
            res[k] = curr_pulse
        else:
            res[k] = np.nan
    return res

cython version:

cpdef np.ndarray[np.int] calc_pulses(np.ndarray time, double pgt, double dt):
    cdef double begin
    cdef int curr_pulse = 1
    cdef Py_ssize_t k, n = len(time)
    cdef np.ndarray[np.int_t, ndim=1] res = np.empty(n, dtype=np.int)
    begin = time[0]
    for k in range(n):
        if time[k] - begin <= pgt:
            res[k] = curr_pulse
        elif time[k] - begin > pgt+dt:
            begin = time[k]
            curr_pulse += 1
            res[k] = curr_pulse
        else:
            res[k] = np.nan
    return res

Before you go through the effort of doing that, I would suggest looking at the annotated output of Cython. It can often identify inefficient things.

Here's what it shows in this case:

The yellower a line is, the more it interacts with Python. There are a bunch of yellow lines outside the loop. Those don't matter. There are also more yellow lines inside the loop. Those do matter.

To give an example of why this function might be slow, let's click on the line

if time[k] - begin <= pgt:

It shows the following C code is used to implement that:

    __pyx_t_6 = __Pyx_GetItemInt(((PyObject *)__pyx_v_time), __pyx_v_k, Py_ssize_t, 1, PyInt_FromSsize_t, 0, 1, 1); if (unlikely(!__pyx_t_6)) __PYX_ERR(0, 11, __pyx_L1_error)
    __Pyx_GOTREF(__pyx_t_6);
    __pyx_t_2 = PyFloat_FromDouble(__pyx_v_begin); if (unlikely(!__pyx_t_2)) __PYX_ERR(0, 11, __pyx_L1_error)
    __Pyx_GOTREF(__pyx_t_2);
    __pyx_t_4 = PyNumber_Subtract(__pyx_t_6, __pyx_t_2); if (unlikely(!__pyx_t_4)) __PYX_ERR(0, 11, __pyx_L1_error)
    __Pyx_GOTREF(__pyx_t_4);
    __Pyx_DECREF(__pyx_t_6); __pyx_t_6 = 0;
    __Pyx_DECREF(__pyx_t_2); __pyx_t_2 = 0;
    __pyx_t_2 = PyFloat_FromDouble(__pyx_v_pgt); if (unlikely(!__pyx_t_2)) __PYX_ERR(0, 11, __pyx_L1_error)
    __Pyx_GOTREF(__pyx_t_2);
    __pyx_t_6 = PyObject_RichCompare(__pyx_t_4, __pyx_t_2, Py_LE); __Pyx_XGOTREF(__pyx_t_6); if (unlikely(!__pyx_t_6)) __PYX_ERR(0, 11, __pyx_L1_error)
    __Pyx_DECREF(__pyx_t_4); __pyx_t_4 = 0;
    __Pyx_DECREF(__pyx_t_2); __pyx_t_2 = 0;
    __pyx_t_11 = __Pyx_PyObject_IsTrue(__pyx_t_6); if (unlikely(__pyx_t_11 < 0)) __PYX_ERR(0, 11, __pyx_L1_error)
    __Pyx_DECREF(__pyx_t_6); __pyx_t_6 = 0;
    if (__pyx_t_11) {
/* … */
      goto __pyx_L5;
    }

What is it doing in this huge block of code?

• It gets an item from the time array. time doesn't have a type for its elements. For that reason, item might be a double, but it could be a Fraction object, or a Decimal object. Cython must cover all these possibilities by calling __Pyx_GetItemInt().
• Because time[k] might not be a double, it must convert begin into a double object with PyFloat_FromDouble(), so that the subtraction can be done.
• Since time[k] could be any object, that object might have overridden how subtraction works. Therefore, it calls PyNumber_Subtract(), which can handle any kind of subtraction.
• Since the result of that subtraction could be any object, it could have overridden how comparisons work. Thefore pgt must also be converted into a double object, and PyObject_RichCompare() must be called, which can handle any kind of comparison.
• Comparison functions normally return True or False, but that is technically not required. Technically, they could return any truthy value as well. Cython must cover this possibility by converting the object to a boolean with __Pyx_PyObject_IsTrue().

Exhausting, right? That's why this is so slow.

Fortunately, there is something you can do to make this much faster: add a type and number of dimensions to the time argument.

cpdef np.ndarray[np.int] calc_pulses(np.ndarray[np.double_t, ndim=1] time, double pgt, double dt):

This gives a 10x performance improvement in my benchmark, because every one of those steps are easier when Cython knows the type.

No types
491 µs ± 6.38 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

With types
39.8 µs ± 274 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

There are other optimizations you can do, but this is the most important.