Is it possible to perform a rolling linear regression with Vaex?

[LamBadaMan]

Hi,

from my understanding, vaex.rolling.Rolling creates a rolling window object and Vaex provides a wrapper for any scikit-learn estimator. However, I do not know how I can get a rolling window regression to work. Would it be possible to provide me with exemplary code to implement a rolling linear regression?

[maartenbreddels]

Interesting question. I don't think this will give stellar performance, but it may be a good first start:

import vaex
import numpy as np
x = np.arange(100)
df = vaex.from_arrays(x=x)
df2 = df.rolling(10, fill_value=0).array()
print(df2)
==== output ====
#    x
0    array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
1    array([0, 0, 0, 0, 0, 0, 0, 0, 0, 1])
2    array([0, 0, 0, 0, 0, 0, 0, 0, 1, 2])
3    array([0, 0, 0, 0, 0, 0, 0, 1, 2, 3])
4    array([0, 0, 0, 0, 0, 0, 1, 2, 3, 4])
5    array([0, 0, 0, 0, 0, 1, 2, 3, 4, 5])
6    array([0, 0, 0, 0, 1, 2, 3, 4, 5, 6])
7    array([0, 0, 0, 1, 2, 3, 4, 5, 6, 7])
8    array([0, 0, 1, 2, 3, 4, 5, 6, 7, 8])
9    array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
...  ...
90   array([81, 82, 83, 84, 85, 86, 87, 88, 89, 90])
91   array([82, 83, 84, 85, 86, 87, 88, 89, 90, 91])
92   array([83, 84, 85, 86, 87, 88, 89, 90, 91, 92])
93   array([84, 85, 86, 87, 88, 89, 90, 91, 92, 93])
94   array([85, 86, 87, 88, 89, 90, 91, 92, 93, 94])
95   array([86, 87, 88, 89, 90, 91, 92, 93, 94, 95])
96   array([87, 88, 89, 90, 91, 92, 93, 94, 95, 96])
97   array([88, 89, 90, 91, 92, 93, 94, 95, 96, 97])
98   array([89, 90, 91, 92, 93, 94, 95, 96, 97, 98])
99   array([90, 91, 92, 93, 94, 95, 96, 97, 98, 99])

debug = True  # this may help to understand what is going on
@vaex.register_function()
def my_func(chunk):
    if debug and len(chunk) != 1:  # 1 means probably a test call of vaex to determine the dtype
        print(chunk.shape, chunk)
    results = []
    for window in chunk:
        # each window is an array of length 10
        # do you linear regression here
        # as placeholer i calculate two statistics
        result1 = window.mean()
        result2 = window.sum()
        if debug and len(chunk) != 1:
            print(result1)
        results.append([result1, result2])
    return np.array(results)
df2.x.my_func()
==== output ====
(5, 10) [[0 0 0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0 0 1]
 [0 0 0 0 0 0 0 0 1 2]
 [0 0 0 0 0 0 0 1 2 3]
 [0 0 0 0 0 0 1 2 3 4]]
0.0
0.1
0.3
0.6
1.0
(5, 10) [[86 87 88 89 90 91 92 93 94 95]
 [87 88 89 90 91 92 93 94 95 96]
 [88 89 90 91 92 93 94 95 96 97]
 [89 90 91 92 93 94 95 96 97 98]
 [90 91 92 93 94 95 96 97 98 99]]
90.5
91.5
92.5
93.5
94.5
Out[19]:
Expression = my_func(x)
Length: 100 dtype: float64 (expression)
---------------------------------------
 0        array([0., 0.])
 1      array([0.1, 1. ])
 2      array([0.3, 3. ])
 3      array([0.6, 6. ])
 4      array([ 1., 10.])
           ...           
95  array([ 90.5, 905. ])
96  array([ 91.5, 915. ])
97  array([ 92.5, 925. ])
98  array([ 93.5, 935. ])
99  array([ 94.5, 945. ])

Working with it

df2['y'] = df2.x.my_func()
df2['y'][:,0].mean()
==== output ====
array(45.165)