How do I include multiple parameters in rolling_apply using Polars Python?

Question

For Polars, I can see that we can apply with multiple column values using pl.struct here, but can't figure out whether you can do the same thing but using rolling_apply.

I am looking to see whether I can performing a rolling regression, or a rolling correlation.

df = pl.DataFrame(
    {
        "time": pl.date_range(
            low=datetime(2021, 1, 1),
            high=datetime(2022, 1, 1),
            interval="1d",
        ),
        "x": pl.arange(0, 366, eager=True).shuffle(seed=1),
        "y": pl.arange(0, 366, eager=True).shuffle(seed=2)
    }
)

If I did the following, I would get a single value

df.select(
     pl.spearman_rank_corr("x", "y").alias("corr")
)

Is it possible to get it on a rolling basis? Thanks

Edit:

I have tried to use the answer suggested below to create rolling regression:

def get_regression_parameter(args: List[pl.Series]) -> pl.Series:
    return pl.Series([sm.OLS(args[0], args[1]).fit().params[0]], dtype=pl.Float64)

df.groupby_rolling(
    index_column='time',
    period='14d',
).agg([
        pl.col('x').alias('x_list'),
        pl.col('y').alias('y_list'),
        pl.apply(
            exprs=["y", "x"], 
            f=get_regression_parameter).alias("corr")
         
])

but I am receiving the following error:

thread '' panicked at 'python apply failed: unrecognized data structures:  / ', src/lazy/apply.rs:200:19

My workaround for this is:

df.groupby_rolling(
    index_column='time',
    period='14d',
).agg([
        pl.col('x').alias('x_list'),
        pl.col('y').alias('y_list'),
      ]
).with_columns(
    pl.struct(['x_list', 'y_list']).apply(lambda x: sm.OLS(x['y_list'], x['x_list']).fit().params[0]).alias('coefficient')
)

But it is outside the .agg function (which is fine - I just wanted to ask what would be considered the best way to perform this).

user_18559875 · Answer

I recommend using groupby_rolling.

df.groupby_rolling(
    index_column='time',
    period='14d',
).agg([
    pl.spearman_rank_corr('x', 'y').alias('sp_rank'),
])

shape: (366, 2)
┌─────────────────────┬───────────┐
│ time                ┆ sp_rank   │
│ ---                 ┆ ---       │
│ datetime[ns]        ┆ f64       │
╞═════════════════════╪═══════════╡
│ 2021-01-01 00:00:00 ┆ NaN       │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-01-02 00:00:00 ┆ -1.0      │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-01-03 00:00:00 ┆ -1.0      │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-01-04 00:00:00 ┆ -0.4      │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ ...                 ┆ ...       │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-12-29 00:00:00 ┆ -0.112088 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-12-30 00:00:00 ┆ -0.257143 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-12-31 00:00:00 ┆ -0.059341 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2022-01-01 00:00:00 ┆ 0.032967  │
└─────────────────────┴───────────┘

One thing I find helpful when first setting up a groupby_rolling is to produce a list of values in each period:

df.groupby_rolling(
    index_column='time',
    period='14d',
).agg([
    pl.col('x').alias('x_list'),
    pl.col('y').alias('y_list'),
    pl.spearman_rank_corr('x', 'y').alias('sp_rank'),
])

shape: (366, 4)
┌─────────────────────┬─────────────────────┬─────────────────────┬───────────┐
│ time                ┆ x_list              ┆ y_list              ┆ sp_rank   │
│ ---                 ┆ ---                 ┆ ---                 ┆ ---       │
│ datetime[ns]        ┆ list[i64]           ┆ list[i64]           ┆ f64       │
╞═════════════════════╪═════════════════════╪═════════════════════╪═══════════╡
│ 2021-01-01 00:00:00 ┆ [245]               ┆ [246]               ┆ NaN       │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-01-02 00:00:00 ┆ [245, 334]          ┆ [246, 128]          ┆ -1.0      │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-01-03 00:00:00 ┆ [245, 334, 74]      ┆ [246, 128, 295]     ┆ -1.0      │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-01-04 00:00:00 ┆ [245, 334, ... 52]  ┆ [246, 128, ... 150] ┆ -0.4      │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ ...                 ┆ ...                 ┆ ...                 ┆ ...       │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-12-29 00:00:00 ┆ [285, 81, ... 124]  ┆ [186, 331, ... 298] ┆ -0.112088 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-12-30 00:00:00 ┆ [81, 174, ... 66]   ┆ [331, 2, ... 214]   ┆ -0.257143 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-12-31 00:00:00 ┆ [174, 221, ... 14]  ┆ [2, 208, ... 82]    ┆ -0.059341 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2022-01-01 00:00:00 ┆ [221, 276, ... 261] ┆ [208, 40, ... 275]  ┆ 0.032967  │
└─────────────────────┴─────────────────────┴─────────────────────┴───────────┘

