Dynamically discounted cumulative sum in Numpy

This is a self-answer to guide future visitors with concrete advice.

TL;DR; prefer Numba, avoid pure Python

It is possible to get a numerically stable version of Riccardo Bucco's answer by performing the calculations in log-domain. Assuming d[0] == 1, the implementation looks like this:

def f_numpy_stable(x, d):
    p = np.cumsum(np.log(d))
    return np.exp(p + np.logaddexp.accumulate(np.log(x) - p))

Comparing it agains the reference implementation in pure Python with random inputs gives outputs that pass an np.allclose test, which I consider to be stable enough. Intuitively, such complicated calculations should be slower than a simple loop that does not involve conversions to and from log-domain. Let's test that intuition.

There are five solutions when collating the above with the question as well as the comments and answer (thank you Riccardo Bucco):

1. Native Python (slow, readable)
2. JIT-compile with Numba (fast, readable)
3. Precompile with Cython (fast)
4. Decompose into cumulative product and sum (pure Numpy, potentially numerically unstable)
5. Decompose but calculate in log-space (pure Numpy)

I coded up and tested all five with fairly large arrays (10k–100M entries), and the result was (on my Intel MacBook pro) that Numba is fastest, followed by Cython (up to 50% slower than Numba, less for larger data sets), followed by Riccardo Bucco's pure Numpy solution (about 3x slower than Numba), followed (surprisingly) by the numerically stable version of the pure Numpy solution (less than 10x slower than the unstable version), with the slowest by far being the pure Python implementation. Here is the full table (all times in seconds).

Here is the code I used for all three implementations. I made sure to trigger Numba's JIT-compilation with a separate call before running timeit.

def f_python(x, d):
    result = np.empty_like(x)
    result[0] = x[0]
    for i in range(1,x.shape[0]):
        result[i] = result[i-1] * d[i] + x[i]
    return result

@numba.jit
def f_numba(x, d):
    result = np.empty_like(x)
    result[0] = x[0]
    for i in range(1,x.shape[0]):
        result[i] = result[i-1] * d[i] + x[i]
    return result

def f_numpy(x, d):
    result = np.cumprod(d)
    return result * np.cumsum(x / result)

def f_numpy_stable(x, d):
    p = np.cumsum(np.log(d))
    return np.exp(p + np.logaddexp.accumulate(np.log(x) - p))

With the Cython implementation being compiled in a separate notebook cell:

%%cython
import numpy as np
cimport numpy as np

cpdef np.ndarray[np.float64_t, ndim=1] f_cython(np.ndarray[np.float64_t, ndim=1] x, np.ndarray[np.float64_t, ndim=1] d):
    cdef:
        int i = 0
        int N = x.shape[0]
        np.ndarray[np.float64_t, ndim=1] result = np.empty_like(x)
    result[0] = x[0]
    for i in range(1, N):
        result[i] = result[i-1] * d[i] + x[i]
    return result

I was a bit surprised that Numba did better than Cython. Given that the decomposition into separate cumulative product and sum is either numerically unstable or fairly slow (although faster than pure Python), my advice would definitely be to avoid it (especially for large data sets) and go with Cython or Numba. Given that the easier solution to just import numba and add a decorator is also faster, it seems like a slam dunk choice to me.

Happy to be shown all the ways I could improve the Cython implementation! :-)