pandas numpy 如何简化多个矢量化函数参数

Question

我有一个 df，你可以通过运行以下代码来获得它：

import numpy as np
import pandas as pd
from io import StringIO
dfs = """
    M0     M1   M2  M3 M4   M5 age
1   1      2    3    4  5    6  3.2        
2   7      5    4    5  8    3  4.5
3   4      8    9    3  5    2  6.7
"""
df = pd.read_csv(StringIO(dfs.strip()), sep='\s+', )

并且基于业务逻辑我有以下功能，输出也是我所期望的：

def func(M0,M1,M2,M3,M4,M5,age):
    newcol=np.prod([M0,M1,M2,M3,M4,M5][0:age])
    return newcol

vfunc = np.frompyfunc(func, 7, 1)
df['newcol']=vfunc(df['M0'].values,df['M1'].values,df['M2'].values,df['M3'].values,df['M4'].values,df['M5'].values,df['age'].values.astype(int))

df

输出是：

M0  M1  M2  M3  M4  M5  age newcol
1   1   2   3   4   5   6   3.2 6
2   7   5   4   5   8   3   4.5 700
3   4   8   9   3   5   2   6.7 8640

问题是这里 def func(M0,M1,M2,M3,M4,M5,age) 中的参数太多了，无论如何我可以将这些参数列一个列表或其他东西，以使函数更干净?

我试过了：

def func(df):
    newcol=np.prod
    ([df['M0'].values,df['M1'].values,df['M2'].values,df['M3'].values,df['M4'].values,df['M5'].values][0:df['age'].values.astype(int)])
    return newcol

vfunc = np.frompyfunc(func,1, 1)
df['newcol']=vfunc(df)

错误：

TypeError: 'float' object is not subscriptable

请注意我不使用 pd.apply() 的原因，是因为在我的实际业务中数据非常大，而 pd.apply() 运行非常缓慢。

Answer 1

这并没有优化，但至少在选择M列方面可以更具可读性，虽然它有一个额外的功能：

M=["M0","M1","M2","M3","M4","M5"]

def func2(df, M):
    return [df[i].values for i in M] 

def func(age,*Ms):
    newcol=np.prod(Ms[0:age])
    return newcol

vfunc = np.frompyfunc(func, len(M)+1, 1)

df['newcol']=vfunc(df['age'].values.astype(int), *func2(df,M))

df