Pandas apply function with additional arguments-CodePudding

I'm trying to use a function "multiply" to create a new column in a dataframe, and I'm using the apply() method to do it. the code currently looks like this:

import pandas as pd

var_a = 10
var_b = 20

def multiply(row):
    if 0.1 in row['Alpha 1']:
        result = row['Alpha 2'] * var_a
        return result
    if 0.12 in row['Alpha 1']:
        result = row['Alpha 2'] * var_b
        return result

data = {
        'Stage Loading':[[0.1], [0.2]],
        'Alpha 1':[[0.1,0.12],[0.2,0.22]],
        'Alpha 2':[[0.1,0.12],[0.2,0.22]]
       }

pdf = pd.DataFrame(data)

pdf['Calc'] = pdf.apply(multiply, axis = 1)

print(pdf)

this works fine, but I want to be able to put the multiply function into another file and pass the variables (var_a, var_b) through as arguments. I have split into "main.py" and "funct.py" as below:

funct.py:

def multiply(row, a, b):
    if 0.1 in row['Alpha 1']:
        result = row['Alpha 2'] * a
        return result
    if 0.2 in row['Alpha 1']:
        result = row['Alpha 2'] * b
        return result

main.py:

import pandas as pd
from funct import multiply

var_a = 10
var_b = 20

data = {
        'Stage Loading':[[0.1], [0.2]],
        'Alpha 1':[[0.1,0.12],[0.2,0.22]],
        'Alpha 2':[[0.1,0.12],[0.2,0.22]]
       }

pdf = pd.DataFrame(data)

pdf['Calc'] = pdf.apply(multiply(pdf.index, var_a, var_b), axis = 1)

print(pdf)

I know index isnt the right way to do this but I cant think of any way to get this to work.

CodePudding user response：

You can use a lambda function or use args argument

pdf['Calc'] = pdf.apply(lambda row: multiply(row, var_a, var_b), axis = 1)

# or

pdf['Calc'] = pdf.apply(multiply, axis = 1, args=(var_a, var_b))