CODEWITHSUNDEEP
pythonpandasvectorizationapplyscipy-optimize
AleVis
AleVis

Reputation: 187

Pandas vectorized root finding

Essentially I'm trying to use the brentq method row-wise on a pandas dataframe to get the root of a function that takes as arguments columns and constants as well.

Something like the following

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

I got this to work with apply and lambda function, on rows

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

but when the dataframe is really big is very slow, as expected.

Any ideas on how to vectorize this?

Many thanks

Upvotes: 0

Views: 448

Answers (1)

Alexander
Alexander

Reputation: 109546

Perhaps you can use a partial. I get a speed-up of almost 4x for a dataframe of 100k rows.

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]

Upvotes: 1

Related Questions