使用两个数据框的熊猫矢量化函数

Question

我有以下操作：

import pandas as pd
import numpy as np

def some_calc(x,y):
    x = x.set_index('Cat')
    y = y.set_index('Cat')
    y = np.sqrt(y['data_point2'])
    vec = pd.DataFrame(x['data_point1'] * y)
    grid = np.random.rand(len(x),len(x))
    result = vec.dot(vec.T).mul(grid).sum().sum()
    return result

sample_size = 100
cats = ['a','b','c','d']

df1 = pd.DataFrame({'Cat':[cats[np.random.randint(4)] for _ in range(sample_size)],
                    'data_point1':np.random.rand(sample_size),
                    'data_point2':np.random.rand(sample_size)})

df2 = df1.groupby('Cat').sum().reset_index()

我想使用来自df1 的相关数据点跨df2 行的每一行运行some_calc。

下面的代码运行良好：

df2['Apply'] = df2.apply(lambda x: some_calc(x=df1[df1['Cat']==x['Cat']][['Cat','data_point1']], 
                                             y=df1[df1['Cat']==x['Cat']][['Cat','data_point2']]),axis=1)

（我重置了df2 中的索引，因为我不知道如何跨索引应用。 另外，我将Cat 作为索引字段和data_point 作为向量传递给some_calc，因为没有索引v.dot(v.T) 会将点积压缩成一个数字。 .mul() 出现此错误，因为我需要完整的 MxM 矩阵而不是浮点值。我可能在这里做错了什么......）

我目前正在探索如何对上述内容进行矢量化处理，以便当sample_size 增长时，我不会因计算速度变慢而受到阻碍。

我看到在以前的线程中，您可以切换raw=True，以便输入处理np.array，而不是pd.Series。

df2['ApplyRaw'] = df2.apply(lambda x: some_calc(x=df1[df1['Cat']==x['Cat']][['Cat','data_point1']], 
                                                y=df1[df1['Cat']==x['Cat']]['Cat','data_point2']),axis=1, raw=True)

但是，它会引发错误：

IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

我尝试从参数中省略 Cat，但仍然是同样的问题。

我可以使用任何代码改进或技巧来对上述内容进行矢量化吗？ 还是我必须修改some_calc？

Answer 1

我不确定是否可以对您的函数进行矢量化，因为它有点复杂。但是，some_calc 本身以及它的调用方式可以进行优化。

什么

df2['Apply'] = df2.apply(lambda x: some_calc(x=df1[df1['Cat']==x['Cat']][['Cat','data_point1']], 
                                             y=df1[df1['Cat']==x['Cat']][['Cat','data_point2']]),axis=1)

does 与groupby 基本相同。因此，不要创建这些组并对它们应用函数，而是使用groupby + apply。同样简化some_calc 函数，我们得到：

def some_calc(df):
    x = df['data_point1'].values
    y = np.sqrt(df['data_point2'].values)
    vec = (x * y).reshape(-1, 1)
    grid = np.random.rand(len(x),len(x))
    result = (vec @ vec.T * grid).sum().sum()
    return result

apply = df1.groupby('Cat').apply(some_calc)
apply.name = 'Apply'
df2.merge(apply, left_on='Cat', right_index=True)

最后的合并只是将结果添加到df2数据框。

时间安排：

# original
20.5 ms ± 2.26 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# above code
3.62 ms ± 668 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)