Calculating a rolling weighted sum using numpy-CodePudding

I am curious to know if there are any more optimal ways to compute this "rolling weighted sum" (unsure what the actual terminology is, but I will provide an example to further clarify). I am asking this because I am certain that my current code snippet is not coded in the most optimal way with respect to memory usage, and there is opportunity to improve its performance by using numpy's more advanced functions.

Example:

import numpy as np

A = np.append(np.linspace(0, 1, 10), np.linspace(1.1, 2, 30))
np.random.seed(0)

B = np.random.randint(3, size=40)   1

# list of [(weight, (lower, upper))]
d = [(1, (-0.25, -0.20)), (0.5, (-0.20, -0.10)), (2, (-0.10, 0.15))]

In Python 3.7:

## A
array([0.        , 0.11111111, 0.22222222, 0.33333333, 0.44444444,
       0.55555556, 0.66666667, 0.77777778, 0.88888889, 1.        ,
       1.1       , 1.13103448, 1.16206897, 1.19310345, 1.22413793,
       1.25517241, 1.2862069 , 1.31724138, 1.34827586, 1.37931034,
       1.41034483, 1.44137931, 1.47241379, 1.50344828, 1.53448276,
       1.56551724, 1.59655172, 1.62758621, 1.65862069, 1.68965517,
       1.72068966, 1.75172414, 1.78275862, 1.8137931 , 1.84482759,
       1.87586207, 1.90689655, 1.93793103, 1.96896552, 2.        ])

## B
array([1, 2, 1, 2, 2, 3, 1, 3, 1, 1, 1, 3, 2, 3, 3, 1, 2, 2, 2, 2, 1, 2,
       1, 1, 2, 3, 1, 3, 1, 2, 2, 3, 1, 2, 2, 2, 1, 3, 1, 3])

Expected Solution:

array([ 6. ,  6.5,  8. , 10.5, 12. , 11. , 11.5, 11.5,  6.5, 13.5, 25. ,
       27.5, 30.5, 34.5, 37.5, 36. , 35. , 35. , 34. , 34.5, 34. , 36.5,
       33. , 34. , 34.5, 34.5, 36. , 39. , 37. , 36. , 37. , 36.5, 37.5,
       39. , 36.5, 37.5, 34. , 31. , 27.5, 23. ])

The logic I want to translate into code:

Let's look at how 10.5 (the fourth element in the expected solution) is computed. d represents a collection of nested tuples with first float element weight, and second tuple element bounds (in the form of (lower, upper)).

We look at the fourth element of A (0.33333333) and apply bounds for each tuple in d. For the first tuple in d:

0.33333333   (-0.25) = 0.08333333
0.33333333   (-0.20) = 0.13333333

We go back to A to see if there are any elements between bounds (0.08333333, 0.1333333). Because the second element of A (0.11111111) falls in this range, we pull the second element of B (2) and multiply it by its weight from d (1) and add it to the second element of the expected output.

After iterating across all tuples in d, the fourth element of the expected output is computed as:

1 * 2   0.5 * 1   2 * (2   2) = 10.5

Here is my attempted code:

D = np.zeros(len(A))
for v in d:
    weight, (_lower, _upper) = v
    lower, upper = A   _lower, A   _upper
    _A = np.tile(A, (len(A), 1))
    __A = np.bitwise_and(_A > lower.reshape(-1, 1), _A < upper.reshape(-1, 1))
    D  = weight * (__A @ B)
D

Hopefully this makes sense. Please feel free to ask clarifying questions. Thanks!

CodePudding user response：

Since intervals (-0.25, -0.20), (-0.20, -0.10) and (-0.10, 0.15) are actually subintervals of partition of an interval (-0.25, 0.15) you could find indices where elements should be inserted in A to maintain order. They specify slices of B to perform addition on. In short:

partition = np.array([-0.25, -0.20, -0.10, 0.15])
weights = np.array([1, 0.5, 2])
out = []
for n in A:
    idx = np.searchsorted(A, n   partition)
    results = np.add.reduceat(B[:idx[-1]], idx[:-1])
    out.append(np.dot(results, weights))
>>> print(out)
[7.5, 7.5, 8.0, 10.5, 12.0, 11.0, 11.5, 11.5, 6.5, 13.5, 27.5, 27.5, 31.5, 35.5, 37.5, 37.0, 36.0, 35.0, 34.0, 34.5, 34.0, 36.5, 33.0, 34.0, 34.5, 34.5, 36.0, 39.0, 37.0, 36.0, 37.0, 36.5, 37.5, 39.0, 36.5, 37.5, 34.0, 31.0, 27.5, 23.0]

Note that results are wrong if there are empty slices of B

CodePudding user response：

Credits to @mathfux for providing me enough guidance. Here's the final code solution that I developed based on conversations here:

partition = np.array([-0.25, -0.20, -0.10, 0.15])
weights = np.array([1, 0.5, 2])
idx = np.searchsorted(A, partition   A[:, None])
_idx = np.lib.stride_tricks.sliding_window_view(idx, 2, axis = 1)
values = np.apply_along_axis(lambda x: B[slice(*(x))].sum(), 2, _idx)
values @ weights