azmyrajab / polars_ols

Polars least squares extension - enables fast linear model polar expressions
MIT License
113 stars 10 forks source link

input expr support wildcard #8

Closed wukan1986 closed 7 months ago

wukan1986 commented 7 months ago

Use Case

import polars as pl
import polars_ols as pls  # noqa
from polars_ols.least_squares import OLSKwargs, RollingKwargs

df = pl.DataFrame({
    "A": [9, 10, 11, 12],
    "B": [1, 2, 3, 4],
}).with_row_index()

df = df.with_columns(df.to_dummies('B'))
df = df.with_columns(pl.col('A').rolling_mean(3).alias('C'))
df = df.with_columns(pls.compute_least_squares(pl.col('A'),
                                               pl.col('B_1'), pl.col('B_2'), pl.col('B_3'), pl.col('B_4'), pl.col('C'),
                                               mode='residuals', ols_kwargs=OLSKwargs(null_policy='drop', solve_method='svd')).alias('resid1'))
df = df.with_columns(pls.compute_least_squares(pl.col('A'),
                                               pl.col(r"^B_\d+$"), pl.col('C'),
                                               mode='residuals', ols_kwargs=OLSKwargs(null_policy='drop', solve_method='svd')).alias('resid2'))
print(df)
"""
shape: (4, 10)
┌───────┬─────┬─────┬─────┬───┬─────┬──────┬─────────────┬─────────────┐
│ index ┆ A   ┆ B   ┆ B_1 ┆ … ┆ B_4 ┆ C    ┆ resid1      ┆ resid2      │
│ ---   ┆ --- ┆ --- ┆ --- ┆   ┆ --- ┆ ---  ┆ ---         ┆ ---         │
│ u32   ┆ i64 ┆ i64 ┆ u8  ┆   ┆ u8  ┆ f64  ┆ f64         ┆ f64         │
╞═══════╪═════╪═════╪═════╪═══╪═════╪══════╪═════════════╪═════════════╡
│ 0     ┆ 9   ┆ 1   ┆ 1   ┆ … ┆ 0   ┆ null ┆ null        ┆ null        │
│ 1     ┆ 10  ┆ 2   ┆ 0   ┆ … ┆ 0   ┆ null ┆ null        ┆ null        │
│ 2     ┆ 11  ┆ 3   ┆ 0   ┆ … ┆ 0   ┆ 10.0 ┆ -1.0658e-14 ┆ -1.0658e-14 │
│ 3     ┆ 12  ┆ 4   ┆ 0   ┆ … ┆ 1   ┆ 11.0 ┆ -8.8818e-15 ┆ -8.8818e-15 │
└───────┴─────┴─────┴─────┴───┴─────┴──────┴─────────────┴─────────────┘
"""
azmyrajab commented 7 months ago

Thank you !