Pandas 矢量化根查找

Question

本质上，我正在尝试在 pandas 数据帧上使用 brentq 方法 row-wise 来获取将列和常量作为参数的函数的根以及。

类似下面的内容

import pandas as pd
import numpy as np

from scipy.optimize import brentq

np.random.seed(0)
df = pd.DataFrame(np.random.rand(10, 3), columns=list('ABC'))

CONST = 0.5

def fcn(x, y, k):
    return x**2 + y - k

def objective_function(x, y, k, delta):
    return fcn(x, y, k) - delta

def get_root(x, k, delta, a=-10.0, b=10.0, xtol=1e-6):

    # avoid mirroring outer scope
    _x, _k, _delta = x, k, delta

    # nested function that takes the target param as the input
    def nfcn(y):

        # get y that makes fcn and delta equal
        return objective_function(_x, y, _k, _delta)

    result = brentq(nfcn, a=a, b=b, xtol=xtol)

    return result

我在行

上使用了 apply 和 lambda 函数

df['D'] = df.apply(lambda x: get_root(x['A'], CONST, x['C'], a=-10., b=10.), axis=1)

但是当 dataframe 真的很大时，正如预期的那样非常慢。

关于如何对其进行矢量化有什么想法吗？

非常感谢

Answer 1

也许你可以使用部分。对于 10 万行的数据帧，我的速度几乎提高了 4 倍。

from functools import partial

# Rearrange parameters (swap `delta` and constant `k`).
def get_root(x, delta, k, a=-10.0, b=10.0, xtol=1e-6):
    ...

p = partial(get_root, k=CONST, a=-10., b=10.)
df['D'] = [p(*vals) for vals in df[['A', 'C']].values]

Pandas 矢量化根查找

Pandas vectorized root finding

python

vectorization

apply

pandas

scipy-optimize