Fast K-step discount in numpy / scipy / python

Question

x has the form [batch_size, n_time], where the parties are independent

If k = 3, d = discount_rate. Pseudocode:

x[:,i] = x[:,i] + x[:,i+1]*(d**1) + x[:,i+2]*(d**2) + x[:,i+3]*(d**3)

Here is the working code, but it is very slow. I will perform this function millions of times, so I hope for a faster implementation

import numpy as np

def k_step_discount(x, k, discount_rate):
    n_time = x.shape[1]
    k_include_cur = k + 1 # k excludes current timestep
    for i in range(n_time):
        k_cur = min(n_time - i, k_include_cur) # prevent out of bounds
        for j in range(1, k_cur):
            x[:, i] += x[:, i+j] * (discount_rate ** j)
    return x

x = np.array([
    [0,0,0,1,0,0],
    [0,1,2,3,4,5.]
])

y = k_step_discount(x+0, k=2, discount_rate=.9)
print('x\n{}\ny\n{}'.format(x, y))


>> x
   [[ 0.  0.  0.  1.  0.  0.]
    [ 0.  1.  2.  3.  4.  5.]]
>> y
   [[  0.     0.81   0.9    1.     0.     0.  ]
    [  2.52   5.23   7.94  10.65   8.5    5.  ]]

Scipy function, which is similar:

import scipy.signal
import numpy as np

x = np.array([[0,0,0,1,0,0.]])
discount_rate = .9

y = np.flip(scipy.signal.lfilter([1], [1, -discount_rate], np.flip(x+0, 1), axis=1), 1)
print('x\n{}\ny\n{}'.format(x, y))
>> x
   [[ 0.  0.  0.  1.  0.  0.]]
>> y
   [[ 0.729  0.81   0.9    1.     0.     0.   ]]

However, it throws off to the end of n_time, and not just for k steps

I am also interested in discounting K-step without lots if it will be easier / faster

import numpy as np

def k_step_discount_no_batch(x, k, discount_rate):
    n_time = x.shape[0]
    k_include_cur = k + 1 # k excludes current timestep
    for i in range(n_time):
        k_cur = min(n_time - i, k_include_cur) # prevent out of bounds
        for j in range(1, k_cur):
            x[i] += x[i+j] * (discount_rate ** j)
    return x

x = np.array([8,0,0,0,1,2.])
y = k_step_discount_no_batch(x+0, k=2, discount_rate=.9)
print('x\n{}\ny\n{}'.format(x, y))

>> x
   [ 8.  0.  0.  0.  1.  2.]
>> y
   [ 8.    0.    0.81  2.52  2.8   2.  ]

Similar function no_batch scipy

import scipy.signal
import numpy as np

x = np.array([8,0,0,0,1,2.])    
discount_rate = .9

y = scipy.signal.lfilter([1], [1, -discount_rate], x[::-1], axis=0)[::-1]
print('x\n{}\ny\n{}'.format(x, y))
>> x
   [ 8.  0.  0.  0.  1.  2.]

>> y
   [ 9.83708  2.0412   2.268    2.52     2.8      2.     ]

Answer 1

2D convolution. , ​​ 2D, discount_rate. convolution, ​​ , , .

, -

from scipy.signal import convolve2d as conv2d
import numpy as np

def k_step_discount(x, k, discount_rate, is_batch=True):
    if is_batch:
        kernel = discount_rate**np.arange(k+1)[::-1][None]
        return conv2d(x,kernel)[:,k:]
    else:
        kernel = discount_rate**np.arange(k+1)[::-1]
        return np.convolve(x, kernel)[k:]

-

In [190]: x
Out[190]: 
array([[ 0.,  0.,  0.,  1.,  0.,  0.],
       [ 0.,  1.,  2.,  3.,  4.,  5.]])

# Proposed method
In [191]: k_step_discount_conv2d(x, k=2, discount_rate=0.9)
Out[191]: 
array([[  0.  ,   0.81,   0.9 ,   1.  ,   0.  ,   0.  ],
       [  2.52,   5.23,   7.94,  10.65,   8.5 ,   5.  ]])

# Original loopy method
In [192]: k_step_discount(x, k=2, discount_rate=.9)
Out[192]: 
array([[  0.  ,   0.81,   0.9 ,   1.  ,   0.  ,   0.  ],
       [  2.52,   5.23,   7.94,  10.65,   8.5 ,   5.  ]])

In [206]: x = np.random.randint(0,9,(100,1000)).astype(float)

In [207]: %timeit k_step_discount_conv2d(x, k=2, discount_rate=0.9)
1000 loops, best of 3: 1.27 ms per loop

In [208]: %timeit k_step_discount(x, k=2, discount_rate=.9)
100 loops, best of 3: 4.83 ms per loop

k's:

In [215]: x = np.random.randint(0,9,(100,1000)).astype(float)

In [216]: %timeit k_step_discount_conv2d(x, k=20, discount_rate=0.9)
100 loops, best of 3: 5.44 ms per loop

In [217]: %timeit k_step_discount(x, k=20, discount_rate=.9)
10 loops, best of 3: 44.8 ms per loop

, k's!

---

@Eric, scipy.ndimage.filters 1D convolution .

Scipy 2D 1D -

from scipy.ndimage.filters import convolve1d as conv1d

def using_conv2d(x, k, discount_rate):
    kernel = discount_rate**np.arange(k+1)[::-1][None]
    return conv2d(x,kernel)[:,k:]

def using_conv1d(x, k, discount_rate):
    kernel = discount_rate**np.arange(k+1)[::-1]
    return conv1d(x,kernel, mode='constant', origin=k//2)

-

In [100]: x = np.random.randint(0,9,(100,1000)).astype(float)

In [101]: out1 = using_conv2d(x, k=20, discount_rate=0.9)
     ...: out2 = using_conv1d(x, k=20, discount_rate=0.9)
     ...: 

In [102]: np.allclose(out1, out2)
Out[102]: True

In [103]: %timeit using_conv2d(x, k=20, discount_rate=0.9)
100 loops, best of 3: 5.27 ms per loop

In [104]: %timeit using_conv1d(x, k=20, discount_rate=0.9)
1000 loops, best of 3: 1.43 ms per loop