Rolling window stats with frequency offset and centered window

rolling.pyx

# cython: profile=False

from cython cimport boundscheck, wraparound

# Py_ssize_t is the proper C type for Python array indices.
from cython cimport Py_ssize_t

import numpy as np
cimport numpy as cnp
from numpy cimport ndarray, double_t, int64_t
cnp.import_array()


cpdef ndarray rolling_centered_offset(ndarray[double_t, ndim=1] data,
                                      ndarray[int64_t, ndim=1] index,
                                      int64_t td,
                                      f):
    """Rolling function with a time-based offset & centered window.

    Parameters
    ----------
    data: 1D ndarray of floats
        Array for which function will be computed.
    index: 1D ndarray of date-ints
        Dates that align with the data where windows are computed.
        If you have a Pandas DatetimeIndex, use `index.asi8`
    td: numpy.timedelta64
        The offset in either direction, in the same precision as `index`.
        If you have a datetime64[ns] Pandas Index as i8 and a timedelta of
        pd.Timedelta('15D'), you would need to pass
        `np.timedelta64(15, 'D') / np.timedelta64(1, 'ns')`
    f:
        The function/callable statistic, i.e. np.nanmean

    Returns
    -------
    out: ndarray

    Example
    -------
    rolling_centered_offset(
        data=mock.values.astype(float),
        index=mock.index.asi8, td=td, f=np.nanmean
    )
    """

    cdef:
        Py_ssize_t i, j, k, n
        ndarray[double_t] out
        ndarray[int64_t] starts, ends

    n = data.shape[0]
    out = np.empty(n, dtype=float)

    starts = index.searchsorted(index - td)
    ends = index.searchsorted(index + td)

    with boundscheck(False), wraparound(False):
        for i, (j, k) in enumerate(zip(starts, ends)):
            out[i] = f(data[j:k])

    return out