I typically only do this for a small subset at first, just to ensure that the results are correct. (For the final run, you won't want to keep the list columns.)

Edit: using `apply`

We can use the apply method in the agg method. Let's perform an OLS using the statsmodel library.

First, let's expand your data so that we can show how to incorporate multiple independent variables.

from datetime import datetime

import polars as pl
import statsmodels.api as sm
import numpy as np

df = pl.DataFrame(
    {
        "time": pl.date_range(
            low=datetime(2021, 1, 1),
            high=datetime(2022, 1, 1),
            interval="1d",
        ),
        "y": pl.arange(0, 366, eager=True).shuffle(seed=4),
        "x1": pl.arange(0, 366, eager=True).shuffle(seed=1),
        "x2": pl.arange(0, 366, eager=True).shuffle(seed=2),
        "x3": pl.arange(0, 366, eager=True).shuffle(seed=3),
    }
)
df

shape: (366, 5)
┌─────────────────────┬─────┬─────┬─────┬─────┐
│ time                ┆ y   ┆ x1  ┆ x2  ┆ x3  │
│ ---                 ┆ --- ┆ --- ┆ --- ┆ --- │
│ datetime[ns]        ┆ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═════════════════════╪═════╪═════╪═════╪═════╡
│ 2021-01-01 00:00:00 ┆ 347 ┆ 245 ┆ 246 ┆ 90  │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 2021-01-02 00:00:00 ┆ 147 ┆ 334 ┆ 128 ┆ 38  │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 2021-01-03 00:00:00 ┆ 85  ┆ 74  ┆ 295 ┆ 24  │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 2021-01-04 00:00:00 ┆ 36  ┆ 52  ┆ 150 ┆ 297 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ ...                 ┆ ... ┆ ... ┆ ... ┆ ... │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 2021-12-29 00:00:00 ┆ 268 ┆ 124 ┆ 298 ┆ 246 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 2021-12-30 00:00:00 ┆ 275 ┆ 66  ┆ 214 ┆ 58  │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 2021-12-31 00:00:00 ┆ 69  ┆ 14  ┆ 82  ┆ 150 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 2022-01-01 00:00:00 ┆ 233 ┆ 261 ┆ 275 ┆ 181 │
└─────────────────────┴─────┴─────┴─────┴─────┘

Rather than use a lambda function, I prefer to create a separate function where I can keep my calculation choices all in one place.

def regress_it(s_list: list[pl.Series]) -> pl.Series:
    np_array = np.column_stack(s_list)
    y = np_array[:, 0]
    X = np_array[:, 1:]
    X = sm.add_constant(X)
    result = sm.OLS(
        endog=y,
        exog=X,
    ).fit().params[0]
    return result

df.groupby_rolling(index_column="time", period="14d").agg(
    [
        pl.apply(
            exprs=["y", "x1", "x2", "x3"],
            f=regress_it)
        .alias('coeff')
    ]
)

The apply method will pass a list of Series to the called function (regress_it in this case).

In the regress_it function, we'll first take this list of Series and column-stack them into a numpy array. We can then easily pass slices of the array to the OLS function.

By the way we ordered the parameters in the apply call, the first column is our dependent variable y (the endog parameter in sm.OLS). The remaining columns (x1, x2, x3) will be our independent variables (the exog parameter in sm.OLS).

I've also added a constant for the regression. Here's the output.

shape: (366, 2)
┌─────────────────────┬─────────────┐
│ time                ┆ coeff       │
│ ---                 ┆ ---         │
│ datetime[ns]        ┆ f64         │
╞═════════════════════╪═════════════╡
│ 2021-01-01 00:00:00 ┆ 0.6608      │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-01-02 00:00:00 ┆ 0.002035    │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-01-03 00:00:00 ┆ -0.013203   │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-01-04 00:00:00 ┆ -798.950143 │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ ...                 ┆ ...         │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-12-29 00:00:00 ┆ 275.147184  │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-12-30 00:00:00 ┆ 285.589173  │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2021-12-31 00:00:00 ┆ 232.471054  │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2022-01-01 00:00:00 ┆ 231.976641  │
└─────────────────────┴─────────────┘

How do I include multiple parameters in rolling_apply using Polars Python?

1 Answers

Edit: using apply

Edit: using `apply`