熊猫矢量化寻根

时间:2019-11-12 18:38:00

标签: python pandas vectorization apply scipy-optimize

本质上,我试图在熊猫数据帧上使用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)

但是,如预期的那样,当数据帧确实很大时,速度非常慢。

关于如何将其向量化的任何想法?

非常感谢

1 个答案:

答案 0 :(得分: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]