如何使用 python 极坐标计算群体的相关性而不贬低？

Question

我正在使用 Python Polars，并且对编写自定义函数感到困惑。

我可以使用下面的代码按组计算皮森相关性。

df.group_by(["UNIVERSE", "datetime"]).agg(corr_xy = pl.corr("ratio_wsm", "y1d_nn_r"))

但是，我想使用不带

demean

的自定义相关函数来替换

pl.corr

。

计算逻辑如下所示。

def rcor(x, y, w=None) -> float:
    if w is not None:
        sxx = np.sum(w*x*x)
        syy = np.sum(w*y*y)
        sxy = np.sum(w*x*y)
        _rcor = sxy / np.sqrt(sxx * syy)
    else:
        sxx = np.sum(x*x)
        syy = np.sum(y*y)
        sxy = np.sum(x*y)
        _rcor = sxy / np.sqrt(sxx * syy)
    return _rcor

如何使用Python Polars 实现这一点。我真的很困惑

map_batches

，

map_elements

，

polars.api.register_expr_namespace

。

Answer 1

您只需稍微重写您的函数即可使用 Polars 表达式：

def rcor(x, y, w=None):
    if w is not None:
        sxx = (w*x*x).sum()
        syy = (w*y*y).sum()
        sxy = (w*x*y).sum()
        r = sxy / (sxx * syy).sqrt()
    else:
        sxx = (x*x).sum()
        syy = (y*y).sum()
        sxy = (x*y).sum()
        r = sxy / (sxx * syy).sqrt()
    return _rcor

然后使用就简单了：

out = (
    df.group_by(["UNIVERSE", "datetime"])
      .agg(corr_xy = rcor(pl.col.ratio_wsm, pl.col.y1d_nn_r))
)